|
| 1 | +#ifndef DUCK_FLOCK_H |
| 2 | +#define DUCK_FLOCK_H |
| 3 | +#include "chsql_extension.hpp" |
| 4 | +namespace duckdb { |
| 5 | + struct DuckFlockData : FunctionData{ |
| 6 | + vector<unique_ptr<Connection>> conn; |
| 7 | + vector<unique_ptr<QueryResult>> results; |
| 8 | + unique_ptr<FunctionData> Copy() const override { |
| 9 | + throw std::runtime_error("not implemented"); |
| 10 | + } |
| 11 | + bool Equals(const FunctionData &other) const override { |
| 12 | + throw std::runtime_error("not implemented"); |
| 13 | + }; |
| 14 | + }; |
| 15 | + |
| 16 | + |
| 17 | + |
| 18 | + unique_ptr<FunctionData> DuckFlockBind(ClientContext &context, TableFunctionBindInput &input, |
| 19 | + vector<LogicalType> &return_types, vector<string> &names) { |
| 20 | + auto data = make_uniq<DuckFlockData>(); |
| 21 | + auto strQuery = input.inputs[0].GetValue<string>(); |
| 22 | + vector<string> flock; |
| 23 | + auto &raw_flock = ListValue::GetChildren(input.inputs[1]); |
| 24 | + for (auto &duck : raw_flock) { |
| 25 | + flock.push_back(duck.ToString()); |
| 26 | + auto conn = make_uniq<Connection>(*context.db); |
| 27 | + conn->Query("SET autoload_known_extensions=1;SET autoinstall_known_extensions=1;"); |
| 28 | + auto req = conn->Prepare("SELECT * FROM read_json($2 || '/?q=' || url_encode($1::VARCHAR))"); |
| 29 | + if (req->HasError()) { |
| 30 | + throw std::runtime_error("duck_flock: error: " + req->GetError()); |
| 31 | + } |
| 32 | + data->conn.push_back(std::move(conn)); |
| 33 | + data->results.push_back(std::move(req->Execute(strQuery.c_str(), duck.ToString()))); |
| 34 | + } |
| 35 | + if (data->results[0]->HasError()) { |
| 36 | + throw std::runtime_error("duck_flock: error: " + data->results[0]->GetError()); |
| 37 | + } |
| 38 | + return_types.clear(); |
| 39 | + copy(data->results[0]->types.begin(), data->results[0]->types.end(), back_inserter(return_types)); |
| 40 | + names.clear(); |
| 41 | + copy(data->results[0]->names.begin(), data->results[0]->names.end(), back_inserter(names)); |
| 42 | + return std::move(data); |
| 43 | + } |
| 44 | + |
| 45 | + void DuckFlockImplementation(ClientContext &context, duckdb::TableFunctionInput &data_p, |
| 46 | + DataChunk &output) { |
| 47 | + auto &data = data_p.bind_data->Cast<DuckFlockData>(); |
| 48 | + for (const auto &res : data.results) { |
| 49 | + ErrorData error_data; |
| 50 | + unique_ptr<DataChunk> data_chunk = make_uniq<DataChunk>(); |
| 51 | + if (res->TryFetch(data_chunk, error_data)) { |
| 52 | + if (data_chunk != nullptr) { |
| 53 | + output.Append(*data_chunk); |
| 54 | + return; |
| 55 | + } |
| 56 | + } |
| 57 | + } |
| 58 | + } |
| 59 | + |
| 60 | + TableFunction DuckFlockTableFunction() { |
| 61 | + TableFunction f( |
| 62 | + "duck_flock", |
| 63 | + {LogicalType::VARCHAR, LogicalType::LIST(LogicalType::VARCHAR)}, |
| 64 | + DuckFlockImplementation, |
| 65 | + DuckFlockBind, |
| 66 | + nullptr, |
| 67 | + nullptr |
| 68 | + ); |
| 69 | + return f; |
| 70 | + } |
| 71 | + |
| 72 | + |
| 73 | +} |
| 74 | + |
| 75 | + |
| 76 | + |
| 77 | + |
| 78 | +#endif |
0 commit comments