Allow generating shard IDs without using the sequence

src/backend/distributed/master/master_node_protocol.c

@@ -68,6 +68,7 @@ int ShardCount = 32;
 int ShardReplicationFactor = 1; /* desired replication factor for shards */
 int ShardMaxSize = 1048576;     /* maximum size in KB one shard can grow to */
 int ShardPlacementPolicy = SHARD_PLACEMENT_ROUND_ROBIN;
+int NextShardId = 0;
 
 static List * GetTableReplicaIdentityCommand(Oid relationId);
 static Datum WorkerNodeGetDatum(WorkerNode *workerNode, TupleDesc tupleDescriptor);
@@ -291,14 +292,26 @@ master_get_new_shardid(PG_FUNCTION_ARGS)
 uint64
 GetNextShardId()
 {
-	text *sequenceName = cstring_to_text(SHARDID_SEQUENCE_NAME);
+	text *sequenceName = NULL;
-	Oid sequenceId = ResolveRelationId(sequenceName);
+	Oid sequenceId = InvalidOid;
-	Datum sequenceIdDatum = ObjectIdGetDatum(sequenceId);
+	Datum sequenceIdDatum = NULL;
 	Oid savedUserId = InvalidOid;
 	int savedSecurityContext = 0;
 	Datum shardIdDatum = 0;
 	uint64 shardId = 0;
 
+	if (NextShardId > 0)
+	{
+		shardId = NextShardId;
+		NextShardId += 1;
+
+		return shardId;
+	}
+
+	sequenceName = cstring_to_text(SHARDID_SEQUENCE_NAME);
+	sequenceId = ResolveRelationId(sequenceName);
+	sequenceIdDatum = ObjectIdGetDatum(sequenceId);
+
 	GetUserIdAndSecContext(&savedUserId, &savedSecurityContext);
 	SetUserIdAndSecContext(CitusExtensionOwner(), SECURITY_LOCAL_USERID_CHANGE);
 

src/backend/distributed/shared_library_init.c

@@ -813,6 +813,21 @@ RegisterCitusConfigVariables(void)
 		GUC_SUPERUSER_ONLY | GUC_NO_SHOW_ALL,
 		NULL, NULL, NULL);
 
+	DefineCustomIntVariable(
+		"citus.next_shard_id",
+		gettext_noop("Set the next shard ID to use in shard creation."),
+		gettext_noop("Shard IDs are normally generated using a sequence. If "
+					 "next_shard_id is set to a non-zero value, shard IDs will "
+					 "instead be generated by incrementing from the value of "
+					 "this GUC and this will be reflected in the GUC. This is "
+					 "mainly useful to ensure consistent shard IDs when running "
+					 "tests in parallel."),
+		&NextShardId,
+		0, 0, INT_MAX,
+		PGC_USERSET,
+		GUC_NO_SHOW_ALL,
+		NULL, NULL, NULL);
+
 	DefineCustomIntVariable(
 		"citus.max_task_string_size",
 		gettext_noop("Sets the maximum size (in bytes) of a worker task call string."),

src/include/distributed/master_protocol.h

@@ -92,6 +92,7 @@ extern int ShardCount;
 extern int ShardReplicationFactor;
 extern int ShardMaxSize;
 extern int ShardPlacementPolicy;
+extern int NextShardId;
 
 
 extern bool IsCoordinator(void);