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)34 get_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)48 tables_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)64 shards_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)82 get_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)113 find_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