Skip to content

Commit

Permalink
Merge branch 'Version_based_connection_reset_using_Queries' of https:…
Browse files Browse the repository at this point in the history
…//github.com/SiddharthBITS/babelfish_extensions into Version_based_connection_reset_using_Queries
  • Loading branch information
SiddharthBITS committed Jan 30, 2025
2 parents 903aa49 + b8824bb commit da16cd1
Show file tree
Hide file tree
Showing 4 changed files with 98 additions and 31 deletions.
58 changes: 57 additions & 1 deletion contrib/babelfishpg_tsql/src/pl_exec.c
Original file line number Diff line number Diff line change
Expand Up @@ -5483,6 +5483,53 @@ pltsql_update_identity_insert_sequence(PLtsql_expr *expr)
Oid seqid = InvalidOid;
SPITupleTable *tuptable = SPI_tuptable;
uint64 n_processed = SPI_processed;
bool is_cross_db = false;
char *schema_name = NULL;
HeapTuple schema_tuple;
Oid current_user_id = InvalidOid;

schema_name = get_namespace_name(get_rel_namespace(tsql_identity_insert.rel_oid));
schema_tuple = SearchSysCache1(SYSNAMESPACENAME, CStringGetDatum(schema_name));

if (HeapTupleIsValid(schema_tuple))
{
Datum datum;
int16 db_id;
bool isnull;

datum = SysCacheGetAttr(SYSNAMESPACENAME, schema_tuple, Anum_namespace_ext_dbid, &isnull);
db_id = DatumGetInt16(datum);

if(!DbidIsValid(db_id))
ereport(ERROR,
(errcode(ERRCODE_INTERNAL_ERROR),
errmsg("dbid in babelfish_namespace_ext catalog doesn't exists in babelfish_sysdatabases catalog")));

if (db_id != get_cur_db_id())
{
char *db_name = get_db_name(db_id);
char *user = get_user_for_database(db_name);

if (user)
{
is_cross_db = true;
pfree(db_name);
pfree(user);
}
else
{
char *login = GetUserNameFromId(GetSessionUserId(), false);
ereport(ERROR,
(errcode(ERRCODE_UNDEFINED_DATABASE),
errmsg("The server principal \"%s\" is not able to access "
"the database \"%s\" under the current security context",
login, db_name)));
}

}

ReleaseSysCache(schema_tuple);
}

/* Get the identity column name */
rel = RelationIdGetRelation(tsql_identity_insert.rel_oid);
Expand Down Expand Up @@ -5566,6 +5613,12 @@ pltsql_update_identity_insert_sequence(PLtsql_expr *expr)

PG_TRY();
{
if (is_cross_db)
{
current_user_id = GetUserId();
SetCurrentRoleId(GetSessionUserId(), false);
}

/*
* We want the T-SQL behavior of setval function.
* Please check the variable definition for
Expand All @@ -5590,6 +5643,9 @@ pltsql_update_identity_insert_sequence(PLtsql_expr *expr)
{
/* reset the value */
pltsql_setval_identity_mode = false;

if (is_cross_db)
SetCurrentRoleId(current_user_id, false);
}
PG_END_TRY();

Expand Down Expand Up @@ -10473,4 +10529,4 @@ pltsql_assign_var(PG_FUNCTION_ARGS)
PG_RETURN_NULL();

PG_RETURN_DATUM(data);
}
}
4 changes: 4 additions & 0 deletions contrib/babelfishpg_tsql/src/pl_handler.c
Original file line number Diff line number Diff line change
Expand Up @@ -400,7 +400,11 @@ assign_identity_insert(const char *newval, void *extra)
* throw permission denied error in that case
*/
if (!curr_user_if_cross_db)
{
/* Get schema name for error message */
logical_schema_name = (char *) lthird(elemlist);
throw_error_for_identity_insert(catalog_name, logical_schema_name, rel_name);
}
}

/* Check the user provided schema value */
Expand Down
54 changes: 29 additions & 25 deletions test/JDBC/expected/BABEL-IDENTITY.out
Original file line number Diff line number Diff line change
Expand Up @@ -763,7 +763,7 @@ Index Scan using test_id_index_pkey on test_id_index

~~START~~
text
Babelfish T-SQL Batch Parsing Time: 9.792 ms
Babelfish T-SQL Batch Parsing Time: 9.564 ms
~~END~~


Expand All @@ -778,7 +778,7 @@ Index Scan using test_id_index_pkey on test_id_index

~~START~~
text
Babelfish T-SQL Batch Parsing Time: 3.872 ms
Babelfish T-SQL Batch Parsing Time: 3.754 ms
~~END~~


Expand All @@ -793,7 +793,7 @@ Index Scan using test_id_index_pkey on test_id_index

~~START~~
text
Babelfish T-SQL Batch Parsing Time: 4.621 ms
Babelfish T-SQL Batch Parsing Time: 4.633 ms
~~END~~


Expand All @@ -808,7 +808,7 @@ Index Scan using test_id_index_pkey on test_id_index

~~START~~
text
Babelfish T-SQL Batch Parsing Time: 42.064 ms
Babelfish T-SQL Batch Parsing Time: 42.276 ms
~~END~~


Expand All @@ -825,7 +825,7 @@ Bitmap Heap Scan on test_id_index

~~START~~
text
Babelfish T-SQL Batch Parsing Time: 9.045 ms
Babelfish T-SQL Batch Parsing Time: 5.772 ms
~~END~~


Expand All @@ -840,7 +840,7 @@ Seq Scan on test_id_index

~~START~~
text
Babelfish T-SQL Batch Parsing Time: 2.038 ms
Babelfish T-SQL Batch Parsing Time: 1.271 ms
~~END~~


Expand All @@ -857,7 +857,7 @@ Bitmap Heap Scan on test_id_index

~~START~~
text
Babelfish T-SQL Batch Parsing Time: 9.868 ms
Babelfish T-SQL Batch Parsing Time: 5.620 ms
~~END~~


Expand All @@ -873,7 +873,7 @@ Index Scan using test_id_index_pkey on test_id_index

~~START~~
text
Babelfish T-SQL Batch Parsing Time: 0.577 ms
Babelfish T-SQL Batch Parsing Time: 0.312 ms
~~END~~


Expand All @@ -888,7 +888,7 @@ Seq Scan on test_id_index

~~START~~
text
Babelfish T-SQL Batch Parsing Time: 2.833 ms
Babelfish T-SQL Batch Parsing Time: 1.654 ms
~~END~~


Expand All @@ -903,7 +903,7 @@ Index Scan using test_id_index_tinyint_pkey on test_id_index_tinyint

~~START~~
text
Babelfish T-SQL Batch Parsing Time: 0.302 ms
Babelfish T-SQL Batch Parsing Time: 0.194 ms
~~END~~


Expand All @@ -918,7 +918,7 @@ Index Scan using test_id_index_smallint_pkey on test_id_index_smallint

~~START~~
text
Babelfish T-SQL Batch Parsing Time: 0.229 ms
Babelfish T-SQL Batch Parsing Time: 0.165 ms
~~END~~


Expand All @@ -933,7 +933,7 @@ Index Scan using test_id_index_bigint_pkey on test_id_index_bigint

~~START~~
text
Babelfish T-SQL Batch Parsing Time: 0.226 ms
Babelfish T-SQL Batch Parsing Time: 0.160 ms
~~END~~


Expand All @@ -948,7 +948,7 @@ Index Scan using test_id_index_numeric_pkey on test_id_index_numeric

~~START~~
text
Babelfish T-SQL Batch Parsing Time: 0.223 ms
Babelfish T-SQL Batch Parsing Time: 0.159 ms
~~END~~


Expand All @@ -974,7 +974,7 @@ Index Scan using test_numeric_index_no_id_pkey on test_numeric_index_no_id

~~START~~
text
Babelfish T-SQL Batch Parsing Time: 0.290 ms
Babelfish T-SQL Batch Parsing Time: 0.191 ms
~~END~~


Expand Down Expand Up @@ -1194,7 +1194,7 @@ Index Only Scan using babel_3384_test_pkey on babel_3384_test

~~START~~
text
Babelfish T-SQL Batch Parsing Time: 0.171 ms
Babelfish T-SQL Batch Parsing Time: 0.165 ms
~~END~~

select id from babel_3384_test WHERE id = @@IDENTITY
Expand All @@ -1208,7 +1208,7 @@ Index Only Scan using babel_3384_test_pkey on babel_3384_test

~~START~~
text
Babelfish T-SQL Batch Parsing Time: 0.164 ms
Babelfish T-SQL Batch Parsing Time: 0.208 ms
~~END~~


Expand Down Expand Up @@ -1318,11 +1318,17 @@ GO
~~ROW COUNT: 1~~


INSERT INTO identity_insert_db.dbo.identity_insert_t1 (a, b) VALUES (3, 103);
GO
~~ROW COUNT: 1~~


SELECT * FROM identity_insert_t1 ORDER BY a
GO
~~START~~
int#!#int
3#!#103
3#!#103
~~END~~


Expand Down Expand Up @@ -1364,15 +1370,13 @@ GO

-- tsql user=identity_insert_l1 password=123
-- Should now have permission
SET IDENTITY_INSERT identity_insert_db.dbo.identity_insert_t1 OFF
SET IDENTITY_INSERT identity_insert_db.dbo.identity_insert_t1 ON
GO

-- Should fail
-- Should not fail
INSERT INTO identity_insert_db.dbo.identity_insert_t1 (a, b) VALUES (5, 105);
GO
~~ERROR (Code: 33557097)~~

~~ERROR (Message: cannot insert a non-DEFAULT value into column "a")~~
~~ROW COUNT: 1~~


-- tsql
Expand Down Expand Up @@ -1473,15 +1477,15 @@ guest#!#0#!#0
SET IDENTITY_INSERT identity_insert_db.identity_insert_sch.identity_insert_t2 ON
GO

USE identity_insert_db
GO

-- Should not fail
INSERT INTO identity_insert_t2 (c, d) VALUES (1, 101);
INSERT INTO identity_insert_db.identity_insert_sch.identity_insert_t2 (c, d) VALUES (1, 101);
GO
~~ROW COUNT: 1~~


USE identity_insert_db
GO

SELECT * FROM identity_insert_t2
GO
~~START~~
Expand Down
13 changes: 8 additions & 5 deletions test/JDBC/input/BABEL-IDENTITY.mix
Original file line number Diff line number Diff line change
Expand Up @@ -613,6 +613,9 @@ GO
INSERT INTO identity_insert_t1 (a, b) VALUES (3, 103);
GO

INSERT INTO identity_insert_db.dbo.identity_insert_t1 (a, b) VALUES (3, 103);
GO

SELECT * FROM identity_insert_t1 ORDER BY a
GO

Expand Down Expand Up @@ -646,10 +649,10 @@ GO

-- tsql user=identity_insert_l1 password=123
-- Should now have permission
SET IDENTITY_INSERT identity_insert_db.dbo.identity_insert_t1 OFF
SET IDENTITY_INSERT identity_insert_db.dbo.identity_insert_t1 ON
GO

-- Should fail
-- Should not fail
INSERT INTO identity_insert_db.dbo.identity_insert_t1 (a, b) VALUES (5, 105);
GO

Expand Down Expand Up @@ -729,11 +732,11 @@ GO
SET IDENTITY_INSERT identity_insert_db.identity_insert_sch.identity_insert_t2 ON
GO

USE identity_insert_db
-- Should not fail
INSERT INTO identity_insert_db.identity_insert_sch.identity_insert_t2 (c, d) VALUES (1, 101);
GO

-- Should not fail
INSERT INTO identity_insert_t2 (c, d) VALUES (1, 101);
USE identity_insert_db
GO

SELECT * FROM identity_insert_t2
Expand Down

0 comments on commit da16cd1

Please sign in to comment.