fediglam/src/sql/lib.zig

const std = @import("std");
const util = @import("util");

const postgres = @import("./postgres.zig");
const sqlite = @import("./sqlite.zig");
const common = @import("./common.zig");
const Allocator = std.mem.Allocator;

pub const QueryOptions = common.QueryOptions;

pub const Engine = enum {
    postgres,
    sqlite,
};

pub const Config = union(Engine) {
    postgres: struct {
        pg_conn_str: [:0]const u8,
    },
    sqlite: struct {
        sqlite_file_path: [:0]const u8,
    },
};

pub const QueryError = error{
    OutOfMemory,
    ConnectionLost,
};

pub fn fieldList(comptime RowType: type) []const u8 {
    comptime {
        const fields = std.meta.fieldNames(RowType);
        const separator = ", ";

        if (fields.len == 0) return "";

        var size: usize = 1; // 1 for null terminator
        for (fields) |f| size += f.len + separator.len;
        size -= separator.len;

        var buf = std.mem.zeroes([size]u8);

        // can't use std.mem.join because of problems with comptime allocation
        // https://github.com/ziglang/zig/issues/5873#issuecomment-1001778218
        //var fba = std.heap.FixedBufferAllocator.init(&buf);
        //return (std.mem.join(fba.allocator(), separator, fields) catch unreachable) ++ " ";

        var buf_idx = 0;
        for (fields) |f, i| {
            std.mem.copy(u8, buf[buf_idx..], f);
            buf_idx += f.len;
            if (i != fields.len - 1) std.mem.copy(u8, buf[buf_idx..], separator);
            buf_idx += separator.len;
        }

        return &buf;
    }
}

//pub const OpenError = sqlite.OpenError | postgres.OpenError;
const RawResults = union(Engine) {
    postgres: postgres.Results,
    sqlite: sqlite.Results,

    fn finish(self: RawResults) void {
        switch (self) {
            .postgres => |pg| pg.finish(),
            .sqlite => |lite| lite.finish(),
        }
    }

    fn columnCount(self: RawResults) u15 {
        return switch (self) {
            .postgres => |pg| pg.columnCount(),
            .sqlite => |lite| lite.columnCount(),
        };
    }

    fn columnNameToIndex(self: RawResults, name: []const u8) !u15 {
        return try switch (self) {
            .postgres => |pg| pg.columnNameToIndex(name),
            .sqlite => |lite| lite.columnNameToIndex(name),
        };
    }

    fn row(self: *RawResults) !?Row {
        return switch (self.*) {
            .postgres => |*pg| if (try pg.row()) |r| Row{ .postgres = r } else null,
            .sqlite => |*lite| if (try lite.row()) |r| Row{ .sqlite = r } else null,
        };
    }
};

// Represents a set of results.
// row() must be called until it returns null, or the query may not complete
// Must be deallocated by a call to finish()
pub fn Results(comptime T: type) type {
    // would normally make this a declaration of the struct, but it causes the compiler to crash
    const fields = std.meta.fields(T);
    return struct {
        const Self = @This();

        underlying: RawResults,
        column_indices: [fields.len]u15,

        fn from(underlying: RawResults) !Self {
            return Self{ .underlying = underlying, .column_indices = blk: {
                var indices: [fields.len]u15 = undefined;
                inline for (fields) |f, i| {
                    indices[i] = if (!std.meta.trait.isTuple(T)) try underlying.columnNameToIndex(f.name) else i;
                }
                break :blk indices;
            } };
        }

        pub fn finish(self: Self) void {
            self.underlying.finish();
        }

        // can be used as an optimization to reduce memory reallocation
        // only works on postgres
        pub fn rowCount(self: Self) ?usize {
            return self.underlying.rowCount();
        }

        // Returns the next row of results, or null if there are no more rows.
        // Caller owns all memory allocated. The entire object can be deallocated with a
        // call to util.deepFree
        pub fn row(self: *Self, alloc: ?Allocator) !?T {
            if (try self.underlying.row()) |row_val| {
                var result: T = undefined;
                var fields_allocated: usize = 0;
                errdefer inline for (fields) |f, i| {
                    // Iteration bounds must be defined at comptime (inline for) but the number of fields we could
                    // successfully allocate is defined at runtime. So we iterate over the entire field array and
                    // conditionally deallocate fields in the loop.
                    if (i < fields_allocated) util.deepFree(alloc, @field(result, f.name));
                };

                inline for (fields) |f, i| {
                    @field(result, f.name) = try row_val.get(f.field_type, self.column_indices[i], alloc);
                    fields_allocated += 1;
                }

                return result;
            } else return null;
        }
    };
}

// Row is invalidated by the next call to result.row()
const Row = union(Engine) {
    postgres: postgres.Row,
    sqlite: sqlite.Row,

    // Returns a value of type T from the zero-indexed column given by idx.
    // Not all types require an allocator to be present. If an allocator is needed but
    // not required, it will return error.AllocatorRequired.
    // The caller is responsible for deallocating T, if relevant.
    fn get(self: Row, comptime T: type, idx: u15, alloc: ?Allocator) anyerror!T {
        return switch (self) {
            .postgres => |pg| pg.get(T, idx, alloc),
            .sqlite => |lite| lite.get(T, idx, alloc),
        };
    }
};

const DbUnion = union(Engine) {
    postgres: postgres.Db,
    sqlite: sqlite.Db,
};

pub const ConstraintMode = enum {
    deferred,
    immediate,
};

pub const Db = struct {
    tx_open: bool = false,
    underlying: DbUnion,

    pub fn open(cfg: Config) !Db {
        return switch (cfg) {
            .postgres => |postgres_cfg| Db{
                .underlying = .{
                    .postgres = try postgres.Db.open(postgres_cfg.pg_conn_str),
                },
            },
            .sqlite => |lite_cfg| Db{
                .underlying = .{
                    .sqlite = try sqlite.Db.open(lite_cfg.sqlite_file_path),
                },
            },
        };
    }

    pub fn close(self: *Db) void {
        switch (self.underlying) {
            .postgres => |pg| pg.close(),
            .sqlite => |lite| lite.close(),
        }
    }

    pub fn queryWithOptions(
        self: *Db,
        comptime RowType: type,
        sql: [:0]const u8,
        args: anytype,
        opt: QueryOptions,
    ) !Results(RowType) {
        if (self.tx_open) return error.TransactionOpen;
        // Create fake transaction to use its functions
        return (Tx{ .db = self }).queryWithOptions(RowType, sql, args, opt);
    }

    pub fn query(
        self: *Db,
        comptime RowType: type,
        sql: [:0]const u8,
        args: anytype,
        alloc: ?Allocator,
    ) !Results(RowType) {
        if (self.tx_open) return error.TransactionOpen;
        // Create fake transaction to use its functions
        return (Tx{ .db = self }).query(RowType, sql, args, alloc);
    }

    pub fn exec(
        self: *Db,
        sql: [:0]const u8,
        args: anytype,
        alloc: ?Allocator,
    ) !void {
        if (self.tx_open) return error.TransactionOpen;
        // Create fake transaction to use its functions
        return (Tx{ .db = self }).exec(sql, args, alloc);
    }

    pub fn queryRow(
        self: *Db,
        comptime RowType: type,
        sql: [:0]const u8,
        args: anytype,
        alloc: ?Allocator,
    ) !?RowType {
        if (self.tx_open) return error.TransactionOpen;
        // Create fake transaction to use its functions
        return (Tx{ .db = self }).queryRow(RowType, sql, args, alloc);
    }

    pub fn insert(
        self: *Db,
        comptime table: []const u8,
        value: anytype,
    ) !void {
        if (self.tx_open) return error.TransactionOpen;
        // Create fake transaction to use its functions
        return (Tx{ .db = self }).insert(table, value);
    }

    pub fn sqlEngine(self: *Db) Engine {
        return self.underlying;
    }

    // Begins a transaction
    pub fn begin(self: *Db) !Tx {
        const tx = Tx{ .db = self };
        try tx.exec("BEGIN", .{}, null);

        return tx;
    }
};

pub const Tx = struct {
    db: *Db,

    // internal helper fn
    fn queryInternal(
        self: Tx,
        sql: [:0]const u8,
        args: anytype,
        opt: QueryOptions,
    ) !RawResults {
        return switch (self.db.underlying) {
            .postgres => |pg| RawResults{ .postgres = try pg.exec(sql, args, opt.prep_allocator) },
            .sqlite => |lite| RawResults{ .sqlite = try lite.exec(sql, args, opt) },
        };
    }

    pub fn queryWithOptions(
        self: Tx,
        comptime RowType: type,
        sql: [:0]const u8,
        args: anytype,
        options: QueryOptions,
    ) !Results(RowType) {
        return Results(RowType).from(try self.queryInternal(sql, args, options));
    }

    // Executes a query and returns the result set
    pub fn query(
        self: Tx,
        comptime RowType: type,
        sql: [:0]const u8,
        args: anytype,
        alloc: ?Allocator,
    ) !Results(RowType) {
        return self.queryWithOptions(RowType, sql, args, .{ .prep_allocator = alloc });
    }

    // Executes a query without returning results
    pub fn exec(
        self: Tx,
        sql: [:0]const u8,
        args: anytype,
        alloc: ?Allocator,
    ) !void {
        _ = try self.queryRow(std.meta.Tuple(&.{}), sql, args, alloc);
    }

    // Runs a query and returns a single row
    pub fn queryRow(
        self: Tx,
        comptime RowType: type,
        q: [:0]const u8,
        args: anytype,
        alloc: ?Allocator,
    ) !?RowType {
        var results = try self.query(RowType, q, args, alloc);
        defer results.finish();
        @compileLog(args);

        const row = (try results.row(alloc)) orelse return null;
        errdefer util.deepFree(alloc, row);

        var more_rows = false;
        while (try results.row(alloc)) |r| {
            util.deepFree(alloc, r);
            more_rows = true;
        }
        if (more_rows) return error.TooManyRows;

        return row;
    }

    // Inserts a single value into a table
    pub fn insert(
        self: Tx,
        comptime table: []const u8,
        value: anytype,
    ) !void {
        const ValueType = comptime @TypeOf(value);

        const fields = std.meta.fields(ValueType);
        comptime var types: [fields.len]type = undefined;
        comptime var table_spec: []const u8 = table ++ "(";
        comptime var value_spec: []const u8 = "(";
        inline for (fields) |field, i| {
            types[i] = field.field_type;
            table_spec = comptime (table_spec ++ field.name ++ ",");
            value_spec = comptime value_spec ++ std.fmt.comptimePrint("${},", .{i + 1});
        }
        table_spec = comptime table_spec[0 .. table_spec.len - 1] ++ ")";
        value_spec = comptime value_spec[0 .. value_spec.len - 1] ++ ")";
        const q = comptime std.fmt.comptimePrint(
            "INSERT INTO {s} VALUES {s}",
            .{ table_spec, value_spec },
        );

        var args_tuple: std.meta.Tuple(&types) = undefined;
        inline for (fields) |field, i| {
            args_tuple[i] = @field(value, field.name);
        }
        try self.exec(q, args_tuple, null);
    }

    pub fn sqlEngine(self: Tx) Engine {
        return self.db.underlying;
    }

    pub fn setConstraintMode(self: Tx, mode: ConstraintMode) !void {
        switch (self.db.underlying) {
            .sqlite => try self.exec(
                switch (mode) {
                    .immediate => "PRAGMA defer_foreign_keys = FALSE",
                    .deferred => "PRAGMA defer_foreign_keys = TRUE",
                },
                .{},
                null,
            ),
            .postgres => try self.exec(
                switch (mode) {
                    .immediate => "SET CONSTRAINTS ALL IMMEDIATE",
                    .deferred => "SET CONSTRAINTS ALL DEFERRED",
                },
                .{},
                null,
            ),
        }
    }

    pub fn rollback(self: Tx) void {
        self.exec("ROLLBACK", .{}, null) catch |err| {
            std.log.err("Error occured during rollback operation: {}", .{err});
        };
    }

    pub fn commit(self: Tx) !void {
        try self.exec("COMMIT", .{}, null);
    }
};