[YQL-14487] Support two versions of PgCall: PgCall("foo", ....) - with String lietral as first arg, and Pg::Foo(....)

ref:192d4d842c70d18b03e01a15be6665ed1638ddf8

5 files changed, 49 insertions, 0 deletions

diff --git a/ydb/library/yql/sql/v1/builtin.cpp b/ydb/library/yql/sql/v1/builtin.cpp
index ae84d22cb52..349af199701 100644
--- a/ydb/library/yql/sql/v1/builtin.cpp
+++ b/ydb/library/yql/sql/v1/builtin.cpp
@@ -664,6 +664,36 @@ public:
     }
 };
 
+class TYqlPgCall : public TCallNode {
+public:
+    TYqlPgCall(TPosition pos, const TVector<TNodePtr>& args)
+        : TCallNode(pos, "PgCall", 1, -1, args)
+    {
+    }
+
+    bool DoInit(TContext& ctx, ISource* src) final {
+        if (!ValidateArguments(ctx)) {
+            return false;
+        }
+
+        if (!Args[0]->Init(ctx, src)) {
+            return false;
+        }
+
+        if (!Args[0]->IsLiteral() || Args[0]->GetLiteralType() != "String") {
+            ctx.Error(Args[0]->GetPos()) << "Expecting string literal as first argument";
+            return false;
+        }
+
+        Args[0] = BuildQuotedAtom(Args[0]->GetPos(), Args[0]->GetLiteralValue());
+        return TCallNode::DoInit(ctx, src);
+    }
+
+    TNodePtr DoClone() const final {
+        return new TYqlPgCall(Pos, CloneContainer(Args));
+    }
+};
+
 template <const char* Name>
 class TYqlSubqueryFor : public TCallNode {
 public:
@@ -2686,6 +2716,7 @@ struct TBuiltinFuncData {
             {"pgtype", BuildSimpleBuiltinFactoryCallback<TYqlPgType>() },
             {"pgconst", BuildSimpleBuiltinFactoryCallback<TYqlPgConst>() },
             {"pgop", BuildSimpleBuiltinFactoryCallback<TYqlPgOp>() },
+            {"pgcall", BuildSimpleBuiltinFactoryCallback<TYqlPgCall>() },
             {"pgcast", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("PgCast", 2, 3) },
             {"frompg", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("FromPg", 1, 1) },
             {"topg", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("ToPg", 1, 1) },
@@ -3092,6 +3123,11 @@ TNodePtr BuildBuiltinFunc(TContext& ctx, TPosition pos, TString name, const TVec
         }
     } else if (ns == "datetime2" && (name == "Format" || name == "Parse")) {
         return BuildUdf(ctx, pos, nameSpace, name, args);
+    } else if (ns == "pg") {
+        TVector<TNodePtr> pgCallArgs;
+        pgCallArgs.push_back(BuildLiteralRawString(pos, name));
+        pgCallArgs.insert(pgCallArgs.end(), args.begin(), args.end());
+        return new TYqlPgCall(pos, pgCallArgs);
     } else if (name == "MakeLibraPreprocessor") {
         if (args.size() != 1) {
             return new TInvalidBuiltin(pos, TStringBuilder() << name << " requires exactly one argument");
diff --git a/ydb/library/yql/sql/v1/context.cpp b/ydb/library/yql/sql/v1/context.cpp
index 85c5001638c..34baf933b4b 100644
--- a/ydb/library/yql/sql/v1/context.cpp
+++ b/ydb/library/yql/sql/v1/context.cpp
@@ -54,6 +54,7 @@ THashMap<TStringBuf, TPragmaField> CTX_PRAGMA_FIELDS = {
     {"UnorderedSubqueries", &TContext::UnorderedSubqueries},
     {"FlexibleTypes", &TContext::FlexibleTypes},
     {"AnsiCurrentRow", &TContext::AnsiCurrentRow},
+    {"PgTypes", &TContext::PgTypes},
 };
 
 typedef TMaybe<bool> TContext::*TPragmaMaybeField;
diff --git a/ydb/library/yql/sql/v1/context.h b/ydb/library/yql/sql/v1/context.h
index 8c09bef194d..0faea0b002f 100644
--- a/ydb/library/yql/sql/v1/context.h
+++ b/ydb/library/yql/sql/v1/context.h
@@ -259,6 +259,7 @@ namespace NSQLTranslationV1 {
         bool FlexibleTypes = false;
         // see YQL-10265
         bool AnsiCurrentRow = false;
+        bool PgTypes = false;
         THashMap<TString, TMaybe<TString>> Libraries; // alias -> optional file
         THashMap<TString, ui32> PackageVersions;
         NYql::TWarningPolicy WarningPolicy;
diff --git a/ydb/library/yql/sql/v1/query.cpp b/ydb/library/yql/sql/v1/query.cpp
index df41d79a877..a4269dbcf40 100644
--- a/ydb/library/yql/sql/v1/query.cpp
+++ b/ydb/library/yql/sql/v1/query.cpp
@@ -1775,6 +1775,11 @@ public:
                             BuildQuotedAtom(Pos, "Attr"), BuildQuotedAtom(Pos, "Dummy_"), BuildQuotedAtom(Pos, "1"))));
                     }
                 }
+
+                if (ctx.PgTypes) {
+                    Add(Y("let", "world", Y(TString(ConfigureName), "world", configSource,
+                        BuildQuotedAtom(Pos, "PgTypes"))));
+                }
             }
         }
 
diff --git a/ydb/library/yql/sql/v1/sql.cpp b/ydb/library/yql/sql/v1/sql.cpp
index 876ed51587d..12188a0260d 100644
--- a/ydb/library/yql/sql/v1/sql.cpp
+++ b/ydb/library/yql/sql/v1/sql.cpp
@@ -9651,6 +9651,12 @@ TNodePtr TSqlQuery::PragmaStatement(const TRule_pragma_stmt& stmt, bool& success
         } else if (normalizedPragma == "disableansicurrentrow") {
             Ctx.AnsiCurrentRow = false;
             Ctx.IncrementMonCounter("sql_pragma", "DisableAnsiCurrentRow");
+        } else if (normalizedPragma == "pgtypes") {
+            Ctx.PgTypes = true;
+            Ctx.IncrementMonCounter("sql_pragma", "PgTypes");
+        } else if (normalizedPragma == "disablepgtypes") {
+            Ctx.PgTypes = false;
+            Ctx.IncrementMonCounter("sql_pragma", "DisablePgTypes");
         } else {
             Error() << "Unknown pragma: " << pragma;
             Ctx.IncrementMonCounter("sql_errors", "UnknownPragma");