fix

Author: vitstn

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)},

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 = {