YQ-2892 supported yt pragmas (ydb-platform#2570)

jepett0 · Mar 11, 2024 · 1fcef17 · 1fcef17
1 parent caf19c3
commit 1fcef17
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 1 deletion.
diff --git a/ydb/library/yql/providers/yt/provider/yql_yt_datasource_constraints.cpp b/ydb/library/yql/providers/yt/provider/yql_yt_datasource_constraints.cpp
@@ -39,6 +39,7 @@ class TYtDataSourceConstraintTransformer : public TVisitorTransformerBase {
         AddHandler({TYtKeyRange::CallableName()}, Hndl(&TYtDataSourceConstraintTransformer::HandleDefault));
         AddHandler({TYtReadTableScheme::CallableName()}, Hndl(&TYtDataSourceConstraintTransformer::HandleDefault));
         AddHandler({TYtLength::CallableName()}, Hndl(&TYtDataSourceConstraintTransformer::HandleDefault));
+        AddHandler({TCoConfigure::CallableName()}, Hndl(&TYtDataSourceConstraintTransformer::HandleDefault));
         AddHandler({TYtConfigure::CallableName()}, Hndl(&TYtDataSourceConstraintTransformer::HandleDefault));
         AddHandler({TYtTablePath::CallableName()}, Hndl(&TYtDataSourceConstraintTransformer::HandleDefault));
         AddHandler({TYtTableRecord::CallableName()}, Hndl(&TYtDataSourceConstraintTransformer::HandleDefault));

diff --git a/ydb/library/yql/providers/yt/provider/yql_yt_datasource_type_ann.cpp b/ydb/library/yql/providers/yt/provider/yql_yt_datasource_type_ann.cpp
@@ -47,7 +47,8 @@ class TYtDataSourceTypeAnnotationTransformer : public TVisitorTransformerBase {
         AddHandler({TYtReadTableScheme::CallableName()}, Hndl(&TYtDataSourceTypeAnnotationTransformer::HandleReadTableScheme));
         AddHandler({TYtTableContent::CallableName()}, Hndl(&TYtDataSourceTypeAnnotationTransformer::HandleTableContent));
         AddHandler({TYtLength::CallableName()}, Hndl(&TYtDataSourceTypeAnnotationTransformer::HandleLength));
-        AddHandler({TYtConfigure::CallableName()}, Hndl(&TYtDataSourceTypeAnnotationTransformer::HandleConfigure));
+        AddHandler({TCoConfigure::CallableName()}, Hndl(&TYtDataSourceTypeAnnotationTransformer::HandleConfigure));
+        AddHandler({TYtConfigure::CallableName()}, Hndl(&TYtDataSourceTypeAnnotationTransformer::HandleYtConfigure));
         AddHandler({TYtTablePath::CallableName()}, Hndl(&TYtDataSourceTypeAnnotationTransformer::HandleTablePath));
         AddHandler({TYtTableRecord::CallableName()}, Hndl(&TYtDataSourceTypeAnnotationTransformer::HandleTableRecord));
         AddHandler({TYtRowNumber::CallableName()}, Hndl(&TYtDataSourceTypeAnnotationTransformer::HandleTableRecord));
@@ -882,6 +883,23 @@ class TYtDataSourceTypeAnnotationTransformer : public TVisitorTransformerBase {
             return TStatus::Error;
         }
 
+        if (!EnsureWorldType(*input.Ptr()->Child(TCoConfigure::idx_World), ctx)) {
+            return TStatus::Error;
+        }
+
+        if (!EnsureSpecificDataSource(*input.Ptr()->Child(TCoConfigure::idx_DataSource), YtProviderName, ctx)) {
+            return TStatus::Error;
+        }
+
+        input.Ptr()->SetTypeAnn(input.Ref().Child(TCoConfigure::idx_World)->GetTypeAnn());
+        return TStatus::Ok;
+    }
+
+    TStatus HandleYtConfigure(TExprBase input, TExprContext& ctx) {
+        if (!EnsureMinArgsCount(input.Ref(), 2, ctx)) {
+            return TStatus::Error;
+        }
+
         if (!EnsureWorldType(*input.Ptr()->Child(TYtConfigure::idx_World), ctx)) {
             return TStatus::Error;
         }