diff --git a/ydb/library/yql/sql/v1/format/sql_format.cpp b/ydb/library/yql/sql/v1/format/sql_format.cpp index c088adc3e753..883417d2ee69 100644 --- a/ydb/library/yql/sql/v1/format/sql_format.cpp +++ b/ydb/library/yql/sql/v1/format/sql_format.cpp @@ -269,6 +269,10 @@ friend struct TStaticData; First = false; } + if (str == "$" && FuncCall) { + FuncCall = false; + } + if (Scopes.back() == EScope::Identifier && !FuncCall) { if (str != "$" && !NYql::LookupSimpleTypeBySqlAlias(str, true)) { SB << "id"; @@ -283,6 +287,25 @@ friend struct TStaticData; } } + void VisitPragmaValue(const TRule_pragma_value& msg) { + switch (msg.Alt_case()) { + case TRule_pragma_value::kAltPragmaValue1: { + NextToken = "0"; + break; + } + case TRule_pragma_value::kAltPragmaValue3: { + NextToken = "'str'"; + break; + } + case TRule_pragma_value::kAltPragmaValue4: { + NextToken = "false"; + break; + } + default:; + } + VisitAllFields(TRule_pragma_value::GetDescriptor(), msg); + } + void VisitLiteralValue(const TRule_literal_value& msg) { switch (msg.Alt_case()) { case TRule_literal_value::kAltLiteralValue1: { @@ -2600,6 +2623,7 @@ TStaticData::TStaticData() , ObfuscatingVisitDispatch({ {TToken::GetDescriptor(), MakeObfuscatingFunctor(&TObfuscatingVisitor::VisitToken)}, {TRule_literal_value::GetDescriptor(), MakeObfuscatingFunctor(&TObfuscatingVisitor::VisitLiteralValue)}, + {TRule_pragma_value::GetDescriptor(), MakeObfuscatingFunctor(&TObfuscatingVisitor::VisitPragmaValue)}, {TRule_atom_expr::GetDescriptor(), MakeObfuscatingFunctor(&TObfuscatingVisitor::VisitAtomExpr)}, {TRule_in_atom_expr::GetDescriptor(), MakeObfuscatingFunctor(&TObfuscatingVisitor::VisitInAtomExpr)}, {TRule_unary_casual_subexpr::GetDescriptor(), MakeObfuscatingFunctor(&TObfuscatingVisitor::VisitUnaryCasualSubexpr)}, diff --git a/ydb/library/yql/sql/v1/format/sql_format_ut.cpp b/ydb/library/yql/sql/v1/format/sql_format_ut.cpp index 8cf05b48157a..8b2dd408821d 100644 --- a/ydb/library/yql/sql/v1/format/sql_format_ut.cpp +++ b/ydb/library/yql/sql/v1/format/sql_format_ut.cpp @@ -1471,7 +1471,7 @@ FROM Input MATCH_RECOGNIZE (PATTERN (A) DEFINE A AS A); setup.Run(cases); } - Y_UNIT_TEST(Obfuscate) { + Y_UNIT_TEST(ObfuscateSelect) { TCases cases = { {"select 1;", "SELECT\n\t0;\n"}, @@ -1495,11 +1495,31 @@ FROM Input MATCH_RECOGNIZE (PATTERN (A) DEFINE A AS A); "DECLARE $id AS int32;\n"}, {"select * from `logs/of/bob` where pwd='foo';", "SELECT\n\t*\nFROM id\nWHERE id = 'str';\n"}, + {"select $f();", + "SELECT\n\t$id();\n"}, }; TSetup setup; setup.Run(cases, NSQLFormat::EFormatMode::Obfuscate); - } + } + + Y_UNIT_TEST(ObfuscatePragma) { + TCases cases = { + {"pragma a=1", + "PRAGMA id = 0;\n"}, + {"pragma a='foo';", + "PRAGMA id = 'str';\n"}, + {"pragma a=true;", + "PRAGMA id = FALSE;\n"}, + {"pragma a=$foo;", + "PRAGMA id = $id;\n"}, + {"pragma a=foo;", + "PRAGMA id = id;\n"}, + }; + + TSetup setup; + setup.Run(cases, NSQLFormat::EFormatMode::Obfuscate); + } Y_UNIT_TEST(CreateView) { TCases cases = {