1 /*------------------------------------------------------------------------- 2 * 3 * test/src/colocations_utils.c 4 * 5 * This file contains functions to test co-location functionality 6 * within Citus. 7 * 8 * Copyright (c) Citus Data, Inc. 9 * 10 *------------------------------------------------------------------------- 11 */ 12 13 #include "postgres.h" 14 #include "fmgr.h" 15 16 #include "catalog/pg_type.h" 17 #include "distributed/colocation_utils.h" 18 #include "distributed/listutils.h" 19 #include "distributed/metadata_cache.h" 20 21 22 /* declarations for dynamic loading */ 23 PG_FUNCTION_INFO_V1(get_table_colocation_id); 24 PG_FUNCTION_INFO_V1(tables_colocated); 25 PG_FUNCTION_INFO_V1(shards_colocated); 26 PG_FUNCTION_INFO_V1(get_colocated_table_array); 27 PG_FUNCTION_INFO_V1(find_shard_interval_index); 28 29 30 /* 31 * get_table_colocation_id returns colocation id of given distributed table. 32 */ 33 Datum get_table_colocation_id(PG_FUNCTION_ARGS)34get_table_colocation_id(PG_FUNCTION_ARGS) 35 { 36 Oid distributedTableId = PG_GETARG_OID(0); 37 uint32 colocationId = TableColocationId(distributedTableId); 38 39 PG_RETURN_INT32(colocationId); 40 } 41 42 43 /* 44 * tables_colocated checks if given two tables are co-located or not. If they are 45 * co-located, this function returns true. 46 */ 47 Datum tables_colocated(PG_FUNCTION_ARGS)48tables_colocated(PG_FUNCTION_ARGS) 49 { 50 Oid leftDistributedTableId = PG_GETARG_OID(0); 51 Oid rightDistributedTableId = PG_GETARG_OID(1); 52 bool tablesColocated = TablesColocated(leftDistributedTableId, 53 rightDistributedTableId); 54 55 PG_RETURN_BOOL(tablesColocated); 56 } 57 58 59 /* 60 * shards_colocated checks if given two shards are co-located or not. If they are 61 * co-located, this function returns true. 62 */ 63 Datum shards_colocated(PG_FUNCTION_ARGS)64shards_colocated(PG_FUNCTION_ARGS) 65 { 66 uint32 leftShardId = PG_GETARG_UINT32(0); 67 uint32 rightShardId = PG_GETARG_UINT32(1); 68 ShardInterval *leftShard = LoadShardInterval(leftShardId); 69 ShardInterval *rightShard = LoadShardInterval(rightShardId); 70 71 bool shardsColocated = ShardsColocated(leftShard, rightShard); 72 73 PG_RETURN_BOOL(shardsColocated); 74 } 75 76 77 /* 78 * get_colocated_tables_array returns array of table oids which are co-located with given 79 * distributed table. 80 */ 81 Datum get_colocated_table_array(PG_FUNCTION_ARGS)82get_colocated_table_array(PG_FUNCTION_ARGS) 83 { 84 Oid distributedTableId = PG_GETARG_OID(0); 85 86 List *colocatedTableList = ColocatedTableList(distributedTableId); 87 int colocatedTableCount = list_length(colocatedTableList); 88 Datum *colocatedTablesDatumArray = palloc0(colocatedTableCount * sizeof(Datum)); 89 Oid arrayTypeId = OIDOID; 90 int colocatedTableIndex = 0; 91 92 Oid colocatedTableId = InvalidOid; 93 foreach_oid(colocatedTableId, colocatedTableList) 94 { 95 Datum colocatedTableDatum = ObjectIdGetDatum(colocatedTableId); 96 97 colocatedTablesDatumArray[colocatedTableIndex] = colocatedTableDatum; 98 colocatedTableIndex++; 99 } 100 101 ArrayType *colocatedTablesArrayType = DatumArrayToArrayType(colocatedTablesDatumArray, 102 colocatedTableCount, 103 arrayTypeId); 104 105 PG_RETURN_ARRAYTYPE_P(colocatedTablesArrayType); 106 } 107 108 109 /* 110 * find_shard_interval_index finds index of given shard in sorted shard interval list. 111 */ 112 Datum find_shard_interval_index(PG_FUNCTION_ARGS)113find_shard_interval_index(PG_FUNCTION_ARGS) 114 { 115 uint32 shardId = PG_GETARG_UINT32(0); 116 ShardInterval *shardInterval = LoadShardInterval(shardId); 117 uint32 shardIndex = ShardIndex(shardInterval); 118 119 PG_RETURN_INT32(shardIndex); 120 } 121