Skip to content

Commit a169e9b

Browse files
authored
support of CREATE TYPE in DDL for extensions (#7157)
1 parent f0a033e commit a169e9b

File tree

7 files changed

+305
-46
lines changed

7 files changed

+305
-46
lines changed

.github/config/muted_ya.txt

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -51,8 +51,6 @@ ydb/core/viewer/ut Viewer.TabletMergingPacked
5151
ydb/library/actors/http/ut HttpProxy.TooLongHeader
5252
ydb/library/actors/http/ut sole*
5353
ydb/library/yql/providers/generic/connector/tests/datasource/ydb* *
54-
ydb/library/yql/tests/sql/yt_native_file/part12/test.py test[pg_catalog-pg_description_pg_syntax-default.txt-Results]
55-
ydb/library/yql/tests/sql/yt_native_file/part7/test.py test[pg_catalog-pg_proc-default.txt-Results]
5654
ydb/public/sdk/cpp/client/ydb_persqueue_core/ut RetryPolicy.TWriteSession_TestBrokenPolicy
5755
ydb/public/sdk/cpp/client/ydb_persqueue_core/ut [*/*]
5856
ydb/public/sdk/cpp/client/ydb_persqueue_public/ut RetryPolicy.TWriteSession_TestBrokenPolicy

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

Lines changed: 77 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1868,7 +1868,10 @@ struct TCatalog : public IExtensionDDLBuilder {
18681868
Conversions = ParseConversions(conversionData, ProcByName);
18691869
Languages = ParseLanguages(languagesData);
18701870

1871-
if (auto exportDir = GetEnv("YQL_EXPORT_PG_FUNCTIONS_DIR")) {
1871+
if (GetEnv("YQL_ALLOW_ALL_PG_FUNCTIONS")) {
1872+
AllowAllFunctions = true;
1873+
} else if (auto exportDir = GetEnv("YQL_EXPORT_PG_FUNCTIONS_DIR")) {
1874+
AllowAllFunctions = true;
18721875
ExportFile.ConstructInPlace(MakeTempName(exportDir.c_str(), "procs"), CreateAlways | RdWr);
18731876
for (const auto& a : Aggregations) {
18741877
const auto& desc = a.second;
@@ -1943,6 +1946,67 @@ struct TCatalog : public IExtensionDDLBuilder {
19431946
ProcByName[newDesc.Name].push_back(newDesc.ProcId);
19441947
}
19451948

1949+
void PrepareType(ui32 extensionIndex, const TString& name) final {
1950+
Y_ENSURE(extensionIndex);
1951+
Y_ENSURE(!TypeByName.contains(name));
1952+
TTypeDesc newDesc;
1953+
newDesc.Name = name;
1954+
newDesc.TypeId = 16000 + Types.size();
1955+
newDesc.ExtensionIndex = extensionIndex;
1956+
newDesc.ArrayTypeId = newDesc.TypeId + 1;
1957+
newDesc.Category = 'U';
1958+
Types[newDesc.TypeId] = newDesc;
1959+
TypeByName[newDesc.Name] = newDesc.TypeId;
1960+
TTypeDesc newArrayDesc = newDesc;
1961+
newArrayDesc.TypeId += 1;
1962+
newArrayDesc.Name = "_" + newArrayDesc.Name;
1963+
newArrayDesc.ElementTypeId = newDesc.TypeId;
1964+
newArrayDesc.ArrayTypeId = newArrayDesc.TypeId;
1965+
newArrayDesc.PassByValue = false;
1966+
newArrayDesc.TypeLen = -1;
1967+
newArrayDesc.SendFuncId = (*ProcByName.FindPtr("array_send"))[0];
1968+
newArrayDesc.ReceiveFuncId = (*ProcByName.FindPtr("array_recv"))[0];
1969+
newArrayDesc.InFuncId = (*ProcByName.FindPtr("array_in"))[0];
1970+
newArrayDesc.OutFuncId = (*ProcByName.FindPtr("array_out"))[0];
1971+
newArrayDesc.Category = 'A';
1972+
Types[newArrayDesc.TypeId] = newArrayDesc;
1973+
TypeByName[newArrayDesc.Name] = newArrayDesc.TypeId;
1974+
}
1975+
1976+
void UpdateType(const TTypeDesc& desc) final {
1977+
auto byIdPtr = Types.FindPtr(desc.TypeId);
1978+
Y_ENSURE(byIdPtr);
1979+
Y_ENSURE(byIdPtr->Name == desc.Name);
1980+
Y_ENSURE(byIdPtr->ArrayTypeId == desc.ArrayTypeId);
1981+
Y_ENSURE(byIdPtr->TypeId == desc.TypeId);
1982+
Y_ENSURE(byIdPtr->ExtensionIndex == desc.ExtensionIndex);
1983+
if (desc.InFuncId) {
1984+
AllowedProcs.insert(Procs.FindPtr(desc.InFuncId)->Name);
1985+
}
1986+
1987+
if (desc.OutFuncId) {
1988+
AllowedProcs.insert(Procs.FindPtr(desc.OutFuncId)->Name);
1989+
}
1990+
1991+
if (desc.SendFuncId) {
1992+
AllowedProcs.insert(Procs.FindPtr(desc.SendFuncId)->Name);
1993+
}
1994+
1995+
if (desc.ReceiveFuncId) {
1996+
AllowedProcs.insert(Procs.FindPtr(desc.ReceiveFuncId)->Name);
1997+
}
1998+
1999+
if (desc.TypeModInFuncId) {
2000+
AllowedProcs.insert(Procs.FindPtr(desc.TypeModInFuncId)->Name);
2001+
}
2002+
2003+
if (desc.TypeModOutFuncId) {
2004+
AllowedProcs.insert(Procs.FindPtr(desc.TypeModOutFuncId)->Name);
2005+
}
2006+
2007+
*byIdPtr = desc;
2008+
}
2009+
19462010
static const TCatalog& Instance() {
19472011
return *Singleton<TCatalog>();
19482012
}
@@ -1981,6 +2045,7 @@ struct TCatalog : public IExtensionDDLBuilder {
19812045
THashMap<TTableInfoKey, TVector<TColumnInfo>> StaticColumns;
19822046

19832047
mutable TMaybe<TFile> ExportFile;
2048+
bool AllowAllFunctions = false;
19842049
TMutex ExportGuard;
19852050

19862051
THashSet<TString> AllowedProcs;
@@ -1998,7 +2063,7 @@ const TProcDesc& LookupProc(ui32 procId, const TVector<ui32>& argTypeIds) {
19982063
throw yexception() << "No such proc: " << procId;
19992064
}
20002065

2001-
if (!catalog.ExportFile && !catalog.AllowedProcs.contains(procPtr->Name)) {
2066+
if (!catalog.AllowAllFunctions && !catalog.AllowedProcs.contains(procPtr->Name)) {
20022067
throw yexception() << "No access to proc: " << procPtr->Name;
20032068
}
20042069

@@ -2022,7 +2087,7 @@ const TProcDesc& LookupProc(const TString& name, const TVector<ui32>& argTypeIds
20222087
for (const auto& id : *procIdPtr) {
20232088
const auto& d = catalog.Procs.FindPtr(id);
20242089
Y_ENSURE(d);
2025-
if (!catalog.ExportFile && !catalog.AllowedProcs.contains(d->Name)) {
2090+
if (!catalog.AllowAllFunctions && !catalog.AllowedProcs.contains(d->Name)) {
20262091
throw yexception() << "No access to proc: " << d->Name;
20272092
}
20282093

@@ -2045,7 +2110,7 @@ const TProcDesc& LookupProc(ui32 procId) {
20452110
throw yexception() << "No such proc: " << procId;
20462111
}
20472112

2048-
if (!catalog.ExportFile && !catalog.AllowedProcs.contains(procPtr->Name)) {
2113+
if (!catalog.AllowAllFunctions && !catalog.AllowedProcs.contains(procPtr->Name)) {
20492114
throw yexception() << "No access to proc: " << procPtr->Name;
20502115
}
20512116

@@ -2056,7 +2121,7 @@ const TProcDesc& LookupProc(ui32 procId) {
20562121
void EnumProc(std::function<void(ui32, const TProcDesc&)> f) {
20572122
const auto& catalog = TCatalog::Instance();
20582123
for (const auto& x : catalog.Procs) {
2059-
if (catalog.ExportFile || catalog.AllowedProcs.contains(x.second.Name)) {
2124+
if (catalog.AllowAllFunctions || catalog.AllowedProcs.contains(x.second.Name)) {
20602125
f(x.first, x.second);
20612126
}
20622127
}
@@ -2673,7 +2738,7 @@ std::variant<const TProcDesc*, const TTypeDesc*> LookupProcWithCasts(const TStri
26732738
const auto& d = catalog.Procs.FindPtr(id);
26742739
Y_ENSURE(d);
26752740

2676-
if (!catalog.ExportFile && !catalog.AllowedProcs.contains(d->Name)) {
2741+
if (!catalog.AllowAllFunctions && !catalog.AllowedProcs.contains(d->Name)) {
26772742
throw yexception() << "No access to proc: " << d->Name;
26782743
}
26792744

@@ -3257,16 +3322,18 @@ const TTableInfo& LookupStaticTable(const TTableInfoKey& tableKey) {
32573322
return *tablePtr;
32583323
}
32593324

3260-
bool IsExportFunctionsEnabled() {
3325+
bool AreAllFunctionsAllowed() {
32613326
const auto& catalog = TCatalog::Instance();
3262-
return catalog.ExportFile.Defined();
3327+
return catalog.AllowAllFunctions;
32633328
}
32643329

32653330
void RegisterExtensions(const TVector<TExtensionDesc>& extensions, bool typesOnly,
32663331
IExtensionDDLParser& parser, IExtensionLoader* loader) {
32673332
auto& catalog = TCatalog::MutableInstance();
32683333
with_lock (catalog.ExtensionsGuard) {
32693334
Y_ENSURE(!catalog.ExtensionsInit);
3335+
auto savedAllowAllFunctions = catalog.AllowAllFunctions;
3336+
catalog.AllowAllFunctions = true;
32703337
for (ui32 i = 0; i < extensions.size(); ++i) {
32713338
auto e = extensions[i];
32723339
e.TypesOnly = e.TypesOnly && typesOnly;
@@ -3284,12 +3351,13 @@ void RegisterExtensions(const TVector<TExtensionDesc>& extensions, bool typesOnl
32843351

32853352
catalog.Extensions.push_back(e);
32863353
TString sql = TFileInput(e.DDLPath).ReadAll();;
3287-
parser.Parse(sql, catalog);
3354+
parser.Parse(i + 1, sql, catalog);
32883355
if (loader && !e.TypesOnly) {
32893356
loader->Load(i + 1, e.Name, e.LibraryPath);
32903357
}
32913358
}
32923359

3360+
catalog.AllowAllFunctions = savedAllowAllFunctions;
32933361
catalog.ExtensionsInit = true;
32943362
}
32953363
}

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

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,8 @@ struct TTypeDesc {
129129

130130
// If TypType is 'c', typrelid is the OID of the class' entry in pg_class.
131131
ETypType TypType = ETypType::Base;
132+
133+
ui32 ExtensionIndex = 0;
132134
};
133135

134136
enum class ECastMethod {
@@ -357,8 +359,7 @@ const TVector<TTableInfo>& GetStaticTables();
357359
const TTableInfo& LookupStaticTable(const TTableInfoKey& tableKey);
358360
const THashMap<TTableInfoKey, TVector<TColumnInfo>>& GetStaticColumns();
359361

360-
void PrepareCatalog();
361-
bool IsExportFunctionsEnabled();
362+
bool AreAllFunctionsAllowed();
362363

363364
struct TExtensionDesc {
364365
TString Name; // postgis
@@ -373,12 +374,16 @@ class IExtensionDDLBuilder {
373374
virtual ~IExtensionDDLBuilder() = default;
374375

375376
virtual void CreateProc(const TProcDesc& desc) = 0;
377+
378+
virtual void PrepareType(ui32 extensionIndex,const TString& name) = 0;
379+
380+
virtual void UpdateType(const TTypeDesc& desc) = 0;
376381
};
377382

378383
class IExtensionDDLParser {
379384
public:
380385
virtual ~IExtensionDDLParser() = default;
381-
virtual void Parse(const TString& sql, IExtensionDDLBuilder& builder) = 0;
386+
virtual void Parse(ui32 extensionIndex, const TString& sql, IExtensionDDLBuilder& builder) = 0;
382387
};
383388

384389
class IExtensionLoader {

ydb/library/yql/parser/pg_wrapper/ut/proc_ut.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ namespace NYql {
1717

1818
Y_UNIT_TEST_SUITE(TProcTests) {
1919
Y_UNIT_TEST(BuiltinsHasRuntimeFuncs) {
20-
if (NPg::IsExportFunctionsEnabled()) {
20+
if (NPg::AreAllFunctionsAllowed()) {
2121
return;
2222
}
2323

0 commit comments

Comments
 (0)