Support IDENTITY() function in SELECT-INTO (babelfish-for-postgresql#166

) Engine changes for Identity support Added hooks for identity Added fields in into clause to pass information Task : BABEL-539 Signed-off-by: Deepakshi Mittal <[email protected]>
amazon-aurora · Nov 22, 2023 · e192a00 · e192a00
1 parent 4984129
commit e192a00
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 2 deletions.
diff --git a/src/backend/commands/createas.c b/src/backend/commands/createas.c
@@ -42,15 +42,19 @@
 #include "nodes/makefuncs.h"
 #include "nodes/nodeFuncs.h"
 #include "parser/parse_clause.h"
+#include "parser/parser.h"
 #include "rewrite/rewriteHandler.h"
 #include "storage/smgr.h"
 #include "tcop/tcopprot.h"
+#include "tcop/utility.h"
 #include "utils/builtins.h"
 #include "utils/lsyscache.h"
 #include "utils/rel.h"
 #include "utils/rls.h"
 #include "utils/snapmgr.h"
 
+bbfSelectIntoAddIdentity_hook_type bbfSelectIntoAddIdentity_hook = NULL;
+
 typedef struct
 {
 	DestReceiver pub;			/* publicly-known function pointers */
@@ -110,6 +114,11 @@ create_ctas_internal(List *attrList, IntoClause *into)
 	create->if_not_exists = false;
 	create->accessMethod = into->accessMethod;
 
+	if (sql_dialect == SQL_DIALECT_TSQL)
+	{
+		if(into->identityName && bbfSelectIntoAddIdentity_hook)
+			(*bbfSelectIntoAddIdentity_hook)(into, create->tableElts);
+	}
 	/*
 	 * Create the relation.  (This will error out if there's an existing view,
 	 * so we don't need more code to complain if "replace" is false.)

diff --git a/src/backend/tcop/utility.c b/src/backend/tcop/utility.c
@@ -77,6 +77,7 @@
 /* Hook for plugins to get control in ProcessUtility() */
 ProcessUtility_hook_type ProcessUtility_hook = NULL;
 bbfCustomProcessUtility_hook_type bbfCustomProcessUtility_hook = NULL;
+bbfSelectIntoUtility_hook_type bbfSelectIntoUtility_hook = NULL;
 
 /* local function declarations */
 static int	ClassifyUtilityCommandAsReadOnly(Node *parsetree);
@@ -1681,8 +1682,14 @@ ProcessUtilitySlow(ParseState *pstate,
 				break;
 
 			case T_CreateTableAsStmt:
-				address = ExecCreateTableAs(pstate, (CreateTableAsStmt *) parsetree,
-											params, queryEnv, qc);
+				{
+					if(sql_dialect == SQL_DIALECT_TSQL && bbfSelectIntoUtility_hook)
+						(*bbfSelectIntoUtility_hook)(pstate, pstmt, queryString, NULL, params, qc);
+					else{
+						address = ExecCreateTableAs(pstate, (CreateTableAsStmt *) parsetree,
+									params, queryEnv, qc);
+					}
+				}
 				break;
 
 			case T_RefreshMatViewStmt:

diff --git a/src/include/nodes/primnodes.h b/src/include/nodes/primnodes.h
@@ -126,6 +126,8 @@ typedef struct IntoClause
 	char	   *tableSpaceName; /* table space to use, or NULL */
 	Node	   *viewQuery;		/* materialized view's SELECT query */
 	bool		skipData;		/* true for WITH NO DATA */
+	char       *identityName;	/* resname for Identity Column*/
+	char       *identityType; 	/* pg type for Identity Column*/
 } IntoClause;
 
 

diff --git a/src/include/tcop/utility.h b/src/include/tcop/utility.h
@@ -111,5 +111,11 @@ extern bool CommandIsReadOnly(PlannedStmt *pstmt);
 typedef bool (*bbfCustomProcessUtility_hook_type)(struct ParseState *pstate, PlannedStmt *pstmt, const char *queryString, ProcessUtilityContext context, 
 						  ParamListInfo params, QueryCompletion *qc);
 extern PGDLLIMPORT bbfCustomProcessUtility_hook_type bbfCustomProcessUtility_hook;
+typedef void (*bbfSelectIntoUtility_hook_type)(struct ParseState *pstate, PlannedStmt *pstmt, const char *queryString, QueryEnvironment *queryEnv, 
+						  ParamListInfo params, QueryCompletion *qc);
+extern PGDLLIMPORT bbfSelectIntoUtility_hook_type bbfSelectIntoUtility_hook;
+
+typedef void (*bbfSelectIntoAddIdentity_hook_type)(IntoClause *into, List *tableElts);
+extern PGDLLIMPORT bbfSelectIntoAddIdentity_hook_type bbfSelectIntoAddIdentity_hook;
 
 #endif							/* UTILITY_H */