1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
|
#include "yql_pg_provider_impl.h"
#include <yql/essentials/providers/common/provider/yql_data_provider_impl.h>
#include <yql/essentials/providers/common/provider/yql_provider_names.h>
#include <yql/essentials/core/expr_nodes/yql_expr_nodes.h>
#include <yql/essentials/providers/pg/expr_nodes/yql_pg_expr_nodes.h>
#include <yql/essentials/core/yql_expr_type_annotation.h>
#include <yql/essentials/utils/log/log.h>
namespace NYql {
using namespace NNodes;
class TPgDataSourceImpl : public TDataProviderBase {
public:
TPgDataSourceImpl(TPgState::TPtr state)
: State_(state)
, TypeAnnotationTransformer_(CreatePgDataSourceTypeAnnotationTransformer(state))
, DqIntegration_(CreatePgDqIntegration(State_))
{}
TStringBuf GetName() const override {
return PgProviderName;
}
IGraphTransformer& GetTypeAnnotationTransformer(bool instantOnly) override {
Y_UNUSED(instantOnly);
return *TypeAnnotationTransformer_;
}
bool CanParse(const TExprNode& node) override {
if (node.IsCallable(TCoRead::CallableName())) {
return TPgDataSource::Match(node.Child(1));
}
return TypeAnnotationTransformer_->CanParse(node);
}
bool GetExecWorld(const TExprNode::TPtr& node, TExprNode::TPtr& root) override {
auto read = TMaybeNode<TPgReadTable>(node);
if (!read) {
return false;
}
root = read.Cast().World().Ptr();
return true;
}
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();
}
auto read = TCoRight(node).Input().Cast<TPgRead>();
auto keyNode = read.FreeArgs().Get(2).Ptr();
if (keyNode->IsCallable("MrTableConcat")) {
if (keyNode->ChildrenSize() != 1) {
ctx.AddError(TIssue(ctx.GetPosition(keyNode->Pos()), TStringBuilder() << "Expected single table name"));
return nullptr;
}
keyNode = keyNode->HeadPtr();
}
const auto maybeKey = TExprBase(keyNode).Maybe<TCoKey>();
if (!maybeKey) {
ctx.AddError(TIssue(ctx.GetPosition(read.FreeArgs().Get(0).Pos()), TStringBuilder() << "Expected Key"));
return nullptr;
}
const auto& keyArg = maybeKey.Cast().Ref().Head();
if (!keyArg.IsList() || keyArg.ChildrenSize() != 2U
|| !keyArg.Head().IsAtom("table") || !keyArg.Tail().IsCallable(TCoString::CallableName())) {
ctx.AddError(TIssue(ctx.GetPosition(keyArg.Pos()), TStringBuilder() << "Expected single table name"));
return nullptr;
}
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)
.Done().Ptr();
}
bool ValidateParameters(TExprNode& node, TExprContext& ctx, TMaybe<TString>& cluster) override {
if (node.IsCallable(TCoDataSource::CallableName())) {
if (!EnsureArgsCount(node, 2, ctx)) {
return false;
}
if (node.Child(0)->Content() == PgProviderName) {
if (!EnsureAtom(*node.Child(1), ctx)) {
return false;
}
if (node.Child(1)->Content().empty()) {
ctx.AddError(TIssue(ctx.GetPosition(node.Child(1)->Pos()), "Empty cluster name"));
return false;
}
cluster = TString(node.Child(1)->Content());
return true;
}
}
ctx.AddError(TIssue(ctx.GetPosition(node.Pos()), "Invalid Pg DataSource parameters"));
return false;
}
IDqIntegration* GetDqIntegration() override {
return DqIntegration_.Get();
}
private:
TPgState::TPtr State_;
const THolder<TVisitorTransformerBase> TypeAnnotationTransformer_;
const THolder<IDqIntegration> DqIntegration_;
};
TIntrusivePtr<IDataProvider> CreatePgDataSource(TPgState::TPtr state) {
return MakeIntrusive<TPgDataSourceImpl>(state);
}
}
|
