diff --git a/src/backend/distributed/master/master_node_protocol.c b/src/backend/distributed/master/master_node_protocol.c index ce2facc4b..035db256d 100644 --- a/src/backend/distributed/master/master_node_protocol.c +++ b/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); - Oid sequenceId = ResolveRelationId(sequenceName); - Datum sequenceIdDatum = ObjectIdGetDatum(sequenceId); + text *sequenceName = NULL; + Oid sequenceId = InvalidOid; + 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); diff --git a/src/backend/distributed/shared_library_init.c b/src/backend/distributed/shared_library_init.c index 3907dc5e2..eef9897bb 100644 --- a/src/backend/distributed/shared_library_init.c +++ b/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."), diff --git a/src/include/distributed/master_protocol.h b/src/include/distributed/master_protocol.h index b4840fa96..9875196dc 100644 --- a/src/include/distributed/master_protocol.h +++ b/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);