From bb6ab4af359e37773108fd1e90864528dce62610 Mon Sep 17 00:00:00 2001
From: Yves <yves@motherduck.com>
Date: Mon, 17 Feb 2025 17:40:41 +0100
Subject: [PATCH] Add notify_ui_connected

---
 src/include/utils/helpers.hpp |  8 ++++++++
 src/ui_extension.cpp          | 30 ++++++++++++++++++++++++++----
 2 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/src/include/utils/helpers.hpp b/src/include/utils/helpers.hpp
index 33c9e7c..7c90af9 100644
--- a/src/include/utils/helpers.hpp
+++ b/src/include/utils/helpers.hpp
@@ -72,8 +72,16 @@ void RegisterTF(DatabaseInstance &instance, const char* name) {
 	ExtensionUtil::RegisterFunction(instance, tf);
 }
 
+template <typename Func, Func func>
+void RegisterTFWithArgs(DatabaseInstance &instance, const char* name, vector<LogicalType> arguments, table_function_bind_t bind) {
+	TableFunction tf(name, arguments, internal::TableFunc<Func, func>, bind, RunOnceTableFunctionState::Init);
+	ExtensionUtil::RegisterFunction(instance, tf);
+}
+
 }
 
 #define RESISTER_TF(name, func) internal::RegisterTF<decltype(&func), &func>(instance, name)
 
+#define RESISTER_TF_ARGS(name, args, func, bind) internal::RegisterTFWithArgs<decltype(&func), &func>(instance, name, args, bind)
+
 } // namespace duckdb
diff --git a/src/ui_extension.cpp b/src/ui_extension.cpp
index 3476d7b..780e71c 100644
--- a/src/ui_extension.cpp
+++ b/src/ui_extension.cpp
@@ -59,10 +59,31 @@ std::string StopUIServerFunction() {
 	return ui::HttpServer::instance()->Stop() ? "UI server stopped" : "UI server already stopped";
 }
 
-// FIXME
-// void HandleConnected(const std::string &token) {
-// 	ui::HttpServer::instance()->SendConnectedEvent(token);
-// }
+// Connected notification
+struct NotifyConnectedFunctionData : public TableFunctionData {
+	NotifyConnectedFunctionData(std::string _token) : token(_token) {}
+
+	std::string token;
+};
+
+static unique_ptr<FunctionData> NotifyConnectedBind(ClientContext &, TableFunctionBindInput &input,
+	vector<LogicalType> &out_types, vector<string> &out_names) {
+	if (input.inputs[0].IsNull()) {
+		throw BinderException("Must provide a token");
+	}
+
+	out_names.emplace_back("result");
+	out_types.emplace_back(LogicalType::VARCHAR);
+	return make_uniq<NotifyConnectedFunctionData>(input.inputs[0].ToString());
+}
+
+std::string NotifyConnectedFunction(ClientContext &context, TableFunctionInput &input) {
+	auto &inputs = input.bind_data->Cast<NotifyConnectedFunctionData>();
+	ui::HttpServer::instance()->SendConnectedEvent(inputs.token);
+	return "OK";
+}
+
+// - connected notification
 
 std::string NotifyCatalogChangedFunction() {
 	ui::HttpServer::instance()->SendCatalogChangedEvent();
@@ -83,6 +104,7 @@ static void LoadInternal(DatabaseInstance &instance) {
 	RESISTER_TF("start_ui_server", StartUIServerFunction);
 	RESISTER_TF("stop_ui_server", StopUIServerFunction);
 	RESISTER_TF("notify_ui_catalog_changed", NotifyCatalogChangedFunction);
+	RESISTER_TF_ARGS("notify_ui_connected", {LogicalType::VARCHAR}, NotifyConnectedFunction, NotifyConnectedBind);
 }
 
 void UiExtension::Load(DuckDB &db) {