diff --git a/src/backend/distributed/commands/drop_distributed_table.c b/src/backend/distributed/commands/drop_distributed_table.c index d4718aab8..1e214adf3 100644 --- a/src/backend/distributed/commands/drop_distributed_table.c +++ b/src/backend/distributed/commands/drop_distributed_table.c @@ -12,6 +12,7 @@ #include "miscadmin.h" +#include "distributed/colocation_utils.h" #include "distributed/commands/utility_hook.h" #include "distributed/commands.h" #include "distributed/metadata_utility.h" @@ -70,6 +71,8 @@ master_remove_partition_metadata(PG_FUNCTION_ARGS) char *schemaName = text_to_cstring(schemaNameText); char *tableName = text_to_cstring(tableNameText); + uint32 colocationId = ColocationIdViaCatalog(relationId); + /* * The SQL_DROP trigger calls this function even for tables that are * not distributed. In that case, silently ignore. This is not very @@ -87,6 +90,8 @@ master_remove_partition_metadata(PG_FUNCTION_ARGS) DeletePartitionRow(relationId); + DeleteColocationGroupIfNoTablesBelong(colocationId); + PG_RETURN_VOID(); } diff --git a/src/backend/distributed/metadata/metadata_cache.c b/src/backend/distributed/metadata/metadata_cache.c index 63c2f8695..70fe44cbe 100644 --- a/src/backend/distributed/metadata/metadata_cache.c +++ b/src/backend/distributed/metadata/metadata_cache.c @@ -594,6 +594,45 @@ PartitionColumnViaCatalog(Oid relationId) } +/* + * ColocationIdViaCatalog gets a relationId and returns the colocation + * id column from pg_dist_partition via reading from catalog. + */ +uint32 +ColocationIdViaCatalog(Oid relationId) +{ + HeapTuple partitionTuple = PgDistPartitionTupleViaCatalog(relationId); + if (!HeapTupleIsValid(partitionTuple)) + { + return INVALID_COLOCATION_ID; + } + + Datum datumArray[Natts_pg_dist_partition]; + bool isNullArray[Natts_pg_dist_partition]; + + Relation pgDistPartition = table_open(DistPartitionRelationId(), AccessShareLock); + + TupleDesc tupleDescriptor = RelationGetDescr(pgDistPartition); + heap_deform_tuple(partitionTuple, tupleDescriptor, datumArray, isNullArray); + + if (isNullArray[Anum_pg_dist_partition_colocationid - 1]) + { + /* colocation id cannot be NULL, still let's make sure */ + heap_freetuple(partitionTuple); + table_close(pgDistPartition, NoLock); + return INVALID_COLOCATION_ID; + } + + Datum colocationIdDatum = datumArray[Anum_pg_dist_partition_colocationid - 1]; + uint32 colocationId = DatumGetUInt32(colocationIdDatum); + + heap_freetuple(partitionTuple); + table_close(pgDistPartition, NoLock); + + return colocationId; +} + + /* * PgDistPartitionTupleViaCatalog is a helper function that searches * pg_dist_partition for the given relationId. The caller is responsible diff --git a/src/include/distributed/metadata_cache.h b/src/include/distributed/metadata_cache.h index e190aef6f..e8b2d6a9c 100644 --- a/src/include/distributed/metadata_cache.h +++ b/src/include/distributed/metadata_cache.h @@ -150,6 +150,7 @@ extern char PgDistPartitionViaCatalog(Oid relationId); extern List * LookupDistShardTuples(Oid relationId); extern char PartitionMethodViaCatalog(Oid relationId); extern Var * PartitionColumnViaCatalog(Oid relationId); +extern uint32 ColocationIdViaCatalog(Oid relationId); extern bool IsCitusLocalTableByDistParams(char partitionMethod, char replicationModel); extern List * CitusTableList(void); extern ShardInterval * LoadShardInterval(uint64 shardId);