From 18a0915486db7af2af6599b831d188b5c39344ab Mon Sep 17 00:00:00 2001
From: Dan Gallagher <d@gllghr.net>
Date: Wed, 14 Feb 2024 08:54:54 -0500
Subject: [PATCH] register as persistent and auto extension

---
 README.md    |  4 +++-
 src/main.zig | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index fe35a69..56da897 100644
--- a/README.md
+++ b/README.md
@@ -65,10 +65,12 @@ The SQLite extension is the dynamic library named `libstanchion` in the `zig-out
 
 ### Load Stanchion
 
-Stanchion is a [Run-Time Loadable Extension](https://sqlite.org/loadext.html) that uses SQLite's virtual table system. Currently, stanchion must be loaded by all connections that access any stanchion virtual tables. This may change in the future when stanchion supports being a [Persistent Loadable Extension](https://sqlite.org/loadext.html#persistent_loadable_extensions). To load an extension from the SQLite CLI, use the `.load` command. Check the documentation of the SQLite bindings you are using to see how to load an extension in your application. Here are some examples for different language bindings: [`sqlite3` for Python](https://docs.python.org/3/library/sqlite3.html#sqlite3.Connection.load_extension), [`rusqlite` for Rust](https://docs.rs/rusqlite/latest/rusqlite/struct.Connection.html#method.load_extension), [`sqlite3` for Ruby](https://rubydoc.info/gems/sqlite3/SQLite3/Database#load_extension-instance_method), and [`go-sqlite3` for Go](https://pkg.go.dev/github.com/mattn/go-sqlite3#SQLiteConn.LoadExtension).
+Stanchion is a [Run-Time Loadable Extension](https://sqlite.org/loadext.html) that uses SQLite's virtual table system. To load an extension from the SQLite CLI, use the `.load` command. Check the documentation of the SQLite bindings you are using to see how to load an extension in your application. Here are some examples for different language bindings: [`sqlite3` for Python](https://docs.python.org/3/library/sqlite3.html#sqlite3.Connection.load_extension), [`rusqlite` for Rust](https://docs.rs/rusqlite/latest/rusqlite/struct.Connection.html#method.load_extension), [`sqlite3` for Ruby](https://rubydoc.info/gems/sqlite3/SQLite3/Database#load_extension-instance_method), and [`go-sqlite3` for Go](https://pkg.go.dev/github.com/mattn/go-sqlite3#SQLiteConn.LoadExtension).
 
 Before loading stanchion (or any extension), you may first need to enable extension loading. Here are some examples for different language bindings: [`sqlite3` for Python](https://docs.python.org/3/library/sqlite3.html#sqlite3.Connection.enable_load_extension), [`rusqlite` for Rust](https://docs.rs/rusqlite/latest/rusqlite/struct.LoadExtensionGuard.html), and [`sqlite3` for Ruby](https://rubydoc.info/gems/sqlite3/SQLite3/Database#enable_load_extension-instance_method). Some bindings enable extension loading by default (e.g. [`go-sqlite3` for Go](https://github.com/mattn/go-sqlite3#feature--extension-list)). For more information, see the [SQLite documentation for the C API](https://sqlite.org/c3ref/enable_load_extension.html).
 
+Stanchion is both a [persistent](https://sqlite.org/loadext.html#persistent_loadable_extensions) and [auto](https://sqlite.org/c3ref/auto_extension.html) extension. It only needs to be loaded by one connection in a process and then it will automatically be loaded by all connections in the same process. Loading it from the other connections anyway is harmless. Connections in other processes connecting to the same SQLite database still need to load stanchion.
+
 ### Create table
 
 Creating a stanchion table works much like creating any table in SQLite:
diff --git a/src/main.zig b/src/main.zig
index 05957db..3405830 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -1,4 +1,5 @@
 const std = @import("std");
+const log = std.log;
 const Allocator = std.mem.Allocator;
 const GeneralPurposeAllocator = std.heap.GeneralPurposeAllocator;
 
@@ -35,6 +36,30 @@ pub export fn sqlite3_stanchion_init(
         return c.SQLITE_ERROR;
     }
 
+    // Register virtual tables and functions for this connection
+    var res = register(db, err_msg, api);
+    if (res != c.SQLITE_OK) {
+        return res;
+    }
+
+    // Automatically register virtual tables and functions on all future connections to this db
+    res = c.sqlite3_auto_extension(
+        @as(?*const fn () callconv(.C) void, @ptrCast(&register)),
+    );
+    if (res != c.SQLITE_OK) {
+        return res;
+    }
+
+    // Must be a persistent extension to be an auto extension
+    return c.SQLITE_OK_LOAD_PERMANENTLY;
+}
+
+fn register(
+    db: *c.sqlite3,
+    err_msg: [*c][*c]u8,
+    api: *c.sqlite3_api_routines,
+) callconv(.C) c_int {
+    _ = api;
     // To store data common to all table instances (global to the module), replace this allocator
     // with a struct containing the common data (see `ModuleContext` in `zig-sqlite`)
     allocator = GeneralPurposeAllocator(.{}){};
@@ -80,6 +105,16 @@ pub export fn sqlite3_stanchion_init(
 
 fn setErrorMsg(err_msg: [*c][*c]u8, msg_text: []const u8) void {
     const msg_buf: [*c]u8 = @ptrCast(c.sqlite3_malloc(@intCast(msg_text.len)));
+
+    // sqlite3_malloc returns a null pointer if it can't allocate the memory
+    if (msg_buf == null) {
+        log.err(
+            "failed to alloc diagnostic message in sqlite: out of memory. original message: {s}",
+            .{msg_text},
+        );
+        return;
+    }
+
     @memcpy(msg_buf[0..msg_text.len], msg_text);
     err_msg.* = msg_buf;
 }