Initial support for PG extensions (#7112)

Author: vitstn

ydb/library/yql/core/pg_ext/ya.make

@@ -0,0 +1,12 @@
+LIBRARY()
+
+SRCS(
+    yql_pg_ext.cpp
+)
+
+PEERDIR(
+    ydb/library/yql/protos
+    ydb/library/yql/parser/pg_catalog
+)
+
+END()

ydb/library/yql/core/pg_ext/yql_pg_ext.cpp

@@ -0,0 +1,19 @@
+#include "yql_pg_ext.h"
+
+namespace NYql {
+
+void PgExtensionsFromProto(const NYql::NProto::TPgExtensions& proto,
+    TVector<NPg::TExtensionDesc>& extensions) {
+    extensions.clear();
+    for (const auto& e: proto.GetExtension()) {
+        NPg::TExtensionDesc desc;
+        desc.Name = e.GetName();
+        desc.InstallName = e.GetInstallName();
+        desc.DDLPath = e.GetDDLPath();
+        desc.LibraryPath = e.GetLibraryPath();
+        desc.TypesOnly = e.GetTypesOnly();
+        extensions.emplace_back(desc);
+    }
+}
+
+} // namespace NYql

ydb/library/yql/core/pg_ext/yql_pg_ext.h

@@ -0,0 +1,11 @@
+#pragma once
+
+#include <ydb/library/yql/protos/pg_ext.pb.h>
+#include <ydb/library/yql/parser/pg_catalog/catalog.h>
+
+namespace NYql {
+
+void PgExtensionsFromProto(const NYql::NProto::TPgExtensions& proto,
+    TVector<NPg::TExtensionDesc>& extensions);
+
+} // namespace NYql

ydb/library/yql/core/ya.make

@@ -108,6 +108,7 @@ RECURSE(
     facade
     file_storage
     issue
+    pg_ext
     peephole_opt
     qplayer
     services

ydb/library/yql/minikql/protobuf_udf/ya.make

@@ -13,6 +13,7 @@ PEERDIR(
     contrib/libs/protobuf
     library/cpp/protobuf/yql
     ydb/library/yql/public/udf
+    ydb/library/yql/minikql
     yt/cpp/mapreduce/interface
     yt/yt_proto/yt/formats
     yt/yt_proto/yt/formats

ydb/library/yql/parser/pg_catalog/catalog.cpp

@@ -6,6 +6,7 @@
 #include <util/string/builder.h>
 #include <util/string/cast.h>
 #include <util/string/split.h>
+#include <util/stream/file.h>
 #include <util/system/env.h>
 #include <util/system/mutex.h>
 #include <util/system/tempfile.h>
@@ -68,6 +69,9 @@ using TConversions = THashMap<std::pair<TString, TString>, TConversionDesc>;
 
 using TLanguages = THashMap<ui32, TLanguageDesc>;
 
+using TExtensions = TVector<TExtensionDesc>;
+using TExtensionsByName = THashMap<TString, ui32>;
+
 bool IsCompatibleTo(ui32 actualTypeId, ui32 expectedTypeId, const TTypes& types) {
     if (actualTypeId == expectedTypeId) {
         return true;
@@ -1592,7 +1596,7 @@ const char* AllowedProcsRaw[] = {
 #include "used_procs.h"
 };
 
-struct TCatalog {
+struct TCatalog : public IExtensionDDLBuilder {
     TCatalog() {
         for (size_t i = 0; i < Y_ARRAY_SIZE(AllStaticTablesRaw); ++i) {
             const auto& raw = AllStaticTablesRaw[i];
@@ -1932,10 +1936,27 @@ struct TCatalog {
         }
     }
 
+    void CreateProc(const TProcDesc& desc) final {
+        TProcDesc newDesc = desc;
+        newDesc.ProcId = 16000 + Procs.size();
+        Procs[newDesc.ProcId] = newDesc;
+        ProcByName[newDesc.Name].push_back(newDesc.ProcId);
+    }
+
     static const TCatalog& Instance() {
         return *Singleton<TCatalog>();
     }
 
+    static TCatalog& MutableInstance() {
+        return *Singleton<TCatalog>();
+    }    
+
+    TMutex ExtensionsGuard;
+    bool ExtensionsInit = false;
+
+    TExtensionsByName ExtensionsByName, ExtensionsByInstallName;
+    TExtensions Extensions;
+
     TOperators Operators;
     TProcs Procs;
     TTypes Types;
@@ -2035,7 +2056,9 @@ const TProcDesc& LookupProc(ui32 procId) {
 void EnumProc(std::function<void(ui32, const TProcDesc&)> f) {
     const auto& catalog = TCatalog::Instance();
     for (const auto& x : catalog.Procs) {
-        f(x.first, x.second);
+        if (catalog.ExportFile || catalog.AllowedProcs.contains(x.second.Name)) {
+            f(x.first, x.second);
+        }
     }
 }
 
@@ -3234,5 +3257,74 @@ const TTableInfo& LookupStaticTable(const TTableInfoKey& tableKey) {
     return *tablePtr;
 }
 
+bool IsExportFunctionsEnabled() {
+    const auto& catalog = TCatalog::Instance();
+    return catalog.ExportFile.Defined();
+}
+
+void RegisterExtensions(const TVector<TExtensionDesc>& extensions, bool typesOnly,
+    IExtensionDDLParser& parser, IExtensionLoader* loader) {
+    auto& catalog = TCatalog::MutableInstance();
+    with_lock (catalog.ExtensionsGuard) {
+        Y_ENSURE(!catalog.ExtensionsInit);
+        for (ui32 i = 0; i < extensions.size(); ++i) {
+            auto e = extensions[i];
+            e.TypesOnly = e.TypesOnly && typesOnly;
+            if (e.Name.empty()) {
+                throw yexception() << "Empty extension name";
+            }
+
+            if (!catalog.ExtensionsByName.insert(std::make_pair(e.Name, i + 1)).second) {
+                throw yexception() << "Duplicated extension name: " << e.Name;
+            }
+
+            if (!catalog.ExtensionsByInstallName.insert(std::make_pair(e.InstallName, i + 1)).second) {
+                throw yexception() << "Duplicated extension install name: " << e.InstallName;
+            }
+
+            catalog.Extensions.push_back(e);
+            TString sql = TFileInput(e.DDLPath).ReadAll();;
+            parser.Parse(sql, catalog);
+            if (loader && !e.TypesOnly) {
+                loader->Load(i + 1, e.Name, e.LibraryPath);
+            }
+        }
+
+        catalog.ExtensionsInit = true;
+    }
+}
+
+void EnumExtensions(std::function<void(ui32, const TExtensionDesc&)> f) {
+    const auto& catalog = TCatalog::Instance();
+    for (ui32 i = 0; i < catalog.Extensions.size(); ++i) {
+        f(i + 1, catalog.Extensions[i]);
+    }
+}
+
+const TExtensionDesc& LookupExtension(ui32 extIndex) {
+    const auto& catalog = TCatalog::Instance();
+    Y_ENSURE(extIndex > 0 && extIndex <= catalog.Extensions.size());
+    return catalog.Extensions[extIndex - 1];
+}
+
+ui32 LookupExtensionByName(const TString& name) {
+    const auto& catalog = TCatalog::Instance();
+    auto indexPtr = catalog.ExtensionsByName.FindPtr(name);
+    if (!indexPtr) {
+        throw yexception() << "Unknown extension name: " << name;
+    }
+
+    return *indexPtr;
+}
+
+ui32 LookupExtensionByInstallName(const TString& installName) {
+    const auto& catalog = TCatalog::Instance();
+    auto indexPtr = catalog.ExtensionsByInstallName.FindPtr(installName);
+    if (!indexPtr) {
+        throw yexception() << "Unknown extension install name: " << installName;
+    }
+
+    return *indexPtr;
+}
 
 }

ydb/library/yql/parser/pg_catalog/catalog.h

@@ -72,6 +72,7 @@ struct TProcDesc {
     ui32 VariadicType = 0;
     ui32 VariadicArgType = 0;
     TString VariadicArgName;
+    ui32 ExtensionIndex = 0;
 };
 
 // Copied from pg_collation_d.h
@@ -355,6 +356,46 @@ struct TColumnInfo {
 const TVector<TTableInfo>& GetStaticTables();
 const TTableInfo& LookupStaticTable(const TTableInfoKey& tableKey);
 const THashMap<TTableInfoKey, TVector<TColumnInfo>>& GetStaticColumns();
+
+void PrepareCatalog();
+bool IsExportFunctionsEnabled();
+
+struct TExtensionDesc {
+    TString Name;           // postgis
+    TString InstallName;    // $libdir/postgis-3
+    TString DDLPath;        // DDL path (CREATE TYPE/CREATE FUNCTION/etc)
+    TString LibraryPath;    // file path
+    bool TypesOnly = false; // Can't be loaded if true
+};
+
+class IExtensionDDLBuilder {
+public:
+    virtual ~IExtensionDDLBuilder() = default;
+
+    virtual void CreateProc(const TProcDesc& desc) = 0;
+};
+
+class IExtensionDDLParser {
+public:
+    virtual ~IExtensionDDLParser() = default;
+    virtual void Parse(const TString& sql, IExtensionDDLBuilder& builder) = 0;
+};
+
+class IExtensionLoader {
+public:
+    virtual ~IExtensionLoader() = default;
+    virtual void Load(ui32 extensionIndex, const TString& name, const TString& path) = 0;
+};
+
+// should be called at most once before other catalog functions
+void RegisterExtensions(const TVector<TExtensionDesc>& extensions, bool typesOnly,
+    IExtensionDDLParser& parser, IExtensionLoader* loader);
+
+void EnumExtensions(std::function<void(ui32 extensionIndex, const TExtensionDesc&)> f);
+const TExtensionDesc& LookupExtension(ui32 extensionIndex);
+ui32 LookupExtensionByName(const TString& name);
+ui32 LookupExtensionByInstallName(const TString& installName);
+
 }
 
 template <>

ydb/library/yql/parser/pg_catalog/safe_procs.h

@@ -18,3 +18,5 @@
 "pg_get_function_arguments",//psql
 "pg_table_is_visible",//psql
 
+// postgis
+"postgis_lib_version",

ydb/library/yql/parser/pg_catalog/used_procs.h

@@ -512,6 +512,7 @@
 "hashinet",
 "hashint2",
 "hashint4",
+"hashint4extended",
 "hashint8",
 "hashmacaddr",
 "hashmacaddr8",

ydb/library/yql/parser/pg_catalog/ya.make

@@ -20,6 +20,7 @@ SRCS(
 PEERDIR(
     library/cpp/resource
     ydb/library/yql/public/issue
+    ydb/library/yql/protos
 )
 
 END()