1 | #include "duckdb/common/string_util.hpp" |
---|---|
2 | #include "duckdb/common/types/value.hpp" |
3 | #include "duckdb/parser/expression/constant_expression.hpp" |
4 | #include "duckdb/parser/statement/copy_statement.hpp" |
5 | #include "duckdb/parser/tableref/basetableref.hpp" |
6 | #include "duckdb/parser/transformer.hpp" |
7 | |
8 | #include <cstring> |
9 | |
10 | namespace duckdb { |
11 | |
12 | void Transformer::TransformCopyOptions(CopyInfo &info, optional_ptr<duckdb_libpgquery::PGList> options) { |
13 | if (!options) { |
14 | return; |
15 | } |
16 | |
17 | // iterate over each option |
18 | duckdb_libpgquery::PGListCell *cell; |
19 | for_each_cell(cell, options->head) { |
20 | auto def_elem = PGPointerCast<duckdb_libpgquery::PGDefElem>(ptr: cell->data.ptr_value); |
21 | if (StringUtil::Lower(str: def_elem->defname) == "format") { |
22 | // format specifier: interpret this option |
23 | auto format_val = PGPointerCast<duckdb_libpgquery::PGValue>(ptr: def_elem->arg); |
24 | if (!format_val || format_val->type != duckdb_libpgquery::T_PGString) { |
25 | throw ParserException("Unsupported parameter type for FORMAT: expected e.g. FORMAT 'csv', 'parquet'"); |
26 | } |
27 | info.format = StringUtil::Lower(str: format_val->val.str); |
28 | continue; |
29 | } |
30 | // otherwise |
31 | if (info.options.find(x: def_elem->defname) != info.options.end()) { |
32 | throw ParserException("Unexpected duplicate option \"%s\"", def_elem->defname); |
33 | } |
34 | if (!def_elem->arg) { |
35 | info.options[def_elem->defname] = vector<Value>(); |
36 | continue; |
37 | } |
38 | switch (def_elem->arg->type) { |
39 | case duckdb_libpgquery::T_PGList: { |
40 | auto column_list = PGPointerCast<duckdb_libpgquery::PGList>(ptr: def_elem->arg); |
41 | for (auto c = column_list->head; c != nullptr; c = lnext(c)) { |
42 | auto target = PGPointerCast<duckdb_libpgquery::PGResTarget>(ptr: c->data.ptr_value); |
43 | info.options[def_elem->defname].push_back(x: Value(target->name)); |
44 | } |
45 | break; |
46 | } |
47 | case duckdb_libpgquery::T_PGAStar: |
48 | info.options[def_elem->defname].push_back(x: Value("*")); |
49 | break; |
50 | default: { |
51 | auto val = PGPointerCast<duckdb_libpgquery::PGValue>(ptr: def_elem->arg); |
52 | info.options[def_elem->defname].push_back(x: TransformValue(val: *val)->value); |
53 | break; |
54 | } |
55 | } |
56 | } |
57 | } |
58 | |
59 | unique_ptr<CopyStatement> Transformer::TransformCopy(duckdb_libpgquery::PGCopyStmt &stmt) { |
60 | auto result = make_uniq<CopyStatement>(); |
61 | auto &info = *result->info; |
62 | |
63 | // get file_path and is_from |
64 | info.is_from = stmt.is_from; |
65 | if (!stmt.filename) { |
66 | // stdin/stdout |
67 | info.file_path = info.is_from ? "/dev/stdin": "/dev/stdout"; |
68 | } else { |
69 | // copy to a file |
70 | info.file_path = stmt.filename; |
71 | } |
72 | if (StringUtil::EndsWith(str: info.file_path, suffix: ".parquet")) { |
73 | info.format = "parquet"; |
74 | } else if (StringUtil::EndsWith(str: info.file_path, suffix: ".json") || StringUtil::EndsWith(str: info.file_path, suffix: ".ndjson")) { |
75 | info.format = "json"; |
76 | } else { |
77 | info.format = "csv"; |
78 | } |
79 | |
80 | // get select_list |
81 | if (stmt.attlist) { |
82 | for (auto n = stmt.attlist->head; n != nullptr; n = n->next) { |
83 | auto target = PGPointerCast<duckdb_libpgquery::PGResTarget>(ptr: n->data.ptr_value); |
84 | if (target->name) { |
85 | info.select_list.emplace_back(args&: target->name); |
86 | } |
87 | } |
88 | } |
89 | |
90 | if (stmt.relation) { |
91 | auto ref = TransformRangeVar(root&: *stmt.relation); |
92 | auto &table = ref->Cast<BaseTableRef>(); |
93 | info.table = table.table_name; |
94 | info.schema = table.schema_name; |
95 | info.catalog = table.catalog_name; |
96 | } else { |
97 | result->select_statement = TransformSelectNode(select&: *PGPointerCast<duckdb_libpgquery::PGSelectStmt>(ptr: stmt.query)); |
98 | } |
99 | |
100 | // handle the different options of the COPY statement |
101 | TransformCopyOptions(info, options: stmt.options); |
102 | |
103 | return result; |
104 | } |
105 | |
106 | } // namespace duckdb |
107 |