diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index fae5360..b0fcaff 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/build.zig b/build.zig index 39bf31b..34a862d 100644 --- a/build.zig +++ b/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() }); diff --git a/src/quack_extension.zig b/src/quack_extension.zig new file mode 100644 index 0000000..7a48e79 --- /dev/null +++ b/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); + } +} diff --git a/src/quack_extension_zig_wrapper.c b/src/quack_extension_zig_wrapper.c new file mode 100644 index 0000000..8cc86dd --- /dev/null +++ b/src/quack_extension_zig_wrapper.c @@ -0,0 +1,21 @@ +#include + +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); +}