Simplify TF registration
This commit is contained in:
@@ -1,9 +1,13 @@
|
|||||||
#pragma once
|
#pragma once
|
||||||
|
|
||||||
#include <duckdb.hpp>
|
#include <duckdb.hpp>
|
||||||
|
#include <duckdb/main/extension_util.hpp>
|
||||||
|
#include <type_traits>
|
||||||
|
|
||||||
namespace duckdb {
|
namespace duckdb {
|
||||||
|
|
||||||
|
typedef std::string (*simple_tf_t) (ClientContext &);
|
||||||
|
|
||||||
struct RunOnceTableFunctionState : GlobalTableFunctionState {
|
struct RunOnceTableFunctionState : GlobalTableFunctionState {
|
||||||
RunOnceTableFunctionState() : run(false) {};
|
RunOnceTableFunctionState() : run(false) {};
|
||||||
std::atomic<bool> run;
|
std::atomic<bool> run;
|
||||||
@@ -20,8 +24,56 @@ T GetSetting(const ClientContext &context, const char *setting_name, const T def
|
|||||||
return context.TryGetCurrentSetting(setting_name, value) ? value.GetValue<T>() : default_value;
|
return context.TryGetCurrentSetting(setting_name, value) ? value.GetValue<T>() : default_value;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
namespace internal {
|
||||||
|
unique_ptr<FunctionData> ResultBind(ClientContext &, TableFunctionBindInput &,
|
||||||
|
vector<LogicalType> &,
|
||||||
|
vector<std::string> &);
|
||||||
|
|
||||||
bool ShouldRun(TableFunctionInput &input);
|
bool ShouldRun(TableFunctionInput &input);
|
||||||
|
|
||||||
void RegisterTF(DatabaseInstance &instance, const char* name, table_function_t func);
|
template <typename Func>
|
||||||
|
struct CallFunctionHelper;
|
||||||
|
|
||||||
|
template <>
|
||||||
|
struct CallFunctionHelper<std::string(*)()> {
|
||||||
|
static std::string call(ClientContext &context, TableFunctionInput &input, std::string(*f)()) {
|
||||||
|
return f();
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
template <>
|
||||||
|
struct CallFunctionHelper<std::string(*)(ClientContext &)> {
|
||||||
|
static std::string call(ClientContext &context, TableFunctionInput &input, std::string(*f)(ClientContext &)) {
|
||||||
|
return f(context);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
template <>
|
||||||
|
struct CallFunctionHelper<std::string(*)(ClientContext &, TableFunctionInput &)> {
|
||||||
|
static std::string call(ClientContext &context, TableFunctionInput &input, std::string(*f)(ClientContext &, TableFunctionInput &)) {
|
||||||
|
return f(context, input);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
template <typename Func, Func func>
|
||||||
|
void TableFunc(ClientContext &context, TableFunctionInput &input, DataChunk &output) {
|
||||||
|
if (!ShouldRun(input)) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
const std::string result = CallFunctionHelper<Func>::call(context, input, func);
|
||||||
|
output.SetCardinality(1);
|
||||||
|
output.SetValue(0, 0, result);
|
||||||
|
}
|
||||||
|
|
||||||
|
template <typename Func, Func func>
|
||||||
|
void RegisterTF(DatabaseInstance &instance, const char* name) {
|
||||||
|
TableFunction tf(name, {}, internal::TableFunc<Func, func>, internal::ResultBind, RunOnceTableFunctionState::Init);
|
||||||
|
ExtensionUtil::RegisterFunction(instance, tf);
|
||||||
|
}
|
||||||
|
|
||||||
|
}
|
||||||
|
|
||||||
|
#define RESISTER_TF(name, func) internal::RegisterTF<decltype(&func), &func>(instance, name)
|
||||||
|
|
||||||
} // namespace duckdb
|
} // namespace duckdb
|
||||||
|
|||||||
@@ -40,47 +40,23 @@ std::string GetHttpServerLocalURL() {
|
|||||||
|
|
||||||
} // namespace internal
|
} // namespace internal
|
||||||
|
|
||||||
void OutputResult(const std::string &result, DataChunk &out_chunk) {
|
std::string StartUIFunction(ClientContext &context) {
|
||||||
out_chunk.SetCardinality(1);
|
|
||||||
out_chunk.SetValue(0, 0, result);
|
|
||||||
}
|
|
||||||
|
|
||||||
void StartUIFunction(ClientContext &context, TableFunctionInput &input,
|
|
||||||
DataChunk &out_chunk) {
|
|
||||||
if (!ShouldRun(input)) {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
|
|
||||||
internal::StartHttpServer(context);
|
internal::StartHttpServer(context);
|
||||||
auto local_url = internal::GetHttpServerLocalURL();
|
auto local_url = internal::GetHttpServerLocalURL();
|
||||||
|
|
||||||
const std::string command = StringUtil::Format("%s %s", OPEN_COMMAND, local_url);
|
const std::string command = StringUtil::Format("%s %s", OPEN_COMMAND, local_url);
|
||||||
std::string result = system(command.c_str()) ?
|
return system(command.c_str()) ?
|
||||||
StringUtil::Format("Navigate browser to %s", local_url) // open command failed
|
StringUtil::Format("Navigate browser to %s", local_url) // open command failed
|
||||||
: StringUtil::Format("MotherDuck UI started at %s", local_url);
|
: StringUtil::Format("MotherDuck UI started at %s", local_url);
|
||||||
OutputResult(result, out_chunk);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
void StartUIServerFunction(ClientContext &context, TableFunctionInput &input,
|
std::string StartUIServerFunction(ClientContext &context) {
|
||||||
DataChunk &out_chunk) {
|
const char* already = internal::StartHttpServer(context) ? "already " : "";
|
||||||
if (!ShouldRun(input)) {
|
return StringUtil::Format("MotherDuck UI server %sstarted at %s", already, internal::GetHttpServerLocalURL());
|
||||||
return;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
const bool already = internal::StartHttpServer(context);
|
std::string StopUIServerFunction() {
|
||||||
const char* already_str = already ? "already " : "";
|
return ui::HttpServer::instance()->Stop() ? "UI server stopped" : "UI server already stopped";
|
||||||
auto result = StringUtil::Format("MotherDuck UI server %sstarted at %s", already_str, internal::GetHttpServerLocalURL());
|
|
||||||
OutputResult(result, out_chunk);
|
|
||||||
}
|
|
||||||
|
|
||||||
void StopUIServerFunction(ClientContext &, TableFunctionInput &input,
|
|
||||||
DataChunk &out_chunk) {
|
|
||||||
if (!ShouldRun(input)) {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
|
|
||||||
auto result = ui::HttpServer::instance()->Stop() ? "UI server stopped" : "UI server already stopped";
|
|
||||||
OutputResult(result, out_chunk);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// FIXME
|
// FIXME
|
||||||
@@ -103,9 +79,9 @@ static void LoadInternal(DatabaseInstance &instance) {
|
|||||||
UI_REMOTE_URL_SETTING_NAME, UI_REMOTE_URL_SETTING_DESCRIPTION, LogicalType::VARCHAR,
|
UI_REMOTE_URL_SETTING_NAME, UI_REMOTE_URL_SETTING_DESCRIPTION, LogicalType::VARCHAR,
|
||||||
Value(GetEnvOrDefault(UI_REMOTE_URL_SETTING_NAME, UI_REMOTE_URL_SETTING_DEFAULT)));
|
Value(GetEnvOrDefault(UI_REMOTE_URL_SETTING_NAME, UI_REMOTE_URL_SETTING_DEFAULT)));
|
||||||
|
|
||||||
RegisterTF(instance, "start_ui", StartUIFunction);
|
RESISTER_TF("start_ui", StartUIFunction);
|
||||||
RegisterTF(instance, "start_ui_server", StartUIServerFunction);
|
RESISTER_TF("start_ui_server", StartUIServerFunction);
|
||||||
RegisterTF(instance, "stop_ui_server", StopUIServerFunction);
|
RESISTER_TF("stop_ui_server", StopUIServerFunction);
|
||||||
}
|
}
|
||||||
|
|
||||||
void UiExtension::Load(DuckDB &db) {
|
void UiExtension::Load(DuckDB &db) {
|
||||||
|
|||||||
@@ -1,7 +1,7 @@
|
|||||||
#include "utils/helpers.hpp"
|
#include "utils/helpers.hpp"
|
||||||
#include <duckdb/main/extension_util.hpp>
|
|
||||||
|
|
||||||
namespace duckdb {
|
namespace duckdb {
|
||||||
|
namespace internal {
|
||||||
|
|
||||||
bool ShouldRun(TableFunctionInput &input) {
|
bool ShouldRun(TableFunctionInput &input) {
|
||||||
auto state = dynamic_cast<RunOnceTableFunctionState *>(input.global_state.get());
|
auto state = dynamic_cast<RunOnceTableFunctionState *>(input.global_state.get());
|
||||||
@@ -22,9 +22,5 @@ unique_ptr<FunctionData> ResultBind(ClientContext &, TableFunctionBindInput &,
|
|||||||
return nullptr;
|
return nullptr;
|
||||||
}
|
}
|
||||||
|
|
||||||
void RegisterTF(DatabaseInstance &instance, const char* name, table_function_t func) {
|
} // namespace internal
|
||||||
TableFunction tf(name, {}, func, ResultBind, RunOnceTableFunctionState::Init);
|
|
||||||
ExtensionUtil::RegisterFunction(instance, tf);
|
|
||||||
}
|
|
||||||
|
|
||||||
} // namespace duckdb
|
} // namespace duckdb
|
||||||
|
|||||||
Reference in New Issue
Block a user