[YQL-15020] Support external declares in TTranslationSettings. Fix declare overrides (last one should win)

ref:34fd8a027e547703ce6c5f880436eaca8c56ac0b

5 files changed, 86 insertions, 1 deletions

diff --git a/ydb/library/yql/sql/settings/translation_settings.h b/ydb/library/yql/sql/settings/translation_settings.h
index 0b9194a249..ca898ce5bf 100644
--- a/ydb/library/yql/sql/settings/translation_settings.h
+++ b/ydb/library/yql/sql/settings/translation_settings.h
@@ -2,6 +2,7 @@
 
 #include <util/generic/hash.h>
 #include <util/generic/hash_set.h>
+#include <util/generic/map.h>
 
 namespace google::protobuf {
     class Arena;
@@ -65,6 +66,10 @@ namespace NSQLTranslation {
         THashMap<TString, TTableBindingSettings> PrivateBindings;
         THashMap<TString, TTableBindingSettings> ScopedBindings;
 
+        // each (name, type) entry in this map is equivalent to
+        // DECLARE $name AS type;
+        TMap<TString, TString> DeclaredNamedExprs;
+
         ESqlMode Mode;
         TString DefaultCluster;
         TIncrementMonCounterFunction IncrementCounter;
diff --git a/ydb/library/yql/sql/sql.cpp b/ydb/library/yql/sql/sql.cpp
index 5a99fbf00f..79b62215b3 100644
--- a/ydb/library/yql/sql/sql.cpp
+++ b/ydb/library/yql/sql/sql.cpp
@@ -32,6 +32,12 @@ namespace NSQLTranslation {
             *actualSyntaxVersion = parsedSettings.SyntaxVersion;
         }
 
+        if (!parsedSettings.DeclaredNamedExprs.empty() && (parsedSettings.PgParser || parsedSettings.SyntaxVersion != 1)) {
+            result.Issues.AddIssue(NYql::YqlIssue(NYql::TPosition(), NYql::TIssuesIds::DEFAULT_ERROR,
+                "Externally declared named expressions are only supported in V1 syntax"));
+            return result;
+        }
+
         if (parsedSettings.PgParser) {
             return NSQLTranslationPG::PGToYql(query, parsedSettings);
         }
diff --git a/ydb/library/yql/sql/v1/context.cpp b/ydb/library/yql/sql/v1/context.cpp
index a46d7e3fcc..6107645e71 100644
--- a/ydb/library/yql/sql/v1/context.cpp
+++ b/ydb/library/yql/sql/v1/context.cpp
@@ -282,7 +282,7 @@ TNodePtr TContext::UniversalAlias(const TString& baseName, TNodePtr&& node) {
 }
 
 void TContext::DeclareVariable(const TString& varName, const TNodePtr& typeNode) {
-    Variables.emplace(varName, typeNode);
+    Variables[varName] = typeNode;
 }
 
 bool TContext::AddExport(TPosition pos, const TString& name) {
diff --git a/ydb/library/yql/sql/v1/sql.cpp b/ydb/library/yql/sql/v1/sql.cpp
index f666435356..bc9bf6b951 100644
--- a/ydb/library/yql/sql/v1/sql.cpp
+++ b/ydb/library/yql/sql/v1/sql.cpp
@@ -10076,6 +10076,31 @@ TSourcePtr TSqlQuery::Build(const TRule_multiple_column_assignment& stmt) {
 }
 
 TNodePtr TSqlQuery::Build(const TSQLv1ParserAST& ast) {
+    if (Mode == NSQLTranslation::ESqlMode::QUERY) {
+        // inject externally declared named expressions
+        for (auto [name, type] : Ctx.Settings.DeclaredNamedExprs) {
+            if (name.empty()) {
+                Error() << "Empty names for externally declared expressions are not allowed";
+                return nullptr;
+            }
+            TString varName = "$" + name;
+            if (IsAnonymousName(varName)) {
+                Error() << "Externally declared name '" << name << "' is anonymous";
+                return nullptr;
+            }
+
+            auto parsed = ParseType(type, *Ctx.Pool, Ctx.Issues, Ctx.Pos());
+            if (!parsed) {
+                Error() << "Failed to parse type for externally declared name '" << name << "'";
+                return nullptr;
+            }
+
+            TNodePtr typeNode = BuildBuiltinFunc(Ctx, Ctx.Pos(), "ParseType", { BuildLiteralRawString(Ctx.Pos(), type) });
+            varName = PushNamedAtom(Ctx.Pos(), varName);
+            Ctx.DeclareVariable(varName, typeNode);
+        }
+    }
+
     const auto& query = ast.GetRule_sql_query();
     TVector<TNodePtr> blocks;
     if (query.Alt_case() == TRule_sql_query::kAltSqlQuery1) {
diff --git a/ydb/library/yql/sql/v1/sql_ut.cpp b/ydb/library/yql/sql/v1/sql_ut.cpp
index c287f696e8..a8300b3624 100644
--- a/ydb/library/yql/sql/v1/sql_ut.cpp
+++ b/ydb/library/yql/sql/v1/sql_ut.cpp
@@ -4276,3 +4276,52 @@ Y_UNIT_TEST_SUITE(FlexibleTypes) {
         UNIT_ASSERT(SqlToYql(q).IsOk());
     }
 }
+
+Y_UNIT_TEST_SUITE(ExternalDeclares) {
+    Y_UNIT_TEST(BasicUsage) {
+        NSQLTranslation::TTranslationSettings settings;
+        settings.DeclaredNamedExprs["foo"] = "String";
+        auto res = SqlToYqlWithSettings("select $foo;", settings);
+        UNIT_ASSERT(res.IsOk());
+        UNIT_ASSERT(res.Issues.Size() == 0);
+
+        TVerifyLineFunc verifyLine = [](const TString& word, const TString& line) {
+            if (word == "declare") {
+                UNIT_ASSERT_VALUES_UNEQUAL(TString::npos, line.find(R"__((declare $foo (DataType 'String)))__"));
+            }
+        };
+
+        TWordCountHive elementStat = {{TString("declare"), 0}};
+        VerifyProgram(res, elementStat, verifyLine);
+
+        UNIT_ASSERT_VALUES_EQUAL(1, elementStat["declare"]);
+    }
+
+    Y_UNIT_TEST(DeclareOverridesExteralDeclares) {
+        NSQLTranslation::TTranslationSettings settings;
+        settings.DeclaredNamedExprs["foo"] = "String";
+        auto res = SqlToYqlWithSettings("declare $foo as Int32; select $foo;", settings);
+        UNIT_ASSERT(res.IsOk());
+
+        TVerifyLineFunc verifyLine = [](const TString& word, const TString& line) {
+            if (word == "declare") {
+                UNIT_ASSERT_VALUES_UNEQUAL(TString::npos, line.find(R"__((declare $foo (DataType 'Int32)))__"));
+            }
+        };
+
+        TWordCountHive elementStat = {{TString("declare"), 0}};
+        VerifyProgram(res, elementStat, verifyLine);
+
+        UNIT_ASSERT_VALUES_EQUAL(1, elementStat["declare"]);
+    }
+
+    Y_UNIT_TEST(DeclaresWithInvalidTypesFails) {
+        NSQLTranslation::TTranslationSettings settings;
+        settings.DeclaredNamedExprs["foo"] = "List<BadType>";
+        auto res = SqlToYqlWithSettings("select 1;", settings);
+        UNIT_ASSERT(!res.Root);
+        UNIT_ASSERT_NO_DIFF(Err2Str(res),
+            "<main>:0:5: Error: Unknown type: 'BadType'\n"
+            "<main>: Error: Failed to parse type for externally declared name 'foo'\n");
+    }
+}