Skip to content

Commit 6055b7b

Browse files
manisha-deshpandemanisha-deshpande
committed
Added serialization support for PLTSQL_STMT_DECL_TABLE
Signed-off-by: Manisha Deshpande <mmdeshp@amazon.com>
1 parent 495ea56 commit 6055b7b

File tree

4 files changed

+194
-3
lines changed

4 files changed

+194
-3
lines changed

‎contrib/babelfishpg_tsql/src/pltsql_deserialize.c‎

Lines changed: 113 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -446,6 +446,10 @@ deserialize_stmt(DeserializeContext *ctx)
446446
stmt_data = palloc0(sizeof(PLtsql_stmt_try_catch));
447447
((PLtsql_stmt_try_catch *)stmt_data)->cmd_type = PLTSQL_STMT_TRY_CATCH;
448448
break;
449+
case PLTSQL_STMT_DECL_TABLE:
450+
stmt_data = palloc0(sizeof(PLtsql_stmt_decl_table));
451+
((PLtsql_stmt_decl_table *)stmt_data)->cmd_type = PLTSQL_STMT_DECL_TABLE;
452+
break;
449453
default:
450454
/* Mark as unsupported instead of throwing error */
451455
*(ctx->status) = PLTSQL_SERIAL_UNSUPPORTED;
@@ -1041,6 +1045,32 @@ deserialize_stmt(DeserializeContext *ctx)
10411045
}
10421046
break;
10431047
}
1048+
1049+
case PLTSQL_STMT_DECL_TABLE:
1050+
{
1051+
PLtsql_stmt_decl_table *decl_table = (PLtsql_stmt_decl_table *)stmt_data;
1052+
if (strcmp(data_key, "lineno") == 0)
1053+
{
1054+
decl_table->lineno = get_int_value(&v);
1055+
}
1056+
else if (strcmp(data_key, "dno") == 0)
1057+
{
1058+
decl_table->dno = get_int_value(&v);
1059+
}
1060+
else if (strcmp(data_key, "tbltypname") == 0)
1061+
{
1062+
decl_table->tbltypname = get_string_value(&v);
1063+
}
1064+
else if (strcmp(data_key, "coldef") == 0)
1065+
{
1066+
decl_table->coldef = get_string_value(&v);
1067+
}
1068+
else
1069+
{
1070+
skip_value(ctx);
1071+
}
1072+
break;
1073+
}
10441074
}
10451075

10461076
pfree(data_key);
@@ -1134,6 +1164,27 @@ deserialize_expr(DeserializeContext *ctx)
11341164
{
11351165
expr->rwparam = get_int_value(&v);
11361166
}
1167+
else if (strcmp(key, "paramnos") == 0)
1168+
{
1169+
/*
1170+
* Deserialize paramnos Bitmapset from JSON array of integers.
1171+
* Each element is a datum dno referenced by this expression.
1172+
*/
1173+
if (tok != WJB_BEGIN_ARRAY)
1174+
elog(ERROR, "Expected array for paramnos");
1175+
1176+
expr->paramnos = NULL;
1177+
while ((tok = JsonbIteratorNext(&ctx->it, &v, false)) != WJB_END_ARRAY)
1178+
{
1179+
int member;
1180+
1181+
if (tok != WJB_ELEM)
1182+
elog(ERROR, "Expected array element in paramnos");
1183+
1184+
member = get_int_value(&v);
1185+
expr->paramnos = bms_add_member(expr->paramnos, member);
1186+
}
1187+
}
11371188
else if (strcmp(key, "expr_simple_mutable") == 0)
11381189
{
11391190
expr->expr_simple_mutable = get_bool_value(&v);
@@ -1363,6 +1414,10 @@ deserialize_variable(DeserializeContext *ctx)
13631414
{
13641415
PLtsql_var *pvar = (PLtsql_var *) palloc0(sizeof(PLtsql_var));
13651416
pvar->dtype = dtype;
1417+
/* Match pltsql_build_variable(): preset to NULL */
1418+
pvar->value = 0;
1419+
pvar->isnull = true;
1420+
pvar->freeval = false;
13661421
var = (PLtsql_variable *) pvar;
13671422
break;
13681423
}
@@ -1945,7 +2000,7 @@ pltsql_deserialize_function_jsonb_internal(Jsonb *jb, PLtsql_function *func, PLt
19452000
pfree(key);
19462001
}
19472002

1948-
if (!parse_tree_found || !datums_found)
2003+
if (!parse_tree_found || !datums_found) // TODO: Check if we should set SERIAL_ERROR and return here
19492004
{
19502005
elog(ERROR, "Missing required fields: parse_tree=%d, datums=%d",
19512006
parse_tree_found, datums_found);
@@ -2006,6 +2061,63 @@ pltsql_deserialize_function_jsonb_internal(Jsonb *jb, PLtsql_function *func, PLt
20062061
*(ref->target_ptr) = found_var;
20072062
}
20082063
}
2064+
2065+
// /* After deserializing all datums, rebuild datatypes to ensure proper initialization */
2066+
// for (i = 0; i < ndatums; i++)
2067+
// {
2068+
// if (datums[i] != NULL)
2069+
// {
2070+
// PLtsql_datum *datum = datums[i];
2071+
2072+
// if (datum->dtype == PLTSQL_DTYPE_VAR || datum->dtype == PLTSQL_DTYPE_PROMISE)
2073+
// {
2074+
// PLtsql_var *var = (PLtsql_var *) datum;
2075+
// if (var->datatype != NULL)
2076+
// {
2077+
// /* Rebuild datatype using the deserialized type info */
2078+
// PLtsql_type *old_type = var->datatype;
2079+
// var->datatype = pltsql_build_datatype(
2080+
// old_type->typoid,
2081+
// old_type->atttypmod,
2082+
// old_type->collation,
2083+
// NULL /* origtypname - not critical for deserialized types */
2084+
// );
2085+
// /* Free the old incomplete type structure */
2086+
// pfree(old_type);
2087+
// }
2088+
// }
2089+
// else if (datum->dtype == PLTSQL_DTYPE_REC)
2090+
// {
2091+
// PLtsql_rec *rec = (PLtsql_rec *) datum;
2092+
// if (rec->datatype != NULL)
2093+
// {
2094+
// PLtsql_type *old_type = rec->datatype;
2095+
// rec->datatype = pltsql_build_datatype(
2096+
// old_type->typoid,
2097+
// old_type->atttypmod,
2098+
// old_type->collation,
2099+
// NULL
2100+
// );
2101+
// pfree(old_type);
2102+
// }
2103+
// }
2104+
// else if (datum->dtype == PLTSQL_DTYPE_TBL)
2105+
// {
2106+
// PLtsql_tbl *tbl = (PLtsql_tbl *) datum;
2107+
// if (tbl->datatype != NULL)
2108+
// {
2109+
// PLtsql_type *old_type = tbl->datatype;
2110+
// tbl->datatype = pltsql_build_datatype(
2111+
// old_type->typoid,
2112+
// old_type->atttypmod,
2113+
// old_type->collation,
2114+
// NULL
2115+
// );
2116+
// pfree(old_type);
2117+
// }
2118+
// }
2119+
// }
2120+
// }
20092121
}
20102122
PG_CATCH();
20112123
{

‎contrib/babelfishpg_tsql/src/pltsql_serialize.c‎

Lines changed: 59 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ pltsql_is_serializable(PLtsql_stmt *stmt)
7272
case PLTSQL_STMT_LABEL:
7373
case PLTSQL_STMT_CASE:
7474
case PLTSQL_STMT_TRY_CATCH:
75+
case PLTSQL_STMT_DECL_TABLE:
7576
return true;
7677

7778
default:
@@ -218,6 +219,8 @@ pltsql_serialize_function_jsonb(PLtsql_function *func, PLtsql_serial_status *sta
218219

219220
if (!pltsql_block_is_serializable(block))
220221
{
222+
elog(LOG, "Parse result caching skipped for function %u: contains unsupported node types",
223+
func->fn_oid);
221224
*status = PLTSQL_SERIAL_UNSUPPORTED;
222225
return NULL;
223226
}
@@ -333,7 +336,8 @@ pltsql_serialize_function_jsonb(PLtsql_function *func, PLtsql_serial_status *sta
333336
MemoryContextSwitchTo(oldcontext);
334337

335338
/* Log the serialization error for debugging */
336-
elog(LOG, "PLtsql function serialization failed: %s", errorData->message);
339+
elog(LOG, "Parse result caching skipped for function %u: serialization error: %s",
340+
func->fn_oid, errorData->message);
337341
FreeErrorData(errorData);
338342

339343
*status = PLTSQL_SERIAL_UNSUPPORTED;
@@ -360,6 +364,7 @@ pltsql_serialize_stmt_block_jsonb(PLtsql_stmt_block *block, PLtsql_serial_status
360364

361365
if (!pltsql_block_is_serializable(block))
362366
{
367+
elog(LOG, "Statement block serialization skipped: contains unsupported node types stmt_type=%d", block->cmd_type);
363368
*status = PLTSQL_SERIAL_UNSUPPORTED;
364369
return NULL;
365370
}
@@ -408,7 +413,8 @@ pltsql_serialize_stmt_block_jsonb(PLtsql_stmt_block *block, PLtsql_serial_status
408413
MemoryContextSwitchTo(oldcontext);
409414

410415
/* Log the serialization error for debugging */
411-
elog(LOG, "PLtsql statement block serialization failed: %s", errorData->message);
416+
elog(LOG, "Statement block serialization skipped: serialization error: %s",
417+
errorData->message);
412418
FreeErrorData(errorData);
413419

414420
*status = PLTSQL_SERIAL_UNSUPPORTED;
@@ -788,8 +794,28 @@ serialize_stmt(SerializeContext *ctx, PLtsql_stmt *stmt)
788794
}
789795
break;
790796

797+
case PLTSQL_STMT_DECL_TABLE:
798+
{
799+
PLtsql_stmt_decl_table *decl_table = (PLtsql_stmt_decl_table *) stmt;
800+
801+
/* Always serialize lineno and dno */
802+
add_int_field(ctx, "lineno", decl_table->lineno);
803+
add_int_field(ctx, "dno", decl_table->dno);
804+
805+
/* Serialize tbltypname if present (one of tbltypname/coldef is used) */
806+
if (decl_table->tbltypname)
807+
add_string_field(ctx, "tbltypname", decl_table->tbltypname);
808+
809+
/* Serialize coldef if present */
810+
if (decl_table->coldef)
811+
add_string_field(ctx, "coldef", decl_table->coldef);
812+
}
813+
break;
814+
791815
default:
792816
/* Mark as unsupported instead of throwing error */
817+
elog(LOG, "Unsupported statement type encountered during serialization: cmd_type=%d",
818+
stmt->cmd_type);
793819
*(ctx->status) = PLTSQL_SERIAL_UNSUPPORTED;
794820
break;
795821
}
@@ -838,6 +864,35 @@ serialize_expr(SerializeContext *ctx, PLtsql_expr *expr)
838864
add_string_field(ctx, "query", expr->query);
839865
add_int_field(ctx, "rwparam", expr->rwparam);
840866

867+
/*
868+
* paramnos: MUST serialize. This is a Bitmapset of all datum dnos
869+
* referenced by this expression. Without it, the executor cannot set up
870+
* the parameter list for query planning/execution, causing crashes in
871+
* pltsql_param_fetch during constant folding (e.g., when the optimizer
872+
* tries to evaluate "WHERE @var = 1" as a constant expression).
873+
*
874+
* Serialized as a JSON array of integer member values.
875+
*/
876+
if (expr->paramnos != NULL)
877+
{
878+
int x;
879+
880+
add_json_key(ctx, "paramnos");
881+
ctx->res = pushJsonbValue(&ctx->jsonstate, WJB_BEGIN_ARRAY, NULL);
882+
883+
x = -1;
884+
while ((x = bms_next_member(expr->paramnos, x)) >= 0)
885+
{
886+
JsonbValue jv;
887+
jv.type = jbvNumeric;
888+
jv.val.numeric = DatumGetNumeric(DirectFunctionCall1(int4_numeric,
889+
Int32GetDatum(x)));
890+
ctx->res = pushJsonbValue(&ctx->jsonstate, WJB_ELEM, &jv);
891+
}
892+
893+
ctx->res = pushJsonbValue(&ctx->jsonstate, WJB_END_ARRAY, NULL);
894+
}
895+
841896
/*
842897
* Simple expression fast-path fields:
843898
*
@@ -1116,6 +1171,8 @@ serialize_variable(SerializeContext *ctx, PLtsql_variable *var)
11161171

11171172
default:
11181173
/* Unsupported variable type */
1174+
elog(LOG, "Unsupported variable type encountered during serialization: dtype=%d",
1175+
var->dtype);
11191176
*(ctx->status) = PLTSQL_SERIAL_UNSUPPORTED;
11201177
break;
11211178
}

‎test/JDBC/expected/BABEL-6037-vu-prepare.out‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,14 @@
11

22
-- BABEL-6037 POC Test - Prepare
33
-- Creates test procedures for parse tree serialization/deserialization testing
4+
SELECT set_config('babelfishpg_tsql.enable_procedure_parse_cache', 'on', false);
5+
GO
6+
~~START~~
7+
text
8+
on
9+
~~END~~
10+
11+
412
-- Test 1: Simple procedure without arguments
513
-- This tests basic serialization with minimal complexity
614
CREATE PROCEDURE dbo.small_proc
@@ -172,5 +180,13 @@ BEGIN
172180
END;
173181
GO
174182

183+
SELECT set_config('babelfishpg_tsql.enable_procedure_parse_cache', 'off', false);
184+
GO
185+
~~START~~
186+
text
187+
off
188+
~~END~~
189+
190+
175191
PRINT 'All test procedures created successfully';
176192
GO

‎test/JDBC/input/BABEL-6037-vu-prepare.sql‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
-- BABEL-6037 POC Test - Prepare
22
-- Creates test procedures for parse tree serialization/deserialization testing
33

4+
SELECT set_config('babelfishpg_tsql.enable_procedure_parse_cache', 'on', false);
5+
GO
6+
47
-- Test 1: Simple procedure without arguments
58
-- This tests basic serialization with minimal complexity
69
CREATE PROCEDURE dbo.small_proc
@@ -172,5 +175,8 @@ BEGIN
172175
END;
173176
GO
174177

178+
SELECT set_config('babelfishpg_tsql.enable_procedure_parse_cache', 'off', false);
179+
GO
180+
175181
PRINT 'All test procedures created successfully';
176182
GO

0 commit comments

Comments
 (0)