Create sequence dependency during metadata syncing

src/backend/distributed/metadata/metadata_sync.c

@@ -24,6 +24,7 @@
 #include "access/xact.h"
 #include "catalog/dependency.h"
 #include "catalog/indexing.h"
+#include "catalog/pg_depend.h"
 #include "catalog/pg_foreign_server.h"
 #include "catalog/pg_namespace.h"
 #include "catalog/pg_type.h"
@@ -58,6 +59,7 @@ static char * LocalGroupIdUpdateCommand(int32 groupId);
 static void UpdateDistNodeBoolAttr(const char *nodeName, int32 nodePort,
 								   int attrNum, bool value);
 static List * SequenceDDLCommandsForTable(Oid relationId);
+static List * SequenceDependencyCommandList(Oid relationId);
 static char * TruncateTriggerCreateCommand(Oid relationId);
 static char * SchemaOwnerName(Oid objectId);
 static bool HasMetadataWorkers(void);
@@ -73,6 +75,7 @@ static char * GenerateSetRoleQuery(Oid roleOid);
 
 PG_FUNCTION_INFO_V1(start_metadata_sync_to_node);
 PG_FUNCTION_INFO_V1(stop_metadata_sync_to_node);
+PG_FUNCTION_INFO_V1(worker_record_sequence_dependency);
 
 
 /*
@@ -389,6 +392,7 @@ MetadataCreateCommands(void)
 		List *workerSequenceDDLCommands = SequenceDDLCommandsForTable(relationId);
 		List *ddlCommandList = GetTableDDLEvents(relationId, includeSequenceDefaults);
 		char *tableOwnerResetCommand = TableOwnerResetCommand(relationId);
+		List *sequenceDependencyCommandList = SequenceDependencyCommandList(relationId);
 
 		/*
 		 * Tables might have dependencies on different objects, since we create shards for
@@ -404,6 +408,8 @@ MetadataCreateCommands(void)
 												  ddlCommandList);
 		metadataSnapshotCommandList = lappend(metadataSnapshotCommandList,
 											  tableOwnerResetCommand);
+		metadataSnapshotCommandList = list_concat(metadataSnapshotCommandList,
+												  sequenceDependencyCommandList);
 	}
 
 	/* construct the foreign key constraints after all tables are created */
@@ -1070,6 +1076,104 @@ SequenceDDLCommandsForTable(Oid relationId)
 }
 
 
+/*
+ * SequenceDependencyCommandList generates commands to record the dependency
+ * of sequences on tables on the worker. This dependency does not exist by
+ * default since the sequences and table are created separately, but it is
+ * necessary to ensure that the sequence is dropped when the table is
+ * dropped.
+ */
+static List *
+SequenceDependencyCommandList(Oid relationId)
+{
+	List *sequenceCommandList = NIL;
+	char *relationName = generate_qualified_relation_name(relationId);
+	HeapTuple tup = NULL;
+	ScanKeyData key[2];
+
+	Relation pgDepend = heap_open(DependRelationId, AccessShareLock);
+
+	ScanKeyInit(&key[0],
+				Anum_pg_depend_refclassid,
+				BTEqualStrategyNumber, F_OIDEQ,
+				ObjectIdGetDatum(RelationRelationId));
+	ScanKeyInit(&key[1],
+				Anum_pg_depend_refobjid,
+				BTEqualStrategyNumber, F_OIDEQ,
+				ObjectIdGetDatum(relationId));
+
+	SysScanDesc scan = systable_beginscan(pgDepend, DependReferenceIndexId, true,
+										  NULL, 2, key);
+
+	while (HeapTupleIsValid(tup = systable_getnext(scan)))
+	{
+		Form_pg_depend dependencyRecord = (Form_pg_depend) GETSTRUCT(tup);
+
+		/*
+		 * We assume any auto or internal dependency of a sequence on a column
+		 * must be what we are looking for.  (We need the relkind test because
+		 * indexes can also have auto dependencies on columns.)
+		 */
+		if (dependencyRecord->classid == RelationRelationId &&
+			dependencyRecord->objsubid == 0 &&
+			dependencyRecord->refobjsubid != 0 &&
+			(dependencyRecord->deptype == DEPENDENCY_AUTO ||
+			 dependencyRecord->deptype == DEPENDENCY_INTERNAL) &&
+			get_rel_relkind(dependencyRecord->objid) == RELKIND_SEQUENCE)
+		{
+			Oid sequenceId = dependencyRecord->objid;
+			char *sequenceName = generate_qualified_relation_name(sequenceId);
+			StringInfo sequenceDependencyCommand = makeStringInfo();
+
+			appendStringInfo(sequenceDependencyCommand,
+							 "SELECT pg_catalog.worker_record_sequence_dependency"
+							 "(%s::regclass,%s::regclass,%d)",
+							 quote_literal_cstr(sequenceName),
+							 quote_literal_cstr(relationName),
+							 dependencyRecord->refobjsubid);
+
+			sequenceCommandList = lappend(sequenceCommandList,
+										  sequenceDependencyCommand->data);
+		}
+	}
+
+	systable_endscan(scan);
+
+	heap_close(pgDepend, AccessShareLock);
+
+	return sequenceCommandList;
+}
+
+
+/*
+ * worker_record_sequence_dependency records the fact that the sequence depends on
+ * the table in pg_depend, such that it will be automatically dropped.
+ */
+Datum
+worker_record_sequence_dependency(PG_FUNCTION_ARGS)
+{
+	Oid sequenceOid = PG_GETARG_OID(0);
+	Oid relationOid = PG_GETARG_OID(1);
+	int columnIndex = PG_GETARG_INT32(2);
+
+	ObjectAddress sequenceAddr = { 0, 0, 0 };
+	ObjectAddress relationAddr = { 0, 0, 0 };
+
+	sequenceAddr.classId = RelationRelationId;
+	sequenceAddr.objectId = sequenceOid;
+	sequenceAddr.objectSubId = 0;
+
+	relationAddr.classId = RelationRelationId;
+	relationAddr.objectId = relationOid;
+	relationAddr.objectSubId = columnIndex;
+
+	/* dependency from sequence to table */
+	recordDependencyOn(&sequenceAddr, &relationAddr, DEPENDENCY_AUTO);
+
+	PG_RETURN_VOID();
+}
+
+
 /*
  * CreateSchemaDDLCommand returns a "CREATE SCHEMA..." SQL string for creating the given
  * schema if not exists and with proper authorization.

src/backend/distributed/sql/citus--9.3-2--9.4-1.sql

@@ -1,3 +1,3 @@
 -- citus--9.3-2--9.4-1
 
--- bump version to 9.4-1
+#include "udfs/worker_record_sequence_dependency/9.4-1.sql"

src/backend/distributed/sql/udfs/worker_record_sequence_dependency/latest.sql

@@ -0,0 +1,9 @@
+CREATE FUNCTION pg_catalog.worker_record_sequence_dependency(seq_name regclass, table_name regclass, column_index int)
+  RETURNS VOID
+  LANGUAGE C STRICT
+  AS 'MODULE_PATHNAME', 'worker_record_sequence_dependency';
+COMMENT ON FUNCTION pg_catalog.worker_record_sequence_dependency(regclass, regclass, int)
+  IS 'record the fact that the sequence depends on the table in pg_depend';
+
+REVOKE ALL ON FUNCTION pg_catalog.worker_record_sequence_dependency(regclass, regclass, int)
+  FROM PUBLIC;