introspect fns ret empty when target doesn't exist

Matches behavior of the table_info pragma
dgllghr · Feb 7, 2024 · f0fc136 · f0fc136
1 parent 1b13440
commit f0fc136
Show file tree

Hide file tree

Showing 4 changed files with 66 additions and 17 deletions.
diff --git a/src/Table.zig b/src/Table.zig
@@ -96,10 +96,14 @@ pub fn create(
     const name = try self.table_static_arena.allocator().dupe(u8, args[2]);
     self.ctx.base = VtabCtxSchemaless.init(conn, name);
 
-    var schema_mgr = try SchemaManager.init(cb_ctx.arena, &self.ctx.base);
+    var schema_mgr = try SchemaManager.create(cb_ctx.arena, &self.ctx.base);
     defer schema_mgr.deinit();
 
-    self.ctx.schema = schema_mgr.create(&self.table_static_arena, cb_ctx.arena, &def) catch |e| {
+    self.ctx.schema = schema_mgr.createSchema(
+        &self.table_static_arena,
+        cb_ctx.arena,
+        &def,
+    ) catch |e| {
         cb_ctx.setErrorMessage("error creating schema: {any}", .{e});
         return e;
     };
@@ -203,7 +207,7 @@ pub fn connect(
     const name = try self.table_static_arena.allocator().dupe(u8, args[2]);
     self.ctx.base = VtabCtxSchemaless.init(conn, name);
 
-    var schema_mgr = try SchemaManager.init(cb_ctx.arena, &self.ctx.base);
+    var schema_mgr = try SchemaManager.open(cb_ctx.arena, &self.ctx.base);
     defer schema_mgr.deinit();
 
     self.ctx.schema = schema_mgr.load(&self.table_static_arena, cb_ctx.arena) catch |e| {

diff --git a/src/functions/SegmentInfo.zig b/src/functions/SegmentInfo.zig
@@ -159,12 +159,20 @@ pub const Cursor = struct {
         try stmt.bind(.Int64, 2, segment_id);
         const exists = try stmt.next();
         if (!exists) {
+            // Ensures that the cursor is eof
+            self.index = 3;
             return;
         }
         const col_rank = stmt.read(.Int64, false, 0);
 
-        // TODO this unnecessarily loads all columns and stores them for the duration of the cursor
-        var schema_manager = try SchemaManager.init(cb_ctx.arena, &self.vtab_ctx);
+        // Match SQLite behavior of `PRAGMA table_info()` and return empty set when the table does
+        // not exist
+        var schema_manager = SchemaManager.open(cb_ctx.arena, &self.vtab_ctx) catch |e| {
+            if (e == SchemaManager.Error.ShadowTableDoesNotExist) {
+                return;
+            }
+            return e;
+        };
         defer schema_manager.deinit();
         const schema = try schema_manager.load(&self.lifetime_arena, cb_ctx.arena);
         if (col_rank == -1) {

diff --git a/src/functions/Segments.zig b/src/functions/Segments.zig
@@ -93,6 +93,7 @@ pub const Cursor = struct {
 
     rg_idx: i64,
     col_idx: i64,
+    begun: bool,
 
     pub fn init(allocator: Allocator, conn: Conn) Cursor {
         return .{
@@ -106,12 +107,15 @@ pub const Cursor = struct {
 
             .rg_idx = 0,
             .col_idx = -1,
+            .begun = false,
         };
     }
 
     pub fn deinit(self: *Cursor) void {
-        self.rg_index_cursor.deinit();
-        self.rg_index.deinit();
+        if (self.begun) {
+            self.rg_index_cursor.deinit();
+            self.rg_index.deinit();
+        }
         self.lifetime_arena.deinit();
     }
 
@@ -127,7 +131,14 @@ pub const Cursor = struct {
         self.vtab_ctx.base.vtab_name = try self.lifetime_arena.allocator()
             .dupe(u8, table_name_ref);
 
-        var schema_manager = try SchemaManager.init(cb_ctx.arena, &self.vtab_ctx.base);
+        // Match SQLite behavior of `PRAGMA table_info()` and return empty set when the table does
+        // not exist
+        var schema_manager = SchemaManager.open(cb_ctx.arena, &self.vtab_ctx.base) catch |e| {
+            if (e == SchemaManager.Error.ShadowTableDoesNotExist) {
+                return;
+            }
+            return e;
+        };
         defer schema_manager.deinit();
         self.vtab_ctx.schema = try schema_manager.load(&self.lifetime_arena, cb_ctx.arena);
 
@@ -145,10 +156,13 @@ pub const Cursor = struct {
             };
             self.rg_index_cursor.readEntry(&self.rg_entry);
         }
+
+        self.begun = true;
     }
 
     pub fn eof(self: *Cursor) bool {
-        return self.rg_index_cursor.eof();
+        // `begun` is false if the table does not exist and `begin` exited early
+        return !self.begun or self.rg_index_cursor.eof();
     }
 
     pub fn next(self: *Cursor) !void {

diff --git a/src/schema/Manager.zig b/src/schema/Manager.zig
@@ -26,21 +26,20 @@ const Self = @This();
 
 const StmtCell = prep_stmt.Cell(VtabCtxSchemaless);
 
-pub const Error = error{ SortKeyColumnNotFound, ExecReturnedData };
-
-pub fn init(tmp_arena: *ArenaAllocator, ctx: *const VtabCtxSchemaless) !Self {
-    try setup(tmp_arena, ctx.*);
+pub const Error = error{ SortKeyColumnNotFound, ExecReturnedData, ShadowTableDoesNotExist };
 
+pub fn create(tmp_arena: *ArenaAllocator, ctx: *const VtabCtxSchemaless) !Self {
+    try createTable(tmp_arena, ctx.*);
     return .{
         .ctx = ctx,
         .load_columns = StmtCell.init(&loadColumnsQuery),
         .create_column = StmtCell.init(&createColumnDml),
     };
 }
 
-fn setup(tmp_arena: *ArenaAllocator, ctx: VtabCtxSchemaless) !void {
+fn createTable(tmp_arena: *ArenaAllocator, ctx: VtabCtxSchemaless) !void {
     const query = try fmt.allocPrintZ(tmp_arena.allocator(),
-        \\CREATE TABLE IF NOT EXISTS "{s}_columns" (
+        \\CREATE TABLE "{s}_columns" (
         \\  rank INTEGER NOT NULL,
         \\  name TEXT NOT NULL COLLATE NOCASE,
         \\  column_type TEXT NOT NULL,
@@ -52,11 +51,35 @@ fn setup(tmp_arena: *ArenaAllocator, ctx: VtabCtxSchemaless) !void {
     try ctx.conn().exec(query);
 }
 
+pub fn open(tmp_arena: *ArenaAllocator, ctx: *const VtabCtxSchemaless) !Self {
+    try verifyExists(ctx.*, tmp_arena);
+    return .{
+        .ctx = ctx,
+        .load_columns = StmtCell.init(&loadColumnsQuery),
+        .create_column = StmtCell.init(&createColumnDml),
+    };
+}
+
 pub fn deinit(self: *Self) void {
     self.load_columns.deinit();
     self.create_column.deinit();
 }
 
+pub fn verifyExists(ctx: VtabCtxSchemaless, tmp_arena: *ArenaAllocator) !void {
+    const query = try fmt.allocPrintZ(tmp_arena.allocator(),
+        \\SELECT 1
+        \\FROM sqlite_master
+        \\WHERE type = 'table' AND name = '{s}_columns'
+    , .{ctx.vtabName()});
+
+    const stmt = try ctx.conn().prepare(query);
+
+    const exists = try stmt.next();
+    if (!exists) {
+        return Error.ShadowTableDoesNotExist;
+    }
+}
+
 pub fn destroy(self: *Self, tmp_arena: *ArenaAllocator) !void {
     const query = try fmt.allocPrintZ(tmp_arena.allocator(),
         \\DROP TABLE "{s}_columns"
@@ -113,7 +136,7 @@ fn readColumn(allocator: Allocator, stmt: Stmt) !Column {
     };
 }
 
-pub fn create(
+pub fn createSchema(
     self: *Self,
     table_static_arena: *ArenaAllocator,
     tmp_arena: *ArenaAllocator,
@@ -212,7 +235,7 @@ test "schema: create and read columns" {
         .conn_ = conn,
         .vtab_name = "test",
     };
-    var mgr = try Self.init(&arena, &ctx);
+    var mgr = try Self.create(&arena, &ctx);
     inline for (&columns) |*col| {
         try mgr.createColumn(&arena, col);
     }