Add pgtle.set_extension_schema() to set an extension's schema

Makefile

@@ -1,5 +1,5 @@
 EXTENSION = pg_tle
-EXTVERSION = 1.5.2
+EXTVERSION = 1.5.3
 
 SCHEMA = pgtle
 MODULE_big = $(EXTENSION)
@@ -9,7 +9,8 @@ OBJS = src/tleextension.o src/guc-file.o src/feature.o src/passcheck.o src/uni_a
 EXTRA_CLEAN	= src/guc-file.c pg_tle.control pg_tle--$(EXTVERSION).sql
 DATA = pg_tle.control pg_tle--1.0.0.sql pg_tle--1.0.0--1.0.1.sql pg_tle--1.0.1--1.0.4.sql pg_tle--1.0.4.sql pg_tle--1.0.4--1.1.1.sql \
 			 pg_tle--1.1.0--1.1.1.sql pg_tle--1.1.1.sql pg_tle--1.1.1--1.2.0.sql pg_tle--1.2.0--1.3.0.sql pg_tle--1.3.0--1.3.3.sql \
-			 pg_tle--1.3.3--1.3.4.sql pg_tle--1.3.4--1.4.0.sql pg_tle--1.4.0--1.5.0.sql pg_tle--1.5.0--1.5.2.sql
+			 pg_tle--1.3.3--1.3.4.sql pg_tle--1.3.4--1.4.0.sql pg_tle--1.4.0--1.5.0.sql pg_tle--1.5.0--1.5.2.sql \
+			 pg_tle--1.5.2--1.5.3.sql
 
 TESTS = $(wildcard test/sql/*.sql)
 REGRESS = $(patsubst test/sql/%.sql,%,$(TESTS))

docs/03_managing_extensions.md

@@ -278,6 +278,33 @@ This functions returns `true` on success.
 * `name`: The name of the extension. This is the value used when calling `CREATE EXTENSION`.
 * `version`: The version of the extension to set the default.
 
+### `pgtle.set_extension_schema(name text, schema text)`
+
+`set_extension_schema` sets, changes, or clears the schema recorded for an installed extension. The schema is the value that PostgreSQL uses to place the extension's objects when `CREATE EXTENSION` is called. This is helpful for pinning an extension that was installed without a schema, for example before schema support was added in `pg_tle` 1.5.0, to a specific schema without uninstalling and reinstalling it (which would drop any data in its tables).
+
+Pass `NULL` as `schema` to clear the recorded schema, reverting the extension to having no fixed schema.
+
+This only affects future `CREATE EXTENSION` calls (including restoring from a `pg_dump`). It does not move an extension that is already created; use `ALTER EXTENSION ... SET SCHEMA` for that.
+
+If the extension in `name` does not already exist, this returns an error.
+
+This function returns `true` on success.
+
+#### Role
+
+`pgtle_admin`
+
+#### Arguments
+
+* `name`: The name of the extension. This is the value used when calling `CREATE EXTENSION`.
+* `schema`: The schema in which the extension's objects are created, or `NULL` to clear the recorded schema.
+
+#### Example
+
+```sql
+SELECT pgtle.set_extension_schema('pg_tle_test', 'tle_schema');
+```
+
 ### `pgtle.uninstall_extension(extname text)`
 
 `uninstall_extension` removes all versions of an extension from a database. This prevents future calls of `CREATE EXTENSION` from installing the extension. If the extension does not exist in the database, then an error is raised.

include/tleextension.h

@@ -25,6 +25,14 @@
 #define PG_TLE_INNER_STR			"$_pgtle_i_$"
 #define PG_TLE_ADMIN				"pgtle_admin"
 
+/*
+ * Characters that cannot be quoted consistently both inside and outside of
+ * string literals.  Identifiers substituted into extension scripts, or stored
+ * where they will be substituted later, are rejected if they contain any of
+ * these.
+ */
+#define PG_TLE_QUOTING_RELEVANT_CHARS	"\"$'\\"
+
 /*
  * creating_extension is only true while running a CREATE EXTENSION or ALTER
  * EXTENSION UPDATE command.  It instructs recordDependencyOnCurrentExtension()

pg_tle--1.5.2--1.5.3.sql

@@ -0,0 +1,40 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+-- complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "CREATE EXTENSION pg_tle" to load this file. \quit
+
+CREATE FUNCTION pgtle.set_extension_schema
+(
+  name text,
+  schema text
+)
+RETURNS boolean
+SET search_path TO 'pgtle'
+AS 'MODULE_PATHNAME', 'pg_tle_set_extension_schema'
+LANGUAGE C;
+
+REVOKE EXECUTE ON FUNCTION pgtle.set_extension_schema
+(
+  name text,
+  schema text
+) FROM PUBLIC;
+
+GRANT EXECUTE ON FUNCTION pgtle.set_extension_schema
+(
+  name text,
+  schema text
+) TO pgtle_admin;

src/tleextension.c

@@ -1411,16 +1411,6 @@ execute_extension_script(Oid extensionOid, ExtensionControlFile *control,
 		char	   *c_sql = read_extension_script_file(control, filename);
 		Datum		t_sql;
 
-		/*
-		 * We filter each substitution through quote_identifier().  When the
-		 * arg contains one of the following characters, no one collection of
-		 * quoting can work inside $$dollar-quoted string literals$$,
-		 * 'single-quoted string literals', and outside of any literal.  To
-		 * avoid a security snare for extension authors, error on substitution
-		 * for arguments containing these.
-		 */
-		const char *quoting_relevant_chars = "\"$'\\";
-
 		/* We use various functions that want to operate on text datums */
 		t_sql = CStringGetTextDatum(c_sql);
 
@@ -1450,11 +1440,20 @@ execute_extension_script(Oid extensionOid, ExtensionControlFile *control,
 											t_sql,
 											CStringGetTextDatum("@extowner@"),
 											CStringGetTextDatum(qUserName));
-			if (strpbrk(userName, quoting_relevant_chars))
+
+			/*
+			 * We filter each substitution through quote_identifier().  When
+			 * the arg contains one of these characters, no one collection of
+			 * quoting can work inside $$dollar-quoted string literals$$,
+			 * 'single-quoted string literals', and outside of any literal. To
+			 * avoid a security snare for extension authors, error on
+			 * substitution for arguments containing these.
+			 */
+			if (strpbrk(userName, PG_TLE_QUOTING_RELEVANT_CHARS))
 				ereport(ERROR,
 						(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
 						 errmsg("invalid character in extension owner: must not contain any of \"%s\"",
-								quoting_relevant_chars)));
+								PG_TLE_QUOTING_RELEVANT_CHARS)));
 		}
 
 		/*
@@ -1474,11 +1473,11 @@ execute_extension_script(Oid extensionOid, ExtensionControlFile *control,
 											t_sql,
 											CStringGetTextDatum("@extschema@"),
 											CStringGetTextDatum(qSchemaName));
-			if (t_sql != old && strpbrk(schemaName, quoting_relevant_chars))
+			if (t_sql != old && strpbrk(schemaName, PG_TLE_QUOTING_RELEVANT_CHARS))
 				ereport(ERROR,
 						(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
 						 errmsg("invalid character in extension \"%s\" schema: must not contain any of \"%s\"",
-								control->name, quoting_relevant_chars)));
+								control->name, PG_TLE_QUOTING_RELEVANT_CHARS)));
 		}
 
 		/*
@@ -5229,6 +5228,139 @@ pg_tle_set_default_version(PG_FUNCTION_ARGS)
 	PG_RETURN_BOOL(true);
 }
 
+Datum		pg_tle_set_extension_schema(PG_FUNCTION_ARGS);
+
+/*
+ * Set (or clear) the target schema recorded in an extension's control
+ * function.  This rewrites just the control function in place, so future
+ * CREATE EXTENSION (or a fresh database restored from pg_dump) honors the new
+ * schema while an already-created extension is left untouched.  A NULL schema
+ * clears it, leaving the extension with no fixed schema.
+ */
+PG_FUNCTION_INFO_V1(pg_tle_set_extension_schema);
+Datum
+pg_tle_set_extension_schema(PG_FUNCTION_ARGS)
+{
+	int			spi_rc;
+	char	   *extname;
+	char	   *schemaName = NULL;
+	char	   *ctlname;
+	Oid			ctlfuncid;
+	StringInfo	ctlstr;
+	char	   *ctlsql;
+	ExtensionControlFile *control;
+	char	   *filename;
+
+	if (PG_ARGISNULL(0))
+		ereport(ERROR,
+				(errcode(ERRCODE_NULL_VALUE_NOT_ALLOWED),
+				 errmsg("\"name\" is a required argument.")));
+
+	extname = text_to_cstring(PG_GETARG_TEXT_PP(0));
+	check_valid_extension_name(extname);
+
+	/*
+	 * Verify that extname does not already exist as a standard file-based
+	 * extension.
+	 */
+	filename = get_extension_control_filename(extname);
+	if (filestat(filename))
+		ereport(ERROR,
+				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+				 errmsg("control file already exists for the %s extension", extname)));
+
+	/*
+	 * A NULL schema clears the recorded schema; otherwise validate it.  The
+	 * schema name is embedded verbatim into the control function and later
+	 * substituted for @extschema@, so reject characters that cannot be quoted
+	 * consistently both inside and outside of string literals.
+	 */
+	if (!PG_ARGISNULL(1))
+	{
+		schemaName = text_to_cstring(PG_GETARG_TEXT_PP(1));
+
+		/*
+		 * Reject the empty string rather than record schema = '', which would
+		 * pin the extension to an unusable zero-length schema.  Use a NULL
+		 * argument to clear the schema instead.
+		 */
+		if (schemaName[0] == '\0')
+			ereport(ERROR,
+					(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+					 errmsg("extension schema must not be empty"),
+					 errhint("Pass NULL as the schema to clear it.")));
+
+		if (strpbrk(schemaName, PG_TLE_QUOTING_RELEVANT_CHARS))
+			ereport(ERROR,
+					(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
+					 errmsg("invalid character in extension schema: must not contain any of \"%s\"",
+							PG_TLE_QUOTING_RELEVANT_CHARS)));
+	}
+
+	/*
+	 * Verify that the extension exists as a TLE extension by looking up its
+	 * control function.
+	 */
+	ctlname = psprintf("%s.control", extname);
+	ctlfuncid = get_tlefunc_oid_if_exists(ctlname);
+	if (ctlfuncid == InvalidOid)
+		ereport(ERROR,
+				(errcode(ERRCODE_UNDEFINED_OBJECT),
+				 errmsg("extension \"%s\" does not exist", extname),
+				 errhint("Try installing the extension with \"%s.install_extension\".", PG_TLE_NSPNAME)));
+
+	/*
+	 * Load the current control state, then replace the schema.
+	 */
+	control = build_default_extension_control_file(extname);
+
+	SET_TLEEXT;
+	parse_extension_control_file(control, NULL);
+	UNSET_TLEEXT;
+
+	control->schema = schemaName;
+
+	ctlstr = build_extension_control_file_string(control);
+
+	/*
+	 * Validate that there are no injections using the dollar-quoted strings
+	 */
+	if (!(validate_tle_sql(ctlstr->data)))
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				 errmsg("invalid character in extension definition"),
+				 errdetail("Use of string delimiters %s and %s are forbidden in extension definitions.",
+						   PG_TLE_OUTER_STR, PG_TLE_INNER_STR)));
+
+	ctlsql = psprintf(
+					  "CREATE OR REPLACE FUNCTION %s.%s() RETURNS TEXT AS %s"
+					  "SELECT %s%s%s%s LANGUAGE SQL",
+					  quote_identifier(PG_TLE_NSPNAME), quote_identifier(ctlname),
+					  PG_TLE_OUTER_STR, PG_TLE_INNER_STR,
+					  ctlstr->data,
+					  PG_TLE_INNER_STR, PG_TLE_OUTER_STR);
+
+	/* flag that we are manipulating pg_tle artifacts */
+	SET_TLEART;
+
+	if (SPI_connect() != SPI_OK_CONNECT)
+		elog(ERROR, "SPI_connect failed");
+
+	spi_rc = SPI_exec(ctlsql, 0);
+	if (spi_rc != SPI_OK_UTILITY)
+		ereport(ERROR,
+				(errcode(ERRCODE_INTERNAL_ERROR),
+				 errmsg("failed to update schema for \"%s\"", extname)));
+
+	if (SPI_finish() != SPI_OK_FINISH)
+		elog(ERROR, "SPI_finish failed");
+
+	/* flag that we are done manipulating pg_tle artifacts */
+	UNSET_TLEART;
+
+	PG_RETURN_BOOL(true);
+}
+
 /*
  * Convert text array to list of strings.
  *