Create the notion of unstable extensions

As a DBaaS provider, Neon needs to provide a stable platform for
customers to build applications upon. At the same time however, we also
need to enable customers to use the latest and greatest technology, so
they can prototype their work, and we can solicit feedback. If all
extensions are treated the same in terms of stability, it is hard to
meet that goal.

There are now two new GUCs created by the Neon extension:

neon.allow_unstable_extensions: This is a session GUC which allows
a session to install and load unstable extensions.

neon.unstable_extensions: This is a comma-separated list of extension
names. We can check if a CREATE EXTENSION statement is attempting to
install an unstable extension, and if so, deny the request if
neon.allow_unstable_extensions is not set to true.

Signed-off-by: Tristan Partin <tristan@neon.tech>

pgxn/neon/Makefile

@@ -16,6 +16,7 @@ OBJS = \
 	neon_walreader.o \
 	pagestore_smgr.o \
 	relsize_cache.o \
+	unstable_extensions.o \
 	walproposer.o \
 	walproposer_pg.o \
 	control_plane_connector.o \

pgxn/neon/neon.c

@@ -30,6 +30,7 @@
 #include "neon.h"
 #include "control_plane_connector.h"
 #include "logical_replication_monitor.h"
+#include "unstable_extensions.h"
 #include "walsender_hooks.h"
 #if PG_MAJORVERSION_NUM >= 16
 #include "storage/ipc.h"
@@ -424,6 +425,7 @@ _PG_init(void)
 	LogicalFuncs_Custom_XLogReaderRoutines = NeonOnDemandXLogReaderRoutines;
 	SlotFuncs_Custom_XLogReaderRoutines = NeonOnDemandXLogReaderRoutines;
 
+	InitUnstableExtensionsSupport();
 	InitLogicalReplicationMonitor();
 	InitControlPlaneConnector();
 

pgxn/neon/neon_pgversioncompat.c

@@ -42,3 +42,59 @@ InitMaterializedSRF(FunctionCallInfo fcinfo, bits32 flags)
 	MemoryContextSwitchTo(old_context);
 }
 #endif
+
+#if PG_MAJORVERSION_NUM < 16
+/*
+ * Some infrastructure for checking malloc/strdup/realloc calls
+ */
+void *
+guc_malloc(int elevel, size_t size)
+{
+	void	   *data;
+
+	/* Avoid unportable behavior of malloc(0) */
+	if (size == 0)
+		size = 1;
+	data = malloc(size);
+	if (data == NULL)
+		ereport(elevel,
+				(errcode(ERRCODE_OUT_OF_MEMORY),
+				 errmsg("out of memory")));
+	return data;
+}
+
+void *
+guc_realloc(int elevel, void *old, size_t size)
+{
+	void	   *data;
+
+	/* Avoid unportable behavior of realloc(NULL, 0) */
+	if (old == NULL && size == 0)
+		size = 1;
+	data = realloc(old, size);
+	if (data == NULL)
+		ereport(elevel,
+				(errcode(ERRCODE_OUT_OF_MEMORY),
+				 errmsg("out of memory")));
+	return data;
+}
+
+char *
+guc_strdup(int elevel, const char *src)
+{
+	char	   *data;
+
+	data = strdup(src);
+	if (data == NULL)
+		ereport(elevel,
+				(errcode(ERRCODE_OUT_OF_MEMORY),
+				 errmsg("out of memory")));
+	return data;
+}
+
+void
+guc_free(void *ptr)
+{
+	free(ptr);
+}
+#endif

pgxn/neon/neon_pgversioncompat.h

@@ -83,6 +83,11 @@ InitBufferTag(BufferTag *tag, const RelFileNode *rnode,
 
 #define DropRelationAllLocalBuffers DropRelFileNodeAllLocalBuffers
 
+void *guc_malloc(int elevel, size_t size);
+void *guc_realloc(int elevel, void *old, size_t size);
+char *guc_strdup(int elevel, const char *src);
+void guc_free(void *ptr);
+
 #else							/* major version >= 16 */
 
 #define USE_RELFILELOCATOR

pgxn/neon/unstable_extensions.c

@@ -0,0 +1,164 @@
+#include <assert.h>
+#include <stdlib.h>
+#include <string.h>
+
+#include "postgres.h"
+
+#include "nodes/plannodes.h"
+#include "nodes/parsenodes.h"
+#include "tcop/utility.h"
+#include "utils/errcodes.h"
+#include "utils/guc.h"
+
+#include "neon_pgversioncompat.h"
+#include "unstable_extensions.h"
+
+typedef struct UnstableExtensions
+{
+	size_t	extc;
+	char   *extv[FLEXIBLE_ARRAY_MEMBER];
+} UnstableExtensions;
+
+static bool					allow_unstable_extensions = false;
+static char				   *unstable_extensions_str = NULL;
+static UnstableExtensions  *unstable_extensions = NULL;
+
+static ProcessUtility_hook_type PreviousProcessUtilityHook = NULL;
+
+static bool
+check_unstable_extensions(char **newval, void **extra, GucSource source)
+{
+	char			   *n;
+	UnstableExtensions *exts;
+	char			   *curr_ext;
+	size_t				max_extc;
+
+	*extra = NULL;
+
+	if (*newval == NULL || (*newval)[0] == '\0')
+		return true;
+
+	/* At this point, we know we have at least 1 extension like: ext */
+	max_extc = 1;
+	for (size_t i = 0; i < strlen(*newval); ++i)
+	{
+		if ((*newval)[i] == ',')
+			++max_extc;
+	}
+
+	exts = guc_malloc(ERROR, sizeof(*exts) + max_extc * sizeof(char *));
+	exts->extc = 0;
+
+	n = guc_strdup(ERROR, *newval);
+	while ((curr_ext = strsep(&n, ",")))
+	{
+		/* In the event, we receive a config like: ",,ext" */
+		if (curr_ext[0] == '\0')
+			continue;
+
+		exts->extv[exts->extc++] = guc_strdup(ERROR, curr_ext);
+	}
+
+	guc_free(n);
+
+	*extra = exts;
+
+	return true;
+}
+
+static void
+assign_unstable_extensions(const char *newval, void *extra)
+{
+	unstable_extensions = extra;
+}
+
+static void
+CheckUnstableExtension(
+	PlannedStmt *pstmt,
+	const char *queryString,
+	bool readOnlyTree,
+	ProcessUtilityContext context,
+	ParamListInfo params,
+	QueryEnvironment *queryEnv,
+	DestReceiver *dest,
+	QueryCompletion *qc)
+{
+	Node	   *parseTree = pstmt->utilityStmt;
+
+	if (allow_unstable_extensions || unstable_extensions == NULL)
+		goto process;
+
+	switch (nodeTag(parseTree))
+	{
+		case T_CreateExtensionStmt:
+		{
+			CreateExtensionStmt *stmt = castNode(CreateExtensionStmt, parseTree);
+
+			for (size_t i = 0; i < unstable_extensions->extc; ++i)
+			{
+				if (strcmp(unstable_extensions->extv[i], stmt->extname) == 0)
+					ereport(ERROR,
+							(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
+							 errmsg("using an unstable extension (%s) is currently prohibited", stmt->extname),
+							 errhint("Set neon.allow_unstable_extensions to true")));
+			}
+
+			break;
+		}
+		default:
+			goto process;
+	}
+
+process:
+	if (PreviousProcessUtilityHook)
+	{
+		PreviousProcessUtilityHook(
+			pstmt,
+			queryString,
+			readOnlyTree,
+			context,
+			params,
+			queryEnv,
+			dest,
+			qc);
+	}
+	else
+	{
+		standard_ProcessUtility(
+			pstmt,
+			queryString,
+			readOnlyTree,
+			context,
+			params,
+			queryEnv,
+			dest,
+			qc);
+	}
+}
+
+void
+InitUnstableExtensionsSupport(void)
+{
+	DefineCustomBoolVariable(
+		"neon.allow_unstable_extensions",
+		"Allow unstable extensions to be installed and used",
+		NULL,
+		&allow_unstable_extensions,
+		false,
+		PGC_USERSET,
+		0,
+		NULL, NULL, NULL);
+
+	DefineCustomStringVariable(
+		"neon.unstable_extensions",
+		"Allow unstable extensions to be installed and used",
+		NULL,
+		&unstable_extensions_str,
+		NULL,
+		PGC_SUSET,
+		0,
+		check_unstable_extensions, assign_unstable_extensions, NULL);
+
+	PreviousProcessUtilityHook = ProcessUtility_hook;
+	ProcessUtility_hook = CheckUnstableExtension;
+}

pgxn/neon/unstable_extensions.h

@@ -0,0 +1,6 @@
+#ifndef __NEON_UNSTABLE_EXTENSIONS_H__
+#define __NEON_UNSTABLE_EXTENSIONS_H__
+
+void InitUnstableExtensionsSupport(void);
+
+#endif

test_runner/regress/test_unstable_extensions.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, cast
+
+import pytest
+from psycopg2.errors import InsufficientPrivilege
+
+if TYPE_CHECKING:
+    from fixtures.neon_fixtures import NeonEnv
+
+
+def test_unstable_extensions_installation(neon_simple_env: NeonEnv):
+    """
+    Test that the unstable extension support within the neon extension can
+    block extension installation.
+    """
+    env = neon_simple_env
+
+    neon_unstable_extensions = "pg_prewarm,amcheck"
+
+    endpoint = env.endpoints.create(
+        "main",
+        config_lines=[
+            "neon.allow_unstable_extensions=false",
+            f"neon.unstable_extensions='{neon_unstable_extensions}'",
+        ],
+    )
+    endpoint.respec(skip_pg_catalog_updates=False)
+    endpoint.start()
+
+    with endpoint.cursor() as cursor:
+        cursor.execute("SELECT current_setting('neon.unstable_extensions')")
+        result = cursor.fetchone()
+        assert result is not None
+        setting = cast("str", result[0])
+        assert setting == neon_unstable_extensions
+
+        with pytest.raises(InsufficientPrivilege):
+            cursor.execute("CREATE EXTENSION pg_prewarm")
+
+        with pytest.raises(InsufficientPrivilege):
+            cursor.execute("CREATE EXTENSION amcheck")
+
+        # Make sure that we can install a "stable" extension
+        cursor.execute("CREATE EXTENSION pageinspect")
+
+        cursor.execute("BEGIN")
+        cursor.execute("SET neon.allow_unstable_extensions TO true")
+        cursor.execute("CREATE EXTENSION pg_prewarm")
+        cursor.execute("COMMIT")