Skip to content

Commit 9d52924

Browse files
authored
Initial support for PG extensions (#7112)
1 parent 7ea1c14 commit 9d52924

File tree

40 files changed

+859
-42
lines changed

40 files changed

+859
-42
lines changed

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
LIBRARY()
2+
3+
SRCS(
4+
yql_pg_ext.cpp
5+
)
6+
7+
PEERDIR(
8+
ydb/library/yql/protos
9+
ydb/library/yql/parser/pg_catalog
10+
)
11+
12+
END()
Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
#include "yql_pg_ext.h"
2+
3+
namespace NYql {
4+
5+
void PgExtensionsFromProto(const NYql::NProto::TPgExtensions& proto,
6+
TVector<NPg::TExtensionDesc>& extensions) {
7+
extensions.clear();
8+
for (const auto& e: proto.GetExtension()) {
9+
NPg::TExtensionDesc desc;
10+
desc.Name = e.GetName();
11+
desc.InstallName = e.GetInstallName();
12+
desc.DDLPath = e.GetDDLPath();
13+
desc.LibraryPath = e.GetLibraryPath();
14+
desc.TypesOnly = e.GetTypesOnly();
15+
extensions.emplace_back(desc);
16+
}
17+
}
18+
19+
} // namespace NYql
Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
#pragma once
2+
3+
#include <ydb/library/yql/protos/pg_ext.pb.h>
4+
#include <ydb/library/yql/parser/pg_catalog/catalog.h>
5+
6+
namespace NYql {
7+
8+
void PgExtensionsFromProto(const NYql::NProto::TPgExtensions& proto,
9+
TVector<NPg::TExtensionDesc>& extensions);
10+
11+
} // namespace NYql

ydb/library/yql/core/ya.make

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,7 @@ RECURSE(
108108
facade
109109
file_storage
110110
issue
111+
pg_ext
111112
peephole_opt
112113
qplayer
113114
services

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ PEERDIR(
1313
contrib/libs/protobuf
1414
library/cpp/protobuf/yql
1515
ydb/library/yql/public/udf
16+
ydb/library/yql/minikql
1617
yt/cpp/mapreduce/interface
1718
yt/yt_proto/yt/formats
1819
yt/yt_proto/yt/formats

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

Lines changed: 94 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
#include <util/string/builder.h>
77
#include <util/string/cast.h>
88
#include <util/string/split.h>
9+
#include <util/stream/file.h>
910
#include <util/system/env.h>
1011
#include <util/system/mutex.h>
1112
#include <util/system/tempfile.h>
@@ -68,6 +69,9 @@ using TConversions = THashMap<std::pair<TString, TString>, TConversionDesc>;
6869

6970
using TLanguages = THashMap<ui32, TLanguageDesc>;
7071

72+
using TExtensions = TVector<TExtensionDesc>;
73+
using TExtensionsByName = THashMap<TString, ui32>;
74+
7175
bool IsCompatibleTo(ui32 actualTypeId, ui32 expectedTypeId, const TTypes& types) {
7276
if (actualTypeId == expectedTypeId) {
7377
return true;
@@ -1592,7 +1596,7 @@ const char* AllowedProcsRaw[] = {
15921596
#include "used_procs.h"
15931597
};
15941598

1595-
struct TCatalog {
1599+
struct TCatalog : public IExtensionDDLBuilder {
15961600
TCatalog() {
15971601
for (size_t i = 0; i < Y_ARRAY_SIZE(AllStaticTablesRaw); ++i) {
15981602
const auto& raw = AllStaticTablesRaw[i];
@@ -1932,10 +1936,27 @@ struct TCatalog {
19321936
}
19331937
}
19341938

1939+
void CreateProc(const TProcDesc& desc) final {
1940+
TProcDesc newDesc = desc;
1941+
newDesc.ProcId = 16000 + Procs.size();
1942+
Procs[newDesc.ProcId] = newDesc;
1943+
ProcByName[newDesc.Name].push_back(newDesc.ProcId);
1944+
}
1945+
19351946
static const TCatalog& Instance() {
19361947
return *Singleton<TCatalog>();
19371948
}
19381949

1950+
static TCatalog& MutableInstance() {
1951+
return *Singleton<TCatalog>();
1952+
}
1953+
1954+
TMutex ExtensionsGuard;
1955+
bool ExtensionsInit = false;
1956+
1957+
TExtensionsByName ExtensionsByName, ExtensionsByInstallName;
1958+
TExtensions Extensions;
1959+
19391960
TOperators Operators;
19401961
TProcs Procs;
19411962
TTypes Types;
@@ -2035,7 +2056,9 @@ const TProcDesc& LookupProc(ui32 procId) {
20352056
void EnumProc(std::function<void(ui32, const TProcDesc&)> f) {
20362057
const auto& catalog = TCatalog::Instance();
20372058
for (const auto& x : catalog.Procs) {
2038-
f(x.first, x.second);
2059+
if (catalog.ExportFile || catalog.AllowedProcs.contains(x.second.Name)) {
2060+
f(x.first, x.second);
2061+
}
20392062
}
20402063
}
20412064

@@ -3234,5 +3257,74 @@ const TTableInfo& LookupStaticTable(const TTableInfoKey& tableKey) {
32343257
return *tablePtr;
32353258
}
32363259

3260+
bool IsExportFunctionsEnabled() {
3261+
const auto& catalog = TCatalog::Instance();
3262+
return catalog.ExportFile.Defined();
3263+
}
3264+
3265+
void RegisterExtensions(const TVector<TExtensionDesc>& extensions, bool typesOnly,
3266+
IExtensionDDLParser& parser, IExtensionLoader* loader) {
3267+
auto& catalog = TCatalog::MutableInstance();
3268+
with_lock (catalog.ExtensionsGuard) {
3269+
Y_ENSURE(!catalog.ExtensionsInit);
3270+
for (ui32 i = 0; i < extensions.size(); ++i) {
3271+
auto e = extensions[i];
3272+
e.TypesOnly = e.TypesOnly && typesOnly;
3273+
if (e.Name.empty()) {
3274+
throw yexception() << "Empty extension name";
3275+
}
3276+
3277+
if (!catalog.ExtensionsByName.insert(std::make_pair(e.Name, i + 1)).second) {
3278+
throw yexception() << "Duplicated extension name: " << e.Name;
3279+
}
3280+
3281+
if (!catalog.ExtensionsByInstallName.insert(std::make_pair(e.InstallName, i + 1)).second) {
3282+
throw yexception() << "Duplicated extension install name: " << e.InstallName;
3283+
}
3284+
3285+
catalog.Extensions.push_back(e);
3286+
TString sql = TFileInput(e.DDLPath).ReadAll();;
3287+
parser.Parse(sql, catalog);
3288+
if (loader && !e.TypesOnly) {
3289+
loader->Load(i + 1, e.Name, e.LibraryPath);
3290+
}
3291+
}
3292+
3293+
catalog.ExtensionsInit = true;
3294+
}
3295+
}
3296+
3297+
void EnumExtensions(std::function<void(ui32, const TExtensionDesc&)> f) {
3298+
const auto& catalog = TCatalog::Instance();
3299+
for (ui32 i = 0; i < catalog.Extensions.size(); ++i) {
3300+
f(i + 1, catalog.Extensions[i]);
3301+
}
3302+
}
3303+
3304+
const TExtensionDesc& LookupExtension(ui32 extIndex) {
3305+
const auto& catalog = TCatalog::Instance();
3306+
Y_ENSURE(extIndex > 0 && extIndex <= catalog.Extensions.size());
3307+
return catalog.Extensions[extIndex - 1];
3308+
}
3309+
3310+
ui32 LookupExtensionByName(const TString& name) {
3311+
const auto& catalog = TCatalog::Instance();
3312+
auto indexPtr = catalog.ExtensionsByName.FindPtr(name);
3313+
if (!indexPtr) {
3314+
throw yexception() << "Unknown extension name: " << name;
3315+
}
3316+
3317+
return *indexPtr;
3318+
}
3319+
3320+
ui32 LookupExtensionByInstallName(const TString& installName) {
3321+
const auto& catalog = TCatalog::Instance();
3322+
auto indexPtr = catalog.ExtensionsByInstallName.FindPtr(installName);
3323+
if (!indexPtr) {
3324+
throw yexception() << "Unknown extension install name: " << installName;
3325+
}
3326+
3327+
return *indexPtr;
3328+
}
32373329

32383330
}

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

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ struct TProcDesc {
7272
ui32 VariadicType = 0;
7373
ui32 VariadicArgType = 0;
7474
TString VariadicArgName;
75+
ui32 ExtensionIndex = 0;
7576
};
7677

7778
// Copied from pg_collation_d.h
@@ -355,6 +356,46 @@ struct TColumnInfo {
355356
const TVector<TTableInfo>& GetStaticTables();
356357
const TTableInfo& LookupStaticTable(const TTableInfoKey& tableKey);
357358
const THashMap<TTableInfoKey, TVector<TColumnInfo>>& GetStaticColumns();
359+
360+
void PrepareCatalog();
361+
bool IsExportFunctionsEnabled();
362+
363+
struct TExtensionDesc {
364+
TString Name; // postgis
365+
TString InstallName; // $libdir/postgis-3
366+
TString DDLPath; // DDL path (CREATE TYPE/CREATE FUNCTION/etc)
367+
TString LibraryPath; // file path
368+
bool TypesOnly = false; // Can't be loaded if true
369+
};
370+
371+
class IExtensionDDLBuilder {
372+
public:
373+
virtual ~IExtensionDDLBuilder() = default;
374+
375+
virtual void CreateProc(const TProcDesc& desc) = 0;
376+
};
377+
378+
class IExtensionDDLParser {
379+
public:
380+
virtual ~IExtensionDDLParser() = default;
381+
virtual void Parse(const TString& sql, IExtensionDDLBuilder& builder) = 0;
382+
};
383+
384+
class IExtensionLoader {
385+
public:
386+
virtual ~IExtensionLoader() = default;
387+
virtual void Load(ui32 extensionIndex, const TString& name, const TString& path) = 0;
388+
};
389+
390+
// should be called at most once before other catalog functions
391+
void RegisterExtensions(const TVector<TExtensionDesc>& extensions, bool typesOnly,
392+
IExtensionDDLParser& parser, IExtensionLoader* loader);
393+
394+
void EnumExtensions(std::function<void(ui32 extensionIndex, const TExtensionDesc&)> f);
395+
const TExtensionDesc& LookupExtension(ui32 extensionIndex);
396+
ui32 LookupExtensionByName(const TString& name);
397+
ui32 LookupExtensionByInstallName(const TString& installName);
398+
358399
}
359400

360401
template <>

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,3 +18,5 @@
1818
"pg_get_function_arguments",//psql
1919
"pg_table_is_visible",//psql
2020

21+
// postgis
22+
"postgis_lib_version",

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -512,6 +512,7 @@
512512
"hashinet",
513513
"hashint2",
514514
"hashint4",
515+
"hashint4extended",
515516
"hashint8",
516517
"hashmacaddr",
517518
"hashmacaddr8",

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ SRCS(
2020
PEERDIR(
2121
library/cpp/resource
2222
ydb/library/yql/public/issue
23+
ydb/library/yql/protos
2324
)
2425

2526
END()

0 commit comments

Comments
 (0)