Merge pull request #1103 from suketa/feature/table-function-phase5-init-info

feat: Add DuckDB::InitInfo and TableFunction#init callback (Phase 5)

Author: suketa

CHANGELOG.md

@@ -28,6 +28,9 @@ All notable changes to this project will be documented in this file.
 - add DuckDB::TableFunction#execute for setting execute callback (Phase 4).
 - add DuckDB::FunctionInfo class for table function execution context (Phase 4).
 - add DuckDB::FunctionInfo#set_error for reporting execution errors.
+- add DuckDB::TableFunction#init for setting init callback (Phase 5).
+- add DuckDB::InitInfo class for table function initialization context (Phase 5).
+- add DuckDB::InitInfo#set_error for reporting initialization errors.
 - bump duckdb to 1.4.4 on CI.
 - add inline style to DuckDB::Connection#register_scalar_function (accepts keyword arguments + block).
 - add DuckDB::ScalarFunction.create class method for declarative API.

ext/duckdb/duckdb.c

@@ -43,6 +43,7 @@ Init_duckdb_native(void) {
     rbduckdb_init_duckdb_value_impl();
     rbduckdb_init_duckdb_scalar_function();
     rbduckdb_init_duckdb_bind_info();
+    rbduckdb_init_duckdb_init_info();
     rbduckdb_init_duckdb_function_info();
     rbduckdb_init_duckdb_vector();
     rbduckdb_init_duckdb_data_chunk();

ext/duckdb/init_info.c

@@ -0,0 +1,65 @@
+#include "ruby-duckdb.h"
+
+VALUE cDuckDBInitInfo;
+
+static void deallocate(void *ctx);
+static VALUE allocate(VALUE klass);
+static size_t memsize(const void *p);
+static VALUE rbduckdb_init_info_set_error(VALUE self, VALUE error);
+
+static const rb_data_type_t init_info_data_type = {
+    "DuckDB/InitInfo",
+    {NULL, deallocate, memsize,},
+    0, 0, RUBY_TYPED_FREE_IMMEDIATELY
+};
+
+static void deallocate(void *ctx) {
+    rubyDuckDBInitInfo *p = (rubyDuckDBInitInfo *)ctx;
+    xfree(p);
+}
+
+static VALUE allocate(VALUE klass) {
+    rubyDuckDBInitInfo *ctx = xcalloc((size_t)1, sizeof(rubyDuckDBInitInfo));
+    return TypedData_Wrap_Struct(klass, &init_info_data_type, ctx);
+}
+
+static size_t memsize(const void *p) {
+    return sizeof(rubyDuckDBInitInfo);
+}
+
+rubyDuckDBInitInfo *get_struct_init_info(VALUE obj) {
+    rubyDuckDBInitInfo *ctx;
+    TypedData_Get_Struct(obj, rubyDuckDBInitInfo, &init_info_data_type, ctx);
+    return ctx;
+}
+
+/*
+ * call-seq:
+ *   init_info.set_error(error_message) -> self
+ *
+ * Sets an error message for the init phase.
+ * This will cause the query to fail with the specified error.
+ *
+ *   init_info.set_error('Invalid initialization')
+ */
+static VALUE rbduckdb_init_info_set_error(VALUE self, VALUE error) {
+    rubyDuckDBInitInfo *ctx;
+    const char *error_msg;
+
+    TypedData_Get_Struct(self, rubyDuckDBInitInfo, &init_info_data_type, ctx);
+
+    error_msg = StringValueCStr(error);
+    duckdb_init_set_error(ctx->info, error_msg);
+
+    return self;
+}
+
+void rbduckdb_init_duckdb_init_info(void) {
+#if 0
+    VALUE mDuckDB = rb_define_module("DuckDB");
+#endif
+    cDuckDBInitInfo = rb_define_class_under(mDuckDB, "InitInfo", rb_cObject);
+    rb_define_alloc_func(cDuckDBInitInfo, allocate);
+
+    rb_define_method(cDuckDBInitInfo, "set_error", rbduckdb_init_info_set_error, 1);
+}

ext/duckdb/init_info.h

@@ -0,0 +1,13 @@
+#ifndef RUBY_DUCKDB_INIT_INFO_H
+#define RUBY_DUCKDB_INIT_INFO_H
+
+struct _rubyDuckDBInitInfo {
+    duckdb_init_info info;
+};
+
+typedef struct _rubyDuckDBInitInfo rubyDuckDBInitInfo;
+
+rubyDuckDBInitInfo *get_struct_init_info(VALUE obj);
+void rbduckdb_init_duckdb_init_info(void);
+
+#endif

ext/duckdb/ruby-duckdb.h

@@ -31,6 +31,7 @@
 #include "./value_impl.h"
 #include "./scalar_function.h"
 #include "./bind_info.h"
+#include "./init_info.h"
 #include "./function_info.h"
 #include "./vector.h"
 #include "./data_chunk.h"

ext/duckdb/table_function.c

@@ -2,6 +2,7 @@
 
 static VALUE cDuckDBTableFunction;
 extern VALUE cDuckDBBindInfo;
+extern VALUE cDuckDBInitInfo;
 extern VALUE cDuckDBFunctionInfo;
 extern VALUE cDuckDBDataChunk;
 
@@ -15,6 +16,8 @@ static VALUE rbduckdb_table_function_add_parameter(VALUE self, VALUE logical_typ
 static VALUE rbduckdb_table_function_add_named_parameter(VALUE self, VALUE name, VALUE logical_type);
 static VALUE rbduckdb_table_function_set_bind(VALUE self);
 static void table_function_bind_callback(duckdb_bind_info info);
+static VALUE rbduckdb_table_function_set_init(VALUE self);
+static void table_function_init_callback(duckdb_init_info info);
 static VALUE rbduckdb_table_function_set_execute(VALUE self);
 static void table_function_execute_callback(duckdb_function_info info, duckdb_data_chunk output);
 
@@ -27,6 +30,7 @@ static const rb_data_type_t table_function_data_type = {
 static void mark(void *ctx) {
     rubyDuckDBTableFunction *p = (rubyDuckDBTableFunction *)ctx;
     rb_gc_mark(p->bind_proc);
+    rb_gc_mark(p->init_proc);
     rb_gc_mark(p->execute_proc);
 }
 
@@ -218,6 +222,68 @@ static void table_function_bind_callback(duckdb_bind_info info) {
     }
 }
 
+/*
+ * call-seq:
+ *   table_function.init { |init_info| ... } -> table_function
+ *
+ * Sets the init callback for the table function.
+ * The callback is invoked once during query initialization to set up execution state.
+ *
+ *   table_function.init do |init_info|
+ *     # Initialize execution state
+ *   end
+ */
+static VALUE rbduckdb_table_function_set_init(VALUE self) {
+    rubyDuckDBTableFunction *ctx;
+
+    if (!rb_block_given_p()) {
+        rb_raise(rb_eArgError, "block is required for init");
+    }
+
+    TypedData_Get_Struct(self, rubyDuckDBTableFunction, &table_function_data_type, ctx);
+
+    if (!ctx->table_function) {
+        rb_raise(eDuckDBError, "Table function is destroyed");
+    }
+
+    ctx->init_proc = rb_block_proc();
+    duckdb_table_function_set_init(ctx->table_function, table_function_init_callback);
+    duckdb_table_function_set_extra_info(ctx->table_function, (void *)self, NULL);
+
+    return self;
+}
+
+static VALUE call_init_proc(VALUE args_val) {
+    VALUE *args = (VALUE *)args_val;
+    return rb_funcall(args[0], rb_intern("call"), 1, args[1]);
+}
+
+static void table_function_init_callback(duckdb_init_info info) {
+    VALUE self = (VALUE)duckdb_init_get_extra_info(info);
+    rubyDuckDBTableFunction *ctx;
+    VALUE init_info_obj;
+    rubyDuckDBInitInfo *init_info_ctx;
+    int state = 0;
+
+    TypedData_Get_Struct(self, rubyDuckDBTableFunction, &table_function_data_type, ctx);
+
+    // Create InitInfo wrapper
+    init_info_obj = rb_funcall(cDuckDBInitInfo, rb_intern("allocate"), 0);
+    init_info_ctx = get_struct_init_info(init_info_obj);
+    init_info_ctx->info = info;
+
+    // Call Ruby block with exception protection
+    VALUE call_args[2] = { ctx->init_proc, init_info_obj };
+    rb_protect(call_init_proc, (VALUE)call_args, &state);
+
+    if (state) {
+        VALUE err = rb_errinfo();
+        VALUE msg = rb_funcall(err, rb_intern("message"), 0);
+        duckdb_init_set_error(info, StringValueCStr(msg));
+        rb_set_errinfo(Qnil); // Clear the error
+    }
+}
+
 /*
  * call-seq:
  *   table_function.execute { |function_info, output| ... } -> table_function
@@ -304,5 +370,6 @@ void rbduckdb_init_duckdb_table_function(void) {
     rb_define_method(cDuckDBTableFunction, "add_parameter", rbduckdb_table_function_add_parameter, 1);
     rb_define_method(cDuckDBTableFunction, "add_named_parameter", rbduckdb_table_function_add_named_parameter, 2);
     rb_define_method(cDuckDBTableFunction, "bind", rbduckdb_table_function_set_bind, 0);
+    rb_define_method(cDuckDBTableFunction, "init", rbduckdb_table_function_set_init, 0);
     rb_define_method(cDuckDBTableFunction, "execute", rbduckdb_table_function_set_execute, 0);
 }

ext/duckdb/table_function.h

@@ -4,6 +4,7 @@
 struct _rubyDuckDBTableFunction {
     duckdb_table_function table_function;
     VALUE bind_proc;
+    VALUE init_proc;
     VALUE execute_proc;
 };
 

lib/duckdb.rb

@@ -16,6 +16,7 @@
 require 'duckdb/logical_type'
 require 'duckdb/scalar_function'
 require 'duckdb/bind_info'
+require 'duckdb/init_info'
 require 'duckdb/function_info'
 require 'duckdb/vector'
 require 'duckdb/data_chunk'

lib/duckdb/init_info.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+module DuckDB
+  #
+  # The DuckDB::InitInfo provides context during table function initialization.
+  #
+  # It is passed to the init callback to set up execution state.
+  #
+  # Example:
+  #
+  #   table_function.init do |init_info|
+  #     # Initialize execution state
+  #     # Can report errors if initialization fails
+  #     init_info.set_error('Initialization failed')
+  #   end
+  #
+  # rubocop:disable Lint/EmptyClass
+  class InitInfo
+    # All methods are defined in C extension (ext/duckdb/init_info.c)
+  end
+  # rubocop:enable Lint/EmptyClass
+end

test/duckdb_test/init_info_test.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+require 'test_helper'
+
+module DuckDBTest
+  class InitInfoTest < Minitest::Test
+    def setup
+      @database = DuckDB::Database.open
+      @connection = @database.connect
+    end
+
+    def teardown
+      @connection.disconnect
+      @database.close
+    end
+
+    # Test 1: Init callback setup
+    def test_init_callback
+      table_function = DuckDB::TableFunction.new
+      table_function.name = 'test_init'
+
+      result = table_function.bind do |bind_info|
+        bind_info.add_result_column('value', DuckDB::LogicalType::BIGINT)
+      end
+
+      result2 = table_function.init do |_init_info|
+        # Will be tested in integration tests (Phase 6)
+      end
+
+      assert_equal table_function, result
+      assert_equal table_function, result2
+    end
+
+    # Test 2: Init without block raises error
+    def test_init_without_block
+      table_function = DuckDB::TableFunction.new
+
+      error = assert_raises(ArgumentError) do
+        table_function.init
+      end
+
+      assert_equal 'block is required for init', error.message
+    end
+
+    # Test 3: InitInfo set_error method exists
+    def test_init_info_set_error
+      table_function = DuckDB::TableFunction.new
+      table_function.name = 'test_init_error'
+
+      result = table_function.init do |init_info|
+        # Verify set_error method exists and returns self
+        # NOTE: This block is stored but not invoked here
+        # Actual execution and error reporting will be tested in Phase 6
+        init_info.set_error('Test error')
+      end
+
+      # Verify init returns table_function for method chaining
+      assert_equal table_function, result
+    end
+  end
+end