/*
* worker_copy_shard_placement implements a internal UDF to copy a table's data from
* a healthy placement into a receiving table on an unhealthy placement. This
* function returns a boolean reflecting success or failure.
*/
Datum
worker_copy_shard_placement(PG_FUNCTION_ARGS)
{
text *shardRelationNameText = PG_GETARG_TEXT_P(0);
text *nodeNameText = PG_GETARG_TEXT_P(1);
int32 nodePort = PG_GETARG_INT32(2);
char *shardRelationName = text_to_cstring(shardRelationNameText);
char *nodeName = text_to_cstring(nodeNameText);
bool fetchSuccessful = false;
Oid shardRelationId = ResolveRelationId(shardRelationNameText);
Relation shardTable = heap_open(shardRelationId, RowExclusiveLock);
TupleDesc tupleDescriptor = RelationGetDescr(shardTable);
Tuplestorestate *tupleStore = tuplestore_begin_heap(false, false, work_mem);
StringInfo selectAllQuery = NULL;
ShardPlacement *placement = NULL;
Task *task = NULL;
selectAllQuery = makeStringInfo();
appendStringInfo(selectAllQuery, SELECT_ALL_QUERY,
quote_identifier(shardRelationName));
placement = (ShardPlacement *) palloc0(sizeof(ShardPlacement));
placement->nodeName = nodeName;
placement->nodePort = nodePort;
task = (Task *) palloc0(sizeof(Task));
task->queryString = selectAllQuery;
task->taskPlacementList = list_make1(placement);
fetchSuccessful = ExecuteTaskAndStoreResults(task, tupleDescriptor, tupleStore);
if (!fetchSuccessful)
{
ereport(ERROR, (errmsg("could not store shard rows from healthy placement"),
errhint("Consult recent messages in the server logs for "
"details.")));
}
CopyDataFromTupleStoreToRelation(tupleStore, shardTable);
tuplestore_end(tupleStore);
heap_close(shardTable, RowExclusiveLock);
PG_RETURN_VOID();
}
/*
* RouterExecutorRun actually executes a single task on a worker.
*/
void
RouterExecutorRun(QueryDesc *queryDesc, ScanDirection direction, long count, Task *task)
{
EState *estate = queryDesc->estate;
CmdType operation = queryDesc->operation;
MemoryContext oldcontext = NULL;
DestReceiver *destination = queryDesc->dest;
MaterialState *routerState = (MaterialState *) queryDesc->planstate;
bool sendTuples = operation == CMD_SELECT || queryDesc->plannedstmt->hasReturning;
Assert(estate != NULL);
Assert(!(estate->es_top_eflags & EXEC_FLAG_EXPLAIN_ONLY));
Assert(task != NULL);
/* we only support default scan direction and row fetch count */
if (!ScanDirectionIsForward(direction))
{
ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
errmsg("scan directions other than forward scans "
"are unsupported")));
}
oldcontext = MemoryContextSwitchTo(estate->es_query_cxt);
if (queryDesc->totaltime != NULL)
{
InstrStartNode(queryDesc->totaltime);
}
estate->es_processed = 0;
/* startup the tuple receiver */
if (sendTuples)
{
(*destination->rStartup)(destination, operation, queryDesc->tupDesc);
}
/*
* If query has not yet been executed, do so now. The main reason why the
* query might already have been executed is cursors.
*/
if (!routerState->eof_underlying)
{
bool resultsOK = false;
bool isModificationQuery = false;
if (operation == CMD_INSERT || operation == CMD_UPDATE ||
operation == CMD_DELETE)
{
isModificationQuery = true;
}
else if (operation != CMD_SELECT)
{
ereport(ERROR, (errmsg("unrecognized operation code: %d",
(int) operation)));
}
resultsOK = ExecuteTaskAndStoreResults(queryDesc, task,
isModificationQuery,
sendTuples);
if (!resultsOK)
{
ereport(ERROR, (errmsg("could not receive query results")));
}
/* mark underlying query as having executed */
routerState->eof_underlying = true;
}
/* if the underlying query produced output, return it */
if (routerState->tuplestorestate != NULL)
{
TupleDesc resultTupleDescriptor = queryDesc->tupDesc;
int64 returnedRows = 0;
/* return rows from the tuplestore */
returnedRows = ReturnRowsFromTuplestore(count, resultTupleDescriptor,
destination,
routerState->tuplestorestate);
/*
* Count tuples processed, if this is a SELECT. (For modifications
* it'll already have been increased, as we want the number of
* modified tuples, not the number of RETURNed tuples.)
*/
if (operation == CMD_SELECT)
{
estate->es_processed += returnedRows;
}
}
/* shutdown tuple receiver, if we started it */
if (sendTuples)
{
(*destination->rShutdown)(destination);
}
if (queryDesc->totaltime != NULL)
{
InstrStopNode(queryDesc->totaltime, estate->es_processed);
}
MemoryContextSwitchTo(oldcontext);
}
