Support IDENTITY() function in SELECT-INTO

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);
+	PG_RETURN_INT64((int64) result);
+}
+
 /*
  * Execute various integrity checks.
  * Returns true if all the checks pass otherwise

contrib/babelfishpg_tsql/sql/sys_functions.sql

@@ -3391,6 +3391,10 @@ 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;
+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;
 GRANT EXECUTE ON FUNCTION sys.degrees(BIGINT) TO PUBLIC;

contrib/babelfishpg_tsql/sql/upgrades/babelfishpg_tsql--3.2.0--3.3.0.sql

@@ -42,6 +42,10 @@ 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;
+GRANT EXECUTE ON FUNCTION sys.identity_into(INT, INT, INT) TO PUBLIC;
+
 CREATE OR REPLACE VIEW sys.sql_expression_dependencies
 AS
 SELECT

contrib/babelfishpg_tsql/src/backend_parser/gram-tsql-epilogue.y.c

@@ -230,6 +230,37 @@ 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, 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, bigint, or numeric")));
+				break;
+
+		}
+
+	return result;
+
+}
+
 /* TsqlFunctionParse -- Implements the PARSE and TRY_PARSE functions.
  * Takes in expression, target type, regional culture, try boolean, location.
  *

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);

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:

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;
+			}
+		}
+	}	
+}

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;