Skip to content

Commit f539014

Browse files
committed
Moved csv parsing in import file cmd to YDB CLI and supported pg-types (only row tables)
1 parent ffbe72e commit f539014

File tree

5 files changed

+136
-61
lines changed

5 files changed

+136
-61
lines changed

ydb/public/lib/ydb_cli/common/csv_parser.cpp

+56-23
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,9 @@ namespace {
1010

1111
class TCsvToYdbConverter {
1212
public:
13-
explicit TCsvToYdbConverter(TTypeParser& parser)
13+
explicit TCsvToYdbConverter(TTypeParser& parser, const TString& nullValue)
1414
: Parser(parser)
15+
, NullValue(nullValue)
1516
{
1617
}
1718

@@ -133,15 +134,15 @@ class TCsvToYdbConverter {
133134

134135
void BuildValue(TStringBuf token) {
135136
switch (Parser.GetKind()) {
136-
case TTypeParser::ETypeKind::Primitive:
137+
case TTypeParser::ETypeKind::Primitive: {
137138
BuildPrimitive(TString(token));
138139
break;
139-
140-
case TTypeParser::ETypeKind::Decimal:
140+
}
141+
case TTypeParser::ETypeKind::Decimal: {
141142
Builder.Decimal(TString(token));
142143
break;
143-
144-
case TTypeParser::ETypeKind::Optional:
144+
}
145+
case TTypeParser::ETypeKind::Optional: {
145146
Parser.OpenOptional();
146147
if (token == NullValue) {
147148
Builder.EmptyOptional(GetType());
@@ -152,23 +153,31 @@ class TCsvToYdbConverter {
152153
}
153154
Parser.CloseOptional();
154155
break;
155-
156-
case TTypeParser::ETypeKind::Null:
156+
}
157+
case TTypeParser::ETypeKind::Null: {
157158
EnsureNull(token);
158159
break;
159-
160-
case TTypeParser::ETypeKind::Void:
160+
}
161+
case TTypeParser::ETypeKind::Void: {
161162
EnsureNull(token);
162163
break;
163-
164-
case TTypeParser::ETypeKind::Tagged:
164+
}
165+
case TTypeParser::ETypeKind::Tagged: {
165166
Parser.OpenTagged();
166167
Builder.BeginTagged(Parser.GetTag());
167168
BuildValue(token);
168169
Builder.EndTagged();
169170
Parser.CloseTagged();
170171
break;
171-
172+
}
173+
case TTypeParser::ETypeKind::Pg: {
174+
if (token == NullValue) {
175+
Builder.Pg(TPgValue(TPgValue::VK_NULL, {}, Parser.GetPg()));
176+
} else {
177+
Builder.Pg(TPgValue(TPgValue::VK_TEXT, TString(token), Parser.GetPg()));
178+
}
179+
break;
180+
}
172181
default:
173182
throw TMisuseException() << "Unsupported type kind: " << Parser.GetKind();
174183
}
@@ -200,6 +209,10 @@ class TCsvToYdbConverter {
200209
Parser.CloseTagged();
201210
break;
202211

212+
case TTypeParser::ETypeKind::Pg:
213+
typeBuilder.Pg(Parser.GetPg());
214+
break;
215+
203216
default:
204217
throw TMisuseException() << "Unsupported type kind: " << Parser.GetKind();
205218
}
@@ -240,23 +253,36 @@ class TCsvToYdbConverter {
240253

241254
}
242255

243-
TCsvParser::TCsvParser(TString&& headerRow, const char delimeter,
244-
const std::map<TString, TType>& paramTypes, const std::map<TString, TString>* paramSources)
256+
TCsvParser::TCsvParser(TString&& headerRow, const char delimeter, const TString& nullValue,
257+
const std::map<TString, TType>* paramTypes,
258+
const std::map<TString, TString>* paramSources)
245259
: HeaderRow(std::move(headerRow))
246260
, Delimeter(delimeter)
261+
, NullValue(nullValue)
247262
, ParamTypes(paramTypes)
248263
, ParamSources(paramSources)
249264
{
250265
NCsvFormat::CsvSplitter splitter(HeaderRow, Delimeter);
251266
Header = static_cast<TVector<TString>>(splitter);
252267
}
253268

254-
TValue TCsvParser::FieldToValue(TTypeParser& parser, TStringBuf token) {
255-
TCsvToYdbConverter converter(parser);
269+
TCsvParser::TCsvParser(TVector<TString>&& header, const char delimeter, const TString& nullValue,
270+
const std::map<TString, TType>* paramTypes,
271+
const std::map<TString, TString>* paramSources)
272+
: Header(std::move(header))
273+
, Delimeter(delimeter)
274+
, NullValue(nullValue)
275+
, ParamTypes(paramTypes)
276+
, ParamSources(paramSources)
277+
{
278+
}
279+
280+
TValue TCsvParser::FieldToValue(TTypeParser& parser, TStringBuf token) const {
281+
TCsvToYdbConverter converter(parser, NullValue);
256282
return converter.Convert(token);
257283
}
258284

259-
void TCsvParser::GetParams(TString&& data, TParamsBuilder& builder) {
285+
void TCsvParser::GetParams(TString&& data, TParamsBuilder& builder) const {
260286
NCsvFormat::CsvSplitter splitter(data, Delimeter);
261287
auto headerIt = Header.begin();
262288
do {
@@ -265,8 +291,8 @@ void TCsvParser::GetParams(TString&& data, TParamsBuilder& builder) {
265291
throw TMisuseException() << "Header contains less fields than data. Header: \"" << HeaderRow << "\", data: \"" << data << "\"";
266292
}
267293
TString fullname = "$" + *headerIt;
268-
auto paramIt = ParamTypes.find(fullname);
269-
if (paramIt == ParamTypes.end()) {
294+
auto paramIt = ParamTypes->find(fullname);
295+
if (paramIt == ParamTypes->end()) {
270296
++headerIt;
271297
continue;
272298
}
@@ -286,7 +312,7 @@ void TCsvParser::GetParams(TString&& data, TParamsBuilder& builder) {
286312
}
287313
}
288314

289-
void TCsvParser::GetValue(TString&& data, TValueBuilder& builder, const TType& type) {
315+
void TCsvParser::GetValue(TString&& data, TValueBuilder& builder, const TType& type) const {
290316
NCsvFormat::CsvSplitter splitter(data, Delimeter);
291317
auto headerIt = Header.cbegin();
292318
std::map<TString, TStringBuf> fields;
@@ -307,6 +333,9 @@ void TCsvParser::GetValue(TString&& data, TValueBuilder& builder, const TType& t
307333
parser.OpenStruct();
308334
while (parser.TryNextMember()) {
309335
TString name = parser.GetMemberName();
336+
if (name == NullValue) {
337+
continue;
338+
}
310339
auto fieldIt = fields.find(name);
311340
if (fieldIt == fields.end()) {
312341
throw TMisuseException() << "No member \"" << name << "\" in csv string for YDB struct type";
@@ -317,11 +346,15 @@ void TCsvParser::GetValue(TString&& data, TValueBuilder& builder, const TType& t
317346
builder.EndStruct();
318347
}
319348

320-
TType TCsvParser::GetColumnsType() {
349+
TType TCsvParser::GetColumnsType() const {
321350
TTypeBuilder builder;
322351
builder.BeginStruct();
323352
for (const auto& colName : Header) {
324-
builder.AddMember(colName, ParamTypes.at(colName));
353+
if (ParamTypes->find(colName) != ParamTypes->end()) {
354+
builder.AddMember(colName, ParamTypes->at(colName));
355+
} else {
356+
builder.AddMember(NullValue, TTypeBuilder().Build());
357+
}
325358
}
326359
builder.EndStruct();
327360
return builder.Build();

ydb/public/lib/ydb_cli/common/csv_parser.h

+21-8
Original file line numberDiff line numberDiff line change
@@ -9,20 +9,33 @@ namespace NConsoleClient {
99

1010
class TCsvParser {
1111
public:
12-
TCsvParser(TString&& headerRow, const char delimeter,
13-
const std::map<TString, TType>& paramTypes, const std::map<TString, TString>* paramSources = nullptr);
12+
TCsvParser() = default;
1413

15-
void GetParams(TString&& data, TParamsBuilder& builder);
16-
void GetValue(TString&& data, TValueBuilder& builder, const TType& type);
17-
TType GetColumnsType();
14+
TCsvParser(const TCsvParser&) = delete;
15+
TCsvParser(TCsvParser&&) = default;
16+
TCsvParser& operator=(const TCsvParser&) = delete;
17+
TCsvParser& operator=(TCsvParser&&) = default;
18+
~TCsvParser() = default;
19+
20+
TCsvParser(TString&& headerRow, const char delimeter, const TString& nullValue,
21+
const std::map<TString, TType>* paramTypes = nullptr,
22+
const std::map<TString, TString>* paramSources = nullptr);
23+
TCsvParser(TVector<TString>&& header, const char delimeter, const TString& nullValue,
24+
const std::map<TString, TType>* paramTypes = nullptr,
25+
const std::map<TString, TString>* paramSources = nullptr);
26+
27+
void GetParams(TString&& data, TParamsBuilder& builder) const;
28+
void GetValue(TString&& data, TValueBuilder& builder, const TType& type) const;
29+
TType GetColumnsType() const;
1830

1931
private:
20-
TValue FieldToValue(TTypeParser& parser, TStringBuf token);
32+
TValue FieldToValue(TTypeParser& parser, TStringBuf token) const;
2133

2234
TVector<TString> Header;
2335
TString HeaderRow;
24-
const char Delimeter;
25-
const std::map<TString, TType>& ParamTypes;
36+
char Delimeter;
37+
TString NullValue;
38+
const std::map<TString, TType>* ParamTypes;
2639
const std::map<TString, TString>* ParamSources;
2740
};
2841

ydb/public/lib/ydb_cli/common/parameters.cpp

+8-8
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,11 @@ void TCommandWithParameters::ParseParameters(TClientCommand::TConfig& config) {
2323

2424
switch (StdinFormat) {
2525
case EOutputFormat::Csv:
26+
Delimiter = ',';
27+
break;
2628
case EOutputFormat::Tsv:
29+
Delimiter = '\t';
30+
break;
2731
case EOutputFormat::Raw:
2832
break;
2933
case EOutputFormat::Default:
@@ -223,11 +227,7 @@ bool TCommandWithParameters::GetNextParams(THolder<TParamsBuilder>& paramBuilder
223227
}
224228
--SkipRows;
225229
}
226-
if (StdinFormat == EOutputFormat::Csv) {
227-
CsvParser = MakeHolder<TCsvParser>(std::move(headerRow), ',', ParamTypes, &ParameterSources);
228-
} else {
229-
CsvParser = MakeHolder<TCsvParser>(std::move(headerRow), '\t', ParamTypes, &ParameterSources);
230-
}
230+
CsvParser = TCsvParser(std::move(headerRow), Delimiter, " ", &ParamTypes, &ParameterSources);
231231
} else {
232232
Input = MakeHolder<TSimpleParamStream>();
233233
}
@@ -257,7 +257,7 @@ bool TCommandWithParameters::GetNextParams(THolder<TParamsBuilder>& paramBuilder
257257
}
258258
case EOutputFormat::Csv:
259259
case EOutputFormat::Tsv: {
260-
CsvParser->GetParams(std::move(*data), *paramBuilder);
260+
CsvParser.GetParams(std::move(*data), *paramBuilder);
261261
break;
262262
}
263263
default:
@@ -300,7 +300,7 @@ bool TCommandWithParameters::GetNextParams(THolder<TParamsBuilder>& paramBuilder
300300
case EOutputFormat::Csv:
301301
case EOutputFormat::Tsv: {
302302
TValueBuilder valueBuilder;
303-
CsvParser->GetValue(std::move(*data), valueBuilder, type);
303+
CsvParser.GetValue(std::move(*data), valueBuilder, type);
304304
paramBuilder->AddParam(fullname, valueBuilder.Build());
305305
break;
306306
}
@@ -379,7 +379,7 @@ bool TCommandWithParameters::GetNextParams(THolder<TParamsBuilder>& paramBuilder
379379
case EOutputFormat::Csv:
380380
case EOutputFormat::Tsv: {
381381
valueBuilder.AddListItem();
382-
CsvParser->GetValue(std::move(*data), valueBuilder, type.GetProto().list_type().item());
382+
CsvParser.GetValue(std::move(*data), valueBuilder, type.GetProto().list_type().item());
383383
break;
384384
}
385385
default:

ydb/public/lib/ydb_cli/common/parameters.h

+2-1
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,8 @@ class TCommandWithParameters : public TCommandWithExamples, public TCommandWithF
4141
THolder<IParamStream> Input;
4242
bool IsFirstEncounter = true;
4343
size_t SkipRows = 0;
44-
THolder<TCsvParser> CsvParser;
44+
char Delimiter;
45+
TCsvParser CsvParser;
4546

4647
protected:
4748
TVector<TString> ParameterOptions, ParameterFiles, StdinParameters;

0 commit comments

Comments
 (0)