Skip to content

Commit 6a0dc99

Browse files
committed
url duck_flock function
1 parent a000d4f commit 6a0dc99

File tree

4 files changed

+85
-2
lines changed

4 files changed

+85
-2
lines changed

chsql/CMakeLists.txt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ include_directories(
2121
../duckdb/third_party/mbedtls
2222
../duckdb/third_party/mbedtls/include
2323
../duckdb/third_party/brotli/include)
24-
set(EXTENSION_SOURCES src/chsql_extension.cpp)
24+
set(EXTENSION_SOURCES src/chsql_extension.cpp src/duck_flock.cpp)
2525
build_static_extension(${TARGET_NAME} ${EXTENSION_SOURCES})
2626
build_loadable_extension(${TARGET_NAME} " " ${EXTENSION_SOURCES})
2727
# Link OpenSSL in both the static library as the loadable extension

chsql/src/chsql_extension.cpp

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -179,7 +179,7 @@ static void LoadInternal(DatabaseInstance &instance) {
179179
ExtensionUtil::RegisterFunction(instance, chsql_openssl_version_scalar_function);
180180

181181
// Macros
182-
for (idx_t index = 0; chsql_macros[index].name != nullptr; index++) {
182+
for (idx_t index = 0; chsql_macros[index].name != nullptr; index++) {
183183
auto info = DefaultFunctionGenerator::CreateInternalMacroInfo(chsql_macros[index]);
184184
ExtensionUtil::RegisterFunction(instance, *info);
185185
}
@@ -189,6 +189,8 @@ static void LoadInternal(DatabaseInstance &instance) {
189189
ExtensionUtil::RegisterFunction(instance, *table_info);
190190
}
191191
ExtensionUtil::RegisterFunction(instance, ReadParquetOrderedFunction());
192+
// Flock
193+
ExtensionUtil::RegisterFunction(instance, DuckFlockTableFunction());
192194
}
193195

194196
void ChsqlExtension::Load(DuckDB &db) {

chsql/src/duck_flock.cpp

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
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

chsql/src/include/chsql_extension.hpp

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,4 +12,7 @@ class ChsqlExtension : public Extension {
1212
};
1313
duckdb::TableFunction ReadParquetOrderedFunction();
1414
static void RegisterSillyBTreeStore(DatabaseInstance &instance);
15+
16+
TableFunction DuckFlockTableFunction();
17+
1518
} // namespace duckdb

0 commit comments

Comments
 (0)