pg query parameters

4 files changed, 90 insertions, 6 deletions

diff --git a/ydb/core/kqp/provider/yql_kikimr_results.cpp b/ydb/core/kqp/provider/yql_kikimr_results.cpp
index 369fa2bb077..81ec7e96344 100644
--- a/ydb/core/kqp/provider/yql_kikimr_results.cpp
+++ b/ydb/core/kqp/provider/yql_kikimr_results.cpp
@@ -460,7 +460,10 @@ const TTypeAnnotationNode* ParseTypeFromKikimrProto(const NKikimrMiniKQL::TType&
 
             return ctx.MakeType<TDictExprType>(keyType, payloadType);
         }
-
+        case NKikimrMiniKQL::ETypeKind::Pg: {
+            const NKikimrMiniKQL::TPgType& protoData = type.GetPg();
+            return ctx.MakeType<TPgExprType>(protoData.Getoid());
+        }
         default: {
             ctx.AddError(TIssue(TPosition(), TStringBuilder() << "Unsupported protobuf type: "
                 << type.ShortDebugString()));
@@ -547,6 +550,13 @@ bool ExportTypeToKikimrProto(const TTypeAnnotationNode& type, NKikimrMiniKQL::TT
             return true;
         }
 
+        case ETypeAnnotationKind::Pg: {
+            protoType.SetKind(NKikimrMiniKQL::ETypeKind::Pg);
+            auto pgTypeId = type.Cast<TPgExprType>()->GetId();
+            auto pgTypeName = type.Cast<TPgExprType>()->GetName();
+            protoType.MutablePg()->Setoid(pgTypeId);
+            return true;
+        }
         default: {
             ctx.AddError(TIssue(TPosition(), TStringBuilder() << "Unsupported protobuf type: " << type));
             return false;
@@ -705,4 +715,3 @@ TExprNode::TPtr ParseKikimrProtoValue(const NKikimrMiniKQL::TType& type, const N
 }
 
 } // namespace NYql
-
diff --git a/ydb/core/kqp/session_actor/kqp_session_actor.cpp b/ydb/core/kqp/session_actor/kqp_session_actor.cpp
index 814220cc8c3..119735f03ec 100644
--- a/ydb/core/kqp/session_actor/kqp_session_actor.cpp
+++ b/ydb/core/kqp/session_actor/kqp_session_actor.cpp
@@ -1043,7 +1043,6 @@ public:
                 QueryState->QueryData->AddMkqlParam(name, type, value);
                 return;
             }
-
             ythrow TRequestFail(Ydb::StatusIds::BAD_REQUEST) << "Missing value for parameter: " << name;
         }
 
diff --git a/ydb/core/kqp/ut/pg/kqp_pg_ut.cpp b/ydb/core/kqp/ut/pg/kqp_pg_ut.cpp
index 3d063c15c09..ade78d357e4 100644
--- a/ydb/core/kqp/ut/pg/kqp_pg_ut.cpp
+++ b/ydb/core/kqp/ut/pg/kqp_pg_ut.cpp
@@ -1163,6 +1163,73 @@ Y_UNIT_TEST_SUITE(KqpPg) {
         UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::INTERNAL_ERROR);
         UNIT_ASSERT(result.GetIssues().begin()->GetMessage().EndsWith("notNull is forbidden for pg types"));
     }
+
+    Y_UNIT_TEST(ValuesInsert) {
+        TKikimrRunner kikimr(NKqp::TKikimrSettings().SetWithSampleTables(false));
+        auto testSingleType = [&kikimr] (const TPgTypeTestSpec& spec) {
+            auto tableClient = kikimr.GetTableClient();
+            auto session = tableClient.CreateSession().GetValueSync().GetSession();
+            auto tableName = createTable(tableClient, session, spec.TypeId, spec.IsKey, true, spec.TextIn, "", 0);
+            auto* typeDesc = NPg::TypeDescFromPgTypeId(spec.TypeId);
+            auto typeName = NPg::PgTypeNameFromTypeDesc(typeDesc);
+            auto keyType = spec.IsKey ? typeName : "pgint2";
+            auto req = Sprintf("\
+            --!syntax_v1\n\
+            DECLARE $key0 as %s;\n\
+            DECLARE $key1 as %s;\n\
+            DECLARE $value0 as %s;\n\
+            DECLARE $value1 as %s;\n\
+            INSERT INTO `%s` (key, value) VALUES ($key0, $value0), ($key1, $value1);\n\
+            ", keyType.c_str(), keyType.c_str(), typeName.c_str(), typeName.c_str(), tableName.c_str());
+            Cerr << req << Endl;
+            auto key0Value = TPgValue(TPgValue::VK_TEXT, spec.IsKey ? spec.TextIn(0) : "0", TPgType(keyType));
+            auto key1Value = TPgValue(TPgValue::VK_TEXT, spec.IsKey ? spec.TextIn(1) : "1", TPgType(keyType));
+            auto params = tableClient.GetParamsBuilder()
+                .AddParam("$key0")
+                    .Pg(key0Value)
+                    .Build()
+                .AddParam("$value0")
+                    .Pg(TPgValue(TPgValue::VK_TEXT, spec.TextIn(0), TPgType(typeName)))
+                    .Build()
+                .AddParam("$key1")
+                    .Pg(key1Value)
+                    .Build()
+                .AddParam("$value1")
+                    .Pg(TPgValue(TPgValue::VK_TEXT, spec.TextIn(1), TPgType(typeName)))
+                    .Build()
+                .Build();
+            auto result = session.ExecuteDataQuery(req, TTxControl::BeginTx().CommitTx(), params).ExtractValueSync();
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+            auto selectResult = ExecutePgSelect(kikimr, tableName);
+            ValidatePgYqlResult(selectResult, spec);
+        };
+        auto testType = [&] (const TPgTypeTestSpec& spec) {
+            auto textInArray = [&spec] (auto i) {
+                auto str = spec.TextIn(i);
+                return spec.ArrayPrint(str);
+            };
+
+            auto textOutArray = [&spec] (auto i) {
+                auto str = spec.TextOut(i);
+                return spec.ArrayPrint(str);
+            };
+
+            auto arrayTypeId = NYql::NPg::LookupType(spec.TypeId).ArrayTypeId;
+            TPgTypeTestSpec arraySpec{arrayTypeId, false, textInArray, textOutArray};
+            testSingleType(spec);
+            testSingleType(arraySpec);
+        };
+
+        for (const auto& spec : typeSpecs) {
+            Cerr << spec.TypeId << Endl;
+            if (spec.TypeId == CHAROID) {
+                continue;
+                // I cant come up with a query with explicit char conversion.
+                // ::char, ::character casts to pg_bpchar
+            }
+            testType(spec);
+        }
+    }
 }
 
 } // namespace NKqp
diff --git a/ydb/library/mkql_proto/mkql_proto.cpp b/ydb/library/mkql_proto/mkql_proto.cpp
index 38eeba2ba82..b95cb6feea2 100644
--- a/ydb/library/mkql_proto/mkql_proto.cpp
+++ b/ydb/library/mkql_proto/mkql_proto.cpp
@@ -1677,9 +1677,18 @@ NUdf::TUnboxedValue TProtoImporter::ImportValueFromProto(const TType* type, cons
         return unboxedValue;
     }
 
-    case TType::EKind::Pg:
-        // TODO: support pg types
-        MKQL_ENSURE(false, "pg types are not supported");
+    case TType::EKind::Pg: {
+        const TPgType* pgType = static_cast<const TPgType*>(type);
+        NYql::NUdf::TUnboxedValue unboxedValue;
+        if (value.Hastext_value()) {
+            unboxedValue = NYql::NCommon::PgValueFromNativeText(value.Gettext_value(), pgType->GetTypeId());
+        } else if (value.Hasbytes_value()) {
+            unboxedValue = NYql::NCommon::PgValueFromNativeBinary(value.Getbytes_value(), pgType->GetTypeId());
+        } else {
+            MKQL_ENSURE(false, "empty pg value proto");
+        }
+        return unboxedValue;
+    }
 
     default:
         MKQL_ENSURE(false, TStringBuilder() << "Unknown kind: " << type->GetKindAsStr());