mirror of https://github.com/citusdata/citus.git
122 lines
3.2 KiB
C
122 lines
3.2 KiB
C
/*-------------------------------------------------------------------------
|
|
*
|
|
* test/src/colocations_utils.c
|
|
*
|
|
* This file contains functions to test co-location functionality
|
|
* within Citus.
|
|
*
|
|
* Copyright (c) 2014-2016, Citus Data, Inc.
|
|
*
|
|
*-------------------------------------------------------------------------
|
|
*/
|
|
|
|
#include "postgres.h"
|
|
#include "fmgr.h"
|
|
|
|
#include "catalog/pg_type.h"
|
|
#include "distributed/colocation_utils.h"
|
|
#include "distributed/listutils.h"
|
|
#include "distributed/metadata_cache.h"
|
|
|
|
|
|
/* declarations for dynamic loading */
|
|
PG_FUNCTION_INFO_V1(get_table_colocation_id);
|
|
PG_FUNCTION_INFO_V1(tables_colocated);
|
|
PG_FUNCTION_INFO_V1(shards_colocated);
|
|
PG_FUNCTION_INFO_V1(get_colocated_table_array);
|
|
PG_FUNCTION_INFO_V1(find_shard_interval_index);
|
|
|
|
|
|
/*
|
|
* get_table_colocation_id returns colocation id of given distributed table.
|
|
*/
|
|
Datum
|
|
get_table_colocation_id(PG_FUNCTION_ARGS)
|
|
{
|
|
Oid distributedTableId = PG_GETARG_OID(0);
|
|
uint32 colocationId = TableColocationId(distributedTableId);
|
|
|
|
PG_RETURN_INT32(colocationId);
|
|
}
|
|
|
|
|
|
/*
|
|
* tables_colocated checks if given two tables are co-located or not. If they are
|
|
* co-located, this function returns true.
|
|
*/
|
|
Datum
|
|
tables_colocated(PG_FUNCTION_ARGS)
|
|
{
|
|
Oid leftDistributedTableId = PG_GETARG_OID(0);
|
|
Oid rightDistributedTableId = PG_GETARG_OID(1);
|
|
bool tablesColocated = TablesColocated(leftDistributedTableId,
|
|
rightDistributedTableId);
|
|
|
|
PG_RETURN_BOOL(tablesColocated);
|
|
}
|
|
|
|
|
|
/*
|
|
* shards_colocated checks if given two shards are co-located or not. If they are
|
|
* co-located, this function returns true.
|
|
*/
|
|
Datum
|
|
shards_colocated(PG_FUNCTION_ARGS)
|
|
{
|
|
uint32 leftShardId = PG_GETARG_UINT32(0);
|
|
uint32 rightShardId = PG_GETARG_UINT32(1);
|
|
ShardInterval *leftShard = LoadShardInterval(leftShardId);
|
|
ShardInterval *rightShard = LoadShardInterval(rightShardId);
|
|
|
|
bool shardsColocated = ShardsColocated(leftShard, rightShard);
|
|
|
|
PG_RETURN_BOOL(shardsColocated);
|
|
}
|
|
|
|
|
|
/*
|
|
* get_colocated_tables_array returns array of table oids which are co-located with given
|
|
* distributed table.
|
|
*/
|
|
Datum
|
|
get_colocated_table_array(PG_FUNCTION_ARGS)
|
|
{
|
|
Oid distributedTableId = PG_GETARG_OID(0);
|
|
|
|
ArrayType *colocatedTablesArrayType = NULL;
|
|
List *colocatedTableList = ColocatedTableList(distributedTableId);
|
|
ListCell *colocatedTableCell = NULL;
|
|
int colocatedTableCount = list_length(colocatedTableList);
|
|
Datum *colocatedTablesDatumArray = palloc0(colocatedTableCount * sizeof(Datum));
|
|
Oid arrayTypeId = OIDOID;
|
|
int colocatedTableIndex = 0;
|
|
|
|
foreach(colocatedTableCell, colocatedTableList)
|
|
{
|
|
Oid colocatedTableId = lfirst_oid(colocatedTableCell);
|
|
Datum colocatedTableDatum = ObjectIdGetDatum(colocatedTableId);
|
|
|
|
colocatedTablesDatumArray[colocatedTableIndex] = colocatedTableDatum;
|
|
colocatedTableIndex++;
|
|
}
|
|
|
|
colocatedTablesArrayType = DatumArrayToArrayType(colocatedTablesDatumArray,
|
|
colocatedTableCount, arrayTypeId);
|
|
|
|
PG_RETURN_ARRAYTYPE_P(colocatedTablesArrayType);
|
|
}
|
|
|
|
|
|
/*
|
|
* find_shard_interval_index finds index of given shard in sorted shard interval list.
|
|
*/
|
|
Datum
|
|
find_shard_interval_index(PG_FUNCTION_ARGS)
|
|
{
|
|
uint32 shardId = PG_GETARG_UINT32(0);
|
|
ShardInterval *shardInterval = LoadShardInterval(shardId);
|
|
uint32 shardIndex = ShardIndex(shardInterval);
|
|
|
|
PG_RETURN_INT32(shardIndex);
|
|
}
|