Add option to compile the extension written in Zig

.github/workflows/ci.yml

@@ -21,5 +21,7 @@ jobs:
         with:
           enable-cache: false
       - run: zig build -Dinstall-headers --verbose --summary new
+      - run: zig build -Dlang=zig -Dinstall-headers --verbose --summary new
       - run: zig build test -Dplatform=linux_amd64_gcc4 --summary new
+      - run: zig build test -Dlang=zig -Dplatform=linux_amd64_gcc4 --summary new
       - run: tree -ash zig-out

build.zig

@@ -7,6 +7,8 @@ pub fn build(b: *Build) void {
     const platforms = b.option([]const Platform, "platform", "DuckDB platform(s) to build for (default: all)") orelse Platform.all;
     const install_headers = b.option(bool, "install-headers", "Install DuckDB C headers") orelse false;
     const flat = b.option(bool, "flat", "Install files without DuckDB version prefix") orelse false;
+    const ExtensionLanguage = enum { c, zig };
+    const lang = b.option(ExtensionLanguage, "lang", "Language to build the extension in (default: c)") orelse .c;
 
     if (flat and duckdb_versions.len > 1) {
         std.zig.fatal("-Dflat requires passing a specific DuckDB version", .{});
@@ -28,31 +30,37 @@ pub fn build(b: *Build) void {
             const platform_string = platform.toString();
             const target = platform.target(b);
 
-            const ext = b.addSharedLibrary(.{
-                .name = "quack",
+            const ext_name = "quack";
+            const ext_mod = b.createModule(.{
+                .link_libc = true,
+                .root_source_file = switch (lang) {
+                    .c => null,
+                    .zig => b.path("src/quack_extension.zig"),
+                },
                 .target = target,
                 .optimize = optimize,
             });
-            ext.addCSourceFiles(.{
-                .files = &.{
-                    "quack_extension.c",
-                },
+            ext_mod.addCSourceFiles(.{
+                .files = &.{switch (lang) {
+                    .c => "quack_extension.c",
+                    .zig => "quack_extension_zig_wrapper.c",
+                }},
                 .root = b.path("src"),
                 .flags = &cflags,
             });
-            ext.addIncludePath(duckdb_headers);
-            ext.linkLibC();
-            ext.root_module.addCMacro("DUCKDB_EXTENSION_NAME", ext.name);
-            ext.root_module.addCMacro("DUCKDB_BUILD_LOADABLE_EXTENSION", "1");
+            ext_mod.addIncludePath(duckdb_headers);
+            ext_mod.addCMacro("DUCKDB_EXTENSION_NAME", ext_name);
+            ext_mod.addCMacro("DUCKDB_BUILD_LOADABLE_EXTENSION", "1");
+            const ext = b.addSharedLibrary(.{ .name = ext_name, .root_module = ext_mod });
 
-            const filename = b.fmt("{s}.duckdb_extension", .{ext.name});
+            const filename = b.fmt("{s}.duckdb_extension", .{ext_name});
             ext.install_name = b.fmt("@rpath/{s}", .{filename}); // macOS only
 
             const ext_path = path: {
                 const cmd = Build.Step.Run.create(b, b.fmt("metadata {s} {s}", .{ version_string, platform_string }));
                 cmd.addArgs(&.{ "uv", "run", "--python=3.12" });
                 cmd.addFileArg(metadata_script);
-                cmd.addArgs(&.{ "--extension-name", ext.name });
+                cmd.addArgs(&.{ "--extension-name", ext_name });
                 cmd.addArgs(&.{ "--extension-version", ext_version });
                 cmd.addArgs(&.{ "--duckdb-platform", platform_string });
                 cmd.addArgs(&.{ "--duckdb-version", duckdb_version.extensionAPIVersion() });

src/quack_extension.zig

@@ -0,0 +1,85 @@
+const std = @import("std");
+const c = @cImport({
+    @cInclude("duckdb_extension.h");
+});
+
+const API = if (c.DUCKDB_EXTENSION_API_VERSION_MAJOR == 0 and
+    c.DUCKDB_EXTENSION_API_VERSION_MINOR == 0 and
+    c.DUCKDB_EXTENSION_API_VERSION_PATCH == 1)
+    c.duckdb_ext_api_v0
+else if (c.DUCKDB_EXTENSION_API_VERSION_MAJOR == 1 and
+    c.DUCKDB_EXTENSION_API_VERSION_MINOR == 2 and
+    c.DUCKDB_EXTENSION_API_VERSION_PATCH == 0)
+    c.duckdb_ext_api_v1
+else
+    @compileError("Unsupported DuckDB API version");
+pub var api: API = undefined;
+
+export fn init(conn: c.duckdb_connection, info: c.duckdb_extension_info, access: *c.duckdb_extension_access) bool {
+    const api_version = std.fmt.comptimePrint("v{d}.{d}.{d}", .{
+        c.DUCKDB_EXTENSION_API_VERSION_MAJOR,
+        c.DUCKDB_EXTENSION_API_VERSION_MINOR,
+        c.DUCKDB_EXTENSION_API_VERSION_PATCH,
+    });
+    const maybe_api: ?*const API = @ptrCast(@alignCast(access.get_api.?(info, api_version)));
+    api = (maybe_api orelse {
+        access.set_error.?(info, "Failed to get API");
+        return false;
+    }).*;
+
+    var func = api.duckdb_create_scalar_function.?(info, "quack").?;
+    api.duckdb_scalar_function_set_name.?(func, "quack");
+
+    var typ = api.duckdb_create_logical_type.?(c.DUCKDB_TYPE_VARCHAR).?;
+    api.duckdb_scalar_function_add_parameter.?(func, typ);
+    api.duckdb_scalar_function_set_return_type.?(func, typ);
+    api.duckdb_destroy_logical_type.?(&typ);
+
+    api.duckdb_scalar_function_set_function.?(func, quack_function);
+    if (api.duckdb_register_scalar_function.?(conn, func) == c.DuckDBError) {
+        access.set_error.?(info, "Failed to register scalar function");
+        return false;
+    }
+    api.duckdb_destroy_scalar_function.?(&func);
+    return true;
+}
+
+const quack_prefix = "Quack ";
+const quack_suffix = " 🐥";
+fn quack_function(
+    info: c.duckdb_function_info,
+    input: c.duckdb_data_chunk,
+    output: c.duckdb_vector,
+) callconv(.c) void {
+    const input_vector = api.duckdb_data_chunk_get_vector.?(input, 0);
+    const input_data: [*]c.duckdb_string_t = @alignCast(@ptrCast(api.duckdb_vector_get_data.?(input_vector)));
+    const input_mask = api.duckdb_vector_get_validity.?(input_vector);
+
+    api.duckdb_vector_ensure_validity_writable.?(output);
+    const result_mask = api.duckdb_vector_get_validity.?(output);
+
+    const num_rows = api.duckdb_data_chunk_get_size.?(input);
+    for (0..@intCast(num_rows)) |row| {
+        if (!api.duckdb_validity_row_is_valid.?(input_mask, row)) {
+            // name is NULL -> set result to NULL
+            api.duckdb_validity_set_row_invalid.?(result_mask, row);
+            continue;
+        }
+
+        var name = input_data[row];
+        const name_slice = api.duckdb_string_t_data.?(&name)[0..api.duckdb_string_t_length.?(name)];
+
+        const res_len = quack_prefix.len + name_slice.len + quack_suffix.len;
+        const res: [*]u8 = @ptrCast(api.duckdb_malloc.?(res_len) orelse {
+            api.duckdb_scalar_function_set_error.?(info, "Failed to allocate memory for result");
+            return;
+        });
+
+        @memcpy(res, quack_prefix);
+        @memcpy(res[quack_prefix.len..], name_slice);
+        @memcpy(res[quack_prefix.len + name_slice.len ..], quack_suffix);
+
+        api.duckdb_vector_assign_string_element_len.?(output, row, res, res_len);
+        api.duckdb_free.?(res);
+    }
+}

src/quack_extension_zig_wrapper.c

@@ -0,0 +1,21 @@
+#include <duckdb_extension.h>
+
+DUCKDB_EXTENSION_EXTERN
+
+// Workaround for missing struct tag in DUCKDB_EXTENSION_ENTRYPOINT (DuckDB 1.1.x)
+typedef struct duckdb_extension_access duckdb_extension_access;
+
+#if DUCKDB_EXTENSION_API_VERSION_MAJOR >= 1
+#define EXTENSION_RETURN(result) return (result)
+#else
+#define EXTENSION_RETURN(result) return
+#endif
+
+extern bool init(duckdb_connection conn, duckdb_extension_info info, duckdb_extension_access *access);
+
+DUCKDB_EXTENSION_ENTRYPOINT(duckdb_connection conn, duckdb_extension_info info, duckdb_extension_access *access) {
+    if (!init(conn, info, access)) {
+        EXTENSION_RETURN(false);
+    }
+    EXTENSION_RETURN(true);
+}