From 2000470f13a3c7f5b117adbe334d811391deda91 Mon Sep 17 00:00:00 2001
From: Deepakshi Mittal <78574784+deepakshi-mittal@users.noreply.github.com>
Date: Sun, 30 Jul 2023 18:37:18 -0700
Subject: [PATCH] Support IDENTITY() function in SELECT-INTO (#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 <depakshi@amazon.com>
---
 src/backend/commands/createas.c |  9 +++++++++
 src/backend/tcop/utility.c      | 11 +++++++++--
 src/include/nodes/primnodes.h   |  2 ++
 src/include/tcop/utility.h      |  6 ++++++
 4 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/src/backend/commands/createas.c b/src/backend/commands/createas.c
index 9abbb6b5552..fe0b380b13e 100644
--- 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
index aa5f81cea4b..61beb35f870 100644
--- 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
index cfa3327805c..d5bdfe861e2 100644
--- 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
index c3016d8484b..3b093e25ed6 100644
--- 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 */