external
/
citus
mirror of https://github.com/citusdata/citus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

mark_aggregate_for_distributed_execution: Add some parameter validation

fix_120_custom_aggregates_distribute_multiarg
Philip Dubé 2019-10-18 18:45:22 +00:00
parent 9b3260f4df
commit 364b33a22d
1 changed files with 45 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

49
src/backend/distributed/commands/function.c
View File

@ -240,6 +240,17 @@ mark_aggregate_for_distributed_execution(PG_FUNCTION_ARGS)
Oid *argtypes = NULL; Oid *argtypes = NULL;
int aggregationStrategy = -1; int aggregationStrategy = -1;
/* if called on NULL input, error out */
if (funcOid == InvalidOid)
{
ereport(ERROR, (errmsg(
"the first parameter for mark_aggregate_for_distributed_execution() "
"should be a single a valid function or procedure name "
"followed by a list of parameters in parantheses")));
}
EnsureFunctionOwner(funcOid);
if (!PG_ARGISNULL(1)) if (!PG_ARGISNULL(1))
{ {
char *strategyParam = TextDatumGetCString(PG_GETARG_TEXT_P(1)); char *strategyParam = TextDatumGetCString(PG_GETARG_TEXT_P(1));
@ -268,22 +279,27 @@ mark_aggregate_for_distributed_execution(PG_FUNCTION_ARGS)
proctup = SearchSysCache1(PROCOID, funcOid); proctup = SearchSysCache1(PROCOID, funcOid);
if (!HeapTupleIsValid(proctup)) if (!HeapTupleIsValid(proctup))
{ {
goto early_exit; elog(ERROR, "citus cache lookup failed for procedure %u", funcOid);
} }
proc = (Form_pg_proc) GETSTRUCT(proctup); proc = (Form_pg_proc) GETSTRUCT(proctup);
if (proc->prokind != PROKIND_AGGREGATE) if (proc->prokind != PROKIND_AGGREGATE)
{ {
goto early_exit; elog(ERROR, "Procedure is not an aggregate");
} }
aggtup = SearchSysCache1(AGGFNOID, funcOid); aggtup = SearchSysCache1(AGGFNOID, funcOid);
if (!HeapTupleIsValid(aggtup)) if (!HeapTupleIsValid(aggtup))
{ {
goto early_exit; elog(ERROR, "citus cache lookup failed for aggregate %u", funcOid);
} }
agg = (Form_pg_aggregate) GETSTRUCT(aggtup); agg = (Form_pg_aggregate) GETSTRUCT(aggtup);
if (agg->aggcombinefn == InvalidOid)
{
elog(ERROR, "Aggregate lacks a combine function.");
}
initStringInfo(&helperSuffix); initStringInfo(&helperSuffix);
appendStringInfoAggregateHelperSuffix(&helperSuffix, proc, agg); appendStringInfoAggregateHelperSuffix(&helperSuffix, proc, agg);
@ -365,11 +381,36 @@ mark_aggregate_for_distributed_execution(PG_FUNCTION_ARGS)
numargs, argtypes, false); numargs, argtypes, false);
} }
} }
else if (aggregationStrategy == AGGREGATION_STRATEGY_COMMUTE)
{
proctup = SearchSysCache1(PROCOID, funcOid);
if (!HeapTupleIsValid(proctup))
{
elog(ERROR, "citus cache lookup failed for procedure %u", funcOid);
}
proc = (Form_pg_proc) GETSTRUCT(proctup);
if (proc->prokind != PROKIND_AGGREGATE)
{
elog(ERROR, "Procedure is not an aggregate");
}
if (proc->pronargs != 1)
{
elog(ERROR,
"Commute aggregation strategy only works for single argument aggregates.");
}
if (proc->proargtypes.values[0] != proc->prorettype)
{
elog(ERROR,
"Commute aggregation strategy requires input type to match result type.");
}
}
/* set strategy column value */ /* set strategy column value */
UpdateDistObjectAggregationStrategy(funcOid, aggregationStrategy); UpdateDistObjectAggregationStrategy(funcOid, aggregationStrategy);
early_exit:
if (aggtup && HeapTupleIsValid(aggtup)) if (aggtup && HeapTupleIsValid(aggtup))
{ {
ReleaseSysCache(aggtup); ReleaseSysCache(aggtup);