[YQL-14487] Support Pg prefixes for types

ref:db66712a7e838104981003931dae54eb901450c7

6 files changed, 47 insertions, 5 deletions

diff --git a/ydb/library/yql/core/type_ann/type_ann_core.cpp b/ydb/library/yql/core/type_ann/type_ann_core.cpp
index 0b5d83497a8..594e581e915 100644
--- a/ydb/library/yql/core/type_ann/type_ann_core.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_core.cpp
@@ -9561,7 +9561,14 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
             return IGraphTransformer::TStatus::Error;
         }
 
-        auto typeId = NPg::LookupType(TString(input->Child(0)->Content())).TypeId;
+        auto type = TString(input->Child(0)->Content());
+        if (!NPg::HasType(type)) {
+            ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Child(0)->Pos()),
+                TStringBuilder() << "Unknown type: '" << type << "'"));
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        auto typeId = NPg::LookupType(type).TypeId;
         input->SetTypeAnn(ctx.Expr.MakeType<TTypeExprType>(ctx.Expr.MakeType<TPgExprType>(typeId)));
         return IGraphTransformer::TStatus::Ok;
     }
diff --git a/ydb/library/yql/sql/v1/CMakeLists.txt b/ydb/library/yql/sql/v1/CMakeLists.txt
index 28dab4b0aec..66e7f9cbf3a 100644
--- a/ydb/library/yql/sql/v1/CMakeLists.txt
+++ b/ydb/library/yql/sql/v1/CMakeLists.txt
@@ -29,6 +29,7 @@ target_link_libraries(yql-sql-v1 PUBLIC
   proto_ast-gen-v1
   proto_ast-gen-v1_ansi
   proto_ast-gen-v1_proto
+  yql-parser-pg_catalog
   tools-enum_parser-enum_serialization_runtime
 )
 target_sources(yql-sql-v1 PRIVATE
diff --git a/ydb/library/yql/sql/v1/node.cpp b/ydb/library/yql/sql/v1/node.cpp
index 2e465373270..8ee87277ef3 100644
--- a/ydb/library/yql/sql/v1/node.cpp
+++ b/ydb/library/yql/sql/v1/node.cpp
@@ -5,6 +5,7 @@
 #include <ydb/library/yql/ast/yql_expr.h>
 #include <ydb/library/yql/core/sql_types/simple_types.h>
 #include <ydb/library/yql/minikql/mkql_type_ops.h>
+#include <ydb/library/yql/parser/pg_catalog/catalog.h>
 #include <ydb/library/yql/utils/yql_panic.h>
 
 #include <library/cpp/containers/stack_vector/stack_vec.h>
@@ -3119,14 +3120,36 @@ TNodePtr BuildDataType(TPosition pos, const TString& typeName) {
     return new TCallNodeImpl(pos, "DataType", {BuildQuotedAtom(pos, typeName, TNodeFlags::Default)});
 }
 
+TMaybe<TString> LookupSimpleType(const TStringBuf& alias, bool flexibleTypes) {
+    if (auto sqlAlias = LookupSimpleTypeBySqlAlias(alias, flexibleTypes)) {
+        return TString(*sqlAlias);
+    }
+
+    TString normalized = to_lower(TString(alias));
+    TString pgType;
+    if (normalized.StartsWith("_pg")) {
+        pgType = normalized.substr(3);
+    } else if (normalized.StartsWith("pg")) {
+        pgType = normalized.substr(2);
+    } else {
+        return {};
+    }
+
+    if (NPg::HasType(pgType)) {
+        return normalized;
+    }
+
+    return {};
+}
+
 TNodePtr BuildSimpleType(TContext& ctx, TPosition pos, const TString& typeName, bool dataOnly) {
-    auto found = LookupSimpleTypeBySqlAlias(typeName, ctx.FlexibleTypes);
+    auto found = LookupSimpleType(typeName, ctx.FlexibleTypes);
     if (!found) {
         ctx.Error(pos) << "Unknown simple type '" << typeName << "'";
         return {};
     }
 
-    auto type = ToString(*found);
+    auto type = *found;
     if (type == "Void" || type == "Unit" || type == "Generic" || type == "EmptyList" || type == "EmptyDict") {
         if (dataOnly) {
             ctx.Error(pos) << "Only data types are allowed here, but got: '" << typeName << "'";
@@ -3136,6 +3159,16 @@ TNodePtr BuildSimpleType(TContext& ctx, TPosition pos, const TString& typeName,
         return new TCallNodeImpl(pos, type, {});
     }
 
+    if (type.StartsWith("_pg") || type.StartsWith("pg")) {
+        TString pgType;
+        if (type.StartsWith("_pg")) {
+            pgType = "_" + type.substr(3);
+        } else {
+            pgType = type.substr(2);
+        }
+        return new TCallNodeImpl(pos, "PgType", { BuildQuotedAtom(pos, pgType, TNodeFlags::Default) });
+    }
+
     return new TCallNodeImpl(pos, "DataType", { BuildQuotedAtom(pos, type, TNodeFlags::Default) });
 }
 
diff --git a/ydb/library/yql/sql/v1/node.h b/ydb/library/yql/sql/v1/node.h
index b228ca92128..c9caf14e48b 100644
--- a/ydb/library/yql/sql/v1/node.h
+++ b/ydb/library/yql/sql/v1/node.h
@@ -1210,6 +1210,7 @@ namespace NSQLTranslationV1 {
     TNodePtr BuildBind(TPosition pos, const TString& module, const TString& alias);
     TNodePtr BuildLambda(TPosition pos, TNodePtr params, TNodePtr body, const TString& resName = TString());
     TNodePtr BuildDataType(TPosition pos, const TString& typeName);
+    TMaybe<TString> LookupSimpleType(const TStringBuf& alias, bool flexibleTypes);
     TNodePtr BuildSimpleType(TContext& ctx, TPosition pos, const TString& typeName, bool dataOnly);
     TNodePtr BuildIsNullOp(TPosition pos, TNodePtr a);
     TNodePtr BuildBinaryOp(TContext& ctx, TPosition pos, const TString& opName, TNodePtr a, TNodePtr b);
diff --git a/ydb/library/yql/sql/v1/sql.cpp b/ydb/library/yql/sql/v1/sql.cpp
index 4850338fdc6..0981adde7a6 100644
--- a/ydb/library/yql/sql/v1/sql.cpp
+++ b/ydb/library/yql/sql/v1/sql.cpp
@@ -9,7 +9,6 @@
 #include <ydb/library/yql/parser/proto_ast/gen/v1/SQLv1Parser.h>
 #include <ydb/library/yql/parser/proto_ast/gen/v1_ansi/SQLv1Lexer.h>
 #include <ydb/library/yql/parser/proto_ast/gen/v1_ansi/SQLv1Parser.h>
-#include <ydb/library/yql/core/sql_types/simple_types.h>
 
 #include <ydb/library/yql/ast/yql_expr.h>
 #include <ydb/library/yql/parser/proto_ast/gen/v1_proto/SQLv1Parser.pb.h>
@@ -4239,7 +4238,7 @@ TNodePtr TSqlExpression::UnaryCasualExpr(const TUnaryCasualExprRule& node, const
     } else {
         const bool flexibleTypes = Ctx.FlexibleTypes;
         bool columnOrType = false;
-        if (auto simpleType = LookupSimpleTypeBySqlAlias(name, flexibleTypes); simpleType && typePossible && suffixIsEmpty) {
+        if (auto simpleType = LookupSimpleType(name, flexibleTypes); simpleType && typePossible && suffixIsEmpty) {
             auto columnRefsState = Ctx.GetColumnReferenceState();
             if (tail.Count > 0 || columnRefsState == EColumnRefState::Deny || !flexibleTypes) {
                 // a type
diff --git a/ydb/library/yql/sql/v1/ya.make b/ydb/library/yql/sql/v1/ya.make
index 432fdefbf71..e7ef5367ea4 100644
--- a/ydb/library/yql/sql/v1/ya.make
+++ b/ydb/library/yql/sql/v1/ya.make
@@ -21,6 +21,7 @@ PEERDIR(
     ydb/library/yql/parser/proto_ast/gen/v1
     ydb/library/yql/parser/proto_ast/gen/v1_ansi
     ydb/library/yql/parser/proto_ast/gen/v1_proto
+    ydb/library/yql/parser/pg_catalog
 )
 
 SRCS(