From 08a459387588131e014c882d727567b3e9410cea Mon Sep 17 00:00:00 2001
From: Deepakshi Mittal <depakshi@amazon.com>
Date: Fri, 28 Jul 2023 20:11:19 +0000
Subject: [PATCH 1/6] Support IDENTITY() function in SELECT-INTO

Changes for Identity in Select Into statement
Added hooks for identity
Added fields in intoclause to pass information about identity column

Task : BABEL-539

Signed-off-by: Deepakshi Mittal <depakshi@amazon.com>
---
 contrib/babelfishpg_tsql/runtime/functions.c  |  10 +
 .../babelfishpg_tsql/sql/sys_functions.sql    |   3 +
 .../babelfishpg_tsql--3.2.0--3.3.0.sql        |   3 +
 .../src/backend_parser/gram-tsql-epilogue.y.c |  30 +++
 .../src/backend_parser/gram-tsql-prologue.y.h |   2 +
 .../src/backend_parser/gram-tsql-rule.y       |  12 +
 contrib/babelfishpg_tsql/src/hooks.c          |  40 +++-
 contrib/babelfishpg_tsql/src/hooks.h          |   3 +
 contrib/babelfishpg_tsql/src/pl_handler.c     | 189 ++++++++++++++++
 contrib/babelfishpg_tsql/src/pltsql.h         |   1 +
 contrib/babelfishpg_tsql/src/tsqlIface.cpp    |  21 ++
 test/JDBC/expected/BABEL_539-vu-cleanup.out   |   8 +
 test/JDBC/expected/BABEL_539-vu-prepare.out   |  12 +
 test/JDBC/expected/BABEL_539-vu-verify.out    | 212 ++++++++++++++++++
 test/JDBC/input/BABEL_539-vu-cleanup.sql      |   8 +
 test/JDBC/input/BABEL_539-vu-prepare.sql      |  10 +
 test/JDBC/input/BABEL_539-vu-verify.sql       | 135 +++++++++++
 test/JDBC/upgrade/latest/schedule             |   1 +
 18 files changed, 699 insertions(+), 1 deletion(-)
 create mode 100644 test/JDBC/expected/BABEL_539-vu-cleanup.out
 create mode 100644 test/JDBC/expected/BABEL_539-vu-prepare.out
 create mode 100644 test/JDBC/expected/BABEL_539-vu-verify.out
 create mode 100644 test/JDBC/input/BABEL_539-vu-cleanup.sql
 create mode 100644 test/JDBC/input/BABEL_539-vu-prepare.sql
 create mode 100644 test/JDBC/input/BABEL_539-vu-verify.sql

diff --git a/contrib/babelfishpg_tsql/runtime/functions.c b/contrib/babelfishpg_tsql/runtime/functions.c
index eb2bc761be..ba9e653299 100644
--- a/contrib/babelfishpg_tsql/runtime/functions.c
+++ b/contrib/babelfishpg_tsql/runtime/functions.c
@@ -87,6 +87,7 @@ PG_FUNCTION_INFO_V1(object_id);
 PG_FUNCTION_INFO_V1(object_name);
 PG_FUNCTION_INFO_V1(sp_datatype_info_helper);
 PG_FUNCTION_INFO_V1(language);
+PG_FUNCTION_INFO_V1(identity_into);
 PG_FUNCTION_INFO_V1(host_name);
 PG_FUNCTION_INFO_V1(host_id);
 PG_FUNCTION_INFO_V1(context_info);
@@ -1898,6 +1899,15 @@ bbf_set_context_info(PG_FUNCTION_ARGS)
 	PG_RETURN_VOID();
 }
 
+Datum
+identity_into(PG_FUNCTION_ARGS)
+{
+	int64 result;
+	Assert(tsql_select_into_seq_oid != InvalidOid);
+	result = nextval_internal(tsql_select_into_seq_oid, false);
+	return result;
+}
+
 /*
  * Execute various integrity checks.
  * Returns true if all the checks pass otherwise
diff --git a/contrib/babelfishpg_tsql/sql/sys_functions.sql b/contrib/babelfishpg_tsql/sql/sys_functions.sql
index 1abb24df00..ddc9986235 100644
--- a/contrib/babelfishpg_tsql/sql/sys_functions.sql
+++ b/contrib/babelfishpg_tsql/sql/sys_functions.sql
@@ -3391,6 +3391,9 @@ CREATE OR REPLACE FUNCTION sys.host_id()
 RETURNS sys.VARCHAR(10)  AS 'babelfishpg_tsql' LANGUAGE C IMMUTABLE PARALLEL SAFE;
 GRANT EXECUTE ON FUNCTION sys.host_id() TO PUBLIC;
 
+CREATE OR REPLACE FUNCTION sys.identity_into(IN typename INT, IN seed INT, IN increment INT)
+RETURNS int AS 'babelfishpg_tsql' LANGUAGE C STABLE;
+
 CREATE OR REPLACE FUNCTION sys.degrees(IN arg1 BIGINT)
 RETURNS bigint  AS 'babelfishpg_tsql','bigint_degrees' LANGUAGE C STRICT IMMUTABLE PARALLEL SAFE;
 GRANT EXECUTE ON FUNCTION sys.degrees(BIGINT) TO PUBLIC;
diff --git a/contrib/babelfishpg_tsql/sql/upgrades/babelfishpg_tsql--3.2.0--3.3.0.sql b/contrib/babelfishpg_tsql/sql/upgrades/babelfishpg_tsql--3.2.0--3.3.0.sql
index 7a0b104f1b..b4ddfef7fd 100644
--- a/contrib/babelfishpg_tsql/sql/upgrades/babelfishpg_tsql--3.2.0--3.3.0.sql
+++ b/contrib/babelfishpg_tsql/sql/upgrades/babelfishpg_tsql--3.2.0--3.3.0.sql
@@ -42,6 +42,9 @@ RETURNS sys.SYSNAME
 AS 'babelfishpg_tsql', 'parsename'
 LANGUAGE C IMMUTABLE STRICT;
 
+CREATE OR REPLACE FUNCTION sys.identity_into(IN typename INT, IN seed INT, IN increment INT)
+RETURNS int AS 'babelfishpg_tsql' LANGUAGE C STABLE;
+
 CREATE OR REPLACE VIEW sys.sql_expression_dependencies
 AS
 SELECT
diff --git a/contrib/babelfishpg_tsql/src/backend_parser/gram-tsql-epilogue.y.c b/contrib/babelfishpg_tsql/src/backend_parser/gram-tsql-epilogue.y.c
index 4c69507f27..5464a78c98 100644
--- a/contrib/babelfishpg_tsql/src/backend_parser/gram-tsql-epilogue.y.c
+++ b/contrib/babelfishpg_tsql/src/backend_parser/gram-tsql-epilogue.y.c
@@ -230,6 +230,36 @@ TsqlFunctionConvert(TypeName *typename, Node *arg, Node *style, bool try, int lo
 	return result;
 }
 
+Node *
+TsqlFunctionIdentityInto(TypeName *typename, Node *seed, Node *increment, int location)
+{
+	Node *result;
+	List *args;
+	int32 typmod;
+	Oid	type_oid;
+	Oid base_oid;
+	typenameTypeIdAndMod(NULL, typename, &type_oid, &typmod);
+	base_oid =  getBaseType(type_oid);
+	switch (base_oid)
+		{
+			case INT2OID:
+			case INT4OID:
+			case INT8OID:
+			case NUMERICOID:				
+				args = list_make3((Node *)makeIntConst((int)type_oid, -1), seed, increment);
+				result = (Node *) makeFuncCall(TsqlSystemFuncName("identity_into"), args, COERCE_EXPLICIT_CALL, location);
+				break;
+			default:
+				ereport(ERROR,
+					(errcode(ERRCODE_SYNTAX_ERROR),
+					 errmsg("identity column type must be smallint, integer, or bigint")));
+
+		}
+
+	return result;
+
+}
+
 /* TsqlFunctionParse -- Implements the PARSE and TRY_PARSE functions.
  * Takes in expression, target type, regional culture, try boolean, location.
  *
diff --git a/contrib/babelfishpg_tsql/src/backend_parser/gram-tsql-prologue.y.h b/contrib/babelfishpg_tsql/src/backend_parser/gram-tsql-prologue.y.h
index a667a3e345..af5ac03e66 100644
--- a/contrib/babelfishpg_tsql/src/backend_parser/gram-tsql-prologue.y.h
+++ b/contrib/babelfishpg_tsql/src/backend_parser/gram-tsql-prologue.y.h
@@ -9,6 +9,7 @@
 #include "parser/parse_type.h"
 #include "parser/scansup.h"
 #include "utils/builtins.h"
+#include "utils/lsyscache.h"
 #include "common/md5.h"
 
 #include "src/backend_parser/gramparse.h"
@@ -44,6 +45,7 @@ static Node *TsqlFunctionConvert(TypeName *typename, Node *arg, Node *style, boo
 static Node *TsqlFunctionParse(Node *arg, TypeName *typename, Node *culture, bool try, int location);
 
 static Node *TsqlFunctionIIF(Node *bool_expr, Node *arg1, Node *arg2, int location);
+static Node *TsqlFunctionIdentityInto(TypeName *typename, Node *seed, Node *increment, int location);
 static Node *TsqlFunctionChoose(Node *int_expr, List *choosable, int location);
 static void tsql_check_param_readonly(const char *paramname, TypeName *typename, bool readonly);
 static ResTarget *TsqlForXMLMakeFuncCall(TSQL_ForClause *forclause);
diff --git a/contrib/babelfishpg_tsql/src/backend_parser/gram-tsql-rule.y b/contrib/babelfishpg_tsql/src/backend_parser/gram-tsql-rule.y
index dbe3d55886..0ee18a0975 100644
--- a/contrib/babelfishpg_tsql/src/backend_parser/gram-tsql-rule.y
+++ b/contrib/babelfishpg_tsql/src/backend_parser/gram-tsql-rule.y
@@ -1884,6 +1884,18 @@ func_expr_common_subexpr:
 				{
 					$$ = (Node *) TsqlJsonModifyMakeFuncCall($3, $5, $7);
 				}
+			| IDENTITY_P '(' Typename ',' a_expr ',' a_expr ')'
+				{
+					$$ = TsqlFunctionIdentityInto($3, $5, $7, @1);	 
+				}
+			| IDENTITY_P '(' Typename ',' a_expr ')'
+				{ 
+					$$ = TsqlFunctionIdentityInto($3, $5, (Node *)makeIntConst(1, -1), @1);
+				}
+			| IDENTITY_P '(' Typename ')'
+				{
+					$$ = TsqlFunctionIdentityInto($3, (Node *)makeIntConst(1, -1), (Node *)makeIntConst(1, -1), @1);
+				}
 		;
 
 target_el:
diff --git a/contrib/babelfishpg_tsql/src/hooks.c b/contrib/babelfishpg_tsql/src/hooks.c
index 8fc889c736..a954bb3221 100644
--- a/contrib/babelfishpg_tsql/src/hooks.c
+++ b/contrib/babelfishpg_tsql/src/hooks.c
@@ -138,6 +138,9 @@ static bool pltsql_bbfCustomProcessUtility(ParseState *pstate,
 									  const char *queryString,
 									  ProcessUtilityContext context,
 									  ParamListInfo params, QueryCompletion *qc);
+static void pltsql_bbfSelectIntoAddIdentity(IntoClause *into,  List *tableElts);
+extern void pltsql_bbfSelectIntoUtility(ParseState *pstate, PlannedStmt *pstmt, const char *queryString, 
+					QueryEnvironment *queryEnv, ParamListInfo params, QueryCompletion *qc);
 
 /*****************************************
  * 			Executor Hooks
@@ -202,6 +205,8 @@ static modify_RangeTblFunction_tupdesc_hook_type prev_modify_RangeTblFunction_tu
 static fill_missing_values_in_copyfrom_hook_type prev_fill_missing_values_in_copyfrom_hook = NULL;
 static check_rowcount_hook_type prev_check_rowcount_hook = NULL;
 static bbfCustomProcessUtility_hook_type prev_bbfCustomProcessUtility_hook = NULL;
+static bbfSelectIntoUtility_hook_type prev_bbfSelectIntoUtility_hook = NULL;
+static bbfSelectIntoAddIdentity_hook_type prev_bbfSelectIntoAddIdentity_hook = NULL;
 static sortby_nulls_hook_type prev_sortby_nulls_hook = NULL;
 static table_variable_satisfies_visibility_hook_type prev_table_variable_satisfies_visibility = NULL;
 static table_variable_satisfies_update_hook_type prev_table_variable_satisfies_update = NULL;
@@ -329,6 +334,12 @@ InstallExtendedHooks(void)
 	prev_bbfCustomProcessUtility_hook = bbfCustomProcessUtility_hook;
 	bbfCustomProcessUtility_hook = pltsql_bbfCustomProcessUtility;
 
+	prev_bbfSelectIntoUtility_hook = bbfSelectIntoUtility_hook;
+	bbfSelectIntoUtility_hook = pltsql_bbfSelectIntoUtility; 
+
+	prev_bbfSelectIntoAddIdentity_hook = bbfSelectIntoAddIdentity_hook;
+	bbfSelectIntoAddIdentity_hook = pltsql_bbfSelectIntoAddIdentity; 
+
 	prev_sortby_nulls_hook = sortby_nulls_hook;
 	sortby_nulls_hook = sort_nulls_first;
 
@@ -400,6 +411,8 @@ UninstallExtendedHooks(void)
 	fill_missing_values_in_copyfrom_hook = prev_fill_missing_values_in_copyfrom_hook;
 	check_rowcount_hook = prev_check_rowcount_hook;
 	bbfCustomProcessUtility_hook = prev_bbfCustomProcessUtility_hook;
+	bbfSelectIntoUtility_hook = prev_bbfSelectIntoUtility_hook;
+	bbfSelectIntoAddIdentity_hook = prev_bbfSelectIntoAddIdentity_hook;
 	sortby_nulls_hook = prev_sortby_nulls_hook;
 	table_variable_satisfies_visibility_hook = prev_table_variable_satisfies_visibility;
 	table_variable_satisfies_update_hook = prev_table_variable_satisfies_update;
@@ -1422,6 +1435,14 @@ pre_transform_target_entry(ResTarget *res, ParseState *pstate,
 			alias_len = strlen(res->name);
 			colname_start = pstate->p_sourcetext + res->name_location;
 		}
+		else if (res->name == NULL && IsA(res->val, FuncCall) ){
+			FuncCall *fc = (FuncCall *) res->val;
+			if (strcasecmp(strVal(llast(fc->funcname)), "identity_into") == 0){
+				ereport(ERROR,
+					(errcode(ERRCODE_SYNTAX_ERROR), errmsg("Incorrect syntax near the keyword 'INTO'"),
+						parser_errposition(pstate, res->location)));
+			}
+		}
 
 		if (alias_len > 0)
 		{
@@ -3789,4 +3810,21 @@ sort_nulls_first(SortGroupClause * sortcl, bool reverse)
 	}
 }
 
-
+static void pltsql_bbfSelectIntoAddIdentity(IntoClause *into, List *tableElts)
+{
+	ListCell   *elements;
+	foreach(elements, tableElts)
+	{
+		Node *element = lfirst(elements);
+		if (nodeTag(element) == T_ColumnDef)
+		{
+			ColumnDef *column = (ColumnDef *) element;
+			if(strcasecmp(column->colname, into->identityName) ==0){
+				column->identity = ATTRIBUTE_IDENTITY_ALWAYS;
+				column->is_not_null = true;
+				column->typeName = typeStringToTypeName(into->identityType);
+				break;
+			}
+		}
+	}	
+}
diff --git a/contrib/babelfishpg_tsql/src/hooks.h b/contrib/babelfishpg_tsql/src/hooks.h
index 0cc74c9c1f..f687140f46 100644
--- a/contrib/babelfishpg_tsql/src/hooks.h
+++ b/contrib/babelfishpg_tsql/src/hooks.h
@@ -3,6 +3,7 @@
 #include "postgres.h"
 #include "catalog/catalog.h"
 #include "parser/analyze.h"
+#include "tcop/cmdtag.h"
 
 extern IsExtendedCatalogHookType PrevIsExtendedCatalogHook;
 extern IsToastRelationHookType PrevIsToastRelationHook;
@@ -21,6 +22,8 @@ extern void pltsql_store_func_default_positions(ObjectAddress address,
 extern Oid	get_tsql_trigger_oid(List *object,
 								 const char *tsql_trigger_name,
 								 bool object_from_input);
+extern void pltsql_bbfSelectIntoUtility(ParseState *pstate, PlannedStmt *pstmt, const char *queryString, 
+					QueryEnvironment *queryEnv, ParamListInfo params, QueryCompletion *qc);
 
 extern char *update_delete_target_alias;
 extern bool sp_describe_first_result_set_inprogress;
diff --git a/contrib/babelfishpg_tsql/src/pl_handler.c b/contrib/babelfishpg_tsql/src/pl_handler.c
index 7bf9a2fdf1..34392ccd6e 100644
--- a/contrib/babelfishpg_tsql/src/pl_handler.c
+++ b/contrib/babelfishpg_tsql/src/pl_handler.c
@@ -27,6 +27,7 @@
 #include "catalog/pg_language.h"
 #include "catalog/pg_proc.h"
 #include "catalog/pg_type.h"
+#include "commands/createas.h"
 #include "commands/dbcommands.h"
 #include "commands/defrem.h"
 #include "commands/sequence.h"
@@ -140,6 +141,8 @@ extern void apply_post_compile_actions(PLtsql_function *func, InlineCodeBlockArg
 Datum		sp_prepare(PG_FUNCTION_ARGS);
 Datum		sp_unprepare(PG_FUNCTION_ARGS);
 static List *transformReturningList(ParseState *pstate, List *returningList);
+static List *transformSelectIntoStmt(CreateTableAsStmt *stmt, const char *queryString);
+static char *parse_type_argument(int type_oid);
 extern char *construct_unique_index_name(char *index_name, char *relation_name);
 extern int	CurrentLineNumber;
 static non_tsql_proc_entry_hook_type prev_non_tsql_proc_entry_hook = NULL;
@@ -175,6 +178,7 @@ static Oid	lang_handler_oid = InvalidOid;	/* Oid of language handler
 											 * function */
 static Oid	lang_validator_oid = InvalidOid;	/* Oid of language validator
 												 * function */
+Oid tsql_select_into_seq_oid = InvalidOid; /* Sequence table oid used by select into*/
 
 PG_MODULE_MAGIC;
 
@@ -5022,6 +5026,191 @@ pltsql_revert_guc(int nest_level)
 
 }
 
+static char *parse_type_argument(int type_oid){
+	char *type_string = NULL;
+	switch(type_oid){
+		case INT2OID:
+			type_string = "pg_catalog.int2";
+			break;
+		case INT4OID:
+			type_string = "pg_catalog.int4";
+			break;
+		case INT8OID:
+			type_string = "pg_catalog.int8";
+			break;
+		case NUMERICOID:
+			type_string = "pg_catalog.numeric";
+			break;
+		default:
+			type_string = "decimal";
+			break;
+	}
+	return type_string;
+}
+
+static List * transformSelectIntoStmt(CreateTableAsStmt *stmt, const char *queryString){
+	List	   *result;
+	ListCell   *elements;
+	CreateSeqStmt *seqstmt;
+	AlterSeqStmt *altseqstmt;
+	List	   *attnamelist;
+	IntoClause *into;
+	Node *n;
+
+	n = stmt->query;
+	into = stmt->into;
+	result = NIL;
+	seqstmt = NULL;
+	altseqstmt = NULL;
+
+	if (n && n->type == T_Query)
+	{
+		Query *q = (Query *) n;
+		bool seen_identity = false;
+		foreach(elements, q->targetList)
+		{	
+			TargetEntry *tle = (TargetEntry *) lfirst(elements);
+			FuncExpr *funcexpr;
+			if (tle->expr && IsA(tle->expr, FuncExpr) && strcasecmp(get_func_name(((FuncExpr *) (tle->expr))->funcid), "identity_into") ==0 ){
+
+				Oid	snamespaceid;
+				char *snamespace;
+				char *sname;
+				List *seqoptions = NIL;
+				ListCell *arg;
+
+				int type_oid;
+				char *type= NULL;
+				TypeName *ofTypename;
+				int64 seed_value;
+				int arg_num;
+
+				if(seen_identity){
+					ereport(ERROR,
+						(errcode(ERRCODE_SYNTAX_ERROR), errmsg("Attempting to add multiple identity columns to table \"%s\" using the SELECT INTO statement.", into->rel->relname )));
+				}
+
+				if(tle->resname == NULL){
+					ereport(ERROR,
+						(errcode(ERRCODE_SYNTAX_ERROR), errmsg("Incorrect syntax near the keyword 'INTO'")));
+				}
+
+				funcexpr = (FuncExpr *)tle->expr;
+				arg_num = 0;
+				foreach(arg, funcexpr->args)
+				{
+					Node *fargNode = (Node *) lfirst(arg);
+					Const *con;
+					arg_num++;
+					switch (arg_num){
+						case 1:
+							con = (Const *) fargNode;
+							type_oid = (int) con->constvalue;
+							type = parse_type_argument(type_oid);
+							ofTypename = typeStringToTypeName(type);
+							seqoptions =  lappend(seqoptions, makeDefElem("as", (Node *) ofTypename, -1));
+							break;
+						case 2:
+							con = (Const *) fargNode;
+							seqoptions = lappend(seqoptions, makeDefElem("start", (Node *)makeInteger(con->constvalue), -1));
+							seed_value = (int64) con->constvalue;
+							break;
+						case 3:
+							con = (Const *) fargNode;
+							seqoptions = lappend(seqoptions,  makeDefElem("increment", (Node *)makeInteger(con->constvalue), -1));
+							if ((int) con->constvalue > 0){
+								seqoptions = lappend(seqoptions, makeDefElem("minvalue", (Node *)makeInteger(seed_value), -1));
+							}else{
+								seqoptions = lappend(seqoptions, makeDefElem("maxvalue", (Node *)makeInteger(seed_value), -1));
+							}
+							break;
+					}
+				}
+
+				into->identityName = tle->resname;
+				into->identityType = TypeNameToString(ofTypename);
+
+				seen_identity = true;
+
+				snamespaceid = RangeVarGetCreationNamespace(into->rel);
+				snamespace = get_namespace_name(snamespaceid);
+				sname = ChooseRelationName(into->rel->relname, tle->resname,"seq",snamespaceid, false);
+				seqstmt = makeNode(CreateSeqStmt);
+				seqstmt->for_identity = true;
+				seqstmt->sequence = makeRangeVar(snamespace, sname, -1);
+				seqstmt->sequence->relpersistence = into->rel->relpersistence;
+				seqstmt->options = seqoptions;
+
+				altseqstmt = makeNode(AlterSeqStmt);
+				altseqstmt->sequence = makeRangeVar(snamespace, sname, -1);
+				attnamelist = list_make3(makeString(snamespace), makeString(into->rel->relname),makeString(tle->resname));
+				altseqstmt->options = list_make1(makeDefElem("owned_by", (Node *) attnamelist, -1));
+				altseqstmt->for_identity = true;
+
+			}
+		}
+	}
+
+	if(seqstmt){
+		result = lappend(result, seqstmt);
+	}
+	result = lappend(result, stmt);
+
+	if(altseqstmt){
+		result = lappend(result, altseqstmt);
+	}
+
+	return result;	
+}
+
+
+void
+pltsql_bbfSelectIntoUtility(ParseState *pstate, PlannedStmt *pstmt, const char *queryString, QueryEnvironment *queryEnv, 
+	ParamListInfo params, QueryCompletion *qc){
+
+	Node *parsetree = pstmt->utilityStmt;	
+	ObjectAddress address;
+	ObjectAddress secondaryObject = InvalidObjectAddress;
+	List *stmts;
+	stmts = transformSelectIntoStmt((CreateTableAsStmt *) parsetree, queryString);
+		while (stmts != NIL)
+		{
+			Node *stmt = (Node *) linitial(stmts);
+			stmts = list_delete_first(stmts);
+			if (IsA(stmt, CreateTableAsStmt))
+			{
+				address = ExecCreateTableAs(pstate, (CreateTableAsStmt *) parsetree,params, queryEnv, qc);
+				EventTriggerCollectSimpleCommand(address,secondaryObject,stmt);
+			}
+			else if(IsA(stmt, CreateSeqStmt)) {
+				address = DefineSequence(pstate, (CreateSeqStmt *) stmt);
+				Assert(address.objectId != InvalidOid);
+				tsql_select_into_seq_oid = address.objectId;
+				EventTriggerCollectSimpleCommand(address,secondaryObject,stmt);
+			}
+			else{
+
+				PlannedStmt *wrapper;
+				wrapper = makeNode(PlannedStmt);
+				wrapper->commandType = CMD_UTILITY;
+				wrapper->canSetTag = false;
+				wrapper->utilityStmt = stmt;
+				wrapper->stmt_location = pstmt->stmt_location;
+				wrapper->stmt_len = pstmt->stmt_len;
+
+				ProcessUtility(wrapper,
+								queryString,
+								false,
+								PROCESS_UTILITY_SUBCOMMAND,
+								params,
+								NULL,
+								None_Receiver,
+								NULL);
+			}
+		if (stmts != NIL)
+			CommandCounterIncrement();
+		}
+}
 void
 set_current_query_is_create_tbl_check_constraint(Node *expr)
 {
diff --git a/contrib/babelfishpg_tsql/src/pltsql.h b/contrib/babelfishpg_tsql/src/pltsql.h
index 0acd92b380..e9cd93e2a7 100644
--- a/contrib/babelfishpg_tsql/src/pltsql.h
+++ b/contrib/babelfishpg_tsql/src/pltsql.h
@@ -1781,6 +1781,7 @@ extern int	fetch_status_var;
 extern int	pltsql_proc_return_code;
 
 extern char *pltsql_version;
+extern Oid tsql_select_into_seq_oid;
 
 typedef struct PLtsqlErrorData
 {
diff --git a/contrib/babelfishpg_tsql/src/tsqlIface.cpp b/contrib/babelfishpg_tsql/src/tsqlIface.cpp
index 298b7bb55f..4c644812ad 100644
--- a/contrib/babelfishpg_tsql/src/tsqlIface.cpp
+++ b/contrib/babelfishpg_tsql/src/tsqlIface.cpp
@@ -244,6 +244,7 @@ static std::string validate_and_stringify_hints();
 static int find_hint_offset(const char * queryTxt);
 
 static bool pltsql_parseonly = false;
+bool has_identity_function = false;
 
 static void
 breakHere()
@@ -1846,6 +1847,16 @@ class tsqlBuilder : public tsqlCommonMutator
 		}
 	}
 
+	void enterFunction_call(TSqlParser::Function_callContext *ctx) override
+	{
+		std::string func_name = ::getFullText(ctx);
+		size_t Lbracket_index = func_name.find('(');
+        func_name = func_name.substr(0, Lbracket_index);
+		if (pg_strcasecmp(func_name.c_str(), "identity") == 0) {
+			has_identity_function = true;
+		}
+	}
+
 	//////////////////////////////////////////////////////////////////////////////
 	// function/procedure call analysis
 	//////////////////////////////////////////////////////////////////////////////
@@ -1961,8 +1972,18 @@ class tsqlBuilder : public tsqlCommonMutator
 	// statement analysis
 	//////////////////////////////////////////////////////////////////////////////
 
+	void enterQuery_specification(TSqlParser::Query_specificationContext *ctx) override
+	{
+		has_identity_function = false;
+	}
+
 	void exitQuery_specification(TSqlParser::Query_specificationContext *ctx) override
 	{
+		// if select doesnt contains into but it contains identity we should throw error
+		if(has_identity_function && !ctx->INTO()){
+			throw PGErrorWrapperException(ERROR, ERRCODE_SYNTAX_ERROR, "Incorrect syntax near ')'", getLineAndPos(ctx));
+		}
+		has_identity_function = false;
 		if (statementMutator)
 			process_query_specification(ctx, statementMutator.get());
 	}
diff --git a/test/JDBC/expected/BABEL_539-vu-cleanup.out b/test/JDBC/expected/BABEL_539-vu-cleanup.out
new file mode 100644
index 0000000000..3dbad8588e
--- /dev/null
+++ b/test/JDBC/expected/BABEL_539-vu-cleanup.out
@@ -0,0 +1,8 @@
+DROP TABLE IF EXISTS babel_539NewTable_proc;
+GO
+
+DROP PROC babel_539_prepare_proc;
+GO
+
+DROP TABLE IF EXISTS babel_539OldTable;
+GO
diff --git a/test/JDBC/expected/BABEL_539-vu-prepare.out b/test/JDBC/expected/BABEL_539-vu-prepare.out
new file mode 100644
index 0000000000..7f24b118ef
--- /dev/null
+++ b/test/JDBC/expected/BABEL_539-vu-prepare.out
@@ -0,0 +1,12 @@
+CREATE TABLE babel_539OldTable(col1 int , name varchar(20));
+GO
+
+INSERT INTO babel_539OldTable VALUES (10, 'user1') , (20, 'user2');
+GO
+~~ROW COUNT: 2~~
+
+
+CREATE PROC babel_539_prepare_proc
+AS
+SELECT col1, IDENTITY(int, 1,2) AS id_num INTO babel_539NewTable_proc FROM babel_539OldTable order by id_num;
+GO
diff --git a/test/JDBC/expected/BABEL_539-vu-verify.out b/test/JDBC/expected/BABEL_539-vu-verify.out
new file mode 100644
index 0000000000..5913a18679
--- /dev/null
+++ b/test/JDBC/expected/BABEL_539-vu-verify.out
@@ -0,0 +1,212 @@
+EXEC babel_539_prepare_proc
+GO
+
+SELECT col1 INTO babel_539NewTable1 FROM babel_539OldTable;
+GO
+
+Select col1 FROM babel_539NewTable1; 
+GO
+~~START~~
+int
+10
+20
+~~END~~
+
+
+DROP TABLE IF EXISTS babel_539NewTable1;
+GO
+
+SELECT col1, IDENTITY(int, 1,1) AS id_num INTO babel_539NewTable1 FROM babel_539OldTable;
+GO
+
+Select col1, id_num FROM babel_539NewTable1 order by id_num;
+GO
+~~START~~
+int#!#int
+10#!#1
+20#!#2
+~~END~~
+
+
+DROP TABLE IF EXISTS babel_539NewTable1;
+GO
+
+SELECT col1, IDENTITY(int, 1) AS id_num INTO #babel_539NewTable1 FROM babel_539OldTable;
+GO
+
+Select col1, id_num FROM #babel_539NewTable1 order by id_num;
+GO
+~~START~~
+int#!#int
+10#!#1
+20#!#2
+~~END~~
+
+
+DROP TABLE IF EXISTS #babel_539NewTable1;
+GO
+
+SELECT col1, IDENTITY(int) AS id_num INTO #babel_539NewTable1 FROM babel_539OldTable;
+GO
+
+Select col1, id_num FROM #babel_539NewTable1 order by id_num;
+GO
+~~START~~
+int#!#int
+10#!#1
+20#!#2
+~~END~~
+
+
+DROP TABLE IF EXISTS #babel_539NewTable1;
+GO
+
+SELECT col1, id_num=IDENTITY(int, 1,1) INTO babel_539NewTable1 FROM babel_539OldTable;
+GO
+
+Select col1, id_num FROM #babel_539NewTable1 order by id_num; 
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: relation "#babel_539newtable1" does not exist)~~
+
+
+DROP TABLE IF EXISTS #babel_539NewTable1;
+GO
+
+SELECT col1, [id_num]=IDENTITY(int, 1,1) INTO #babel_539NewTable1 FROM babel_539OldTable;
+GO
+
+Select col1, id_num FROM #babel_539NewTable1 order by id_num; 
+GO
+~~START~~
+int#!#int
+10#!#1
+20#!#2
+~~END~~
+
+
+DROP TABLE IF EXISTS #babel_539NewTable1;
+GO
+
+SELECT col1, IDENTITY(int, 1,1) AS [id_num] INTO #babel_539NewTempTable1 FROM babel_539OldTable;
+GO
+
+Select col1, id_num FROM #babel_539NewTempTable1 order by id_num; 
+GO
+~~START~~
+int#!#int
+10#!#1
+20#!#2
+~~END~~
+
+
+DROP TABLE IF EXISTS #babel_539NewTempTable1;
+GO
+
+SELECT *, IDENTITY(int) AS [id_num] INTO #babel_539NewTempTable1 FROM babel_539OldTable;
+GO
+
+Select col1, name, id_num FROM #babel_539NewTempTable1 order by id_num; 
+GO
+~~START~~
+int#!#varchar#!#int
+10#!#user1#!#1
+20#!#user2#!#2
+~~END~~
+
+
+DROP TABLE IF EXISTS #babel_539NewTempTable1;
+GO
+
+SELECT col1, IDENTITY(char, 1,1) AS id_num INTO babel_539NewTable1 FROM babel_539OldTable;
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: identity column type must be smallint, integer, or bigint)~~
+
+
+DROP TABLE IF EXISTS babel_539NewTable1;
+GO
+
+SELECT col1, IDENTITY() AS id_num INTO babel_539NewTable1 FROM babel_539OldTable;
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: syntax error at or near ")")~~
+
+
+-- impact to other similar queries and functions
+-- normal create table cases
+CREATE TABLE babel_539OldTable2 (col1 int NOT NULL, id_num INT IDENTITY(1, 2));
+GO
+
+INSERT INTO babel_539OldTable2 VALUES (10), (20), (30), (40);
+GO
+~~ROW COUNT: 4~~
+
+
+SELECT col1, id_num INTO babel_539NewTable2 FROM babel_539OldTable2 order by id_num;
+GO
+
+DROP TABLE IF EXISTS babel_539OldTable2;
+GO
+
+DROP TABLE IF EXISTS babel_539NewTable2;
+GO
+
+-- create table as temp table
+CREATE TABLE #babel_539NewTempTable2 (col1 int, id_num int IDENTITY(-1, 2));
+GO
+
+INSERT INTO #babel_539NewTempTable2(col1) values (10),(20),(30);
+GO
+~~ROW COUNT: 3~~
+
+
+Select col1, id_num FROM #babel_539NewTempTable2 order by id_num; 
+GO
+~~START~~
+int#!#int
+10#!#-1
+20#!#1
+30#!#3
+~~END~~
+
+
+DROP TABLE IF EXISTS #babel_539NewTempTable2;
+GO
+
+CREATE TABLE #babel_539NewTempTable2 (col1 int);
+GO
+
+Select col1 FROM #babel_539NewTempTable2 order by col1; 
+GO
+~~START~~
+int
+~~END~~
+
+
+
+ALTER TABLE #babel_539NewTempTable2 ADD id_num int IDENTITY(-1, 2);
+-- try altering table and check other columns, sequence should drop and any constraints also
+Select col1, id_num FROM #babel_539NewTempTable2 order by id_num; 
+GO
+~~START~~
+int#!#int
+~~END~~
+
+
+DROP TABLE IF EXISTS #babel_539NewTempTable2;
+GO
+
+-- Two identity columns in a query
+SELECT col1, IDENTITY(int, 1,1) as id_num, IDENTITY(int, 1,1) as id_num2 INTO #babel_539NewTempTable2 FROM babel_539OldTable;
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: Attempting to add multiple identity columns to table "#babel_539newtemptable2" using the SELECT INTO statement.)~~
+
+
+DROP TABLE IF EXISTS #babel_539NewTempTable2;
+GO
diff --git a/test/JDBC/input/BABEL_539-vu-cleanup.sql b/test/JDBC/input/BABEL_539-vu-cleanup.sql
new file mode 100644
index 0000000000..5fefbe0105
--- /dev/null
+++ b/test/JDBC/input/BABEL_539-vu-cleanup.sql
@@ -0,0 +1,8 @@
+DROP TABLE IF EXISTS babel_539NewTable_proc;
+GO
+
+DROP PROC babel_539_prepare_proc;
+GO
+
+DROP TABLE IF EXISTS babel_539OldTable;
+GO
\ No newline at end of file
diff --git a/test/JDBC/input/BABEL_539-vu-prepare.sql b/test/JDBC/input/BABEL_539-vu-prepare.sql
new file mode 100644
index 0000000000..70c79e3272
--- /dev/null
+++ b/test/JDBC/input/BABEL_539-vu-prepare.sql
@@ -0,0 +1,10 @@
+CREATE TABLE babel_539OldTable(col1 int , name varchar(20));
+GO
+
+INSERT INTO babel_539OldTable VALUES (10, 'user1') , (20, 'user2');
+GO
+
+CREATE PROC babel_539_prepare_proc
+AS
+SELECT col1, IDENTITY(int, 1,2) AS id_num INTO babel_539NewTable_proc FROM babel_539OldTable order by id_num;
+GO
\ No newline at end of file
diff --git a/test/JDBC/input/BABEL_539-vu-verify.sql b/test/JDBC/input/BABEL_539-vu-verify.sql
new file mode 100644
index 0000000000..d2b0be9402
--- /dev/null
+++ b/test/JDBC/input/BABEL_539-vu-verify.sql
@@ -0,0 +1,135 @@
+EXEC babel_539_prepare_proc
+GO
+
+SELECT col1 INTO babel_539NewTable1 FROM babel_539OldTable;
+GO
+
+Select col1 FROM babel_539NewTable1; 
+GO
+
+DROP TABLE IF EXISTS babel_539NewTable1;
+GO
+
+SELECT col1, IDENTITY(int, 1,1) AS id_num INTO babel_539NewTable1 FROM babel_539OldTable;
+GO
+
+Select col1, id_num FROM babel_539NewTable1 order by id_num;
+GO
+
+DROP TABLE IF EXISTS babel_539NewTable1;
+GO
+
+SELECT col1, IDENTITY(int, 1) AS id_num INTO #babel_539NewTable1 FROM babel_539OldTable;
+GO
+
+Select col1, id_num FROM #babel_539NewTable1 order by id_num;
+GO
+
+DROP TABLE IF EXISTS #babel_539NewTable1;
+GO
+
+SELECT col1, IDENTITY(int) AS id_num INTO #babel_539NewTable1 FROM babel_539OldTable;
+GO
+
+Select col1, id_num FROM #babel_539NewTable1 order by id_num;
+GO
+
+DROP TABLE IF EXISTS #babel_539NewTable1;
+GO
+
+SELECT col1, id_num=IDENTITY(int, 1,1) INTO babel_539NewTable1 FROM babel_539OldTable;
+GO
+
+Select col1, id_num FROM #babel_539NewTable1 order by id_num; 
+GO
+
+DROP TABLE IF EXISTS #babel_539NewTable1;
+GO
+
+SELECT col1, [id_num]=IDENTITY(int, 1,1) INTO #babel_539NewTable1 FROM babel_539OldTable;
+GO
+
+Select col1, id_num FROM #babel_539NewTable1 order by id_num; 
+GO
+
+DROP TABLE IF EXISTS #babel_539NewTable1;
+GO
+
+SELECT col1, IDENTITY(int, 1,1) AS [id_num] INTO #babel_539NewTempTable1 FROM babel_539OldTable;
+GO
+
+Select col1, id_num FROM #babel_539NewTempTable1 order by id_num; 
+GO
+
+DROP TABLE IF EXISTS #babel_539NewTempTable1;
+GO
+
+SELECT *, IDENTITY(int) AS [id_num] INTO #babel_539NewTempTable1 FROM babel_539OldTable;
+GO
+
+Select col1, name, id_num FROM #babel_539NewTempTable1 order by id_num; 
+GO
+
+DROP TABLE IF EXISTS #babel_539NewTempTable1;
+GO
+
+SELECT col1, IDENTITY(char, 1,1) AS id_num INTO babel_539NewTable1 FROM babel_539OldTable;
+GO
+
+DROP TABLE IF EXISTS babel_539NewTable1;
+GO
+
+SELECT col1, IDENTITY() AS id_num INTO babel_539NewTable1 FROM babel_539OldTable;
+GO
+
+-- impact to other similar queries and functions
+-- normal create table cases
+CREATE TABLE babel_539OldTable2 (col1 int NOT NULL, id_num INT IDENTITY(1, 2));
+GO
+
+INSERT INTO babel_539OldTable2 VALUES (10), (20), (30), (40);
+GO
+
+SELECT col1, id_num INTO babel_539NewTable2 FROM babel_539OldTable2 order by id_num;
+GO
+
+DROP TABLE IF EXISTS babel_539OldTable2;
+GO
+
+DROP TABLE IF EXISTS babel_539NewTable2;
+GO
+
+-- create table as temp table
+CREATE TABLE #babel_539NewTempTable2 (col1 int, id_num int IDENTITY(-1, 2));
+GO
+
+INSERT INTO #babel_539NewTempTable2(col1) values (10),(20),(30);
+GO
+
+Select col1, id_num FROM #babel_539NewTempTable2 order by id_num; 
+GO
+
+DROP TABLE IF EXISTS #babel_539NewTempTable2;
+GO
+
+CREATE TABLE #babel_539NewTempTable2 (col1 int);
+GO
+
+Select col1 FROM #babel_539NewTempTable2 order by col1; 
+GO
+
+ALTER TABLE #babel_539NewTempTable2 ADD id_num int IDENTITY(-1, 2);
+-- try altering table and check other columns, sequence should drop and any constraints also
+
+Select col1, id_num FROM #babel_539NewTempTable2 order by id_num; 
+GO
+
+DROP TABLE IF EXISTS #babel_539NewTempTable2;
+GO
+
+-- Two identity columns in a query
+SELECT col1, IDENTITY(int, 1,1) as id_num, IDENTITY(int, 1,1) as id_num2 INTO #babel_539NewTempTable2 FROM babel_539OldTable;
+GO
+
+DROP TABLE IF EXISTS #babel_539NewTempTable2;
+GO
\ No newline at end of file
diff --git a/test/JDBC/upgrade/latest/schedule b/test/JDBC/upgrade/latest/schedule
index 9cad8338ac..24d3d20d69 100644
--- a/test/JDBC/upgrade/latest/schedule
+++ b/test/JDBC/upgrade/latest/schedule
@@ -210,6 +210,7 @@ fulltextserviceproperty
 get_tds_id
 HAS_DBACCESS
 identity_function
+BABEL_539
 indexproperty
 insteadof_nested_trigger_inside_proc
 insteadof_nested_trigger_with_dml

From 9af3466b4c106b8341006aa9d87f152c1e98d8f7 Mon Sep 17 00:00:00 2001
From: Deepakshi Mittal <depakshi@amazon.com>
Date: Fri, 28 Jul 2023 21:30:40 +0000
Subject: [PATCH 2/6] Added test to all upgrade files

Added syntax error if query tries to calls internal function

Signed-off-by: Deepakshi Mittal <depakshi@amazon.com>
---
 contrib/babelfishpg_tsql/src/tsqlIface.cpp | 3 ++-
 test/JDBC/upgrade/13_4/schedule            | 1 +
 test/JDBC/upgrade/13_5/schedule            | 1 +
 test/JDBC/upgrade/13_6/schedule            | 1 +
 test/JDBC/upgrade/13_7/schedule            | 1 +
 test/JDBC/upgrade/13_8/schedule            | 1 +
 test/JDBC/upgrade/13_9/schedule            | 1 +
 test/JDBC/upgrade/14_3/schedule            | 1 +
 test/JDBC/upgrade/14_5/schedule            | 1 +
 test/JDBC/upgrade/14_6/schedule            | 1 +
 test/JDBC/upgrade/14_7/schedule            | 1 +
 test/JDBC/upgrade/14_8/schedule            | 1 +
 test/JDBC/upgrade/14_9/schedule            | 1 +
 test/JDBC/upgrade/15_1/schedule            | 1 +
 test/JDBC/upgrade/15_2/schedule            | 1 +
 test/JDBC/upgrade/15_3/schedule            | 1 +
 test/JDBC/upgrade/latest/schedule          | 2 +-
 17 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/contrib/babelfishpg_tsql/src/tsqlIface.cpp b/contrib/babelfishpg_tsql/src/tsqlIface.cpp
index 4c644812ad..0d2bfb1c39 100644
--- a/contrib/babelfishpg_tsql/src/tsqlIface.cpp
+++ b/contrib/babelfishpg_tsql/src/tsqlIface.cpp
@@ -1852,7 +1852,8 @@ class tsqlBuilder : public tsqlCommonMutator
 		std::string func_name = ::getFullText(ctx);
 		size_t Lbracket_index = func_name.find('(');
         func_name = func_name.substr(0, Lbracket_index);
-		if (pg_strcasecmp(func_name.c_str(), "identity") == 0) {
+		if (pg_strcasecmp(func_name.c_str(), "identity") == 0 || pg_strcasecmp(func_name.c_str(), "identity_into") == 0
+			|| pg_strcasecmp(func_name.c_str(), "sys.identity_into") == 0) {
 			has_identity_function = true;
 		}
 	}
diff --git a/test/JDBC/upgrade/13_4/schedule b/test/JDBC/upgrade/13_4/schedule
index c7be099e66..ed11eb933d 100644
--- a/test/JDBC/upgrade/13_4/schedule
+++ b/test/JDBC/upgrade/13_4/schedule
@@ -61,6 +61,7 @@ BABEL-404
 BABEL-4078-before-14_8-or-15_3
 babel_417
 BABEL-493
+BABEL_539
 BABEL-621
 BABEL-728
 BABEL-741
diff --git a/test/JDBC/upgrade/13_5/schedule b/test/JDBC/upgrade/13_5/schedule
index 6c4e12bec3..1e224e0b43 100644
--- a/test/JDBC/upgrade/13_5/schedule
+++ b/test/JDBC/upgrade/13_5/schedule
@@ -67,6 +67,7 @@ BABEL-404
 BABEL-4078-before-14_8-or-15_3
 babel_417
 BABEL-493
+BABEL_539
 BABEL-621
 BABEL-728
 BABEL-741
diff --git a/test/JDBC/upgrade/13_6/schedule b/test/JDBC/upgrade/13_6/schedule
index ba0280cf44..4f0b98bba0 100644
--- a/test/JDBC/upgrade/13_6/schedule
+++ b/test/JDBC/upgrade/13_6/schedule
@@ -87,6 +87,7 @@ BABEL-405
 BABEL-4078-before-14_8-or-15_3
 babel_417
 BABEL-493
+BABEL_539
 BABEL-621
 BABEL-728
 BABEL-741
diff --git a/test/JDBC/upgrade/13_7/schedule b/test/JDBC/upgrade/13_7/schedule
index 5631caa980..474e6900d4 100644
--- a/test/JDBC/upgrade/13_7/schedule
+++ b/test/JDBC/upgrade/13_7/schedule
@@ -85,6 +85,7 @@ BABEL-405
 BABEL-4078-before-14_8-or-15_3
 babel_417
 BABEL-493
+BABEL_539
 BABEL-621
 BABEL-728
 BABEL-741
diff --git a/test/JDBC/upgrade/13_8/schedule b/test/JDBC/upgrade/13_8/schedule
index 5631caa980..474e6900d4 100644
--- a/test/JDBC/upgrade/13_8/schedule
+++ b/test/JDBC/upgrade/13_8/schedule
@@ -85,6 +85,7 @@ BABEL-405
 BABEL-4078-before-14_8-or-15_3
 babel_417
 BABEL-493
+BABEL_539
 BABEL-621
 BABEL-728
 BABEL-741
diff --git a/test/JDBC/upgrade/13_9/schedule b/test/JDBC/upgrade/13_9/schedule
index d9f9545674..ee4988f5f6 100644
--- a/test/JDBC/upgrade/13_9/schedule
+++ b/test/JDBC/upgrade/13_9/schedule
@@ -84,6 +84,7 @@ BABEL-405
 BABEL-4078-before-14_8-or-15_3
 babel_417
 BABEL-493
+BABEL_539
 BABEL-621
 BABEL-728
 BABEL-741
diff --git a/test/JDBC/upgrade/14_3/schedule b/test/JDBC/upgrade/14_3/schedule
index 9e13e0d3c6..0c43f378f1 100644
--- a/test/JDBC/upgrade/14_3/schedule
+++ b/test/JDBC/upgrade/14_3/schedule
@@ -89,6 +89,7 @@ BABEL-405
 BABEL-4078-before-14_8-or-15_3
 babel_417
 BABEL-493
+BABEL_539
 BABEL-621
 BABEL-728
 BABEL-741
diff --git a/test/JDBC/upgrade/14_5/schedule b/test/JDBC/upgrade/14_5/schedule
index 79649590da..bd8f56d93e 100644
--- a/test/JDBC/upgrade/14_5/schedule
+++ b/test/JDBC/upgrade/14_5/schedule
@@ -87,6 +87,7 @@ BABEL-405
 BABEL-4078-before-14_8-or-15_3
 babel_417
 BABEL-493
+BABEL_539
 BABEL-621
 BABEL-728
 BABEL-741
diff --git a/test/JDBC/upgrade/14_6/schedule b/test/JDBC/upgrade/14_6/schedule
index 333593cf49..2718651c3c 100644
--- a/test/JDBC/upgrade/14_6/schedule
+++ b/test/JDBC/upgrade/14_6/schedule
@@ -100,6 +100,7 @@ BABEL-405
 BABEL-4078-before-14_8-or-15_3
 babel_417
 BABEL-493
+BABEL_539
 BABEL-621
 BABEL-728
 BABEL-741
diff --git a/test/JDBC/upgrade/14_7/schedule b/test/JDBC/upgrade/14_7/schedule
index 3697af7547..ac2ca1e61c 100644
--- a/test/JDBC/upgrade/14_7/schedule
+++ b/test/JDBC/upgrade/14_7/schedule
@@ -107,6 +107,7 @@ BABEL-4078-before-14_8-or-15_3
 BABEL-4098
 babel_417
 BABEL-493
+BABEL_539
 BABEL-621
 BABEL-728
 BABEL-733
diff --git a/test/JDBC/upgrade/14_8/schedule b/test/JDBC/upgrade/14_8/schedule
index 79c461bb3b..5b32ccd913 100644
--- a/test/JDBC/upgrade/14_8/schedule
+++ b/test/JDBC/upgrade/14_8/schedule
@@ -107,6 +107,7 @@ BABEL-4078
 BABEL-4098
 babel_417
 BABEL-493
+BABEL_539
 BABEL-621
 BABEL-728
 BABEL-733
diff --git a/test/JDBC/upgrade/14_9/schedule b/test/JDBC/upgrade/14_9/schedule
index d8e10fbe7d..e9682f9ca4 100644
--- a/test/JDBC/upgrade/14_9/schedule
+++ b/test/JDBC/upgrade/14_9/schedule
@@ -86,6 +86,7 @@ sys-trigger_nestlevel
 schema_resolution_proc
 BABEL-404
 BABEL-493
+BABEL_539
 BABEL-621
 BABEL-775
 BABEL-1206
diff --git a/test/JDBC/upgrade/15_1/schedule b/test/JDBC/upgrade/15_1/schedule
index a1b2d06dac..61bd5db0f9 100644
--- a/test/JDBC/upgrade/15_1/schedule
+++ b/test/JDBC/upgrade/15_1/schedule
@@ -100,6 +100,7 @@ BABEL-405
 BABEL-4078-before-14_8-or-15_3
 babel_417
 BABEL-493
+BABEL_539
 BABEL-621
 BABEL-728
 BABEL-741
diff --git a/test/JDBC/upgrade/15_2/schedule b/test/JDBC/upgrade/15_2/schedule
index 775fb97c34..8980ac31de 100644
--- a/test/JDBC/upgrade/15_2/schedule
+++ b/test/JDBC/upgrade/15_2/schedule
@@ -109,6 +109,7 @@ BABEL-4098
 BABEL-4168
 babel_417
 BABEL-493
+BABEL_539
 BABEL-621
 BABEL-728
 BABEL-733
diff --git a/test/JDBC/upgrade/15_3/schedule b/test/JDBC/upgrade/15_3/schedule
index 43b356e331..ba3d839a2e 100644
--- a/test/JDBC/upgrade/15_3/schedule
+++ b/test/JDBC/upgrade/15_3/schedule
@@ -115,6 +115,7 @@ BABEL-4078
 BABEL-4098
 babel_417
 BABEL-493
+BABEL_539
 BABEL-621
 BABEL-728
 BABEL-733
diff --git a/test/JDBC/upgrade/latest/schedule b/test/JDBC/upgrade/latest/schedule
index 24d3d20d69..e5ebc61230 100644
--- a/test/JDBC/upgrade/latest/schedule
+++ b/test/JDBC/upgrade/latest/schedule
@@ -115,6 +115,7 @@ BABEL-4078
 BABEL-4098
 babel_417
 BABEL-493
+BABEL_539
 BABEL-621
 BABEL-728
 BABEL-733
@@ -210,7 +211,6 @@ fulltextserviceproperty
 get_tds_id
 HAS_DBACCESS
 identity_function
-BABEL_539
 indexproperty
 insteadof_nested_trigger_inside_proc
 insteadof_nested_trigger_with_dml

From 7d3255ca3b5b67faa19a4a231bc7d256b19098ba Mon Sep 17 00:00:00 2001
From: Deepakshi Mittal <depakshi@amazon.com>
Date: Fri, 28 Jul 2023 22:04:31 +0000
Subject: [PATCH 3/6] Added missing function for failing upgrade test

Task: BABEL-539

Signed-off-by: Deepakshi Mittal <depakshi@amazon.com>
---
 .../python/expected/sql_validation_framework/expected_create.out | 1 +
 test/python/expected/upgrade_validation/expected_dependency.out  | 1 +
 2 files changed, 2 insertions(+)

diff --git a/test/python/expected/sql_validation_framework/expected_create.out b/test/python/expected/sql_validation_framework/expected_create.out
index c4c0022283..d61b60e7df 100644
--- a/test/python/expected/sql_validation_framework/expected_create.out
+++ b/test/python/expected/sql_validation_framework/expected_create.out
@@ -124,6 +124,7 @@ Could not find upgrade tests for function sys.get_current_full_xact_id
 Could not find upgrade tests for function sys.get_host_os
 Could not find upgrade tests for function sys.ident_incr
 Could not find upgrade tests for function sys.ident_seed
+Could not find upgrade tests for function sys.identity_into
 Could not find upgrade tests for function sys.is_collated_ci_as
 Could not find upgrade tests for function sys.is_collated_ci_as_internal
 Could not find upgrade tests for function sys.is_rolemember_internal
diff --git a/test/python/expected/upgrade_validation/expected_dependency.out b/test/python/expected/upgrade_validation/expected_dependency.out
index 0e25b4c6dd..63be382a45 100644
--- a/test/python/expected/upgrade_validation/expected_dependency.out
+++ b/test/python/expected/upgrade_validation/expected_dependency.out
@@ -429,6 +429,7 @@ Function sys.host_name()
 Function sys.ident_current(text)
 Function sys.ident_incr(text)
 Function sys.ident_seed(text)
+Function sys.identity_into(integer,integer,integer)
 Function sys.indexproperty(integer,sys.nvarchar,sys."varchar")
 Function sys.inject_fault(text)
 Function sys.inject_fault(text,integer)

From 52937897903605784ea6634e69f04a34029e4f1e Mon Sep 17 00:00:00 2001
From: Deepakshi Mittal <depakshi@amazon.com>
Date: Fri, 28 Jul 2023 22:18:38 +0000
Subject: [PATCH 4/6] Added missing function for failing Validate test

Task: BABEL-539

Signed-off-by: Deepakshi Mittal <depakshi@amazon.com>
---
 .../python/expected/sql_validation_framework/expected_create.out | 1 +
 1 file changed, 1 insertion(+)

diff --git a/test/python/expected/sql_validation_framework/expected_create.out b/test/python/expected/sql_validation_framework/expected_create.out
index d61b60e7df..5c604e31f9 100644
--- a/test/python/expected/sql_validation_framework/expected_create.out
+++ b/test/python/expected/sql_validation_framework/expected_create.out
@@ -17,6 +17,7 @@ Could not find tests for function sys.format_datetime
 Could not find tests for function sys.format_numeric
 Could not find tests for function sys.get_current_full_xact_id
 Could not find tests for function sys.get_host_os
+Could not find tests for function sys.identity_into
 Could not find tests for function sys.is_collated_ci_as
 Could not find tests for function sys.is_collated_ci_as_internal
 Could not find tests for function sys.is_rolemember_internal

From e0774d605680902d778cb6de8c557a9db2856409 Mon Sep 17 00:00:00 2001
From: Deepakshi Mittal <depakshi@amazon.com>
Date: Sun, 30 Jul 2023 18:07:59 +0000
Subject: [PATCH 5/6] Code review comments

Corrected formatting
Added case for decimal
used inbuilt function for getting procedure name
Added syntax error if identity_into function is called directly

Task: BABEL-539

Signed-off-by: Deepakshi Mittal <depakshi@amazon.com>
---
 contrib/babelfishpg_tsql/runtime/functions.c  |   2 +-
 .../babelfishpg_tsql/sql/sys_functions.sql    |   1 +
 .../babelfishpg_tsql--3.2.0--3.3.0.sql        |   1 +
 .../src/backend_parser/gram-tsql-epilogue.y.c |   5 +-
 contrib/babelfishpg_tsql/src/pl_handler.c     | 182 +++++++++---------
 contrib/babelfishpg_tsql/src/tsqlIface.cpp    |  26 +--
 test/JDBC/expected/BABEL_539-vu-verify.out    |  30 ++-
 test/JDBC/input/BABEL_539-vu-verify.sql       |  12 +-
 8 files changed, 156 insertions(+), 103 deletions(-)

diff --git a/contrib/babelfishpg_tsql/runtime/functions.c b/contrib/babelfishpg_tsql/runtime/functions.c
index ba9e653299..a5c91f242f 100644
--- a/contrib/babelfishpg_tsql/runtime/functions.c
+++ b/contrib/babelfishpg_tsql/runtime/functions.c
@@ -1905,7 +1905,7 @@ identity_into(PG_FUNCTION_ARGS)
 	int64 result;
 	Assert(tsql_select_into_seq_oid != InvalidOid);
 	result = nextval_internal(tsql_select_into_seq_oid, false);
-	return result;
+	PG_RETURN_INT64((int64) result);
 }
 
 /*
diff --git a/contrib/babelfishpg_tsql/sql/sys_functions.sql b/contrib/babelfishpg_tsql/sql/sys_functions.sql
index ddc9986235..e6f942a271 100644
--- a/contrib/babelfishpg_tsql/sql/sys_functions.sql
+++ b/contrib/babelfishpg_tsql/sql/sys_functions.sql
@@ -3393,6 +3393,7 @@ GRANT EXECUTE ON FUNCTION sys.host_id() TO PUBLIC;
 
 CREATE OR REPLACE FUNCTION sys.identity_into(IN typename INT, IN seed INT, IN increment INT)
 RETURNS int AS 'babelfishpg_tsql' LANGUAGE C STABLE;
+GRANT EXECUTE ON FUNCTION sys.identity_into(INT, INT, INT) TO PUBLIC;
 
 CREATE OR REPLACE FUNCTION sys.degrees(IN arg1 BIGINT)
 RETURNS bigint  AS 'babelfishpg_tsql','bigint_degrees' LANGUAGE C STRICT IMMUTABLE PARALLEL SAFE;
diff --git a/contrib/babelfishpg_tsql/sql/upgrades/babelfishpg_tsql--3.2.0--3.3.0.sql b/contrib/babelfishpg_tsql/sql/upgrades/babelfishpg_tsql--3.2.0--3.3.0.sql
index b4ddfef7fd..e841005455 100644
--- a/contrib/babelfishpg_tsql/sql/upgrades/babelfishpg_tsql--3.2.0--3.3.0.sql
+++ b/contrib/babelfishpg_tsql/sql/upgrades/babelfishpg_tsql--3.2.0--3.3.0.sql
@@ -44,6 +44,7 @@ LANGUAGE C IMMUTABLE STRICT;
 
 CREATE OR REPLACE FUNCTION sys.identity_into(IN typename INT, IN seed INT, IN increment INT)
 RETURNS int AS 'babelfishpg_tsql' LANGUAGE C STABLE;
+GRANT EXECUTE ON FUNCTION sys.identity_into(INT, INT, INT) TO PUBLIC;
 
 CREATE OR REPLACE VIEW sys.sql_expression_dependencies
 AS
diff --git a/contrib/babelfishpg_tsql/src/backend_parser/gram-tsql-epilogue.y.c b/contrib/babelfishpg_tsql/src/backend_parser/gram-tsql-epilogue.y.c
index 5464a78c98..fc870925db 100644
--- a/contrib/babelfishpg_tsql/src/backend_parser/gram-tsql-epilogue.y.c
+++ b/contrib/babelfishpg_tsql/src/backend_parser/gram-tsql-epilogue.y.c
@@ -246,13 +246,14 @@ TsqlFunctionIdentityInto(TypeName *typename, Node *seed, Node *increment, int lo
 			case INT4OID:
 			case INT8OID:
 			case NUMERICOID:				
-				args = list_make3((Node *)makeIntConst((int)type_oid, -1), seed, increment);
+				args = list_make3((Node *)makeIntConst((int)type_oid, location), seed, increment);
 				result = (Node *) makeFuncCall(TsqlSystemFuncName("identity_into"), args, COERCE_EXPLICIT_CALL, location);
 				break;
 			default:
 				ereport(ERROR,
 					(errcode(ERRCODE_SYNTAX_ERROR),
-					 errmsg("identity column type must be smallint, integer, or bigint")));
+					 errmsg("identity column type must be smallint, integer, bigint, or numeric")));
+				break;
 
 		}
 
diff --git a/contrib/babelfishpg_tsql/src/pl_handler.c b/contrib/babelfishpg_tsql/src/pl_handler.c
index 34392ccd6e..b5cd391d10 100644
--- a/contrib/babelfishpg_tsql/src/pl_handler.c
+++ b/contrib/babelfishpg_tsql/src/pl_handler.c
@@ -142,7 +142,7 @@ Datum		sp_prepare(PG_FUNCTION_ARGS);
 Datum		sp_unprepare(PG_FUNCTION_ARGS);
 static List *transformReturningList(ParseState *pstate, List *returningList);
 static List *transformSelectIntoStmt(CreateTableAsStmt *stmt, const char *queryString);
-static char *parse_type_argument(int type_oid);
+static char *get_oid_type_string(int type_oid);
 extern char *construct_unique_index_name(char *index_name, char *relation_name);
 extern int	CurrentLineNumber;
 static non_tsql_proc_entry_hook_type prev_non_tsql_proc_entry_hook = NULL;
@@ -5026,9 +5026,16 @@ pltsql_revert_guc(int nest_level)
 
 }
 
-static char *parse_type_argument(int type_oid){
+static char *get_oid_type_string(int type_oid){
 	char *type_string = NULL;
-	switch(type_oid){
+	if ((*common_utility_plugin_ptr->is_tsql_decimal_datatype) (type_oid))
+	{
+		type_string = "decimal";
+		return type_string;
+	}
+	
+	switch(type_oid)
+	{
 		case INT2OID:
 			type_string = "pg_catalog.int2";
 			break;
@@ -5042,18 +5049,20 @@ static char *parse_type_argument(int type_oid){
 			type_string = "pg_catalog.numeric";
 			break;
 		default:
-			type_string = "decimal";
-			break;
+			ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED), 
+				errmsg("Identity column type must be smallint, integer, bigint, or numeric")));
+			break;	
 	}
 	return type_string;
 }
 
-static List * transformSelectIntoStmt(CreateTableAsStmt *stmt, const char *queryString){
-	List	   *result;
-	ListCell   *elements;
+static List *transformSelectIntoStmt(CreateTableAsStmt *stmt, const char *queryString)
+{
+	List *result;
+	ListCell *elements;
 	CreateSeqStmt *seqstmt;
 	AlterSeqStmt *altseqstmt;
-	List	   *attnamelist;
+	List *attnamelist;
 	IntoClause *into;
 	Node *n;
 
@@ -5065,62 +5074,69 @@ static List * transformSelectIntoStmt(CreateTableAsStmt *stmt, const char *query
 
 	if (n && n->type == T_Query)
 	{
-		Query *q = (Query *) n;
+		Query *q = (Query *)n;
 		bool seen_identity = false;
-		foreach(elements, q->targetList)
-		{	
-			TargetEntry *tle = (TargetEntry *) lfirst(elements);
-			FuncExpr *funcexpr;
-			if (tle->expr && IsA(tle->expr, FuncExpr) && strcasecmp(get_func_name(((FuncExpr *) (tle->expr))->funcid), "identity_into") ==0 ){
-
-				Oid	snamespaceid;
+		foreach (elements, q->targetList)
+		{
+			TargetEntry *tle = (TargetEntry *)lfirst(elements);
+			
+			if (tle->expr && IsA(tle->expr, FuncExpr) && strcasecmp(get_func_name(((FuncExpr *)(tle->expr))->funcid), "identity_into") == 0)
+			{
+				FuncExpr *funcexpr;
+				Oid snamespaceid;
 				char *snamespace;
 				char *sname;
 				List *seqoptions = NIL;
 				ListCell *arg;
 
 				int type_oid;
-				char *type= NULL;
+				char *type = NULL;
 				TypeName *ofTypename;
 				int64 seed_value;
 				int arg_num;
 
-				if(seen_identity){
+				if (seen_identity)
+				{
 					ereport(ERROR,
-						(errcode(ERRCODE_SYNTAX_ERROR), errmsg("Attempting to add multiple identity columns to table \"%s\" using the SELECT INTO statement.", into->rel->relname )));
+							(errcode(ERRCODE_SYNTAX_ERROR), errmsg("Attempting to add multiple identity columns to table \"%s\" using the SELECT INTO statement.", into->rel->relname)));
 				}
 
-				if(tle->resname == NULL){
+				if (tle->resname == NULL)
+				{
 					ereport(ERROR,
-						(errcode(ERRCODE_SYNTAX_ERROR), errmsg("Incorrect syntax near the keyword 'INTO'")));
+							(errcode(ERRCODE_SYNTAX_ERROR), errmsg("Incorrect syntax near the keyword 'INTO'")));
 				}
 
 				funcexpr = (FuncExpr *)tle->expr;
 				arg_num = 0;
-				foreach(arg, funcexpr->args)
+				foreach (arg, funcexpr->args)
 				{
-					Node *fargNode = (Node *) lfirst(arg);
+					Node *fargNode = (Node *)lfirst(arg);
 					Const *con;
 					arg_num++;
-					switch (arg_num){
+					switch (arg_num)
+					{
 						case 1:
-							con = (Const *) fargNode;
-							type_oid = (int) con->constvalue;
-							type = parse_type_argument(type_oid);
+							con = (Const *)fargNode;
+							type_oid = (int)con->constvalue;
+							type = get_oid_type_string(type_oid);
 							ofTypename = typeStringToTypeName(type);
-							seqoptions =  lappend(seqoptions, makeDefElem("as", (Node *) ofTypename, -1));
+							seqoptions = lappend(seqoptions, makeDefElem("as", (Node *)ofTypename, -1));
 							break;
 						case 2:
-							con = (Const *) fargNode;
-							seqoptions = lappend(seqoptions, makeDefElem("start", (Node *)makeInteger(con->constvalue), -1));
-							seed_value = (int64) con->constvalue;
+							con = (Const *)fargNode;
+							seqoptions = lappend(seqoptions, makeDefElem("start", (Node *)makeInteger((int64)con->constvalue), -1));
+							seed_value = (int64)con->constvalue;
 							break;
 						case 3:
-							con = (Const *) fargNode;
-							seqoptions = lappend(seqoptions,  makeDefElem("increment", (Node *)makeInteger(con->constvalue), -1));
-							if ((int) con->constvalue > 0){
+							con = (Const *)fargNode;
+							seqoptions = lappend(seqoptions, makeDefElem("increment", (Node *)makeInteger((int64)con->constvalue), -1));
+							if ((int)con->constvalue > 0)
+							{
 								seqoptions = lappend(seqoptions, makeDefElem("minvalue", (Node *)makeInteger(seed_value), -1));
-							}else{
+							}
+							else
+							{
 								seqoptions = lappend(seqoptions, makeDefElem("maxvalue", (Node *)makeInteger(seed_value), -1));
 							}
 							break;
@@ -5134,7 +5150,7 @@ static List * transformSelectIntoStmt(CreateTableAsStmt *stmt, const char *query
 
 				snamespaceid = RangeVarGetCreationNamespace(into->rel);
 				snamespace = get_namespace_name(snamespaceid);
-				sname = ChooseRelationName(into->rel->relname, tle->resname,"seq",snamespaceid, false);
+				sname = ChooseRelationName(into->rel->relname, tle->resname, "seq", snamespaceid, false);
 				seqstmt = makeNode(CreateSeqStmt);
 				seqstmt->for_identity = true;
 				seqstmt->sequence = makeRangeVar(snamespace, sname, -1);
@@ -5143,74 +5159,68 @@ static List * transformSelectIntoStmt(CreateTableAsStmt *stmt, const char *query
 
 				altseqstmt = makeNode(AlterSeqStmt);
 				altseqstmt->sequence = makeRangeVar(snamespace, sname, -1);
-				attnamelist = list_make3(makeString(snamespace), makeString(into->rel->relname),makeString(tle->resname));
-				altseqstmt->options = list_make1(makeDefElem("owned_by", (Node *) attnamelist, -1));
+				attnamelist = list_make3(makeString(snamespace), makeString(into->rel->relname), makeString(tle->resname));
+				altseqstmt->options = list_make1(makeDefElem("owned_by", (Node *)attnamelist, -1));
 				altseqstmt->for_identity = true;
-
 			}
 		}
 	}
 
-	if(seqstmt){
-		result = lappend(result, seqstmt);
+	if (seqstmt)
+	{
+			result = lappend(result, seqstmt);
 	}
 	result = lappend(result, stmt);
-
-	if(altseqstmt){
-		result = lappend(result, altseqstmt);
+	if (altseqstmt)
+	{
+			result = lappend(result, altseqstmt);
 	}
 
-	return result;	
+	return result;
 }
 
+void pltsql_bbfSelectIntoUtility(ParseState *pstate, PlannedStmt *pstmt, const char *queryString, QueryEnvironment *queryEnv,
+								 ParamListInfo params, QueryCompletion *qc)
+{
 
-void
-pltsql_bbfSelectIntoUtility(ParseState *pstate, PlannedStmt *pstmt, const char *queryString, QueryEnvironment *queryEnv, 
-	ParamListInfo params, QueryCompletion *qc){
-
-	Node *parsetree = pstmt->utilityStmt;	
+	Node *parsetree = pstmt->utilityStmt;
 	ObjectAddress address;
 	ObjectAddress secondaryObject = InvalidObjectAddress;
 	List *stmts;
-	stmts = transformSelectIntoStmt((CreateTableAsStmt *) parsetree, queryString);
-		while (stmts != NIL)
+	stmts = transformSelectIntoStmt((CreateTableAsStmt *)parsetree, queryString);
+	while (stmts != NIL)
+	{
+		Node *stmt = (Node *)linitial(stmts);
+		stmts = list_delete_first(stmts);
+		if (IsA(stmt, CreateTableAsStmt))
 		{
-			Node *stmt = (Node *) linitial(stmts);
-			stmts = list_delete_first(stmts);
-			if (IsA(stmt, CreateTableAsStmt))
-			{
-				address = ExecCreateTableAs(pstate, (CreateTableAsStmt *) parsetree,params, queryEnv, qc);
-				EventTriggerCollectSimpleCommand(address,secondaryObject,stmt);
-			}
-			else if(IsA(stmt, CreateSeqStmt)) {
-				address = DefineSequence(pstate, (CreateSeqStmt *) stmt);
-				Assert(address.objectId != InvalidOid);
-				tsql_select_into_seq_oid = address.objectId;
-				EventTriggerCollectSimpleCommand(address,secondaryObject,stmt);
-			}
-			else{
-
-				PlannedStmt *wrapper;
-				wrapper = makeNode(PlannedStmt);
-				wrapper->commandType = CMD_UTILITY;
-				wrapper->canSetTag = false;
-				wrapper->utilityStmt = stmt;
-				wrapper->stmt_location = pstmt->stmt_location;
-				wrapper->stmt_len = pstmt->stmt_len;
-
-				ProcessUtility(wrapper,
-								queryString,
-								false,
-								PROCESS_UTILITY_SUBCOMMAND,
-								params,
-								NULL,
-								None_Receiver,
-								NULL);
-			}
+			address = ExecCreateTableAs(pstate, (CreateTableAsStmt *)parsetree, params, queryEnv, qc);
+			EventTriggerCollectSimpleCommand(address, secondaryObject, stmt);
+		}
+		else if (IsA(stmt, CreateSeqStmt))
+		{
+			address = DefineSequence(pstate, (CreateSeqStmt *)stmt);
+			Assert(address.objectId != InvalidOid);
+			tsql_select_into_seq_oid = address.objectId;
+			EventTriggerCollectSimpleCommand(address, secondaryObject, stmt);
+		}
+		else
+		{
+			PlannedStmt *wrapper;
+			wrapper = makeNode(PlannedStmt);
+			wrapper->commandType = CMD_UTILITY;
+			wrapper->canSetTag = false;
+			wrapper->utilityStmt = stmt;
+			wrapper->stmt_location = pstmt->stmt_location;
+			wrapper->stmt_len = pstmt->stmt_len;
+
+			ProcessUtility(wrapper, queryString, false, PROCESS_UTILITY_SUBCOMMAND, params, NULL, None_Receiver, NULL);
+		}
 		if (stmts != NIL)
 			CommandCounterIncrement();
-		}
+	}
 }
+
 void
 set_current_query_is_create_tbl_check_constraint(Node *expr)
 {
diff --git a/contrib/babelfishpg_tsql/src/tsqlIface.cpp b/contrib/babelfishpg_tsql/src/tsqlIface.cpp
index 0d2bfb1c39..219c609750 100644
--- a/contrib/babelfishpg_tsql/src/tsqlIface.cpp
+++ b/contrib/babelfishpg_tsql/src/tsqlIface.cpp
@@ -1847,17 +1847,6 @@ class tsqlBuilder : public tsqlCommonMutator
 		}
 	}
 
-	void enterFunction_call(TSqlParser::Function_callContext *ctx) override
-	{
-		std::string func_name = ::getFullText(ctx);
-		size_t Lbracket_index = func_name.find('(');
-        func_name = func_name.substr(0, Lbracket_index);
-		if (pg_strcasecmp(func_name.c_str(), "identity") == 0 || pg_strcasecmp(func_name.c_str(), "identity_into") == 0
-			|| pg_strcasecmp(func_name.c_str(), "sys.identity_into") == 0) {
-			has_identity_function = true;
-		}
-	}
-
 	//////////////////////////////////////////////////////////////////////////////
 	// function/procedure call analysis
 	//////////////////////////////////////////////////////////////////////////////
@@ -1966,6 +1955,21 @@ class tsqlBuilder : public tsqlCommonMutator
                               }
 
 			}
+
+			if (ctx->func_proc_name_server_database_schema()->procedure)
+			{
+				std::string proc_name = stripQuoteFromId(ctx->func_proc_name_server_database_schema()->procedure);
+				if (pg_strcasecmp(proc_name.c_str(), "identity") == 0) 
+				{
+					has_identity_function = true;
+				}
+			
+				if (pg_strcasecmp(proc_name.c_str(), "identity_into") == 0)
+				{
+					throw PGErrorWrapperException(ERROR, ERRCODE_FEATURE_NOT_SUPPORTED, 
+						format_errmsg("function %s does not exist", proc_name.c_str()), getLineAndPos(ctx));
+				}
+			}
 		}
 	}
 
diff --git a/test/JDBC/expected/BABEL_539-vu-verify.out b/test/JDBC/expected/BABEL_539-vu-verify.out
index 5913a18679..fc728d1a45 100644
--- a/test/JDBC/expected/BABEL_539-vu-verify.out
+++ b/test/JDBC/expected/BABEL_539-vu-verify.out
@@ -123,7 +123,7 @@ SELECT col1, IDENTITY(char, 1,1) AS id_num INTO babel_539NewTable1 FROM babel_53
 GO
 ~~ERROR (Code: 33557097)~~
 
-~~ERROR (Message: identity column type must be smallint, integer, or bigint)~~
+~~ERROR (Message: identity column type must be smallint, integer, bigint, or numeric)~~
 
 
 DROP TABLE IF EXISTS babel_539NewTable1;
@@ -208,5 +208,31 @@ GO
 ~~ERROR (Message: Attempting to add multiple identity columns to table "#babel_539newtemptable2" using the SELECT INTO statement.)~~
 
 
-DROP TABLE IF EXISTS #babel_539NewTempTable2;
+--calling internal function directly
+SELECT col1, IDENTITY_INTO(1, 1,1) as id_num INTO #babel_539NewTempTable2 FROM babel_539OldTable;
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: function IDENTITY_INTO does not exist)~~
+
+
+SELECT sys.IDENTITY(23, 1);
 GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: Incorrect syntax near ')')~~
+
+
+SELECT IDENTITY(int, 21);
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: Incorrect syntax near ')')~~
+
+
+SELECT sys.IDENTITY_INTO(23, 1, 1);
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: function IDENTITY_INTO does not exist)~~
+
diff --git a/test/JDBC/input/BABEL_539-vu-verify.sql b/test/JDBC/input/BABEL_539-vu-verify.sql
index d2b0be9402..c6154da9d1 100644
--- a/test/JDBC/input/BABEL_539-vu-verify.sql
+++ b/test/JDBC/input/BABEL_539-vu-verify.sql
@@ -131,5 +131,15 @@ GO
 SELECT col1, IDENTITY(int, 1,1) as id_num, IDENTITY(int, 1,1) as id_num2 INTO #babel_539NewTempTable2 FROM babel_539OldTable;
 GO
 
-DROP TABLE IF EXISTS #babel_539NewTempTable2;
+--calling internal function directly
+SELECT col1, IDENTITY_INTO(1, 1,1) as id_num INTO #babel_539NewTempTable2 FROM babel_539OldTable;
+GO
+
+SELECT sys.IDENTITY(23, 1);
+GO
+
+SELECT IDENTITY(int, 21);
+GO
+
+SELECT sys.IDENTITY_INTO(23, 1, 1);
 GO
\ No newline at end of file

From d33a11d6d5cad14344fd7c7549548015e5f8349f Mon Sep 17 00:00:00 2001
From: Deepakshi Mittal <depakshi@amazon.com>
Date: Sun, 30 Jul 2023 18:45:33 +0000
Subject: [PATCH 6/6] Revert changes for failing test

Not required because no added test cases for callin these functions

Task: BABEL-539

Signed-off-by: Deepakshi Mittal <depakshi@amazon.com>
---
 .../expected/sql_validation_framework/expected_create.out       | 2 --
 1 file changed, 2 deletions(-)

diff --git a/test/python/expected/sql_validation_framework/expected_create.out b/test/python/expected/sql_validation_framework/expected_create.out
index 5c604e31f9..c4c0022283 100644
--- a/test/python/expected/sql_validation_framework/expected_create.out
+++ b/test/python/expected/sql_validation_framework/expected_create.out
@@ -17,7 +17,6 @@ Could not find tests for function sys.format_datetime
 Could not find tests for function sys.format_numeric
 Could not find tests for function sys.get_current_full_xact_id
 Could not find tests for function sys.get_host_os
-Could not find tests for function sys.identity_into
 Could not find tests for function sys.is_collated_ci_as
 Could not find tests for function sys.is_collated_ci_as_internal
 Could not find tests for function sys.is_rolemember_internal
@@ -125,7 +124,6 @@ Could not find upgrade tests for function sys.get_current_full_xact_id
 Could not find upgrade tests for function sys.get_host_os
 Could not find upgrade tests for function sys.ident_incr
 Could not find upgrade tests for function sys.ident_seed
-Could not find upgrade tests for function sys.identity_into
 Could not find upgrade tests for function sys.is_collated_ci_as
 Could not find upgrade tests for function sys.is_collated_ci_as_internal
 Could not find upgrade tests for function sys.is_rolemember_internal