url duck_flock function

quackscience · Oct 18, 2024 · 6a0dc99 · 6a0dc99
1 parent a000d4f
commit 6a0dc99
Show file tree

Hide file tree

Showing 4 changed files with 85 additions and 2 deletions.
diff --git a/chsql/CMakeLists.txt b/chsql/CMakeLists.txt
@@ -21,7 +21,7 @@ include_directories(
         ../duckdb/third_party/mbedtls
         ../duckdb/third_party/mbedtls/include
         ../duckdb/third_party/brotli/include)
-set(EXTENSION_SOURCES src/chsql_extension.cpp)
+set(EXTENSION_SOURCES src/chsql_extension.cpp src/duck_flock.cpp)
 build_static_extension(${TARGET_NAME} ${EXTENSION_SOURCES})
 build_loadable_extension(${TARGET_NAME} " " ${EXTENSION_SOURCES})
 # Link OpenSSL in both the static library as the loadable extension

diff --git a/chsql/src/chsql_extension.cpp b/chsql/src/chsql_extension.cpp
@@ -179,7 +179,7 @@ static void LoadInternal(DatabaseInstance &instance) {
     ExtensionUtil::RegisterFunction(instance, chsql_openssl_version_scalar_function);
 
     // Macros
-	for (idx_t index = 0; chsql_macros[index].name != nullptr; index++) {
+    for (idx_t index = 0; chsql_macros[index].name != nullptr; index++) {
 		auto info = DefaultFunctionGenerator::CreateInternalMacroInfo(chsql_macros[index]);
 		ExtensionUtil::RegisterFunction(instance, *info);
 	}
@@ -189,6 +189,8 @@ static void LoadInternal(DatabaseInstance &instance) {
         ExtensionUtil::RegisterFunction(instance, *table_info);
 	}
 	ExtensionUtil::RegisterFunction(instance, ReadParquetOrderedFunction());
+    // Flock
+    ExtensionUtil::RegisterFunction(instance, DuckFlockTableFunction());
 }
 
 void ChsqlExtension::Load(DuckDB &db) {

diff --git a/chsql/src/duck_flock.cpp b/chsql/src/duck_flock.cpp
@@ -0,0 +1,78 @@
+#ifndef DUCK_FLOCK_H
+#define DUCK_FLOCK_H
+#include "chsql_extension.hpp"
+namespace duckdb {
+    struct DuckFlockData : FunctionData{
+        vector<unique_ptr<Connection>> conn;
+        vector<unique_ptr<QueryResult>> results;
+        unique_ptr<FunctionData> Copy() const override {
+            throw std::runtime_error("not implemented");
+        }
+        bool Equals(const FunctionData &other) const override {
+            throw std::runtime_error("not implemented");
+        };
+    };
+
+
+
+    unique_ptr<FunctionData> DuckFlockBind(ClientContext &context, TableFunctionBindInput &input,
+                                                        vector<LogicalType> &return_types, vector<string> &names) {
+        auto data = make_uniq<DuckFlockData>();
+        auto strQuery = input.inputs[0].GetValue<string>();
+        vector<string> flock;
+        auto &raw_flock = ListValue::GetChildren(input.inputs[1]);
+        for (auto &duck : raw_flock) {
+            flock.push_back(duck.ToString());
+            auto conn = make_uniq<Connection>(*context.db);
+            conn->Query("SET autoload_known_extensions=1;SET autoinstall_known_extensions=1;");
+            auto req = conn->Prepare("SELECT * FROM read_json($2 || '/?q=' || url_encode($1::VARCHAR))");
+            if (req->HasError()) {
+                throw std::runtime_error("duck_flock: error: " + req->GetError());
+            }
+            data->conn.push_back(std::move(conn));
+            data->results.push_back(std::move(req->Execute(strQuery.c_str(), duck.ToString())));
+        }
+        if (data->results[0]->HasError()) {
+            throw std::runtime_error("duck_flock: error: " + data->results[0]->GetError());
+        }
+        return_types.clear();
+        copy(data->results[0]->types.begin(), data->results[0]->types.end(), back_inserter(return_types));
+        names.clear();
+        copy(data->results[0]->names.begin(), data->results[0]->names.end(), back_inserter(names));
+        return std::move(data);
+    }
+
+    void DuckFlockImplementation(ClientContext &context, duckdb::TableFunctionInput &data_p,
+                                      DataChunk &output) {
+        auto &data = data_p.bind_data->Cast<DuckFlockData>();
+        for (const auto &res : data.results) {
+            ErrorData error_data;
+            unique_ptr<DataChunk> data_chunk = make_uniq<DataChunk>();
+            if (res->TryFetch(data_chunk, error_data)) {
+                if (data_chunk != nullptr) {
+                    output.Append(*data_chunk);
+                    return;
+                }
+            }
+        }
+    }
+
+    TableFunction DuckFlockTableFunction() {
+      TableFunction f(
+          "duck_flock",
+          {LogicalType::VARCHAR, LogicalType::LIST(LogicalType::VARCHAR)},
+          DuckFlockImplementation,
+          DuckFlockBind,
+          nullptr,
+          nullptr
+      );
+      return f;
+    }
+
+
+}
+
+
+
+
+#endif
diff --git a/chsql/src/include/chsql_extension.hpp b/chsql/src/include/chsql_extension.hpp
@@ -12,4 +12,7 @@ class ChsqlExtension : public Extension {
 };
 duckdb::TableFunction ReadParquetOrderedFunction();
 static void RegisterSillyBTreeStore(DatabaseInstance &instance);
+
+TableFunction DuckFlockTableFunction();
+
 } // namespace duckdb