merge to ydb stable YQ-4070 supported insert values for s3 (#14915)

GrigoriyPA · evanevanevanevannnn · web-flow · commit 7fb005fc9fff · 2025-02-23T12:38:27.000+03:00
Co-authored-by: Ivan Sukhov &lt;evanevannnn@ydb.tech&gt;
diff --git a/ydb/core/external_sources/object_storage.cpp b/ydb/core/external_sources/object_storage.cpp
@@ -150,6 +150,7 @@ struct TObjectStorageExternalSource : public IExternalSource {
         }
         const bool hasPartitioning = objectStorage.projection_size() || objectStorage.partitioned_by_size();
         issues.AddIssues(ValidateFormatSetting(objectStorage.format(), objectStorage.format_setting(), location, hasPartitioning));
+        issues.AddIssues(ValidateSchema(schema));
         issues.AddIssues(ValidateJsonListFormat(objectStorage.format(), schema, objectStorage.partitioned_by()));
         issues.AddIssues(ValidateRawFormat(objectStorage.format(), schema, objectStorage.partitioned_by()));
         if (hasPartitioning) {
@@ -268,6 +269,22 @@ struct TObjectStorageExternalSource : public IExternalSource {
         return issues;
     }
 
+    template<typename TScheme>
+    static NYql::TIssues ValidateSchema(const TScheme& schema) {
+        NYql::TIssues issues;
+        for (const auto& column: schema.column()) {
+            const auto type = column.type();
+            if (type.has_optional_type() && type.optional_type().item().has_optional_type()) {
+                issues.AddIssue(MakeErrorIssue(
+                    Ydb::StatusIds::BAD_REQUEST,
+                    TStringBuilder{} << "Double optional types are not supported (you have '" 
+                        << column.name() << " " << NYdb::TType(column.type()).ToString() << "' field)"));
+            }
+        }
+
+        return issues;
+    }
+
     template<typename TScheme>
     static NYql::TIssues ValidateJsonListFormat(const TString& format, const TScheme& schema, const google::protobuf::RepeatedPtrField<TString>& partitionedBy) {
         NYql::TIssues issues;
diff --git a/ydb/core/external_sources/object_storage_ut.cpp b/ydb/core/external_sources/object_storage_ut.cpp
@@ -55,6 +55,15 @@ Y_UNIT_TEST_SUITE(ObjectStorageTest) {
         UNIT_ASSERT_EXCEPTION_CONTAINS(source->Pack(schema, general), NExternalSource::TExternalSourceException, "Date, Timestamp and Interval types are not allowed in json_list format");
     }
 
+    Y_UNIT_TEST(FailedOptionalTypeValidation) {
+        auto source = NExternalSource::CreateObjectStorageExternalSource({}, nullptr, 1000, nullptr, false, false);
+        NKikimrExternalSources::TSchema schema;
+        NKikimrExternalSources::TGeneral general;
+        auto newColumn = schema.add_column();
+        newColumn->mutable_type()->mutable_optional_type()->mutable_item()->mutable_optional_type()->mutable_item()->set_type_id(Ydb::Type::INT32);
+        UNIT_ASSERT_EXCEPTION_CONTAINS(source->Pack(schema, general), NExternalSource::TExternalSourceException, "Double optional types are not supported");
+    }
+
     Y_UNIT_TEST(WildcardsValidation) {
         auto source = NExternalSource::CreateObjectStorageExternalSource({}, nullptr, 1000, nullptr, false, false);
         NKikimrExternalSources::TSchema schema;
diff --git a/ydb/core/kqp/ut/federated_query/s3/kqp_federated_query_ut.cpp b/ydb/core/kqp/ut/federated_query/s3/kqp_federated_query_ut.cpp
@@ -805,6 +805,82 @@ Y_UNIT_TEST_SUITE(KqpFederatedQuery) {
         UNIT_ASSERT_EQUAL(GetObjectKeys(writeBucket).size(), 3);
     }
 
+    Y_UNIT_TEST(InsertIntoBucketValuesCast) {
+        const TString writeDataSourceName = "/Root/write_data_source";
+        const TString writeTableName = "/Root/write_binding";
+        const TString writeBucket = "test_bucket_values_cast";
+        const TString writeObject = "test_object_write/";
+        {
+            Aws::S3::S3Client s3Client = MakeS3Client();
+            CreateBucket(writeBucket, s3Client);
+        }
+
+        auto kikimr = NTestUtils::MakeKikimrRunner();
+
+        auto tc = kikimr->GetTableClient();
+        auto session = tc.CreateSession().GetValueSync().GetSession();
+        {
+            const TString query = fmt::format(R"(
+                CREATE EXTERNAL DATA SOURCE `{write_source}` WITH (
+                    SOURCE_TYPE="ObjectStorage",
+                    LOCATION="{write_location}",
+                    AUTH_METHOD="NONE"
+                );
+                CREATE EXTERNAL TABLE `{write_table}` (
+                    key Uint64 NOT NULL,
+                    value String NOT NULL
+                ) WITH (
+                    DATA_SOURCE="{write_source}",
+                    LOCATION="{write_object}",
+                    FORMAT="tsv_with_names"
+                );
+                )",
+                "write_source"_a = writeDataSourceName,
+                "write_table"_a = writeTableName,
+                "write_location"_a = GetBucketLocation(writeBucket),
+                "write_object"_a = writeObject);
+
+            const auto result = session.ExecuteSchemeQuery(query).GetValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToOneLineString());
+        }
+
+        auto db = kikimr->GetQueryClient();
+        {
+            const TString query = fmt::format(R"(
+                INSERT INTO `{write_table}`
+                    (key, value)
+                VALUES
+                    (1, "#######"),
+                    (4294967295u, "#######");
+
+                INSERT INTO `{write_source}`.`{write_object}` WITH (FORMAT = "tsv_with_names")
+                    (key, value)
+                VALUES
+                    (1, "#######"),
+                    (4294967295u, "#######");
+
+                INSERT INTO `{write_table}` SELECT * FROM AS_TABLE([
+                    <|key: 1, value: "#####"|>,
+                    <|key: 4294967295u, value: "#####"|>
+                ]);
+
+                INSERT INTO `{write_source}`.`{write_object}` WITH (FORMAT = "tsv_with_names")
+                SELECT * FROM AS_TABLE([
+                    <|key: 1, value: "#####"|>,
+                    <|key: 4294967295u, value: "#####"|>
+                ]);
+                )",
+                "write_source"_a = writeDataSourceName,
+                "write_table"_a = writeTableName,
+                "write_object"_a = writeObject);
+
+            const auto result = db.ExecuteQuery(query, NYdb::NQuery::TTxControl::BeginTx().CommitTx()).GetValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToOneLineString());
+        }
+
+        UNIT_ASSERT_EQUAL(GetObjectKeys(writeBucket).size(), 4);
+    }
+
     Y_UNIT_TEST(UpdateExternalTable) {
         const TString readDataSourceName = "/Root/read_data_source";
         const TString readTableName = "/Root/read_binding";
diff --git a/ydb/library/yql/providers/s3/common/util.cpp b/ydb/library/yql/providers/s3/common/util.cpp
@@ -78,4 +78,20 @@ TString TUrlBuilder::Build() const {
     return std::move(result);
 }
 
+bool ValidateS3ReadWriteSchema(const TStructExprType* schemaStructRowType, TExprContext& ctx) {
+    for (const TItemExprType* item : schemaStructRowType->GetItems()) {
+        const TTypeAnnotationNode* rowType = item->GetItemType();
+        if (rowType->GetKind() == ETypeAnnotationKind::Optional) {
+            rowType = rowType->Cast<TOptionalExprType>()->GetItemType();
+        }
+
+        if (rowType->GetKind() == ETypeAnnotationKind::Optional) {
+            ctx.AddError(TIssue(TStringBuilder() << "Double optional types are not supported (you have '"
+                << item->GetName() << " " << FormatType(item->GetItemType()) << "' field)"));
+            return false;
+        }
+    }
+    return true;
+}
+
 }
diff --git a/ydb/library/yql/providers/s3/common/util.h b/ydb/library/yql/providers/s3/common/util.h
@@ -2,6 +2,7 @@
 
 #include <util/string/builder.h>
 #include <ydb/library/yql/public/issue/yql_issue.h>
+#include <ydb/library/yql/ast/yql_expr.h>
 
 namespace NYql::NS3Util {
 
@@ -30,4 +31,6 @@ class TUrlBuilder {
     TString MainUri;
 };
 
+bool ValidateS3ReadWriteSchema(const TStructExprType* schemaStructRowType, TExprContext& ctx);
+
 }
diff --git a/ydb/library/yql/providers/s3/common/ya.make b/ydb/library/yql/providers/s3/common/ya.make
@@ -19,6 +19,7 @@ PEERDIR(
     ydb/library/yql/providers/s3/events
     ydb/library/yql/public/issue
     ydb/library/yql/public/issue/protos
+    ydb/library/yql/ast
 )
 
 IF (CLANG AND NOT WITH_VALGRIND)
diff --git a/ydb/library/yql/providers/s3/provider/yql_s3_datasink_type_ann.cpp b/ydb/library/yql/providers/s3/provider/yql_s3_datasink_type_ann.cpp
@@ -2,6 +2,7 @@
 
 #include <ydb/library/yql/core/expr_nodes/yql_expr_nodes.h>
 #include <ydb/library/yql/core/yql_opt_utils.h>
+#include <ydb/library/yql/providers/s3/common/util.h>
 #include <ydb/library/yql/providers/s3/expr_nodes/yql_s3_expr_nodes.h>
 
 #include <ydb/library/yql/providers/common/provider/yql_provider.h>
@@ -26,6 +27,7 @@ TExprNode::TListType GetPartitionKeys(const TExprNode::TPtr& partBy) {
 
     return {};
 }
+
 }
 
 namespace {
@@ -64,7 +66,47 @@ class TS3DataSinkTypeAnnotationTransformer : public TVisitorTransformerBase {
             return TStatus::Error;
         }
 
-        auto source = input->Child(TS3WriteObject::idx_Input);
+        const auto targetNode = input->Child(TS3WriteObject::idx_Target);
+        if (!TS3Target::Match(targetNode)) {
+            ctx.AddError(TIssue(ctx.GetPosition(targetNode->Pos()), "Expected S3 target."));
+            return TStatus::Error;
+        }
+
+        const TTypeAnnotationNode* targetType = nullptr;
+        if (const TS3Target target(targetNode); const auto settings = target.Settings()) {
+            if (const auto userschema = GetSetting(settings.Cast().Ref(), "userschema")) {
+                targetType = userschema->Child(1)->GetTypeAnn()->Cast<TTypeExprType>()->GetType();
+            }
+        }
+
+        const auto source = input->ChildPtr(TS3WriteObject::idx_Input);
+        if (const auto maybeTuple = TMaybeNode<TExprList>(source)) {
+            const auto tuple = maybeTuple.Cast();
+
+            TVector<TExprBase> convertedValues;
+            convertedValues.reserve(tuple.Size());
+            for (const auto& value : tuple) {
+                if (!EnsureStructType(input->Pos(), *value.Ref().GetTypeAnn(), ctx)) {
+                    return TStatus::Error;
+                }
+
+                TExprNode::TPtr node = value.Ptr();
+                if (targetType && TryConvertTo(node, *targetType, ctx) == TStatus::Error) {
+                    ctx.AddError(TIssue(ctx.GetPosition(source->Pos()), "Failed to convert input columns types to scheme types"));
+                    return TStatus::Error;
+                }
+
+                convertedValues.emplace_back(std::move(node));
+            }
+
+            const auto list = Build<TCoAsList>(ctx, input->Pos())
+                .Add(std::move(convertedValues))
+                .Done();
+
+            input->ChildRef(TS3WriteObject::idx_Input) = list.Ptr();
+            return TStatus::Repeat;
+        }
+
         if (!EnsureListType(*source, ctx)) {
             return TStatus::Error;
         }
@@ -74,23 +116,17 @@ class TS3DataSinkTypeAnnotationTransformer : public TVisitorTransformerBase {
             return TStatus::Error;
         }
 
-        auto target = input->Child(TS3WriteObject::idx_Target);
-        if (!TS3Target::Match(target)) {
-            ctx.AddError(TIssue(ctx.GetPosition(target->Pos()), "Expected S3 target."));
+        if (!NS3Util::ValidateS3ReadWriteSchema(sourceType->Cast<TStructExprType>(), ctx)) {
             return TStatus::Error;
         }
 
-        TS3Target tgt(target);
-        if (auto settings = tgt.Settings()) {
-            if (auto userschema = GetSetting(settings.Cast().Ref(), "userschema")) {
-                const TTypeAnnotationNode* targetType = userschema->Child(1)->GetTypeAnn()->Cast<TTypeExprType>()->GetType();
-                if (!IsSameAnnotation(*targetType, *sourceType)) {
-                    ctx.AddError(TIssue(ctx.GetPosition(source->Pos()),
-                                        TStringBuilder() << "Type mismatch between schema type: " << *targetType
-                                                         << " and actual data type: " << *sourceType << ", diff is: "
-                                                         << GetTypeDiff(*targetType, *sourceType)));
-                    return TStatus::Error;
-                }
+        if (targetType) {
+            const auto status = TryConvertTo(input->ChildRef(TS3WriteObject::idx_Input), *ctx.MakeType<TListExprType>(targetType), ctx);
+            if (status == TStatus::Error) {
+                ctx.AddError(TIssue(ctx.GetPosition(source->Pos()), "Row type mismatch for S3 external table"));
+                return TStatus::Error;
+            } else if (status != TStatus::Ok) {
+                return status;
             }
         }
 
diff --git a/ydb/library/yql/providers/s3/provider/yql_s3_datasource_type_ann.cpp b/ydb/library/yql/providers/s3/provider/yql_s3_datasource_type_ann.cpp
@@ -1,6 +1,7 @@
 #include "yql_s3_provider_impl.h"
 
 #include <ydb/library/yql/core/expr_nodes/yql_expr_nodes.h>
+#include <ydb/library/yql/providers/s3/common/util.h>
 #include <ydb/library/yql/providers/s3/expr_nodes/yql_s3_expr_nodes.h>
 #include <ydb/library/yql/providers/s3/path_generator/yql_s3_path_generator.h>
 #include <ydb/library/yql/providers/s3/range_helpers/path_list_reader.h>
@@ -491,6 +492,10 @@ class TS3DataSourceTypeAnnotationTransformer : public TVisitorTransformerBase {
             auto format = s3Object.Format().Ref().Content();
             const TStructExprType* structRowType = rowType->Cast<TStructExprType>();
 
+            if (!NS3Util::ValidateS3ReadWriteSchema(structRowType, ctx)) {
+                return TStatus::Error;
+            }
+
             THashSet<TStringBuf> columns;
             for (const TItemExprType* item : structRowType->GetItems()) {
                 columns.emplace(item->GetName());
diff --git a/ydb/tests/fq/s3/test_bindings_1.py b/ydb/tests/fq/s3/test_bindings_1.py
@@ -113,7 +113,7 @@ def test_s3_format_mismatch(self, kikimr, s3, client, unique_prefix):
 
         describe_result = client.describe_query(query_id).result
         describe_string = "{}".format(describe_result)
-        assert "Type mismatch between schema type" in describe_string, describe_string
+        assert "Row type mismatch for S3 external table" in describe_string, describe_string
 
     @yq_all
     @pytest.mark.parametrize("client", [{"folder_id": "my_folder"}], indirect=True)
diff --git a/ydb/tests/fq/s3/test_s3_0.py b/ydb/tests/fq/s3/test_s3_0.py
@@ -1114,3 +1114,56 @@ def wait_checkpoints(require_query_is_on=False):
 
         client.abort_query(query_id)
         client.wait_query(query_id)
+
+    @yq_v2
+    @pytest.mark.parametrize("client", [{"folder_id": "my_folder"}], indirect=True)
+    def test_double_optional_types_validation(self, kikimr, s3, client, unique_prefix):
+        resource = boto3.resource(
+            "s3", endpoint_url=s3.s3_url, aws_access_key_id="key", aws_secret_access_key="secret_key"
+        )
+
+        bucket = resource.Bucket("fbucket")
+        bucket.create(ACL='public-read')
+        bucket.objects.all().delete()
+
+        s3_client = boto3.client(
+            "s3", endpoint_url=s3.s3_url, aws_access_key_id="key", aws_secret_access_key="secret_key"
+        )
+
+        fruits = '''Fruit,Price,Weight
+Banana,3,100
+Apple,2,22
+Pear,15,33'''
+        s3_client.put_object(Body=fruits, Bucket='fbucket', Key='fruits.csv', ContentType='text/plain')
+
+        kikimr.control_plane.wait_bootstrap(1)
+        storage_connection_name = unique_prefix + "fruitbucket"
+        client.create_storage_connection(storage_connection_name, "fbucket")
+
+        sql = f'''
+            SELECT *
+            FROM `{storage_connection_name}`.`fruits.csv`
+            WITH (format='csv_with_names', SCHEMA (
+                Name Int32??,
+            ));
+            '''
+
+        query_id = client.create_query("simple", sql, type=fq.QueryContent.QueryType.ANALYTICS).result.query_id
+        client.wait_query_status(query_id, fq.QueryMeta.FAILED)
+        issues = str(client.describe_query(query_id).result.query.issue)
+
+        assert "Double optional types are not supported" in issues, "Incorrect issues: " + issues
+
+        sql = f'''
+            INSERT INTO `{storage_connection_name}`.`insert/`
+            WITH
+            (
+                FORMAT="csv_with_names"
+            )
+            SELECT CAST(42 AS Int32??) as Weight;'''
+
+        query_id = client.create_query("simple", sql, type=fq.QueryContent.QueryType.ANALYTICS).result.query_id
+        client.wait_query_status(query_id, fq.QueryMeta.FAILED)
+        issues = str(client.describe_query(query_id).result.query.issue)
+
+        assert "Double optional types are not supported" in issues, "Incorrect issues: " + issues

Original file line number	Diff line number	Diff line change
`@@ -19,6 +19,7 @@ PEERDIR(`
`19`	`19`	`ydb/library/yql/providers/s3/events`
`20`	`20`	`ydb/library/yql/public/issue`
`21`	`21`	`ydb/library/yql/public/issue/protos`
	`22`	`+ ydb/library/yql/ast`
`22`	`23`	`)`
`23`	`24`
`24`	`25`	`IF (CLANG AND NOT WITH_VALGRIND)`