Allow using generated identity column based on int/smallint when creating a distributed table (#7008)

Allow using generated identity column based on int/smallint when
creating a distributed table so that applications that rely on
those data types don't break.

Inserting into / modifying such columns from workers is not allowed
but it's better than not allowing such columns altogether.

src/backend/distributed/commands/create_distributed_table.c

@@ -1835,8 +1835,6 @@ EnsureRelationCanBeDistributed(Oid relationId, Var *distributionColumn,
 {
 	Oid parentRelationId = InvalidOid;
 
-	ErrorIfTableHasUnsupportedIdentityColumn(relationId);
-
 	EnsureLocalTableEmptyIfNecessary(relationId, distributionMethod);
 
 	/* user really wants triggers? */

src/backend/distributed/commands/table.c

@@ -4049,36 +4049,6 @@ MakeNameListFromRangeVar(const RangeVar *rel)
 }
 
 
-/*
- * ErrorIfTableHasUnsupportedIdentityColumn errors out if the given table has any identity column other than bigint identity column.
- */
-void
-ErrorIfTableHasUnsupportedIdentityColumn(Oid relationId)
-{
-	Relation relation = relation_open(relationId, AccessShareLock);
-	TupleDesc tupleDescriptor = RelationGetDescr(relation);
-
-	for (int attributeIndex = 0; attributeIndex < tupleDescriptor->natts;
-		 attributeIndex++)
-	{
-		Form_pg_attribute attributeForm = TupleDescAttr(tupleDescriptor, attributeIndex);
-
-		if (attributeForm->attidentity && attributeForm->atttypid != INT8OID)
-		{
-			char *qualifiedRelationName = generate_qualified_relation_name(relationId);
-			ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-							errmsg(
-								"cannot complete operation on %s with smallint/int identity column",
-								qualifiedRelationName),
-							errhint(
-								"Use bigint identity column instead.")));
-		}
-	}
-
-	relation_close(relation, NoLock);
-}
-
-
 /*
  * ErrorIfTableHasIdentityColumn errors out if the given table has identity column
  */

src/backend/distributed/shared_library_init.c

@@ -185,6 +185,7 @@ static void CitusObjectAccessHook(ObjectAccessType access, Oid classId, Oid obje
 static void DoInitialCleanup(void);
 static void ResizeStackToMaximumDepth(void);
 static void multi_log_hook(ErrorData *edata);
+static bool IsSequenceOverflowError(ErrorData *edata);
 static void RegisterConnectionCleanup(void);
 static void RegisterExternalClientBackendCounterDecrement(void);
 static void CitusCleanupConnectionsAtExit(int code, Datum arg);
@@ -683,6 +684,15 @@ multi_log_hook(ErrorData *edata)
 	 * Show the user a meaningful error message when a backend is cancelled
 	 * by the distributed deadlock detection. Also reset the state for this,
 	 * since the next cancelation of the backend might have another reason.
+	 *
+	 * We also want to provide a useful hint for sequence overflow errors
+	 * because they're likely to be caused by the way Citus handles smallint/int
+	 * based sequences on worker nodes. Note that we add the hint without checking
+	 * whether we're on a worker node or the sequence was used on a distributed
+	 * table because catalog might not be available at this point. And given
+	 * that this hint might be shown for regular Postgres tables too, we inject
+	 * the hint only when EnableUnsupportedFeatureMessages is set to true.
+	 * Otherwise, vanilla tests would fail.
 	 */
 	bool clearState = true;
 	if (edata->elevel == ERROR && edata->sqlerrcode == ERRCODE_QUERY_CANCELED &&
@@ -700,6 +710,16 @@ multi_log_hook(ErrorData *edata)
 		edata->message = pstrdup("canceling the transaction since it was "
 								 "involved in a distributed deadlock");
 	}
+	else if (EnableUnsupportedFeatureMessages &&
+			 IsSequenceOverflowError(edata))
+	{
+		edata->detail = pstrdup("nextval(sequence) calls in worker nodes "
+								"are not supported for column defaults of "
+								"type int or smallint");
+		edata->hint = pstrdup("If the command was issued from a worker node, "
+							  "try issuing it from the coordinator node "
+							  "instead.");
+	}
 
 	if (original_emit_log_hook)
 	{
@@ -708,6 +728,25 @@ multi_log_hook(ErrorData *edata)
 }
 
 
+/*
+ * IsSequenceOverflowError returns true if the given error is a sequence
+ * overflow error.
+ */
+static bool
+IsSequenceOverflowError(ErrorData *edata)
+{
+	static const char *sequenceOverflowedMsgPrefix =
+		"nextval: reached maximum value of sequence";
+	static const int sequenceOverflowedMsgPrefixLen = 42;
+
+	return edata->elevel == ERROR &&
+		   edata->sqlerrcode == ERRCODE_SEQUENCE_GENERATOR_LIMIT_EXCEEDED &&
+		   edata->message != NULL &&
+		   strncmp(edata->message, sequenceOverflowedMsgPrefix,
+				   sequenceOverflowedMsgPrefixLen) == 0;
+}
+
+
 /*
  * StartupCitusBackend initializes per-backend infrastructure, and is called
  * the first time citus is used in a database.

src/include/distributed/commands.h

@@ -591,7 +591,6 @@ extern bool ConstrTypeCitusCanDefaultName(ConstrType constrType);
 extern char * GetAlterColumnWithNextvalDefaultCmd(Oid sequenceOid, Oid relationId,
 												  char *colname, bool missingTableOk);
 
-extern void ErrorIfTableHasUnsupportedIdentityColumn(Oid relationId);
 extern void ErrorIfTableHasIdentityColumn(Oid relationId);
 extern void ConvertNewTableIfNecessary(Node *createStmt);
 

src/test/regress/citus_tests/run_test.py

@@ -158,6 +158,14 @@ DEPS = {
         None, ["isolation_setup", "isolation_add_remove_node"]
     ),
     "schema_based_sharding": TestDeps("minimal_schedule"),
+    "multi_sequence_default": TestDeps(
+        None,
+        [
+            "multi_test_helpers",
+            "multi_cluster_management",
+            "multi_table_ddl",
+        ],
+    ),
 }
 
 

src/test/regress/expected/create_single_shard_table.out

@@ -1040,12 +1040,27 @@ CREATE TABLE identity_test (
     c bigint GENERATED BY DEFAULT AS IDENTITY (START WITH 1000 INCREMENT BY 1000)
 );
 SELECT create_distributed_table('identity_test', NULL, distribution_type=>null);
-ERROR:  cannot complete operation on create_single_shard_table.identity_test with smallint/int identity column
+ create_distributed_table
-HINT:  Use bigint identity column instead.
+---------------------------------------------------------------------
+
+(1 row)
+
+INSERT INTO identity_test (a) VALUES (DEFAULT) RETURNING a;
+ a
+---------------------------------------------------------------------
+ 10
+(1 row)
+
+SELECT result FROM run_command_on_workers($$
+    INSERT INTO create_single_shard_table.identity_test (a) VALUES (DEFAULT)
+$$);
+                                        result
+---------------------------------------------------------------------
+ ERROR:  nextval: reached maximum value of sequence "identity_test_a_seq" (2147483647)
+ ERROR:  nextval: reached maximum value of sequence "identity_test_a_seq" (2147483647)
+(2 rows)
+
 DROP TABLE identity_test;
--- Above failed because we don't support using a data type other than BIGINT
--- for identity columns, so drop the table and create a new one with BIGINT
--- identity columns.
 CREATE TABLE identity_test (
     a bigint GENERATED BY DEFAULT AS IDENTITY (START WITH 10 INCREMENT BY 10),
     b bigint GENERATED ALWAYS AS IDENTITY (START WITH 100 INCREMENT BY 100),
@@ -1131,7 +1146,7 @@ BEGIN;
   INSERT INTO referencing_table VALUES (1, 2);
   -- fails
   INSERT INTO referencing_table VALUES (2, 2);
-ERROR:  insert or update on table "referencing_table_xxxxxxx" violates foreign key constraint "referencing_table_a_fkey_1730099"
+ERROR:  insert or update on table "referencing_table_xxxxxxx" violates foreign key constraint "referencing_table_a_fkey_1730100"
 DETAIL:  Key (a)=(2) is not present in table "referenced_table_xxxxxxx".
 CONTEXT:  while executing command on localhost:xxxxx
 ROLLBACK;
@@ -1177,7 +1192,7 @@ BEGIN;
   INSERT INTO referencing_table VALUES (1, 2);
   -- fails
   INSERT INTO referencing_table VALUES (2, 2);
-ERROR:  insert or update on table "referencing_table_xxxxxxx" violates foreign key constraint "referencing_table_a_fkey_1730135"
+ERROR:  insert or update on table "referencing_table_xxxxxxx" violates foreign key constraint "referencing_table_a_fkey_1730136"
 DETAIL:  Key (a)=(2) is not present in table "referenced_table_xxxxxxx".
 CONTEXT:  while executing command on localhost:xxxxx
 ROLLBACK;
@@ -1295,8 +1310,8 @@ SELECT result, success FROM run_command_on_workers($$
 $$);
                                                               result                                                              | success
 ---------------------------------------------------------------------
- ERROR:  insert or update on table "referencing_table_xxxxxxx" violates foreign key constraint "referencing_table_a_fkey_1730152" | f
+ ERROR:  insert or update on table "referencing_table_xxxxxxx" violates foreign key constraint "referencing_table_a_fkey_1730153" | f
- ERROR:  insert or update on table "referencing_table_xxxxxxx" violates foreign key constraint "referencing_table_a_fkey_1730152" | f
+ ERROR:  insert or update on table "referencing_table_xxxxxxx" violates foreign key constraint "referencing_table_a_fkey_1730153" | f
 (2 rows)
 
 DROP TABLE referencing_table, referenced_table;
@@ -1311,8 +1326,8 @@ SELECT create_distributed_table('self_fkey_test', NULL, distribution_type=>null)
 
 INSERT INTO self_fkey_test VALUES (1, 1); -- ok
 INSERT INTO self_fkey_test VALUES (2, 3); -- fails
-ERROR:  insert or update on table "self_fkey_test_1730153" violates foreign key constraint "self_fkey_test_b_fkey_1730153"
+ERROR:  insert or update on table "self_fkey_test_1730154" violates foreign key constraint "self_fkey_test_b_fkey_1730154"
-DETAIL:  Key (b)=(3) is not present in table "self_fkey_test_1730153".
+DETAIL:  Key (b)=(3) is not present in table "self_fkey_test_1730154".
 CONTEXT:  while executing command on localhost:xxxxx
 -- similar foreign key tests but this time create the referencing table later on
 -- referencing table is a single-shard table
@@ -1336,7 +1351,7 @@ BEGIN;
   INSERT INTO referencing_table VALUES (1, 2);
   -- fails
   INSERT INTO referencing_table VALUES (2, 2);
-ERROR:  insert or update on table "referencing_table_xxxxxxx" violates foreign key constraint "referencing_table_a_fkey_1730155"
+ERROR:  insert or update on table "referencing_table_xxxxxxx" violates foreign key constraint "referencing_table_a_fkey_1730156"
 DETAIL:  Key (a)=(2) is not present in table "referenced_table_xxxxxxx".
 CONTEXT:  while executing command on localhost:xxxxx
 ROLLBACK;
@@ -1359,7 +1374,7 @@ BEGIN;
   INSERT INTO referencing_table VALUES (2, 1);
   -- fails
   INSERT INTO referencing_table VALUES (1, 2);
-ERROR:  insert or update on table "referencing_table_xxxxxxx" violates foreign key constraint "referencing_table_a_b_fkey_1730157"
+ERROR:  insert or update on table "referencing_table_xxxxxxx" violates foreign key constraint "referencing_table_a_b_fkey_1730158"
 DETAIL:  Key (a, b)=(1, 2) is not present in table "referenced_table_xxxxxxx".
 CONTEXT:  while executing command on localhost:xxxxx
 ROLLBACK;
@@ -1466,7 +1481,7 @@ BEGIN;
   INSERT INTO referencing_table VALUES (1, 2);
   -- fails
   INSERT INTO referencing_table VALUES (2, 2);
-ERROR:  insert or update on table "referencing_table_xxxxxxx" violates foreign key constraint "referencing_table_a_fkey_1730198"
+ERROR:  insert or update on table "referencing_table_xxxxxxx" violates foreign key constraint "referencing_table_a_fkey_1730199"
 DETAIL:  Key (a)=(2) is not present in table "referenced_table_xxxxxxx".
 CONTEXT:  while executing command on localhost:xxxxx
 ROLLBACK;

src/test/regress/expected/generated_identity.out

@@ -16,44 +16,185 @@ SELECT 1 from citus_add_node('localhost', :master_port, groupId=>0);
         1
 (1 row)
 
--- smallint identity column can not be distributed
 CREATE TABLE smallint_identity_column (
     a smallint GENERATED BY DEFAULT AS IDENTITY
 );
-SELECT create_distributed_table('smallint_identity_column', 'a');
+CREATE VIEW verify_smallint_identity_column AS
-ERROR:  cannot complete operation on generated_identities.smallint_identity_column with smallint/int identity column
+SELECT attidentity, attgenerated FROM pg_attribute WHERE attrelid = 'smallint_identity_column'::regclass AND attname = 'a';
-HINT:  Use bigint identity column instead.
+BEGIN;
-SELECT create_distributed_table_concurrently('smallint_identity_column', 'a');
+    SELECT create_distributed_table('smallint_identity_column', 'a');
-ERROR:  cannot complete operation on generated_identities.smallint_identity_column with smallint/int identity column
+ create_distributed_table
-HINT:  Use bigint identity column instead.
+---------------------------------------------------------------------
-SELECT create_reference_table('smallint_identity_column');
+
-ERROR:  cannot complete operation on a table with identity column
+(1 row)
-SELECT citus_add_local_table_to_metadata('smallint_identity_column');
+
+    SELECT * FROM verify_smallint_identity_column;
+ attidentity | attgenerated
+---------------------------------------------------------------------
+ d           |
+(1 row)
+
+ROLLBACK;
+BEGIN;
+    SELECT create_reference_table('smallint_identity_column');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+    SELECT * FROM verify_smallint_identity_column;
+ attidentity | attgenerated
+---------------------------------------------------------------------
+ d           |
+(1 row)
+
+ROLLBACK;
+BEGIN;
+    SELECT citus_add_local_table_to_metadata('smallint_identity_column');
  citus_add_local_table_to_metadata
 ---------------------------------------------------------------------
 
 (1 row)
 
-DROP TABLE smallint_identity_column;
+    SELECT * FROM verify_smallint_identity_column;
--- int identity column can not be distributed
+ attidentity | attgenerated
+---------------------------------------------------------------------
+ d           |
+(1 row)
+
+ROLLBACK;
+SELECT create_distributed_table_concurrently('smallint_identity_column', 'a');
+ create_distributed_table_concurrently
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT * FROM verify_smallint_identity_column;
+ attidentity | attgenerated
+---------------------------------------------------------------------
+ d           |
+(1 row)
+
+SELECT result FROM run_command_on_workers('INSERT INTO generated_identities.smallint_identity_column (a) VALUES (DEFAULT);');
+                                           result
+---------------------------------------------------------------------
+ ERROR:  nextval: reached maximum value of sequence "smallint_identity_column_a_seq" (32767)
+ ERROR:  nextval: reached maximum value of sequence "smallint_identity_column_a_seq" (32767)
+(2 rows)
+
+DROP TABLE smallint_identity_column CASCADE;
 CREATE TABLE int_identity_column (
     a int GENERATED BY DEFAULT AS IDENTITY
 );
-SELECT create_distributed_table('int_identity_column', 'a');
+CREATE VIEW verify_int_identity_column AS
-ERROR:  cannot complete operation on generated_identities.int_identity_column with smallint/int identity column
+SELECT attidentity, attgenerated FROM pg_attribute WHERE attrelid = 'int_identity_column'::regclass AND attname = 'a';
-HINT:  Use bigint identity column instead.
+BEGIN;
-SELECT create_distributed_table_concurrently('int_identity_column', 'a');
+    SELECT create_distributed_table('int_identity_column', 'a');
-ERROR:  cannot complete operation on generated_identities.int_identity_column with smallint/int identity column
+ create_distributed_table
-HINT:  Use bigint identity column instead.
+---------------------------------------------------------------------
-SELECT create_reference_table('int_identity_column');
+
-ERROR:  cannot complete operation on a table with identity column
+(1 row)
-SELECT citus_add_local_table_to_metadata('int_identity_column');
+
+    SELECT * FROM verify_int_identity_column;
+ attidentity | attgenerated
+---------------------------------------------------------------------
+ d           |
+(1 row)
+
+ROLLBACK;
+BEGIN;
+    SELECT create_reference_table('int_identity_column');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+    SELECT * FROM verify_int_identity_column;
+ attidentity | attgenerated
+---------------------------------------------------------------------
+ d           |
+(1 row)
+
+ROLLBACK;
+BEGIN;
+    SELECT citus_add_local_table_to_metadata('int_identity_column');
  citus_add_local_table_to_metadata
 ---------------------------------------------------------------------
 
 (1 row)
 
-DROP TABLE int_identity_column;
+    SELECT * FROM verify_int_identity_column;
+ attidentity | attgenerated
+---------------------------------------------------------------------
+ d           |
+(1 row)
+
+ROLLBACK;
+SELECT create_distributed_table_concurrently('int_identity_column', 'a');
+ create_distributed_table_concurrently
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT * FROM verify_int_identity_column;
+ attidentity | attgenerated
+---------------------------------------------------------------------
+ d           |
+(1 row)
+
+SELECT result FROM run_command_on_workers('INSERT INTO generated_identities.int_identity_column (a) VALUES (DEFAULT);');
+                                           result
+---------------------------------------------------------------------
+ ERROR:  nextval: reached maximum value of sequence "int_identity_column_a_seq" (2147483647)
+ ERROR:  nextval: reached maximum value of sequence "int_identity_column_a_seq" (2147483647)
+(2 rows)
+
+DROP TABLE int_identity_column CASCADE;
+CREATE TABLE reference_int_identity_column (
+    a int GENERATED BY DEFAULT AS IDENTITY
+);
+SELECT create_reference_table('reference_int_identity_column');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+INSERT INTO generated_identities.reference_int_identity_column (a) VALUES (DEFAULT) RETURNING a;
+ a
+---------------------------------------------------------------------
+ 1
+(1 row)
+
+SELECT result FROM run_command_on_workers('INSERT INTO generated_identities.reference_int_identity_column (a) VALUES (DEFAULT);');
+                                                result
+---------------------------------------------------------------------
+ ERROR:  nextval: reached maximum value of sequence "reference_int_identity_column_a_seq" (2147483647)
+ ERROR:  nextval: reached maximum value of sequence "reference_int_identity_column_a_seq" (2147483647)
+(2 rows)
+
+CREATE TABLE citus_local_int_identity_column (
+    a int GENERATED BY DEFAULT AS IDENTITY
+);
+SELECT citus_add_local_table_to_metadata('citus_local_int_identity_column');
+ citus_add_local_table_to_metadata
+---------------------------------------------------------------------
+
+(1 row)
+
+INSERT INTO generated_identities.citus_local_int_identity_column (a) VALUES (DEFAULT) RETURNING a;
+ a
+---------------------------------------------------------------------
+ 1
+(1 row)
+
+SELECT result FROM run_command_on_workers('INSERT INTO generated_identities.citus_local_int_identity_column (a) VALUES (DEFAULT);');
+                                                 result
+---------------------------------------------------------------------
+ ERROR:  nextval: reached maximum value of sequence "citus_local_int_identity_column_a_seq" (2147483647)
+ ERROR:  nextval: reached maximum value of sequence "citus_local_int_identity_column_a_seq" (2147483647)
+(2 rows)
+
+DROP TABLE reference_int_identity_column, citus_local_int_identity_column;
 RESET citus.shard_replication_factor;
 CREATE TABLE bigint_identity_column (
     a bigint GENERATED BY DEFAULT AS IDENTITY,

src/test/regress/expected/multi_sequence_default.out

@@ -668,8 +668,12 @@ ERROR:  nextval(sequence) calls in worker nodes are not supported for column def
 -- nextval from worker should fail for int and smallint sequences
 SELECT nextval('seq_12');
 ERROR:  nextval: reached maximum value of sequence "seq_12" (32767)
+DETAIL:  nextval(sequence) calls in worker nodes are not supported for column defaults of type int or smallint
+HINT:  If the command was issued from a worker node, try issuing it from the coordinator node instead.
 SELECT nextval('seq_13');
 ERROR:  nextval: reached maximum value of sequence "seq_13" (2147483647)
+DETAIL:  nextval(sequence) calls in worker nodes are not supported for column defaults of type int or smallint
+HINT:  If the command was issued from a worker node, try issuing it from the coordinator node instead.
 -- nextval from worker should work for bigint sequences
 SELECT nextval('seq_14');
      nextval
@@ -738,8 +742,12 @@ SELECT nextval('seq_12');
 -- nextval from worker should fail for int and smallint sequences
 SELECT nextval('seq_13');
 ERROR:  nextval: reached maximum value of sequence "seq_13" (2147483647)
+DETAIL:  nextval(sequence) calls in worker nodes are not supported for column defaults of type int or smallint
+HINT:  If the command was issued from a worker node, try issuing it from the coordinator node instead.
 SELECT nextval('seq_14');
 ERROR:  nextval: reached maximum value of sequence "seq_14" (32767)
+DETAIL:  nextval(sequence) calls in worker nodes are not supported for column defaults of type int or smallint
+HINT:  If the command was issued from a worker node, try issuing it from the coordinator node instead.
 \c - - - :master_port
 SET citus.shard_replication_factor TO 1;
 SET search_path = sequence_default, public;
@@ -804,8 +812,12 @@ SELECT nextval('seq_12');
 -- nextval from worker should fail for int and smallint sequences
 SELECT nextval('seq_13');
 ERROR:  nextval: reached maximum value of sequence "seq_13" (2147483647)
+DETAIL:  nextval(sequence) calls in worker nodes are not supported for column defaults of type int or smallint
+HINT:  If the command was issued from a worker node, try issuing it from the coordinator node instead.
 SELECT nextval('seq_14');
 ERROR:  nextval: reached maximum value of sequence "seq_14" (32767)
+DETAIL:  nextval(sequence) calls in worker nodes are not supported for column defaults of type int or smallint
+HINT:  If the command was issued from a worker node, try issuing it from the coordinator node instead.
 \c - - - :master_port
 -- Show that sequence and its dependency schema will be propagated if a distributed
 -- table with default column is added

src/test/regress/sql/create_single_shard_table.sql

@@ -734,11 +734,13 @@ CREATE TABLE identity_test (
 
 SELECT create_distributed_table('identity_test', NULL, distribution_type=>null);
 
+INSERT INTO identity_test (a) VALUES (DEFAULT) RETURNING a;
+SELECT result FROM run_command_on_workers($$
+    INSERT INTO create_single_shard_table.identity_test (a) VALUES (DEFAULT)
+$$);
+
 DROP TABLE identity_test;
 
--- Above failed because we don't support using a data type other than BIGINT
--- for identity columns, so drop the table and create a new one with BIGINT
--- identity columns.
 CREATE TABLE identity_test (
     a bigint GENERATED BY DEFAULT AS IDENTITY (START WITH 10 INCREMENT BY 10),
     b bigint GENERATED ALWAYS AS IDENTITY (START WITH 100 INCREMENT BY 100),

src/test/regress/sql/generated_identity.sql

@@ -9,26 +9,78 @@ SET citus.shard_replication_factor TO 1;
 
 SELECT 1 from citus_add_node('localhost', :master_port, groupId=>0);
 
--- smallint identity column can not be distributed
 CREATE TABLE smallint_identity_column (
     a smallint GENERATED BY DEFAULT AS IDENTITY
 );
-SELECT create_distributed_table('smallint_identity_column', 'a');
+
+CREATE VIEW verify_smallint_identity_column AS
+SELECT attidentity, attgenerated FROM pg_attribute WHERE attrelid = 'smallint_identity_column'::regclass AND attname = 'a';
+
+BEGIN;
+    SELECT create_distributed_table('smallint_identity_column', 'a');
+    SELECT * FROM verify_smallint_identity_column;
+ROLLBACK;
+
+BEGIN;
+    SELECT create_reference_table('smallint_identity_column');
+    SELECT * FROM verify_smallint_identity_column;
+ROLLBACK;
+
+BEGIN;
+    SELECT citus_add_local_table_to_metadata('smallint_identity_column');
+    SELECT * FROM verify_smallint_identity_column;
+ROLLBACK;
+
 SELECT create_distributed_table_concurrently('smallint_identity_column', 'a');
-SELECT create_reference_table('smallint_identity_column');
+SELECT * FROM verify_smallint_identity_column;
-SELECT citus_add_local_table_to_metadata('smallint_identity_column');
+SELECT result FROM run_command_on_workers('INSERT INTO generated_identities.smallint_identity_column (a) VALUES (DEFAULT);');
 
-DROP TABLE smallint_identity_column;
+DROP TABLE smallint_identity_column CASCADE;
 
--- int identity column can not be distributed
 CREATE TABLE int_identity_column (
     a int GENERATED BY DEFAULT AS IDENTITY
 );
-SELECT create_distributed_table('int_identity_column', 'a');
+
+CREATE VIEW verify_int_identity_column AS
+SELECT attidentity, attgenerated FROM pg_attribute WHERE attrelid = 'int_identity_column'::regclass AND attname = 'a';
+
+BEGIN;
+    SELECT create_distributed_table('int_identity_column', 'a');
+    SELECT * FROM verify_int_identity_column;
+ROLLBACK;
+
+BEGIN;
+    SELECT create_reference_table('int_identity_column');
+    SELECT * FROM verify_int_identity_column;
+ROLLBACK;
+
+BEGIN;
+    SELECT citus_add_local_table_to_metadata('int_identity_column');
+    SELECT * FROM verify_int_identity_column;
+ROLLBACK;
+
 SELECT create_distributed_table_concurrently('int_identity_column', 'a');
-SELECT create_reference_table('int_identity_column');
+SELECT * FROM verify_int_identity_column;
-SELECT citus_add_local_table_to_metadata('int_identity_column');
+SELECT result FROM run_command_on_workers('INSERT INTO generated_identities.int_identity_column (a) VALUES (DEFAULT);');
-DROP TABLE int_identity_column;
+
+DROP TABLE int_identity_column CASCADE;
+
+CREATE TABLE reference_int_identity_column (
+    a int GENERATED BY DEFAULT AS IDENTITY
+);
+SELECT create_reference_table('reference_int_identity_column');
+INSERT INTO generated_identities.reference_int_identity_column (a) VALUES (DEFAULT) RETURNING a;
+SELECT result FROM run_command_on_workers('INSERT INTO generated_identities.reference_int_identity_column (a) VALUES (DEFAULT);');
+
+CREATE TABLE citus_local_int_identity_column (
+    a int GENERATED BY DEFAULT AS IDENTITY
+);
+SELECT citus_add_local_table_to_metadata('citus_local_int_identity_column');
+INSERT INTO generated_identities.citus_local_int_identity_column (a) VALUES (DEFAULT) RETURNING a;
+SELECT result FROM run_command_on_workers('INSERT INTO generated_identities.citus_local_int_identity_column (a) VALUES (DEFAULT);');
+
+DROP TABLE reference_int_identity_column, citus_local_int_identity_column;
+
 RESET citus.shard_replication_factor;