[YQL-14488] Support Pg literals

ref:8773ba61fb7d6001dd6402daf399a466b17eaf56

4 files changed, 66 insertions, 17 deletions

diff --git a/ydb/library/yql/sql/v1/SQLv1.g.in b/ydb/library/yql/sql/v1/SQLv1.g.in
index 8ac069aaa3..bcf648137d 100644
--- a/ydb/library/yql/sql/v1/SQLv1.g.in
+++ b/ydb/library/yql/sql/v1/SQLv1.g.in
@@ -1341,7 +1341,7 @@ fragment STRING_SINGLE: (QUOTE_SINGLE STRING_CORE_SINGLE* QUOTE_SINGLE);
 fragment STRING_DOUBLE: (QUOTE_DOUBLE STRING_CORE_DOUBLE* QUOTE_DOUBLE);
 fragment STRING_MULTILINE: (DOUBLE_AT .* DOUBLE_AT)+ AT?;
 
-STRING_VALUE: ((STRING_SINGLE | STRING_DOUBLE | STRING_MULTILINE) (U | Y | J)?);
+STRING_VALUE: ((STRING_SINGLE | STRING_DOUBLE | STRING_MULTILINE) (U | Y | J | P (T | B | V)?)?);
 
 ID_PLAIN: ('a'..'z' | 'A'..'Z' | '_') ('a'..'z' | 'A'..'Z' | '_' | DIGIT)*;
 
@@ -1356,11 +1356,12 @@ fragment BINDIGITS: '0' B ('0' | '1')+;
 fragment DECDIGITS: DIGIT+;
 DIGITS: DECDIGITS | HEXDIGITS | OCTDIGITS | BINDIGITS;
 
-INTEGER_VALUE: DIGITS (U? (L | S | T)?);
+// not all combinations of P/U with L/S/T/I/B/N are actually valid - this is resolved in sql.cpp
+INTEGER_VALUE: DIGITS ((P | U)? (L | S | T | I | B | N)?);
 
 fragment FLOAT_EXP : E (PLUS | MINUS)? DECDIGITS ;
 REAL:
-      DECDIGITS DOT DIGIT* FLOAT_EXP? F?
+      DECDIGITS DOT DIGIT* FLOAT_EXP? (F | P F ('4'|'8'))?
     |   DECDIGITS FLOAT_EXP F?
 //    |   DOT DECDIGITS FLOAT_EXP?    // Conflicts with tuple element access through DOT
     ;
diff --git a/ydb/library/yql/sql/v1/node.cpp b/ydb/library/yql/sql/v1/node.cpp
index fccb14af00..9b089d2eaa 100644
--- a/ydb/library/yql/sql/v1/node.cpp
+++ b/ydb/library/yql/sql/v1/node.cpp
@@ -2184,19 +2184,32 @@ StringContentInternal(TContext& ctx, TPosition pos, const TString& input, EStrin
 
     TString str = input;
     if (mode == EStringContentMode::TypedStringLiteral) {
-        if (str.EndsWith("u") || str.EndsWith("U")) {
+        auto lower = to_lower(str);
+        if (lower.EndsWith("u")) {
             str = str.substr(0, str.Size() - 1);
             result.Type = NKikimr::NUdf::EDataSlot::Utf8;
-        } else if (str.EndsWith("y")) {
+        } else if (lower.EndsWith("y")) {
             str = str.substr(0, str.Size() - 1);
             result.Type = NKikimr::NUdf::EDataSlot::Yson;
-        } else if (str.EndsWith("j")) {
+        } else if (lower.EndsWith("j")) {
             str = str.substr(0, str.Size() - 1);
             result.Type = NKikimr::NUdf::EDataSlot::Json;
+        } else if (lower.EndsWith("p")) {
+            str = str.substr(0, str.Size() - 1);
+            result.PgType = "PgText";
+        } else if (lower.EndsWith("pt")) {
+            str = str.substr(0, str.Size() - 2);
+            result.PgType = "PgText";
+        } else if (lower.EndsWith("pb")) {
+            str = str.substr(0, str.Size() - 2);
+            result.PgType = "PgBytea";
+        } else if (lower.EndsWith("pv")) {
+            str = str.substr(0, str.Size() - 2);
+            result.PgType = "PgVarchar";
         }
     }
 
-    if (mode == EStringContentMode::Default && result.Type != NKikimr::NUdf::EDataSlot::String) {
+    if (mode == EStringContentMode::Default && (result.Type != NKikimr::NUdf::EDataSlot::String || result.PgType)) {
         ctx.Error(pos) << "Type suffix is not allowed here";
         return {};
     }
@@ -2226,7 +2239,7 @@ StringContentInternal(TContext& ctx, TPosition pos, const TString& input, EStrin
         return {};
     }
 
-    if (!NKikimr::NMiniKQL::IsValidStringValue(result.Type, result.Content)) {
+    if (!result.PgType.Defined() && !NKikimr::NMiniKQL::IsValidStringValue(result.Type, result.Content)) {
         ctx.Error() << "Invalid value " << result.Content.Quote() << " for type " << result.Type;
         return {};
     }
@@ -2347,7 +2360,11 @@ TLiteralNode::TLiteralNode(TPosition pos, const TString& type, const TString& va
     , Type(type)
     , Value(value)
 {
-    Add(Type, BuildQuotedAtom(Pos, Value));
+    if (Type.StartsWith("Pg")) {
+        Add("PgConst", Y("PgType", Q(to_lower(Type.substr(2)))), BuildQuotedAtom(Pos, Value));
+    } else {
+        Add(Type, BuildQuotedAtom(Pos, Value));
+    }
 }
 
 TLiteralNode::TLiteralNode(TPosition pos, const TString& value, ui32 nodeFlags)
@@ -2367,7 +2384,11 @@ TLiteralNode::TLiteralNode(TPosition pos, const TString& value, ui32 nodeFlags,
     , Type(type)
     , Value(value)
 {
-    Add(Type, BuildQuotedAtom(pos, Value, nodeFlags));
+    if (Type.StartsWith("Pg")) {
+        Add("PgConst", Y("PgType", Q(to_lower(Type.substr(2)))), BuildQuotedAtom(Pos, Value, nodeFlags));
+    } else {
+        Add(Type, BuildQuotedAtom(pos, Value, nodeFlags));
+    }
 }
 
 bool TLiteralNode::IsNull() const {
@@ -2438,13 +2459,13 @@ TNodePtr TLiteralNumberNode<T>::ApplyUnaryOp(TContext& ctx, TPosition pos, const
                 // negated value fits in Int32
                 i32 v;
                 YQL_ENSURE(TryFromString(negated, v));
-                return new TLiteralNumberNode<i32>(pos, "Int32", negated);
+                return new TLiteralNumberNode<i32>(pos, Type.StartsWith("Pg") ? "PgInt4" : "Int32", negated);
             }
             if (val <= ui64(std::numeric_limits<i64>::max()) + 1) {
                 // negated value fits in Int64
                 i64 v;
                 YQL_ENSURE(TryFromString(negated, v));
-                return new TLiteralNumberNode<i64>(pos, "Int64", negated);
+                return new TLiteralNumberNode<i64>(pos, Type.StartsWith("Pg") ? "PgInt8" : "Int64", negated);
             }
 
             ctx.Error(pos) << "Failed to parse negative integer: " << negated << ", number limit overflow";
@@ -2503,7 +2524,7 @@ TMaybe<TExprOrIdent> BuildLiteralTypedSmartStringOrId(TContext& ctx, const TStri
         return {};
     }
 
-    TString type = ToString(unescaped->Type);
+    TString type = unescaped->PgType ? *unescaped->PgType : ToString(unescaped->Type);
     result.Expr = new TLiteralNode(ctx.Pos(), unescaped->Content, unescaped->Flags, type);
     return result;
 }
diff --git a/ydb/library/yql/sql/v1/node.h b/ydb/library/yql/sql/v1/node.h
index ceb36e1991..9ad0097b1c 100644
--- a/ydb/library/yql/sql/v1/node.h
+++ b/ydb/library/yql/sql/v1/node.h
@@ -1046,6 +1046,7 @@ namespace NSQLTranslationV1 {
     struct TStringContent {
         TString Content;
         NYql::NUdf::EDataSlot Type = NYql::NUdf::EDataSlot::String;
+        TMaybe<TString> PgType;
         ui32 Flags = NYql::TNodeFlags::Default;
     };
 
diff --git a/ydb/library/yql/sql/v1/sql.cpp b/ydb/library/yql/sql/v1/sql.cpp
index 876ed51587..1d1dc7a797 100644
--- a/ydb/library/yql/sql/v1/sql.cpp
+++ b/ydb/library/yql/sql/v1/sql.cpp
@@ -3618,10 +3618,10 @@ bool ParseNumbers(TContext& ctx, const TString& strOrig, ui64& value, TString& s
     }
     if (strLen > 1) {
         auto iter = str.cend() - 1;
-        if (*iter == 'l' || *iter == 's' || *iter == 't' || /* deprecated */ *iter == 'b') {
+        if (*iter == 'l' || *iter == 's' || *iter == 't' || *iter == 's' || *iter == 'i' ||  *iter == 'b' || *iter == 'n') {
             --iter;
         }
-        if (*iter == 'u') {
+        if (*iter == 'u' || *iter == 'p') {
             --iter;
         }
         suffix = TString(++iter, str.cend());
@@ -3655,6 +3655,11 @@ bool ParseNumbers(TContext& ctx, const TString& strOrig, ui64& value, TString& s
 
 TNodePtr LiteralNumber(TContext& ctx, const TRule_integer& node) {
     const TString intergerString = ctx.Token(node.GetToken1());
+    if (to_lower(intergerString).EndsWith("pn")) {
+        // TODO: add validation
+        return new TLiteralNode(ctx.Pos(), "PgNumeric", intergerString.substr(0, intergerString.size() - 2));
+    }
+
     ui64 value;
     TString suffix;
     if (!ParseNumbers(ctx, intergerString, value, suffix)) {
@@ -3671,6 +3676,15 @@ TNodePtr LiteralNumber(TContext& ctx, const TRule_integer& node) {
             return new TLiteralNumberNode<i64>(ctx.Pos(), "Int64", ToString(value), implicitType);
         }
         return new TLiteralNumberNode<i32>(ctx.Pos(), "Int32", ToString(value), implicitType);
+    } else if (suffix == "p") {
+        bool implicitType = true;
+        if (noSpaceForInt64) {
+            ctx.Error(ctx.Pos()) << "Failed to parse number from string: " << intergerString << ", 64 bit signed integer overflow";
+            return {};
+        } else if (noSpaceForInt32) {
+            return new TLiteralNumberNode<i64>(ctx.Pos(), "PgInt8", ToString(value), implicitType);
+        }
+        return new TLiteralNumberNode<i32>(ctx.Pos(), "PgInt4", ToString(value), implicitType);
     } else if (suffix == "u") {
         return new TLiteralNumberNode<ui32>(ctx.Pos(), "Uint32", ToString(value));
     } else if (suffix == "ul") {
@@ -3685,6 +3699,14 @@ TNodePtr LiteralNumber(TContext& ctx, const TRule_integer& node) {
         return new TLiteralNumberNode<ui16>(ctx.Pos(), "Uint16", ToString(value));
     } else if (suffix == "s") {
         return new TLiteralNumberNode<i16>(ctx.Pos(), "Int16", ToString(value));
+    } else if (suffix == "ps") {
+        return new TLiteralNumberNode<i16>(ctx.Pos(), "PgInt2", ToString(value));
+    } else if (suffix == "pi") {
+        return new TLiteralNumberNode<i32>(ctx.Pos(), "PgInt4", ToString(value));
+    } else if (suffix == "pb") {
+        return new TLiteralNumberNode<i64>(ctx.Pos(), "PgInt8", ToString(value));
+    } else if (suffix == "pn") {
+        return new TLiteralNode(ctx.Pos(), "PgNumeric", ToString(value));
     } else {
         ctx.Error(ctx.Pos()) << "Failed to parse number from string: " << intergerString << ", invalid suffix: " << suffix;
         return {};
@@ -3694,9 +3716,13 @@ TNodePtr LiteralNumber(TContext& ctx, const TRule_integer& node) {
 TNodePtr LiteralReal(TContext& ctx, const TRule_real& node) {
     const TString value(ctx.Token(node.GetToken1()));
     YQL_ENSURE(!value.empty());
-    const auto lastValue = value[value.size() - 1];
-    if (lastValue == 'f' || lastValue == 'F') {
+    auto lower = to_lower(value);
+    if (lower.EndsWith("f")) {
         return new TLiteralNumberNode<float>(ctx.Pos(), "Float", value.substr(0, value.size()-1));
+    } else if (lower.EndsWith("pf4")) {
+        return new TLiteralNumberNode<float>(ctx.Pos(), "PgFloat4", value.substr(0, value.size()-3));
+    } else if (lower.EndsWith("pf8")) {
+        return new TLiteralNumberNode<float>(ctx.Pos(), "PgFloat8", value.substr(0, value.size()-3));
     } else {
         return new TLiteralNumberNode<double>(ctx.Pos(), "Double", value);
     }