Redirect pg_tables reads to system table. Resolves #1918

qrort · qrort · commit 36805439a349 · 2024-06-12T15:56:09.000Z
diff --git a/ydb/core/kqp/executer_actor/kqp_partition_helper.cpp b/ydb/core/kqp/executer_actor/kqp_partition_helper.cpp
@@ -441,6 +441,7 @@ TSerializedTableRange MakeKeyRange(const TVector<NScheme::TTypeInfo>& keyColumnT
 {
     YQL_ENSURE(range.HasFrom());
     YQL_ENSURE(range.HasTo());
+    auto guard = typeEnv.BindAllocator();
 
     auto fromValues = FillKeyValues(keyColumnTypes, range.GetFrom(), stageInfo, holderFactory, typeEnv);
     if (range.GetFrom().GetIsInclusive()) {
diff --git a/ydb/core/kqp/host/kqp_host.cpp b/ydb/core/kqp/host/kqp_host.cpp
@@ -1749,7 +1749,7 @@ class TKqpHost : public IKqpHost {
     void InitPgProvider() {
         auto state = MakeIntrusive<NYql::TPgState>();
         state->Types = TypesCtx.Get();
-
+        state->InitializedInKqpHost = true;
         TypesCtx->AddDataSource(NYql::PgProviderName, NYql::CreatePgDataSource(state));
         TypesCtx->AddDataSink(NYql::PgProviderName, NYql::CreatePgDataSink(state));
     }
diff --git a/ydb/core/kqp/ut/pg/pg_catalog_ut.cpp b/ydb/core/kqp/ut/pg/pg_catalog_ut.cpp
@@ -404,6 +404,60 @@ Y_UNIT_TEST_SUITE(PgCatalog) {
             ])", FormatResultSetYson(result.GetResultSet(0)));
         }
     }
+
+    Y_UNIT_TEST(PgTables) {
+        TKikimrRunner kikimr(NKqp::TKikimrSettings().SetWithSampleTables(false));
+        auto db = kikimr.GetQueryClient();
+        auto settings = NYdb::NQuery::TExecuteQuerySettings().Syntax(NYdb::NQuery::ESyntax::Pg);
+        {
+            auto result = db.ExecuteQuery(R"(
+                CREATE TABLE table1 (
+                    id int4 primary key
+                );
+                CREATE TABLE table2 (
+                    id varchar primary key
+                );
+            )", NYdb::NQuery::TTxControl::NoTx(), settings).ExtractValueSync();
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+        }
+        {
+            auto result = db.ExecuteQuery(R"(
+                DROP TABLE table1;
+            )", NYdb::NQuery::TTxControl::NoTx(), settings).ExtractValueSync();
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+        }
+        {
+            auto result = db.ExecuteQuery(R"(
+                SELECT * FROM pg_tables WHERE schemaname = 'public' ORDER BY tablename;
+            )", NYdb::NQuery::TTxControl::BeginTx().CommitTx(), settings).ExtractValueSync();
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+            UNIT_ASSERT_C(!result.GetResultSets().empty(), "no result sets");
+            CompareYson(R"([
+                ["public";"config";#;#;#;#;#;#];
+                ["public";"dbversion";#;#;#;#;#;#];
+                ["public";"migration_log";#;#;#;#;#;#];
+                ["public";"table2";"root@builtin";#;"t";"f";"f";"f"]
+            ])", FormatResultSetYson(result.GetResultSet(0)));
+        }
+        {
+            auto result = db.ExecuteQuery(R"(
+                --!syntax_pg
+                select
+                count(*) n,
+                min(schemaname) min_s,
+                min(tablename) min_t,
+                max(schemaname) max_s,
+                max(tablename) max_t
+                from pg_catalog.pg_tables;
+            )", NYdb::NQuery::TTxControl::BeginTx().CommitTx()).ExtractValueSync();
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+            UNIT_ASSERT_C(!result.GetResultSets().empty(), "no result sets");
+            CompareYson(R"([
+                ["207";"information_schema";"_pg_foreign_data_wrappers";"public";"views"]
+            ])", FormatResultSetYson(result.GetResultSet(0)));
+
+        }
+    }
 }
 
 } // namespace NKqp
diff --git a/ydb/core/sys_view/common/schema.cpp b/ydb/core/sys_view/common/schema.cpp
@@ -5,15 +5,15 @@
 namespace NKikimr {
 namespace NSysView {
 
-TVector<Schema::PgColumn> Schema::PgTables::Columns = { //lexicographical order
-    Schema::PgColumn(5, "pgbool", "hasindexes"),
-    Schema::PgColumn(6, "pgbool", "hasrules"),
-    Schema::PgColumn(7, "pgbool", "hastriggers"),
-    Schema::PgColumn(8, "pgbool", "rowsecurity"),
+const TVector<Schema::PgColumn> Schema::PgTables::Columns = {
     Schema::PgColumn(1, "pgname", "schemaname"),
     Schema::PgColumn(2, "pgname", "tablename"),
     Schema::PgColumn(3, "pgname", "tableowner"),
-    Schema::PgColumn(4, "pgname", "tablespace")
+    Schema::PgColumn(4, "pgname", "tablespace"),
+    Schema::PgColumn(5, "pgbool", "hasindexes"),
+    Schema::PgColumn(6, "pgbool", "hasrules"),
+    Schema::PgColumn(7, "pgbool", "hastriggers"),
+    Schema::PgColumn(8, "pgbool", "rowsecurity")
 };
 
 bool MaybeSystemViewPath(const TVector<TString>& path) {
@@ -182,24 +182,14 @@ class TSystemViewResolver : public ISystemViewResolver {
     void RegisterPgTablesSystemView() {
         auto& dsv  = DomainSystemViews[PgTablesName];
         auto& sdsv = SubDomainSystemViews[PgTablesName];
-        auto PgTablesSchema = Schema::PgTables();
-        for (const auto& column : PgTablesSchema.Columns) {
-            dsv.Columns[column._ColumnId - 1] = TSysTables::TTableColumnInfo(
-                column.GetColumnName(), column._ColumnId, column._ColumnTypeInfo, "", -1
+        for (const auto& column : Schema::PgTables::Columns) {
+            dsv.Columns[column._ColumnId] = TSysTables::TTableColumnInfo(
+                column._ColumnName, column._ColumnId, column._ColumnTypeInfo, "", -1
             );
-            sdsv.Columns[column._ColumnId - 1] = TSysTables::TTableColumnInfo(
-                column.GetColumnName(), column._ColumnId, column._ColumnTypeInfo, "", -1
+            sdsv.Columns[column._ColumnId] = TSysTables::TTableColumnInfo(
+                column._ColumnName, column._ColumnId, column._ColumnTypeInfo, "", -1
             );
         }
-        auto fillKey = [&](TSchema& schema, i32 index) -> void {
-            auto& column = schema.Columns[index];
-            column.KeyOrder = index;
-            schema.KeyColumnTypes.push_back(column.PType);
-        };
-        for (size_t i = 0; i < 2; i++) {
-            fillKey(dsv, i);
-            fillKey(sdsv, i);
-        }
     }
 
     template <typename Table>
diff --git a/ydb/core/sys_view/common/schema.h b/ydb/core/sys_view/common/schema.h
@@ -596,15 +596,12 @@ struct Schema : NIceDb::Schema {
         NScheme::TTypeInfo _ColumnTypeInfo;
         TString _ColumnName;
         PgColumn(NIceDb::TColumnId columnId, TStringBuf columnTypeName, TStringBuf columnName) 
-            : _ColumnId(columnId), _ColumnTypeInfo(NScheme::TTypeInfo(NScheme::NTypeIds::Pg, NPg::TypeDescFromPgTypeName(columnTypeName))), _ColumnName(columnName) {} 
-        
-        TString GetColumnName() const {
-            return _ColumnName;
-        }
+            : _ColumnId(columnId), _ColumnTypeInfo(NScheme::NTypeIds::Pg, NPg::TypeDescFromPgTypeName(columnTypeName)), _ColumnName(columnName) 
+        {}
     };
 
     struct PgTables {
-        static TVector<PgColumn> Columns;
+        const static TVector<PgColumn> Columns;
     };
 };
 
diff --git a/ydb/core/sys_view/pg_tables/pg_tables.cpp b/ydb/core/sys_view/pg_tables/pg_tables.cpp
@@ -22,30 +22,85 @@ class TPgTablesScan : public NKikimr::NSysView::TScanActorBase<TPgTablesScan> {
 private:
     TCell MakePgCell(const Schema::PgColumn& column, const TString& value, TVector<TString>& cellData) {
         NYql::NUdf::TStringRef ref;
-        auto typeDesc = column._ColumnTypeInfo.GetTypeDesc();
-        auto convert = NPg::PgNativeBinaryFromNativeText(value, NPg::PgTypeIdFromTypeDesc(typeDesc));
+        auto convert = NPg::PgNativeBinaryFromNativeText(value, NPg::PgTypeIdFromTypeDesc(column._ColumnTypeInfo.GetTypeDesc()));
         if (convert.Error) {
             ConvertError_ = *convert.Error;
             return TCell();
         }
         cellData.emplace_back(convert.Str);
         ref = NYql::NUdf::TStringRef(cellData.back());
-    
+        Y_ENSURE(ref.Size() > 0);
         return TCell(ref.Data(), ref.Size());
     }
 
     TVector<TCell> MakePgTablesRow(const TString& tableName, const TString& tableOwner, TVector<TString>& cellData) {
-        const auto &Columns =  Schema::PgTables::Columns;
-        return {
-            MakePgCell(Columns[0], "true", cellData), //hasindexes
-            MakePgCell(Columns[1], "false", cellData), //hasrules
-            MakePgCell(Columns[2], "false", cellData), //hastriggers
-            MakePgCell(Columns[3], "false", cellData), //rowsecurity
-            MakePgCell(Columns[4], "public", cellData), //schemaname
-            MakePgCell(Columns[5], tableName, cellData), //tablename
-            MakePgCell(Columns[6], tableOwner, cellData), //tableowner
-            TCell() //tablespace
-        };
+        TVector<TCell> res;
+        res.reserve(Columns.size());
+        for (const auto& column : Columns) {
+            TCell cell;
+            switch (column.Tag) {
+                case 1: {
+                    cell = MakePgCell(Schema::PgTables::Columns[0], "public", cellData);
+                    break;
+                }
+                case 2: {
+                    cell = MakePgCell(Schema::PgTables::Columns[1], tableName, cellData);
+                    break;
+                }
+                case 3: {
+                    cell = MakePgCell(Schema::PgTables::Columns[2], tableOwner, cellData);
+                    break;
+                }
+                case 4: {
+                    cell = TCell();
+                    break;
+                }
+                case 5: {
+                    cell = MakePgCell(Schema::PgTables::Columns[4], "true", cellData);
+                    break;
+                }
+                case 6: {
+                    cell = MakePgCell(Schema::PgTables::Columns[5], "false", cellData);
+                    break;
+                }
+                case 7: {
+                    cell = MakePgCell(Schema::PgTables::Columns[6], "false", cellData);
+                    break;
+                }
+                case 8: {
+                    cell = MakePgCell(Schema::PgTables::Columns[7], "false", cellData);
+                    break;
+                }
+
+            }
+            res.emplace_back(std::move(cell));
+        }
+        return res;
+    }
+
+    TVector<TCell> MakePgTablesStaticRow(const NYql::NPg::TTableInfo& tableInfo, TVector<TString>& cellData) {
+        TVector<TCell> res;
+        res.reserve(Columns.size());
+        for (const auto& column : Columns) {
+            TCell cell;
+            switch (column.Tag) {
+                case 1: {
+                    cell = MakePgCell(Schema::PgTables::Columns[0], tableInfo.Schema, cellData);
+                    break;
+                }
+                case 2: {
+                    cell = MakePgCell(Schema::PgTables::Columns[1], tableInfo.Name, cellData);
+                    break;
+                }
+                default: {
+                    cell = TCell();
+                    break;
+                }
+
+            }
+            res.emplace_back(std::move(cell));
+        }
+        return res;
     }
 public:
     using TBase = NKikimr::NSysView::TScanActorBase<TPgTablesScan>;
@@ -71,6 +126,19 @@ class TPgTablesScan : public NKikimr::NSysView::TScanActorBase<TPgTablesScan> {
         Become(&TPgTablesScan::StateWork);
     }
 
+    void ExpandBatchWithStaticTables(const THolder<NKqp::TEvKqpCompute::TEvScanData>& batch) {
+        for (const auto& tableDesc : NYql::NPg::GetStaticTables()) {
+            TVector<TString> cellData;
+            TVector<TCell> cells = MakePgTablesStaticRow(tableDesc, cellData);
+            if (!ConvertError_.Empty()) {
+                ReplyErrorAndDie(Ydb::StatusIds::INTERNAL_ERROR, ConvertError_);
+                return;
+            }
+            TArrayRef<const TCell> ref(cells);
+            batch->Rows.emplace_back(TOwnedCellVec::Make(ref));
+        }
+    }
+
     void Handle(NSchemeShard::TEvSchemeShard::TEvDescribeSchemeResult::TPtr& ev, const TActorContext& ctx) {
         Y_UNUSED(ctx);
         const auto& record = ev->Get()->GetRecord();
@@ -102,6 +170,7 @@ class TPgTablesScan : public NKikimr::NSysView::TScanActorBase<TPgTablesScan> {
     
         auto batch = MakeHolder<NKqp::TEvKqpCompute::TEvScanData>(ScanId);
 
+        ExpandBatchWithStaticTables(batch);
 
         for (size_t i = 0; i < record.GetPathDescription().ChildrenSize(); ++i) {
             TVector<TString> cellData;
diff --git a/ydb/core/sys_view/ut_kqp.cpp b/ydb/core/sys_view/ut_kqp.cpp
@@ -326,7 +326,7 @@ Y_UNIT_TEST_SUITE(SystemView) {
         auto session = client.CreateSession().GetValueSync().GetSession();
         {
             auto result = session.ExecuteDataQuery(R"(
-                SELECT schemaname, tablename, tableowner, tablespace, hasindexes, hasrules, hastriggers, rowsecurity FROM `Root/.sys/pg_tables`;
+                SELECT schemaname, tablename, tableowner, tablespace, hasindexes, hasrules, hastriggers, rowsecurity FROM `Root/.sys/pg_tables` WHERE tablename = PgName("Table0") OR tablename = PgName("Table1") ORDER BY tablename;
             )", TTxControl::BeginTx().CommitTx()).ExtractValueSync();
 
             UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
diff --git a/ydb/library/yql/providers/pg/provider/yql_pg_datasource.cpp b/ydb/library/yql/providers/pg/provider/yql_pg_datasource.cpp
@@ -1,3 +1,4 @@
+#include "ydb/library/yql/providers/common/provider/yql_provider.h"
 #include "yql_pg_provider_impl.h"
 
 #include <ydb/library/yql/providers/common/provider/yql_data_provider_impl.h>
@@ -48,11 +49,16 @@ class TPgDataSourceImpl : public TDataProviderBase {
 
     TExprNode::TPtr RewriteIO(const TExprNode::TPtr& node, TExprContext& ctx) override {
         YQL_CLOG(INFO, ProviderPg) << "RewriteIO";
-        if (auto left = TMaybeNode<TCoLeft>(node)) {
-            return left.Input().Maybe<TPgRead>().World().Cast().Ptr();
+
+        TMaybeNode<TPgRead> maybeRead;
+        auto left = TMaybeNode<TCoLeft>(node);
+        if (left) {
+            maybeRead = left.Input().Cast<TPgRead>();
+        } else {
+            maybeRead = TCoRight(node).Input().Cast<TPgRead>();
         }
 
-        auto read = TCoRight(node).Input().Cast<TPgRead>();
+        auto read = maybeRead.Cast();
         auto keyNode = read.FreeArgs().Get(2).Ptr();
         if (keyNode->IsCallable("MrTableConcat")) {
             if (keyNode->ChildrenSize() != 1) {
@@ -77,13 +83,54 @@ class TPgDataSourceImpl : public TDataProviderBase {
         }
 
         const auto tableName = TString(keyArg.Tail().Head().Content());
-        auto childrenList = read.Ref().ChildrenList();
-        childrenList[2] = ctx.NewAtom(childrenList[2]->Pos(), tableName);
-        auto newRead = ctx.NewCallable(read.Ref().Pos(), TPgReadTable::CallableName(), std::move(childrenList));
-
-        return Build<TCoRight>(ctx, read.Pos())
-            .Input(newRead)
+        if (tableName == "pg_tables" && State_->InitializedInKqpHost) {
+            TNodeOnNodeOwnedMap replaces;
+            replaces[keyArg.TailPtr()->TailPtr().Get()] = ctx.NewAtom(read.Pos(), "/Root/.sys/pg_tables");
+
+            auto ydbSysTableRead = Build<TCoRead>(ctx, read.Pos())
+                .InitFrom(read.Cast<TCoRead>())
+                .DataSource<TCoDataSource>()
+                    .Category(ctx.NewAtom(read.Pos(), KikimrProviderName))
+                    .FreeArgs()
+                        .Add(ctx.NewAtom(read.Pos(), "db"))
+                    .Build()
+                .Build()
+                .FreeArgs()
+                    .Add(ctx.ReplaceNodes(maybeKey.Cast().Ptr(), replaces))
+                    .Add(ctx.NewCallable(read.Pos(), "Void", {}))
+                    .Add(ctx.NewList(read.Pos(), {}))
+                .Build()
             .Done().Ptr();
+
+            ctx.Step
+                .Repeat(TExprStep::ExprEval)
+                .Repeat(TExprStep::DiscoveryIO)
+                .Repeat(TExprStep::Epochs)
+                .Repeat(TExprStep::Intents)
+                .Repeat(TExprStep::LoadTablesMetadata)
+                .Repeat(TExprStep::RewriteIO);
+
+            if (left) {
+                return Build<TCoLeft>(ctx, ydbSysTableRead->Pos())
+                    .Input(ydbSysTableRead)
+                .Done().Ptr();
+            } else {
+                return Build<TCoRight>(ctx, ydbSysTableRead->Pos())
+                    .Input(ydbSysTableRead)
+                .Done().Ptr();
+            }
+        } else {
+            if (left) {
+                return left.Input().Maybe<TPgRead>().World().Cast().Ptr();
+            } else {
+                auto childrenList = read.Ref().ChildrenList();
+                childrenList[2] = ctx.NewAtom(childrenList[2]->Pos(), tableName);
+                auto newRead = ctx.NewCallable(read.Ref().Pos(), TPgReadTable::CallableName(), std::move(childrenList));
+                return Build<TCoRight>(ctx, read.Pos())
+                    .Input(newRead)
+                .Done().Ptr();
+            }
+        }
     }
 
     bool ValidateParameters(TExprNode& node, TExprContext& ctx, TMaybe<TString>& cluster) override {
diff --git a/ydb/library/yql/providers/pg/provider/yql_pg_provider.cpp b/ydb/library/yql/providers/pg/provider/yql_pg_provider.cpp
@@ -62,4 +62,4 @@ TDataProviderInitializer GetPgDataProviderInitializer() {
     };
 }
 
-}
+}
diff --git a/ydb/library/yql/providers/pg/provider/yql_pg_provider.h b/ydb/library/yql/providers/pg/provider/yql_pg_provider.h
@@ -9,6 +9,7 @@ struct TPgState : public TThrRefBase
     using TPtr = TIntrusivePtr<TPgState>;
 
     TTypeAnnotationContext* Types = nullptr;
+    bool InitializedInKqpHost = false;
 };
 
 TDataProviderInitializer GetPgDataProviderInitializer();

Original file line number	Diff line number	Diff line change
`@@ -441,6 +441,7 @@ TSerializedTableRange MakeKeyRange(const TVector<NScheme::TTypeInfo>& keyColumnT`
`441`	`441`	`{`
`442`	`442`	`YQL_ENSURE(range.HasFrom());`
`443`	`443`	`YQL_ENSURE(range.HasTo());`
	`444`	`+ auto guard = typeEnv.BindAllocator();`
`444`	`445`
`445`	`446`	`auto fromValues = FillKeyValues(keyColumnTypes, range.GetFrom(), stageInfo, holderFactory, typeEnv);`
`446`	`447`	`if (range.GetFrom().GetIsInclusive()) {`
Original file line number	Diff line number	Diff line change
`@@ -1749,7 +1749,7 @@ class TKqpHost : public IKqpHost {`
`1749`	`1749`	`void InitPgProvider() {`
`1750`	`1750`	`auto state = MakeIntrusive<NYql::TPgState>();`
`1751`	`1751`	`state->Types = TypesCtx.Get();`
`1752`		`-`
	`1752`	`+ state->InitializedInKqpHost = true;`
`1753`	`1753`	`TypesCtx->AddDataSource(NYql::PgProviderName, NYql::CreatePgDataSource(state));`
`1754`	`1754`	`TypesCtx->AddDataSink(NYql::PgProviderName, NYql::CreatePgDataSink(state));`
`1755`	`1755`	`}`
Original file line number	Diff line number	Diff line change
`@@ -326,7 +326,7 @@ Y_UNIT_TEST_SUITE(SystemView) {`
`326`	`326`	`auto session = client.CreateSession().GetValueSync().GetSession();`
`327`	`327`	`{`
`328`	`328`	`auto result = session.ExecuteDataQuery(R"(`
`329`		- SELECT schemaname, tablename, tableowner, tablespace, hasindexes, hasrules, hastriggers, rowsecurity FROM `Root/.sys/pg_tables`;
	`329`	+ SELECT schemaname, tablename, tableowner, tablespace, hasindexes, hasrules, hastriggers, rowsecurity FROM `Root/.sys/pg_tables` WHERE tablename = PgName("Table0") OR tablename = PgName("Table1") ORDER BY tablename;
`330`	`330`	`)", TTxControl::BeginTx().CommitTx()).ExtractValueSync();`
`331`	`331`
`332`	`332`	`UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());`
Original file line number	Diff line number	Diff line change
`@@ -62,4 +62,4 @@ TDataProviderInitializer GetPgDataProviderInitializer() {`
`62`	`62`	`};`
`63`	`63`	`}`
`64`	`64`
`65`		`-}`
	`65`	`+}`