diff --git a/src/backend/distributed/commands/non_main_db_distribute_object_ops.c b/src/backend/distributed/commands/non_main_db_distribute_object_ops.c index 2cbf3486a..c2fc3891b 100644 --- a/src/backend/distributed/commands/non_main_db_distribute_object_ops.c +++ b/src/backend/distributed/commands/non_main_db_distribute_object_ops.c @@ -143,25 +143,42 @@ IsCommandToCreateOrDropMainDB(Node *parsetree) /* - * RunPreprocessMainDBCommand runs the necessary commands for a query, in main - * database before query is run on the local node with PrevProcessUtility + * RunPreprocessNonMainDBCommand runs the necessary commands for a query, in main + * database before query is run on the local node with PrevProcessUtility. + * + * Returns true if previous utility hook needs to be skipped after completing + * preprocess phase. */ -void -RunPreprocessMainDBCommand(Node *parsetree) +bool +RunPreprocessNonMainDBCommand(Node *parsetree) { - if (!IsStatementSupportedFromNonMainDb(parsetree)) + if (IsMainDB || !IsStatementSupportedFromNonMainDb(parsetree)) { - return; + return false; + } + + if (IsCommandToCreateOrDropMainDB(parsetree)) + { + /* + * We don't try to send the query to the main database if the CREATE/DROP DATABASE + * command is for the main database itself, this is a very rare case but it's + * exercised by our test suite. + */ + return false; } char *queryString = DeparseTreeNode(parsetree); - if (IsA(parsetree, CreatedbStmt) || - IsA(parsetree, DropdbStmt)) + if (IsA(parsetree, CreatedbStmt) || IsA(parsetree, DropdbStmt)) { + /* + * We always execute CREATE/DROP DATABASE from the main database. There are no + * transactional visibility issues, since these commands are non-transactional. + * And this way we only have to consider one codepath when creating databases. + */ IsMainDBCommandInXact = false; RunCitusMainDBQuery((char *) queryString); - return; + return true; } IsMainDBCommandInXact = true; @@ -184,18 +201,24 @@ RunPreprocessMainDBCommand(Node *parsetree) List *unmarkParams = GetDistObjectOperationParams(parsetree); UnMarkObjectDistributedLocallyFromNonMainDb(unmarkParams); } + + return false; } /* - * RunPostprocessMainDBCommand runs the necessary commands for a query, in main - * database after query is run on the local node with PrevProcessUtility + * RunPostprocessNonMainDBCommand runs the necessary commands for a query, in main + * database after query is run on the local node with PrevProcessUtility. */ void -RunPostprocessMainDBCommand(Node *parsetree) +RunPostprocessNonMainDBCommand(Node *parsetree) { - if (IsStatementSupportedFromNonMainDb(parsetree) && - StatementRequiresMarkDistributedGloballyFromNonMainDb(parsetree)) + if (IsMainDB || !IsStatementSupportedFromNonMainDb(parsetree)) + { + return; + } + + if (StatementRequiresMarkDistributedGloballyFromNonMainDb(parsetree)) { MarkObjectDistributedGloballyFromNonMainDb(parsetree); } diff --git a/src/backend/distributed/commands/utility_hook.c b/src/backend/distributed/commands/utility_hook.c index 96c515fe8..9426e13c0 100644 --- a/src/backend/distributed/commands/utility_hook.c +++ b/src/backend/distributed/commands/utility_hook.c @@ -247,36 +247,25 @@ citus_ProcessUtility(PlannedStmt *pstmt, if (!CitusHasBeenLoaded()) { /* - * We always execute CREATE/DROP DATABASE from the main database. There are no - * transactional visibility issues, since these commands are non-transactional. - * And this way we only have to consider one codepath when creating databases. - * We don't try to send the query to the main database if the CREATE/DROP DATABASE - * command is for the main database itself, this is a very rare case but it's - * exercised by our test suite. + * Process the command via RunPreprocessNonMainDBCommand and + * RunPostprocessNonMainDBCommand hooks if we're in a non-main database + * and if the command is a node-wide object management command that we + * support from non-main databases. */ - if (!IsMainDB && - !IsCommandToCreateOrDropMainDB(parsetree)) - { - RunPreprocessMainDBCommand(parsetree); - if (IsA(parsetree, CreatedbStmt) || - IsA(parsetree, DropdbStmt)) - { - return; - } + bool shouldSkipPrevUtilityHook = RunPreprocessNonMainDBCommand(parsetree); + + if (!shouldSkipPrevUtilityHook) + { + /* + * Ensure that utility commands do not behave any differently until CREATE + * EXTENSION is invoked. + */ + PrevProcessUtility(pstmt, queryString, false, context, + params, queryEnv, dest, completionTag); } - /* - * Ensure that utility commands do not behave any differently until CREATE - * EXTENSION is invoked. - */ - PrevProcessUtility(pstmt, queryString, false, context, - params, queryEnv, dest, completionTag); - - if (!IsMainDB) - { - RunPostprocessMainDBCommand(parsetree); - } + RunPostprocessNonMainDBCommand(parsetree); return; } diff --git a/src/include/distributed/commands.h b/src/include/distributed/commands.h index 37043a91e..e9468521f 100644 --- a/src/include/distributed/commands.h +++ b/src/include/distributed/commands.h @@ -106,8 +106,8 @@ const DistributeObjectOps * GetDistributeObjectOps(Node *node); /* functions to support node-wide object management commands from non-main dbs */ extern bool IsCommandToCreateOrDropMainDB(Node *parsetree); -extern void RunPreprocessMainDBCommand(Node *parsetree); -extern void RunPostprocessMainDBCommand(Node *parsetree); +extern bool RunPreprocessNonMainDBCommand(Node *parsetree); +extern void RunPostprocessNonMainDBCommand(Node *parsetree); /* * Flags that can be passed to GetForeignKeyOids to indicate