Fix setval() behaviour in T-SQL mode (babelfish-for-postgresql#803)

We should not change the setval() behaviour in T-SQL dialect as some of the users may use this function as an alternative to DBCC CHECKIDENT for resetting the identity. Task: BABEL-3506 Signed-off-by: Kuntal Ghosh kuntalgh@amazon.com
Bit-Quill · Nov 9, 2022 · 428ecf7 · 428ecf7
1 parent e17df72
commit 428ecf7
Show file tree

Hide file tree

Showing 19 changed files with 346 additions and 15 deletions.
diff --git a/contrib/babelfishpg_tsql/src/pl_exec.c b/contrib/babelfishpg_tsql/src/pl_exec.c
@@ -5508,9 +5508,6 @@ pltsql_update_identity_insert_sequence(PLtsql_expr *expr)
 								max_identity = last_identity;
 						}
 
-						/* update last used identity */
-						pltsql_update_last_identity(seqid, last_identity);
-
 						/*
 						 * We also need to reset the seed.  If the increment
 						 * is positive, we need to find the max identity that
@@ -5528,18 +5525,35 @@ pltsql_update_identity_insert_sequence(PLtsql_expr *expr)
 
 						}
 
-						if (seq_incr > 0)
-							DirectFunctionCall2(setval_oid,
-												ObjectIdGetDatum(seqid),
-												Int64GetDatum(max_identity));
-						else if (seq_incr < 0)
-							DirectFunctionCall2(setval_oid,
-												ObjectIdGetDatum(seqid),
-												Int64GetDatum(min_identity));
-						else {
-							/* increment can't be zero */
-							Assert(0);
+						PG_TRY();
+						{
+							/*
+							 * We want the T-SQL behavior of setval function.
+							 * Please check the variable definition for details.
+							 */
+							pltsql_setval_identity_mode = true;
+							if (seq_incr > 0)
+								DirectFunctionCall2(setval_oid,
+													ObjectIdGetDatum(seqid),
+													Int64GetDatum(max_identity));
+							else if (seq_incr < 0)
+								DirectFunctionCall2(setval_oid,
+													ObjectIdGetDatum(seqid),
+													Int64GetDatum(min_identity));
+							else {
+								/* increment can't be zero */
+								Assert(0);
+							}
 						}
+						PG_FINALLY();
+						{
+							/* reset the value */
+							pltsql_setval_identity_mode = false;
+						}
+						PG_END_TRY();
+
+						/* update last used identity if setval is successful */
+						pltsql_update_last_identity(seqid, last_identity);
 
 						/* more than one identity column isn't allowed */
 						break;

diff --git a/contrib/babelfishpg_tsql/src/pltsql.h b/contrib/babelfishpg_tsql/src/pltsql.h
@@ -2032,6 +2032,14 @@ void pltsql_function_probin_writer(CreateFunctionStmt *stmt, Oid languageOid, ch
 void pltsql_function_probin_reader(ParseState *pstate,
 						List *fargs, Oid *actual_arg_types, Oid *declared_arg_types, Oid funcid);
 
+/*
+ * This variable is set to true, if setval should behave in T-SQL way, i.e.,
+ * setval sets the max/min(current identity value, new identity value to be
+ * inserted.  By default, it is set to fale which means setval should behave
+ * PG way irrespective of the dialect - reset identity seed.
+ */
+extern bool pltsql_setval_identity_mode;
+
 /*
  * Functions in pltsql_identity.c
  */

diff --git a/contrib/babelfishpg_tsql/src/pltsql_identity.c b/contrib/babelfishpg_tsql/src/pltsql_identity.c
@@ -37,6 +37,12 @@ typedef struct SeqTableIdentityData
 	int64		last_identity;	/* sequence identity value */
 } SeqTableIdentityData;
 
+/*
+ * By default, it is set to false.  This is set to true only when we want setval
+ * to set the max/min(current identity value, new identity value to be inserted.
+ */
+bool pltsql_setval_identity_mode = false;
+
 static HTAB *seqhashtabidentity = NULL;
 
 static SeqTableIdentityData *last_used_seq_identity = NULL;
@@ -319,7 +325,7 @@ void pltsql_resetcache_identity()
 int64
 pltsql_setval_identity(Oid seqid, int64 val, int64 last_val)
 {
-	if (sql_dialect == SQL_DIALECT_TSQL)
+	if (sql_dialect == SQL_DIALECT_TSQL && pltsql_setval_identity_mode)
 	{
 		ListCell *seq_lc;
 		List *seq_options;

diff --git a/test/JDBC/expected/BABEL-PG-SYSTEM-FUNCTIONS-vu-cleanup.out b/test/JDBC/expected/BABEL-PG-SYSTEM-FUNCTIONS-vu-cleanup.out
@@ -0,0 +1,2 @@
+DROP SEQUENCE BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1
+go
diff --git a/test/JDBC/expected/BABEL-PG-SYSTEM-FUNCTIONS-vu-prepare.out b/test/JDBC/expected/BABEL-PG-SYSTEM-FUNCTIONS-vu-prepare.out
@@ -0,0 +1,78 @@
+-- For DMS, we've suggested the following PG functions related to identity
+-- feature that can be called from TDS endpoint and they should get the exact
+-- same behaviour as PG endpoint.  So, let's add some tests.
+-- basic sequence operations for setval, nextval, currentval (tests are taken
+-- src/test/regress/sql/sequence.sql PG regression test suite)
+CREATE SEQUENCE BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1;
+go
+
+SELECT nextval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1');
+go
+~~START~~
+bigint
+1
+~~END~~
+
+SELECT currval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1');
+go
+~~START~~
+bigint
+1
+~~END~~
+
+SELECT setval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1', 32);
+go
+~~START~~
+bigint
+32
+~~END~~
+
+SELECT nextval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1');
+go
+~~START~~
+bigint
+33
+~~END~~
+
+SELECT setval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1', 99, false);
+go
+~~START~~
+bigint
+99
+~~END~~
+
+SELECT nextval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1');
+go
+~~START~~
+bigint
+99
+~~END~~
+
+SELECT setval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1', 32);
+go
+~~START~~
+bigint
+32
+~~END~~
+
+SELECT nextval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1');
+go
+~~START~~
+bigint
+33
+~~END~~
+
+SELECT setval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1', 99, false);
+go
+~~START~~
+bigint
+99
+~~END~~
+
+SELECT nextval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1');
+go
+~~START~~
+bigint
+99
+~~END~~
+
diff --git a/test/JDBC/expected/BABEL-PG-SYSTEM-FUNCTIONS-vu-verify.out b/test/JDBC/expected/BABEL-PG-SYSTEM-FUNCTIONS-vu-verify.out
@@ -0,0 +1,70 @@
+SELECT nextval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1');
+go
+~~START~~
+bigint
+100
+~~END~~
+
+SELECT currval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1');
+go
+~~START~~
+bigint
+100
+~~END~~
+
+SELECT setval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1', 32);
+go
+~~START~~
+bigint
+32
+~~END~~
+
+SELECT nextval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1');
+go
+~~START~~
+bigint
+33
+~~END~~
+
+SELECT setval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1', 99, false);
+go
+~~START~~
+bigint
+99
+~~END~~
+
+SELECT nextval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1');
+go
+~~START~~
+bigint
+99
+~~END~~
+
+SELECT setval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1', 32);
+go
+~~START~~
+bigint
+32
+~~END~~
+
+SELECT nextval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1');
+go
+~~START~~
+bigint
+33
+~~END~~
+
+SELECT setval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1', 99, false);
+go
+~~START~~
+bigint
+99
+~~END~~
+
+SELECT nextval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1');
+go
+~~START~~
+bigint
+99
+~~END~~
+
diff --git a/test/JDBC/input/BABEL-IDENTITY.sql b/test/JDBC/input/BABEL-IDENTITY.sql
@@ -254,6 +254,18 @@ SET IDENTITY_INSERT dbo.t1_identity_1 ON;
 INSERT INTO dbo.t1_identity_1 (a,b) VALUES (1,1);
 go
 
+-- Test with an error in setval
+ALTER SEQUENCE t1_identity_1_a_seq MAXVALUE 700
+INSERT INTO dbo.t1_identity_1 (a,b) VALUES (800,2);
+SELECT @@IDENTITY; SELECT IDENT_CURRENT('dbo.t1_identity_1'); SELECT SCOPE_IDENTITY();
+go
+
+-- Test with setval after an error
+-- It should update on IDENT_CURRENT, not other identity variables
+SELECT setval('t1_identity_1_a_seq', 32);
+SELECT @@IDENTITY; SELECT IDENT_CURRENT('dbo.t1_identity_1'); SELECT SCOPE_IDENTITY();
+go
+
 -- Check transaction rollback should increase identity
 BEGIN TRAN t1; INSERT INTO dbo.t1_identity_1 (a,b) VALUES (300,2); ROLLBACK TRAN t1;
 SELECT @@IDENTITY; SELECT IDENT_CURRENT('dbo.t1_identity_1'); SELECT SCOPE_IDENTITY();
@@ -269,9 +281,19 @@ INSERT INTO dbo.t1_identity_1 (a,b) VALUES (100,3);
 SELECT @@IDENTITY; SELECT IDENT_CURRENT('dbo.t1_identity_1'); SELECT SCOPE_IDENTITY();
 go
 
+-- After identity insert off, the insert should start from the next seed that
+-- setval sets
+SELECT setval('t1_identity_1_a_seq', 500);
+go
+
 SET IDENTITY_INSERT dbo.t1_identity_1 OFF;
 go
 
+INSERT INTO dbo.t1_identity_1 (b) VALUES (4);
+SELECT a FROM dbo.t1_identity_1 where b = 4;
+SELECT @@IDENTITY; SELECT IDENT_CURRENT('dbo.t1_identity_1'); SELECT SCOPE_IDENTITY();
+go
+
 CREATE TABLE dbo.t1_identity_2(a int identity(1, -1) primary key, b int unique not null);
 SET IDENTITY_INSERT dbo.t1_identity_2 ON;
 INSERT INTO dbo.t1_identity_2 (a,b) VALUES (1,1);

diff --git a/test/JDBC/input/BABEL-PG-SYSTEM-FUNCTIONS-vu-cleanup.sql b/test/JDBC/input/BABEL-PG-SYSTEM-FUNCTIONS-vu-cleanup.sql
@@ -0,0 +1,2 @@
+DROP SEQUENCE BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1
+go
diff --git a/test/JDBC/input/BABEL-PG-SYSTEM-FUNCTIONS-vu-prepare.sql b/test/JDBC/input/BABEL-PG-SYSTEM-FUNCTIONS-vu-prepare.sql
@@ -0,0 +1,28 @@
+-- For DMS, we've suggested the following PG functions related to identity
+-- feature that can be called from TDS endpoint and they should get the exact
+-- same behaviour as PG endpoint.  So, let's add some tests.
+-- basic sequence operations for setval, nextval, currentval (tests are taken
+-- src/test/regress/sql/sequence.sql PG regression test suite)
+CREATE SEQUENCE BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1;
+go
+
+SELECT nextval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1');
+go
+SELECT currval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1');
+go
+SELECT setval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1', 32);
+go
+SELECT nextval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1');
+go
+SELECT setval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1', 99, false);
+go
+SELECT nextval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1');
+go
+SELECT setval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1', 32);
+go
+SELECT nextval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1');
+go
+SELECT setval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1', 99, false);
+go
+SELECT nextval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1');
+go
diff --git a/test/JDBC/input/BABEL-PG-SYSTEM-FUNCTIONS-vu-verify.sql b/test/JDBC/input/BABEL-PG-SYSTEM-FUNCTIONS-vu-verify.sql
@@ -0,0 +1,20 @@
+SELECT nextval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1');
+go
+SELECT currval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1');
+go
+SELECT setval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1', 32);
+go
+SELECT nextval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1');
+go
+SELECT setval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1', 99, false);
+go
+SELECT nextval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1');
+go
+SELECT setval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1', 32);
+go
+SELECT nextval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1');
+go
+SELECT setval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1', 99, false);
+go
+SELECT nextval('BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1');
+go
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		DROP SEQUENCE BABEL_PG_SYSTEM_FUNCTIONS_vu_prepare_seq1
		go