Skip to content

Commit

Permalink
support of CREATE TYPE in DDL for extensions (#7157)
Browse files Browse the repository at this point in the history
  • Loading branch information
@vitstn
vitstn authored Jul 29, 2024
1 parent f0a033e commit a169e9b
Show file tree
Hide file tree
Showing 7 changed files with 305 additions and 46 deletions.
2 changes: 0 additions & 2 deletions .github/config/muted_ya.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,8 +51,6 @@ ydb/core/viewer/ut Viewer.TabletMergingPacked
ydb/library/actors/http/ut HttpProxy.TooLongHeader
ydb/library/actors/http/ut sole*
ydb/library/yql/providers/generic/connector/tests/datasource/ydb* *
ydb/library/yql/tests/sql/yt_native_file/part12/test.py test[pg_catalog-pg_description_pg_syntax-default.txt-Results]
ydb/library/yql/tests/sql/yt_native_file/part7/test.py test[pg_catalog-pg_proc-default.txt-Results]
ydb/public/sdk/cpp/client/ydb_persqueue_core/ut RetryPolicy.TWriteSession_TestBrokenPolicy
ydb/public/sdk/cpp/client/ydb_persqueue_core/ut [*/*]
ydb/public/sdk/cpp/client/ydb_persqueue_public/ut RetryPolicy.TWriteSession_TestBrokenPolicy
Expand Down
86 changes: 77 additions & 9 deletions ydb/library/yql/parser/pg_catalog/catalog.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1868,7 +1868,10 @@ struct TCatalog : public IExtensionDDLBuilder {
Conversions = ParseConversions(conversionData, ProcByName);
Languages = ParseLanguages(languagesData);

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

void PrepareType(ui32 extensionIndex, const TString& name) final {
Y_ENSURE(extensionIndex);
Y_ENSURE(!TypeByName.contains(name));
TTypeDesc newDesc;
newDesc.Name = name;
newDesc.TypeId = 16000 + Types.size();
newDesc.ExtensionIndex = extensionIndex;
newDesc.ArrayTypeId = newDesc.TypeId + 1;
newDesc.Category = 'U';
Types[newDesc.TypeId] = newDesc;
TypeByName[newDesc.Name] = newDesc.TypeId;
TTypeDesc newArrayDesc = newDesc;
newArrayDesc.TypeId += 1;
newArrayDesc.Name = "_" + newArrayDesc.Name;
newArrayDesc.ElementTypeId = newDesc.TypeId;
newArrayDesc.ArrayTypeId = newArrayDesc.TypeId;
newArrayDesc.PassByValue = false;
newArrayDesc.TypeLen = -1;
newArrayDesc.SendFuncId = (*ProcByName.FindPtr("array_send"))[0];
newArrayDesc.ReceiveFuncId = (*ProcByName.FindPtr("array_recv"))[0];
newArrayDesc.InFuncId = (*ProcByName.FindPtr("array_in"))[0];
newArrayDesc.OutFuncId = (*ProcByName.FindPtr("array_out"))[0];
newArrayDesc.Category = 'A';
Types[newArrayDesc.TypeId] = newArrayDesc;
TypeByName[newArrayDesc.Name] = newArrayDesc.TypeId;
}

void UpdateType(const TTypeDesc& desc) final {
auto byIdPtr = Types.FindPtr(desc.TypeId);
Y_ENSURE(byIdPtr);
Y_ENSURE(byIdPtr->Name == desc.Name);
Y_ENSURE(byIdPtr->ArrayTypeId == desc.ArrayTypeId);
Y_ENSURE(byIdPtr->TypeId == desc.TypeId);
Y_ENSURE(byIdPtr->ExtensionIndex == desc.ExtensionIndex);
if (desc.InFuncId) {
AllowedProcs.insert(Procs.FindPtr(desc.InFuncId)->Name);
}

if (desc.OutFuncId) {
AllowedProcs.insert(Procs.FindPtr(desc.OutFuncId)->Name);
}

if (desc.SendFuncId) {
AllowedProcs.insert(Procs.FindPtr(desc.SendFuncId)->Name);
}

if (desc.ReceiveFuncId) {
AllowedProcs.insert(Procs.FindPtr(desc.ReceiveFuncId)->Name);
}

if (desc.TypeModInFuncId) {
AllowedProcs.insert(Procs.FindPtr(desc.TypeModInFuncId)->Name);
}

if (desc.TypeModOutFuncId) {
AllowedProcs.insert(Procs.FindPtr(desc.TypeModOutFuncId)->Name);
}

*byIdPtr = desc;
}

static const TCatalog& Instance() {
return *Singleton<TCatalog>();
}
Expand Down Expand Up @@ -1981,6 +2045,7 @@ struct TCatalog : public IExtensionDDLBuilder {
THashMap<TTableInfoKey, TVector<TColumnInfo>> StaticColumns;

mutable TMaybe<TFile> ExportFile;
bool AllowAllFunctions = false;
TMutex ExportGuard;

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

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

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

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

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

Expand All @@ -2056,7 +2121,7 @@ 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) {
if (catalog.ExportFile || catalog.AllowedProcs.contains(x.second.Name)) {
if (catalog.AllowAllFunctions || catalog.AllowedProcs.contains(x.second.Name)) {
f(x.first, x.second);
}
}
Expand Down Expand Up @@ -2673,7 +2738,7 @@ std::variant<const TProcDesc*, const TTypeDesc*> LookupProcWithCasts(const TStri
const auto& d = catalog.Procs.FindPtr(id);
Y_ENSURE(d);

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

Expand Down Expand Up @@ -3257,16 +3322,18 @@ const TTableInfo& LookupStaticTable(const TTableInfoKey& tableKey) {
return *tablePtr;
}

bool IsExportFunctionsEnabled() {
bool AreAllFunctionsAllowed() {
const auto& catalog = TCatalog::Instance();
return catalog.ExportFile.Defined();
return catalog.AllowAllFunctions;
}

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

catalog.Extensions.push_back(e);
TString sql = TFileInput(e.DDLPath).ReadAll();;
parser.Parse(sql, catalog);
parser.Parse(i + 1, sql, catalog);
if (loader && !e.TypesOnly) {
loader->Load(i + 1, e.Name, e.LibraryPath);
}
}

catalog.AllowAllFunctions = savedAllowAllFunctions;
catalog.ExtensionsInit = true;
}
}
Expand Down
11 changes: 8 additions & 3 deletions ydb/library/yql/parser/pg_catalog/catalog.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,8 @@ struct TTypeDesc {

// If TypType is 'c', typrelid is the OID of the class' entry in pg_class.
ETypType TypType = ETypType::Base;

ui32 ExtensionIndex = 0;
};

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

void PrepareCatalog();
bool IsExportFunctionsEnabled();
bool AreAllFunctionsAllowed();

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

virtual void CreateProc(const TProcDesc& desc) = 0;

virtual void PrepareType(ui32 extensionIndex,const TString& name) = 0;

virtual void UpdateType(const TTypeDesc& desc) = 0;
};

class IExtensionDDLParser {
public:
virtual ~IExtensionDDLParser() = default;
virtual void Parse(const TString& sql, IExtensionDDLBuilder& builder) = 0;
virtual void Parse(ui32 extensionIndex, const TString& sql, IExtensionDDLBuilder& builder) = 0;
};

class IExtensionLoader {
Expand Down
2 changes: 1 addition & 1 deletion ydb/library/yql/parser/pg_wrapper/ut/proc_ut.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ namespace NYql {

Y_UNIT_TEST_SUITE(TProcTests) {
Y_UNIT_TEST(BuiltinsHasRuntimeFuncs) {
if (NPg::IsExportFunctionsEnabled()) {
if (NPg::AreAllFunctionsAllowed()) {
return;
}

Expand Down
Loading

0 comments on commit a169e9b

Please sign in to comment.