Add chunks to schema-based publications

Chunks live in _timescaledb_internal, so FOR TABLES IN SCHEMA
publications miss them and replication silently skips chunk data.

Backfill chunk membership by reading the schema set straight from the
CREATE/ALTER PUBLICATION parse tree (no before/after snapshots):

- CREATE PUBLICATION is reconciled in the ddl_command_end event trigger
  (the publication Oid exists by then); add CREATE PUBLICATION to the
  timescaledb_ddl_command_end WHEN TAG list.
- ALTER PUBLICATION ADD/DROP/SET TABLES IN SCHEMA is reconciled in
  process_ddl_command_start after prev_ProcessUtility. SET rebuilds the
  full new schema set (PG drops every tracked chunk row on SET); only
  DROP removes them. Options-only ALTER is a no-op.

ALTER TABLE ... SET SCHEMA on a hypertable now reconciles its chunks:
remove from the old schema publications, add to the new ones. Runs
before ts_hypertable_set_schema flips the catalog schema.

Skip adding a redundant explicit pg_publication_rel row when a chunk
already lives in a schema this publication covers natively (chunks
placed via associated_schema_name); the explicit row would survive a
later DROP TABLES IN SCHEMA and orphan. OSM (foreign-table) chunks are
skipped throughout; publication_add_relation would fail on them.

Tests 6b-6i cover CREATE/ADD/SET/DROP/CURRENT_SCHEMA, schema moves,
the redundancy guard (counts of pg_publication_rel rows), options-only
ALTER, and GUC-off on the ALTER/schema-change paths.

Signed-off-by: Arunprasad Rajkumar <ar.arunprasad@gmail.com>

Author: arajkumar

.unreleased/pr_10128

@@ -0,0 +1 @@
+Fixes: #10128 Include hypertable chunks in publications created with FOR TABLES IN SCHEMA

sql/ddl_triggers.sql

@@ -9,7 +9,7 @@ AS '@MODULE_PATHNAME@', 'ts_timescaledb_process_ddl_event' LANGUAGE C;
 
 --EVENT TRIGGER MUST exclude the ALTER EXTENSION tag.
 CREATE EVENT TRIGGER timescaledb_ddl_command_end ON ddl_command_end
-WHEN TAG IN ('ALTER TABLE','CREATE TRIGGER','CREATE TABLE','CREATE INDEX','ALTER INDEX', 'DROP TABLE', 'DROP INDEX', 'DROP SCHEMA')
+WHEN TAG IN ('ALTER TABLE','CREATE TRIGGER','CREATE TABLE','CREATE INDEX','ALTER INDEX', 'DROP TABLE', 'DROP INDEX', 'DROP SCHEMA', 'CREATE PUBLICATION')
 EXECUTE FUNCTION _timescaledb_functions.process_ddl_event();
 
 DROP EVENT TRIGGER IF EXISTS timescaledb_ddl_sql_drop;

src/chunk.c

@@ -11,6 +11,7 @@
 #include <access/tableam.h>
 #include <access/tupdesc.h>
 #include <access/xact.h>
+#include <catalog/dependency.h>
 #include <catalog/indexing.h>
 #include <catalog/namespace.h>
 #include <catalog/pg_class.h>
@@ -1048,14 +1049,170 @@ chunk_add_to_publication(Oid puboid, const Chunk *chunk)
 static void
 chunk_add_to_publications(const Chunk *chunk)
 {
-	List *puboids;
+	Oid ht_nspid = get_rel_namespace(chunk->hypertable_relid);
+	Oid chunk_nspid = get_rel_namespace(chunk->fd.relid);
+	List *puboids = list_concat_unique_oid(GetRelationIncludedPublications(chunk->hypertable_relid),
+										   GetSchemaPublications(ht_nspid));
+	/* If the chunk lives in a published schema, PostgreSQL already covers it
+	 * via the schema publication; skip the redundant explicit row. */
+	List *chunk_schema_pubs = GetSchemaPublications(chunk_nspid);
 	ListCell *lc;
-
-	puboids = GetRelationIncludedPublications(chunk->hypertable_relid);
 	foreach (lc, puboids)
 	{
-		Oid puboid = lfirst_oid(lc);
-		chunk_add_to_publication(puboid, chunk);
+		if (list_member_oid(chunk_schema_pubs, lfirst_oid(lc)))
+		{
+			continue;
+		}
+		chunk_add_to_publication(lfirst_oid(lc), chunk);
+	}
+}
+
+static void
+chunk_remove_from_publication(Oid puboid, const Chunk *chunk)
+{
+	ObjectAddress obj;
+	Oid prid = GetSysCacheOid2(PUBLICATIONRELMAP,
+							   Anum_pg_publication_rel_oid,
+							   ObjectIdGetDatum(chunk->fd.relid),
+							   ObjectIdGetDatum(puboid));
+	if (!OidIsValid(prid))
+	{
+		return;
+	}
+
+	/* No core PG object depends on a pg_publication_rel row, so RESTRICT is
+	 * sufficient and matches PG's own AlterPublication convention. */
+	ObjectAddressSet(obj, PublicationRelRelationId, prid);
+	performDeletion(&obj, DROP_RESTRICT, 0);
+}
+
+/*
+ * Add or remove the chunks of every hypertable whose root table lives in the
+ * given schemas to/from a publication. Schema publications miss chunks
+ * (chunks live in _timescaledb_internal, not the hypertable's schema);
+ * FOR TABLE / FOR ALL TABLES are handled by PostgreSQL, which expands
+ * inheritance children itself.
+ *
+ * Honors enable_chunk_auto_publication: when off, the whole reconciliation is
+ * skipped. This is consistent with the opt-in design of the feature - if the
+ * user disables the GUC after chunks were auto-published, they own any cleanup
+ * (a subsequent ALTER PUBLICATION ... DROP TABLES IN SCHEMA will not remove
+ * the previously backfilled chunk rows in that case). Drop the publication or
+ * re-enable the GUC before such an ALTER to have the rows removed here.
+ */
+void
+ts_chunk_publication_reconcile_schema_chunks(Oid pubid, List *schema_oids, bool add)
+{
+	if (!ts_guc_enable_chunk_auto_publication)
+	{
+		return;
+	}
+
+	Cache *hcache = ts_hypertable_cache_pin();
+	ListCell *sc;
+	foreach (sc, schema_oids)
+	{
+		Oid published_schema = lfirst_oid(sc);
+		ListCell *rc;
+		foreach (rc, GetSchemaPublicationRelations(published_schema, PUBLICATION_PART_ROOT))
+		{
+			Hypertable *ht =
+				ts_hypertable_cache_get_entry(hcache, lfirst_oid(rc), CACHE_FLAG_MISSING_OK);
+			if (ht == NULL)
+			{
+				continue;
+			}
+			ListCell *cc;
+			foreach (cc, ts_chunk_get_by_hypertable_id(ht->fd.id))
+			{
+				Chunk *chunk = lfirst(cc);
+				/* OSM chunks are foreign tables; publication_add_relation would fail. */
+				if (IS_OSM_CHUNK(chunk))
+				{
+					continue;
+				}
+				/*
+				 * Skip adding an explicit row for a chunk that lives in a schema this
+				 * publication already covers natively: the explicit row would be
+				 * redundant and would survive a DROP that removes the schema mapping.
+				 * (Mirrors chunk_add_to_publications' guard.)
+				 */
+				if (add)
+				{
+					if (list_member_oid(GetSchemaPublications(get_rel_namespace(chunk->fd.relid)),
+										pubid))
+					{
+						continue;
+					}
+					chunk_add_to_publication(pubid, chunk);
+				}
+				else
+				{
+					chunk_remove_from_publication(pubid, chunk);
+				}
+			}
+		}
+	}
+	ts_cache_release(&hcache);
+}
+
+/*
+ * Reconcile chunk publication membership when a hypertable's schema changes
+ * (ALTER TABLE ... SET SCHEMA). Chunks live in _timescaledb_internal and do not
+ * move with the root table, so PostgreSQL's schema publication membership
+ * (which follows the root table's schema) leaves chunk rows behind: chunks that
+ * were backfilled for the old schema's publications must be removed, and they
+ * must be added for the new schema's publications. Existing chunks are handled
+ * here; chunks created later are covered by the per-chunk create hook.
+ */
+void
+ts_chunk_publication_reconcile_ht_schema_change(int32 hypertable_id, Oid old_schema, Oid new_schema)
+{
+	ListCell *cc;
+	List *old_pubs;
+	List *new_pubs;
+
+	if (!ts_guc_enable_chunk_auto_publication)
+	{
+		return;
+	}
+
+	if (old_schema == new_schema)
+	{
+		return;
+	}
+
+	/* The old/new schema publication lists are the same for every chunk. */
+	old_pubs = GetSchemaPublications(old_schema);
+	new_pubs = GetSchemaPublications(new_schema);
+
+	foreach (cc, ts_chunk_get_by_hypertable_id(hypertable_id))
+	{
+		Chunk *chunk = lfirst(cc);
+		ListCell *pc;
+
+		/* OSM chunks are foreign tables; publication_add_relation would fail. */
+		if (IS_OSM_CHUNK(chunk))
+		{
+			continue;
+		}
+
+		foreach (pc, old_pubs)
+		{
+			chunk_remove_from_publication(lfirst_oid(pc), chunk);
+		}
+
+		/* A chunk whose schema is already covered natively by a new-schema
+		 * publication needs no explicit row (it would survive a later DROP
+		 * that removes the schema mapping). */
+		List *chunk_schema_pubs = GetSchemaPublications(get_rel_namespace(chunk->fd.relid));
+		foreach (pc, new_pubs)
+		{
+			if (!list_member_oid(chunk_schema_pubs, lfirst_oid(pc)))
+			{
+				chunk_add_to_publication(lfirst_oid(pc), chunk);
+			}
+		}
 	}
 }
 

src/chunk.h

@@ -230,6 +230,11 @@ extern TSDLLEXPORT Datum ts_chunk_id_from_relid(PG_FUNCTION_ARGS);
 extern TSDLLEXPORT Datum ts_chunk_status_text(PG_FUNCTION_ARGS);
 extern TSDLLEXPORT List *ts_chunk_get_chunk_ids_by_hypertable_id(int32 hypertable_id);
 extern TSDLLEXPORT List *ts_chunk_get_by_hypertable_id(int32 hypertable_id);
+extern TSDLLEXPORT void ts_chunk_publication_reconcile_schema_chunks(Oid pubid, List *schema_oids,
+												bool add);
+extern TSDLLEXPORT void ts_chunk_publication_reconcile_ht_schema_change(int32 hypertable_id,
+																		Oid old_schema,
+																		Oid new_schema);
 
 extern TSDLLEXPORT int64 ts_chunk_primary_dimension_start(const Chunk *chunk);
 

src/process_utility.c

@@ -16,6 +16,7 @@
 #include <catalog/pg_class_d.h>
 #include <catalog/pg_constraint.h>
 #include <catalog/pg_inherits.h>
+#include <catalog/pg_publication.h>
 #include <catalog/pg_trigger.h>
 #include <commands/alter.h>
 #include <commands/copy.h>
@@ -805,6 +806,24 @@ process_altertableschema(ProcessUtilityArgs *args)
 
 	if (ht)
 	{
+		Oid old_schema = get_rel_namespace(relid);
+		Oid new_schema = get_namespace_oid(alterstmt->newschema, true);
+
+		/*
+		 * Chunks live in _timescaledb_internal and don't move with the root
+		 * table, so a schema publication (which follows the root table's
+		 * schema) leaves chunk rows behind. Reconcile them for the schema the
+		 * hypertable is leaving and the one it is entering. This must run
+		 * before ts_hypertable_set_schema flips the catalog schema, since
+		 * ts_chunk_publication_reconcile_ht_schema_change resolves the
+		 * hypertable's relid from the catalog (which still points at the old
+		 * schema's physical table); the chunk rows themselves are stable.
+		 */
+		if (OidIsValid(new_schema))
+		{
+			ts_chunk_publication_reconcile_ht_schema_change(ht->fd.id, old_schema, new_schema);
+		}
+
 		ts_hypertable_set_schema(ht, alterstmt->newschema);
 	}
 	else
@@ -6030,6 +6049,97 @@ preprocess_execute(ProcessUtilityArgs *args)
 	return DDL_CONTINUE;
 }
 
+/*
+ * Collect the Oids of schemas named by FOR TABLES IN SCHEMA in a
+ * CREATE/ALTER PUBLICATION object list. Mirrors PostgreSQL's schema handling
+ * in ObjectsInPublicationToOids; FOR TABLE / EXCEPT TABLE objects are left to
+ * PostgreSQL, which already expands inheritance children (chunks) itself.
+ */
+static List *
+publication_schema_oids(List *pubobjects)
+{
+	List *schemas = NIL;
+	ListCell *lc;
+
+	foreach (lc, pubobjects)
+	{
+		PublicationObjSpec *pubobj = lfirst(lc);
+		Oid schemaid;
+
+		switch (pubobj->pubobjtype)
+		{
+			case PUBLICATIONOBJ_TABLES_IN_SCHEMA:
+				schemaid = get_namespace_oid(pubobj->name, false);
+				break;
+			case PUBLICATIONOBJ_TABLES_IN_CUR_SCHEMA:
+			{
+				List *search_path = fetch_search_path(false);
+
+				if (search_path == NIL)
+				{
+					ereport(ERROR,
+							(errcode(ERRCODE_UNDEFINED_SCHEMA),
+							 errmsg("no schema has been selected for CURRENT_SCHEMA")));
+				}
+				schemaid = linitial_oid(search_path);
+				list_free(search_path);
+				break;
+			}
+			default:
+				continue;
+		}
+		schemas = list_append_unique_oid(schemas, schemaid);
+	}
+	return schemas;
+}
+
+/*
+ * Reconcile chunk membership for the schemas named in a publication command.
+ * Chunks live in _timescaledb_internal, so FOR TABLES IN SCHEMA never adds them
+ * (PostgreSQL only expands inheritance children for FOR TABLE / FOR ALL
+ * TABLES). For SET, PostgreSQL has already dropped every chunk row it
+ * tracked, so we rebuild for the full new schema set rather than a difference;
+ * only DROP removes chunks, because PostgreSQL leaves our backfilled chunk
+ * rows untouched when a schema is dropped from the publication. We read the
+ * schema set straight from the parse tree, so only the schemas the command
+ * touched are reconciled -- no before/after snapshot is needed.
+ */
+static void
+reconcile_publication_schemas(const char *pubname, List *schema_oids, bool add)
+{
+	Publication *pub = GetPublicationByName(pubname, true);
+	if (pub)
+	{
+		ts_chunk_publication_reconcile_schema_chunks(pub->oid, schema_oids, add);
+	}
+}
+
+/*
+ * Keep chunk membership of a schema publication in sync with the schemas named
+ * by ALTER PUBLICATION ... ADD/DROP/SET TABLES IN SCHEMA. CREATE PUBLICATION is
+ * handled later in process_ddl_event_command_end, where the publication's Oid
+ * is already available.
+ */
+static DDLResult
+process_publication(ProcessUtilityArgs *args)
+{
+	AlterPublicationStmt *stmt = castNode(AlterPublicationStmt, args->parsetree);
+	List *schema_oids;
+
+	/* Let PostgreSQL alter the publication first. */
+	prev_ProcessUtility(args);
+
+	schema_oids = publication_schema_oids(stmt->pubobjects);
+	if (schema_oids == NIL)
+	{
+		return DDL_DONE;
+	}
+
+	reconcile_publication_schemas(stmt->pubname, schema_oids, stmt->action != AP_DropObjects);
+
+	return DDL_DONE;
+}
+
 /*
  * Handle DDL commands before they have been processed by PostgreSQL.
  */
@@ -6123,6 +6233,15 @@ process_ddl_command_start(ProcessUtilityArgs *args)
 			handler = preprocess_execute;
 			break;
 
+		/* ALTER PUBLICATION is reconciled here; CREATE PUBLICATION is reconciled
+		 * later in process_ddl_event_command_end (where the publication's Oid is
+		 * available). check_read_only stays off to match prior behavior; PG
+		 * rejects publication DDL on read-only standbys itself. */
+		case T_AlterPublicationStmt:
+			check_read_only = false;
+			handler = process_publication;
+			break;
+
 		default:
 			handler = NULL;
 			break;
@@ -6141,6 +6260,25 @@ process_ddl_command_start(ProcessUtilityArgs *args)
 	return handler(args);
 }
 
+/*
+ * Backfill chunk membership for a publication created with FOR TABLES IN
+ * SCHEMA. Runs in the ddl_command_end event trigger, so the publication (and
+ * its Oid) already exists.
+ */
+static void
+process_create_publication_end(Node *parsetree)
+{
+	CreatePublicationStmt *stmt = castNode(CreatePublicationStmt, parsetree);
+	List *schema_oids = publication_schema_oids(stmt->pubobjects);
+
+	if (schema_oids == NIL)
+	{
+		return;
+	}
+
+	reconcile_publication_schemas(stmt->pubname, schema_oids, true);
+}
+
 /*
  * Handle DDL commands after they've been processed by PostgreSQL.
  */
@@ -6155,6 +6293,9 @@ process_ddl_command_end(CollectedCommand *cmd)
 		case T_AlterTableStmt:
 			process_altertable_end(cmd->parsetree, cmd);
 			break;
+		case T_CreatePublicationStmt:
+			process_create_publication_end(cmd->parsetree);
+			break;
 		default:
 			break;
 	}
@@ -6418,6 +6559,7 @@ process_ddl_event_command_end(EventTriggerData *trigdata)
 		case T_CreateTrigStmt:
 		case T_CreateStmt:
 		case T_IndexStmt:
+		case T_CreatePublicationStmt:
 			foreach (lc, ts_event_trigger_ddl_commands())
 			{
 				process_ddl_command_end(lfirst(lc));