sorm/sorm.py

import ctypes
import csv
from csv import DictReader
import io
import tempfile
import time


class DictReader:

    def get_column_types(self):
        types = []
        for column in self.columns:
            name = column.split("(")[0]
            type = column.split("(")[1]
            type = type.split(")")[0]
            if type == "integer":
                types.append(int)
            if type == "text":
                types.append(str)
        return types

    def get_column_names(self):
        names = []
        for column in self.columns:
            name = column.split("(")[0]
            names.append(name)
        return names

    def __init__(self, data):
        self.result = 0
        if type(data) == int:
            self.data = ""
        else:
            self.data = data.decode()
        self.rows = [row.split(";")[:-1] for row in self.data.split("\n")]

        self.columns = self.rows[0]
        self.rows.pop(0)
        if type(data) == str:
            self.result = len(self.rows)
        else:
            self.result = data
        self.column_types = self.get_column_types()
        for row in self.rows:
            for index, field in enumerate(row):
                row[index] = self.column_types[index](field)
        self.column_names = self.get_column_names()

    def __iter__(self):
        return self.rows.__iter__()


libc = ctypes.CDLL("libc.so.6")


class Sorm:

    def __init__(self):
        self.lib = ctypes.CDLL("./sorm.so")
        self.sormq = self.lib.sormq
        self.sormq.argtypes = [ctypes.c_int, ctypes.c_char_p]
        self.sormq.restype = ctypes.c_char_p

        self.sormc = self.lib.sormc
        self.sormc.argtypes = [ctypes.c_char_p]
        self.sormc.restype = ctypes.c_int

        self.sormd = self.lib.sormd
        self.sormd.argtypes = [ctypes.c_int]
        self.sormd.restype = None

        self.sormm = self.lib.sormm
        self.sormm.argtypes = [ctypes.c_int]
        self.sormm.restype = ctypes.c_char_p


class SormDb(Sorm):

    def __init__(self, path):
        super().__init__()
        self.path = path
        self.conn = None

    def c(self):
        if not self.conn:
            self.conn = self.sormc(self.path.encode())
        return self.conn

    def __enter__(self):
        self.c()
        return self

    def __exit__(self, *args, **kwargs):
        self.d()

    def q(self, sql, *args) -> DictReader:
        ctypes_list = []
        for arg in args:
            if type(arg) == int:
                ctypes_list.append(ctypes.c_int)
            if type(arg) == str:
                ctypes_list.append(ctypes.c_char_p)
        self.sormq.argtypes = [ctypes.c_int, ctypes.c_char_p] + ctypes_list
        if not sql.lower().startswith("select"):
            self.sormq.restype = ctypes.c_int
        else:
            self.sormq.restype = ctypes.c_char_p
        params = tuple(
            [self.conn, sql.encode()]
            + list(arg.encode() if type(arg) == str else arg for arg in args)
        )
        result = DictReader(self.sormq(*params))
        self.m = self.sormm(self.conn).decode()
        return result

    def d(self):
        if not self.conn:
            return
        self.sormd(self.conn)
        self.conn = None


# Load the shared library
lib = ctypes.CDLL("./sorm.so")

free = libc.free
free.argtypes = [ctypes.c_void_p]
free.restype = None


rsomm = lib.sormm
rsomm.argtypes = [ctypes.c_int]
rsomm.restype = ctypes.c_char_p

disconnect = lib.sormd
disconnect.argtypes = [ctypes.c_int]

start = time.time()
for x in range(1):

    with SormDb("db.sqlite3") as db:

        for x in range(1):
            # db.q("BEGIN TRANSACTION")
            # for x in range(100000):
            #    db.q("INSERT INTO pony (name,age) VALUES (?s,?d);","Python Pony",1337)
            # db.q("COMMIT")
            result = db.q(
                "SELECT * FROM pony WHERE id > ?d AND name like ?s ORDER BY id",
                1337,
                "%hon Pon%",
            )
            # for row in result:
            #    print(row)
            print(result.column_names)
            print(len(result.rows), "records")
            print(db.m)
end = time.time()
duration = end - start
print("Duration: {}s".format(duration))