lib.mys

from os.path import Path

c"""header-before-namespace
#include <sqlite3.h>
#include <string.h>
"""

enum _Result:
    Ok = c"SQLITE_OK"
    Row = c"SQLITE_ROW"
    Done = c"SQLITE_DONE"

enum Type(u8):
    """Value type.

    """

    Integer = c"SQLITE_INTEGER"
    Float = c"SQLITE_FLOAT"
    String = c"SQLITE_TEXT"
    Null = c"SQLITE_NULL"
    Bytes = c"SQLITE_BLOB"

class SqlError(Error):
    message: string

func _to_utf8(value: string) -> bytes:
    value_utf8 = value.to_utf8()
    value_utf8 += 0

    return value_utf8

class Database:
    c"sqlite3 *database_p;"

    func __init__(self, path: Path):
        """Create or open a database.

        """

        path_utf8 = _to_utf8(str(path))
        res = _Result.Ok

        c"""
        res = sqlite3_open((const char *)path_utf8.m_bytes->data(), &this->database_p);
        """

        if res != _Result.Ok:
            raise SqlError("failed to open the database")

    func __del__(self):
        c"sqlite3_close(this->database_p);"

    func execute(self, sql: string):
        """Execute given statement.

        """

        sql_utf8 = _to_utf8(sql)
        res = _Result.Ok

        c"""
        res = sqlite3_exec(this->database_p,
                           (const char *)sql_utf8.m_bytes->data(),
                           NULL,
                           NULL,
                           NULL);
        """

        if res != _Result.Ok:
            raise SqlError(_create_error_message(self))

    func prepare(self, sql: string) -> Statement:
        """Prepare a statement. Safer than execute(), and faster if used more
        than once.

        """

        return Statement(sql, self)

func _create_error_message(database: Database) -> string:
    message: string? = None

    c"message = String(sqlite3_errmsg(database->database_p));"

    return message

class Statement:
    database: Database
    _number_of_columns: u32
    c"sqlite3_stmt *stmt_p;"

    func __init__(self, sql: string, database: Database):
        self.database = database
        self._number_of_columns = 0
        sql_utf8 = _to_utf8(sql)
        res = _Result.Ok

        c"""
        res = sqlite3_prepare(database->database_p,
                              (const char *)sql_utf8.m_bytes->data(),
                              -1,
                              &this->stmt_p,
                              NULL);
        """

        if res != _Result.Ok:
            raise SqlError(_create_error_message(database))

    func __del__(self):
        c"sqlite3_finalize(this->stmt_p);"

    func bind_int(self, column: u32, value: i64):
        """Bind given integer to given column.

        """

        res = _Result.Ok

        c"res = sqlite3_bind_int64(this->stmt_p, column, value);"

        if res != _Result.Ok:
            raise SqlError(_create_error_message(self.database))

    func bind_float(self, column: u32, value: f64):
        """Bind given float to given column.

        """

        res = _Result.Ok

        c"res = sqlite3_bind_double(this->stmt_p, column, value);"

        if res != _Result.Ok:
            raise SqlError(_create_error_message(self.database))

    func bind_string(self, column: u32, value: string):
        """Bind given string to given column.

        """

        value_utf8 = _to_utf8(value)
        res = _Result.Ok

        c"""
        res = sqlite3_bind_text(this->stmt_p,
                                column,
                                (const char *)value_utf8.m_bytes->data(),
                                -1,
                                SQLITE_TRANSIENT);
        """

        if res != _Result.Ok:
            raise SqlError(_create_error_message(self.database))

    func bind_bytes(self, column: u32, value: bytes):
        """Bind given bytes to given column.

        """

        res = _Result.Ok

        c"""
        res = sqlite3_bind_blob(this->stmt_p,
                                column,
                                (const char *)value.m_bytes->data(),
                                value.m_bytes->size(),
                                SQLITE_TRANSIENT);
        """

        if res != _Result.Ok:
            raise SqlError(_create_error_message(self.database))

    func bind_null(self, column: u32):
        """Bind null to given column.

        """

        res = _Result.Ok

        c"res = sqlite3_bind_null(this->stmt_p, column);"

        if res != _Result.Ok:
            raise SqlError(_create_error_message(self.database))

    func execute(self):
        """Execute the statement. Bind any values to columns before executing
        it. Calls reset() once complete.

        """

        result = self._step()

        try:
            if result != _Result.Done:
                message = _create_error_message(self.database)

                raise SqlError(_create_error_message(self.database))
        finally:
            self.reset()

    func fetch(self) -> bool:
        """Fetch the next row from the database. Returns True if a row was
        fetched, or calls reset() and returns False when there are no
        more rows available. If fetched, get column values with
        column_*() methods.

        """

        result = self._step()

        if result == _Result.Row:
            c"this->_number_of_columns = sqlite3_data_count(this->stmt_p); "

            return True
        else:
            self.reset()

            if result == _Result.Done:
                return False
            else:
                raise SqlError("fetch")

    func _step(self) -> _Result:
        res = 0

        c"res = sqlite3_step(this->stmt_p);"

        return _Result(res)

    func reset(self):
        """Reset the statement so it can be used again.

        """

        self._number_of_columns = 0

        c"sqlite3_reset(this->stmt_p);"

    func column_type(self, column: u32) -> Type:
        """Get the type of given column.

        """

        if column >= self._number_of_columns:
            raise SqlError(f"bad column {column}")

        ctype: u8 = 0

        c"ctype = sqlite3_column_type(this->stmt_p, column);"

        return Type(ctype)

    func column_int(self, column: u32) -> i64:
        """Get the value of given column as an integer.

        """

        if column >= self._number_of_columns:
            raise SqlError(f"bad column {column}")

        value = 0

        c"value = sqlite3_column_int64(this->stmt_p, column);"

        return value

    func column_float(self, column: u32) -> f64:
        """Get the value of given column as a float.

        """

        if column >= self._number_of_columns:
            raise SqlError(f"bad column {column}")

        value = 0.0

        c"value = sqlite3_column_double(this->stmt_p, column);"

        return value

    func column_string(self, column: u32) -> string?:
        """Get the value of given column as a string.

        """

        if column >= self._number_of_columns:
            raise SqlError(f"bad column {column}")

        value: bytes? = b""

        c"""
        const unsigned char *value_p = sqlite3_column_text(this->stmt_p, column);

        if (value_p != NULL) {
            for (size_t i = 0; i < strlen((const char *)value_p); i++) {
                value += value_p[i];
            }
        } else {
            value = nullptr;
        }
        """

        if value is None:
            return None

        return string(value)

    func column_bytes(self, column: u32) -> bytes:
        """Get the value of given column as bytes.

        """

        if column >= self._number_of_columns:
            raise SqlError(f"bad column {column}")

        value: bytes? = None

        c"""
        value = Bytes(sqlite3_column_bytes(this->stmt_p, column));
        memcpy(value.m_bytes->data(),
               sqlite3_column_blob(this->stmt_p, column),
               value.m_bytes->size());
        """

        return value

    func column_value_string(self, column: u32) -> string:
        """Get given columns value as a string.

        """

        column_type = self.column_type(column)

        if column_type == Type.Integer:
            return str(self.column_int(column))
        elif column_type == Type.Float:
            return str(self.column_float(column))
        elif column_type == Type.String:
            return f"\"{self.column_string(column)}\""
        elif column_type == Type.Bytes:
            return str(self.column_bytes(column))
        elif column_type == Type.Null:
            return "null"
        else:
            raise SqlError(f"invalid column type {column_type}")

test basics():
    Path("the.db").rm(force=True)

    database = Database(Path("the.db"))

    database.execute("CREATE TABLE tab(foo, bar, baz)")
    database.execute("INSERT INTO tab VALUES(1, 'one', null)")
    database.execute("INSERT INTO tab VALUES(2,  2.2, 'two')")
    database.execute("INSERT INTO tab VALUES(3,  'three', null)")
    database.execute("INSERT INTO tab VALUES(4,  X'89', null)")

    statement = database.prepare("SELECT * FROM tab WHERE foo >= ? ORDER BY foo")
    statement.bind_int(1, 2)

    assert statement.fetch()
    assert statement.column_type(0) == Type.Integer
    assert statement.column_int(0) == 2
    assert statement.column_type(1) == Type.Float
    assert statement.column_float(1) == 2.2
    assert statement.column_type(2) == Type.String
    assert statement.column_string(2) == "two"

    assert statement.fetch()
    assert statement.column_type(0) == Type.Integer
    assert statement.column_int(0) == 3
    assert statement.column_type(1) == Type.String
    assert statement.column_string(1) == "three"
    assert statement.column_type(2) == Type.Null

    assert statement.fetch()
    assert statement.column_type(0) == Type.Integer
    assert statement.column_int(0) == 4
    assert statement.column_type(1) == Type.Bytes
    assert statement.column_bytes(1) == b"\x89"
    assert statement.column_type(2) == Type.Null

    assert not statement.fetch()

    Path("the.db").rm(force=True)

test advanced():
    Path("the.db").rm(force=True)

    database = Database(Path("the.db"))

    database.execute("CREATE TABLE tab(foo, bar, baz)")
    database.execute("INSERT INTO tab VALUES(1, 'one', null)")
    database.execute("INSERT INTO tab VALUES(2,  2.2, 'two')")

    statement = database.prepare("INSERT INTO tab VALUES(?, ?, ?)")

    statement.bind_int(1, 3)
    statement.bind_string(2, "three")
    statement.bind_int(3, 333)
    statement.execute()

    statement.bind_int(1, 4)
    statement.bind_string(2, "four")
    statement.bind_null(3)
    statement.execute()

    statement.bind_int(1, 5)
    statement.bind_bytes(2, b"\x12\x34")
    statement.bind_null(3)
    statement.execute()

    statement = database.prepare("SELECT * FROM tab WHERE foo >= ? ORDER BY foo")
    statement.bind_int(1, 2)

    assert statement.fetch()
    assert statement.column_type(0) == Type.Integer
    assert statement.column_int(0) == 2
    assert statement.column_type(1) == Type.Float
    assert statement.column_float(1) == 2.2
    assert statement.column_type(2) == Type.String
    assert statement.column_string(2) == "two"

    assert statement.fetch()
    assert statement.column_type(0) == Type.Integer
    assert statement.column_int(0) == 3
    assert statement.column_type(1) == Type.String
    assert statement.column_string(1) == "three"
    assert statement.column_type(2) == Type.Integer
    assert statement.column_int(2) == 333

    assert statement.fetch()
    assert statement.column_type(0) == Type.Integer
    assert statement.column_int(0) == 4
    assert statement.column_type(1) == Type.String
    assert statement.column_string(1) == "four"
    assert statement.column_type(2) == Type.Null

    assert statement.fetch()
    assert statement.column_type(0) == Type.Integer
    assert statement.column_int(0) == 5
    assert statement.column_type(1) == Type.Bytes
    assert statement.column_bytes(1) == b"\x12\x34"
    assert statement.column_type(2) == Type.Null

    assert not statement.fetch()

    Path("the.db").rm(force=True)

test try_to_create_existing_table():
    Path("the.db").rm(force=True)

    database = Database(Path("the.db"))

    database.execute("CREATE TABLE tab(foo, bar, baz)")
    message: string? = None

    try:
        database.execute("CREATE TABLE tab(foo, bar, baz)")
    except SqlError as e:
        message = e.message

    assert message == "table tab already exists"

    Path("the.db").rm(force=True)

test prepare_bad_statement():
    Path("the.db").rm(force=True)

    database = Database(Path("the.db"))

    try:
        message = ""
        database.prepare("FOOBAR 123")
    except SqlError as e:
        message = e.message

    assert message == "near \"FOOBAR\": syntax error"

    Path("the.db").rm(force=True)

test bad_column():
    Path("the.db").rm(force=True)

    database = Database(Path("the.db"))

    database.execute("CREATE TABLE tab(foo, bar, baz)")
    database.execute("INSERT INTO tab VALUES(1, 'one', null)")

    statement = database.prepare("SELECT * FROM tab")

    assert statement.fetch()
    assert statement.column_type(2) == Type.Null

    try:
        message = ""
        statement.column_type(3)
    except SqlError as e:
        message = e.message

    assert message == "bad column 3"

    try:
        message = ""
        statement.column_int(4)
    except SqlError as e:
        message = e.message

    assert message == "bad column 4"

    try:
        message = ""
        statement.column_string(10)
    except SqlError as e:
        message = e.message

    assert message == "bad column 10"

    try:
        message = ""
        statement.column_float(3)
    except SqlError as e:
        message = e.message

    assert message == "bad column 3"

    try:
        message = ""
        statement.column_bytes(3)
    except SqlError as e:
        message = e.message

    assert message == "bad column 3"

    Path("the.db").rm(force=True)

test column_value_string():
    Path("the.db").rm(force=True)

    database = Database(Path("the.db"))

    database.execute("CREATE TABLE tab(foo, bar, baz, bak, bat)")
    database.execute("INSERT INTO tab VALUES(1, 'one', null, 1.0, X'012345')")

    statement = database.prepare("SELECT * FROM tab")

    assert statement.fetch()
    assert statement.column_value_string(0) == "1"
    assert statement.column_value_string(1) == "\"one\""
    assert statement.column_value_string(2) == "null"
    assert statement.column_value_string(3) == "1.000000"
    assert statement.column_value_string(4) == "b\"\\x01#E\""

test statement():
    Path("the.db").rm(force=True)

    database = Database(Path("the.db"))

    database.execute("CREATE TABLE tab(x)")
    database.execute("INSERT INTO tab VALUES(5)")
    database.execute("INSERT INTO tab VALUES(6)")
    database.execute("INSERT INTO tab VALUES(7)")

    statement = database.prepare("SELECT * FROM tab WHERE x >= ? ORDER BY x")

    # Get all.
    statement.bind_int(1, 2)
    assert statement.fetch()
    assert statement.column_int(0) == 5
    assert statement.fetch()
    assert statement.column_int(0) == 6
    assert statement.fetch()
    assert not statement.fetch()

    # Get one then reset.
    statement.bind_int(1, 2)
    assert statement.fetch()
    assert statement.column_int(0) == 5
    statement.reset()
    assert statement.fetch()
    assert statement.column_int(0) == 5

    # Bind before reset.
    try:
        ok = False
        statement.bind_int(1, 2)
    except SqlError:
        ok = True

    assert ok
    statement.reset()

    # Fetch from beginning once all rows read.
    statement.bind_int(1, 2)
    assert statement.fetch()
    assert statement.column_int(0) == 5
    assert statement.fetch()
    assert statement.fetch()
    assert not statement.fetch()
    assert statement.fetch()
    assert statement.column_int(0) == 5

    Path("the.db").rm(force=True)

test two_database_connections():
    Path("the.db").rm(force=True)

    database_1 = Database(Path("the.db"))
    database_2 = Database(Path("the.db"))

    database_1.execute("CREATE TABLE tab(x)")
    database_1.execute("INSERT INTO tab VALUES(5)")
    database_1.execute("INSERT INTO tab VALUES(6)")

    statement_1 = database_1.prepare("SELECT * FROM tab WHERE x = ?")
    statement_2 = database_2.prepare("SELECT * FROM tab WHERE x = ?")

    # Fetch interleaved.
    statement_1.bind_int(1, 5)
    statement_2.bind_int(1, 6)
    assert statement_2.fetch()
    assert statement_2.column_int(0) == 6
    assert statement_1.fetch()
    assert statement_1.column_int(0) == 5
    assert not statement_1.fetch()
    assert not statement_2.fetch()

    Path("the.db").rm(force=True)

test string():
    path = Path("the.db")
    path.rm(force=True)
    database = Database(path)
    database.execute("CREATE TABLE tab(x)")
    statement = database.prepare("INSERT INTO tab VALUES(?)")
    statement.bind_string(1, "📦")
    statement.execute()
    statement = database.prepare("SELECT * FROM tab")
    assert statement.fetch()
    assert statement.column_string(0) == "📦"