dataclassdb 0.1.0__py3-none-any.whl

dataclassdb/dataclass_field_codec.py

@@ -0,0 +1,216 @@
+import dataclasses
+import json
+import logging
+import pickle
+from datetime import datetime, timezone
+from typing import Any, TypeVar, Union, get_args, get_origin
+
+import dacite
+
+from dataclassdb.dataclass_types import Codec, IsDataclass
+from dataclassdb.utils import is_enum_class, is_strenum
+
+logger = logging.getLogger(__name__)
+
+DataclassT = TypeVar("DataclassT", bound=IsDataclass)
+T = TypeVar("T")
+def create_field_codec(field_class: T, sql_type: str) -> Codec:
+    if origin := get_origin(field_class):
+        if origin is Union:
+            field_class = get_args(field_class)[0]
+        else:
+            field_class = origin
+
+    if sql_type == "BLOB":
+        if dataclasses.is_dataclass(field_class):
+            return DataclassPickleCodec(field_class)
+        else:
+            return pickle_codec
+    elif dataclasses.is_dataclass(field_class):
+        return DataclassJsonCodec(field_class)
+    elif is_strenum(field_class):
+        if sql_type != "TEXT":
+            raise ValueError("The SQL type for StrEnum must be TEXT or BLOB")
+        return StrEnumCodec(field_class=field_class, sql_type=sql_type)
+    elif is_enum_class(field_class):
+        if sql_type != "INTEGER":
+            raise ValueError("The SQL type for Enums must be INTEGER or BLOB")
+        return EnumCodec(field_class=field_class, sql_type=sql_type)
+    elif issubclass(field_class, datetime):
+        if sql_type == "INTEGER":
+            return datetime_int_codec
+        else:
+            return datetime_text_codec
+
+    elif sql_type == "TEXT":
+        if field_class is str:
+            return identity_codec
+        else:
+            return json_codec
+    elif field_class is bool:
+        if sql_type != "INTEGER":
+            raise ValueError("bool fields must be BLOB, TEXT, or INTEGER")
+        return bool_codec
+    return identity_codec
+
+
+class DataclassJsonCodec(Codec):
+    def __init__(self, field_class, *args, **kwargs):
+        self.field_class = field_class
+
+    def encode(self, data: any) -> str:
+        if data is None:
+            return ""
+        data_dict = None
+        if isinstance(data, dict):
+            data_dict = data
+        else:
+            data_dict = dataclasses.asdict(data)
+        return json.dumps(data_dict)
+
+    def decode(self, data: str | dict):
+        if data is None or data == "":
+            return None
+        if isinstance(data, dict):
+            data_dict = data
+        else:
+            data_dict = json.loads(data)
+        return dacite.from_dict(self.field_class, data_dict)
+
+
+class DataclassPickleCodec(Codec):
+    def __init__(self, field_class, *args, **kwargs):
+        self.field_class = field_class
+
+    def encode(self, data):
+        if data is None:
+            return None
+        data_dict = None
+        if isinstance(data, dict):
+            data_dict = data
+        else:
+            data_dict = dataclasses.asdict(data)
+        return pickle.dumps(data_dict)
+
+    def decode(self, data):
+        if data is None:
+            return None
+        data_dict = None
+        if isinstance(data, dict):
+            data_dict = data
+        else:
+            data_dict = pickle.loads(data)
+        return dacite.from_dict(self.field_class, data_dict)
+
+
+class StrEnumCodec(Codec):
+    def __init__(self, field_class: T, *args, **kwargs):
+        self.field_class = field_class
+
+    def encode(self, data: T) -> str | int:
+        if data is None:
+            return None
+        return data.value
+
+    def decode(self, data: str | int) -> T:
+        if data is None:
+            return None
+        return self.field_class(data)
+
+
+class EnumCodec(Codec):
+    def __init__(self, field_class: T, *args, **kwargs):
+        self.field_class = field_class
+
+    def encode(self, data: T) -> str | int:
+        if data is None:
+            return None
+        return data.value
+
+    def decode(self, data: str | int) -> T:
+        if data is None:
+            return None
+        return self.field_class(data)
+
+
+class JsonCodec(Codec):
+    def encode(self, data: any) -> str:
+        if data is None:
+            return None
+        if dataclasses.is_dataclass(data):
+            return json.dumps(dataclasses.asdict(data))
+        return json.dumps(data)
+
+    def decode(self, data: str) -> any:
+        if data is None:
+            return None
+        return json.loads(data)
+
+
+class PickleCodec(Codec):
+    def encode(self, data: Any) -> bytes:
+        if data is None:
+            return None
+        return pickle.dumps(data)
+
+    def decode(self, data: bytes) -> Any:
+        if data is None:
+            return None
+        return pickle.loads(data)
+
+
+class IdentityCodec(Codec):
+    def encode(self, data: T) -> T:
+        return data
+
+    def decode(self, data: T) -> T:
+        return data
+
+
+class BoolCodec(Codec):
+    def encode(self, data: bool) -> int:
+        if data is None:
+            return None
+        return 1 if data is True else 0
+
+    def decode(self, data: int) -> bool:
+        if data is None:
+            return None
+        return data >= 1
+
+
+class DatetimeIntCodec(Codec):
+    """Converts between python datetime(utc) to sqlite3 INTEGER"""
+
+    def encode(self, data: datetime) -> int:
+        if data is None:
+            return None
+        return int(data.timestamp())
+
+    def decode(self, data: int) -> datetime:
+        if data is None:
+            return None
+        return datetime.fromtimestamp(timestamp=data, tz=timezone.utc)
+
+
+class DatetimeTextCodec(Codec):
+    """Converts between python datetime to sqlite3 TEXT(ISO 8601)"""
+
+    def encode(self, data: datetime) -> int:
+        if data is None:
+            return None
+        return data.isoformat()
+
+    def decode(self, data: int) -> datetime:
+        if data is None:
+            return None
+        return datetime.fromisoformat(data)
+
+
+# Reusable codecs:
+bool_codec = BoolCodec()
+identity_codec = IdentityCodec()
+json_codec = JsonCodec()
+pickle_codec = PickleCodec()
+datetime_int_codec = DatetimeIntCodec()
+datetime_text_codec = DatetimeTextCodec()

dataclassdb/dataclass_sqlite_field.py

@@ -0,0 +1,193 @@
+import inspect
+import logging
+from dataclasses import MISSING, Field, is_dataclass
+from datetime import datetime
+from enum import Enum, StrEnum
+from types import NoneType, UnionType
+from typing import Annotated, Union, get_args, get_origin
+
+from dataclassdb.builders.query_builder import QueryBuilder
+from dataclassdb.dataclass_field_codec import create_field_codec
+from dataclassdb.dataclass_types import Codec
+
+logger = logging.getLogger(__name__)
+
+
+class DataclassSqliteField:
+    """This class reads annotations and determines the provided type, sqlite constraints, and sqlite type
+    It also provides an encoder/decoder to and from sqlite.
+    """
+
+    _type_map = {
+        int: "INTEGER",
+        StrEnum: "TEXT",
+        Enum: "INTEGER",
+        str: "TEXT",
+        float: "REAL",
+        bool: "INTEGER",
+        list: "TEXT",
+        dict: "TEXT",
+        datetime: "TEXT",
+    }
+
+    @property
+    def unique(self) -> bool:
+        return "UNIQUE" in self.constraints
+
+    @property
+    def notnull(self) -> bool:
+        return "NOT" in self.constraints
+
+    @property
+    def primary_key(self) -> bool:
+        return "PRIMARY" in self.constraints
+
+    def encode(self, value):
+        return self.codec.encode(value)
+
+    def decode(self, value):
+        return self.codec.decode(value)
+
+    def __init__(self, dc_type: type, dc_field: Field, default_sql_type="TEXT"):
+        self.name = dc_field.name
+        self.type = None
+        self.sql_type = default_sql_type
+        self.sql_create_type = ""
+        self.constraints = {}
+        self.codec: Codec = None
+        self.type_override = False
+
+        # Handle field outer wrapper
+        if is_dataclass(dc_type):
+            field_type = dc_type
+        if get_origin(dc_type) is Annotated:
+            field_type = self.parse_annotated(dc_type)
+        else:
+            field_type = dc_type
+
+        self.parse_type(field_type)
+        if self.codec is None:
+            self.codec = create_field_codec(field_type, self.sql_type)
+
+        self.parse_default(dc_field)
+
+    def parse_annotated(self, dataclass_type: type):
+        args = get_args(dataclass_type)
+        field_type = args[0]
+        field_args = args[1:]
+        for i, arg in enumerate(field_args):
+            if isinstance(arg, Codec):
+                logger.info("Overloading codec for field %s", self.name)
+                self.codec = arg
+                continue
+            elif not isinstance(arg, str):
+                logger.warning(
+                    "Skipping constraint field %i constraint field for %s",
+                    i,
+                    self.name,
+                )
+                continue
+            split_arg = arg.split(" ")
+            key = split_arg[0]
+            if key == "CONSTRAINT":
+                key = split_arg[1]
+                arg = " ".join(split_arg[1:])
+
+            if i == 0:
+                no_par_arg = arg.split("(")[0]
+                if no_par_arg in QueryBuilder._affinity_map:
+                    affinity = QueryBuilder.get_sql_col_affinity(arg)
+                    self.type_override = True
+                    self.affinity = affinity
+                    self.sql_type = affinity
+                    self.sql_create_type = arg
+                    continue
+            if key in [
+                "PRIMARY",
+                "NOT",
+                "UNIQUE",
+                "CHECK",
+                "DEFAULT",
+                "COLLATE",
+                "REFERENCES",
+                "GENERATED",
+                "AS",
+                "ON",
+            ]:
+                self.constraints[key] = arg
+            else:
+                logger.warning("Skipping annotation field %s", key)
+        return field_type
+
+    def parse_type(self, field_type: type):
+        # Handle field inner wrapper (e.g. int | None)
+        types = (
+            get_args(field_type)
+            if isinstance(field_type, UnionType)
+            else [
+                field_type,
+            ]
+        )
+        if "NOT" not in self.constraints and NoneType not in types:
+            self.constraints["NOT"] = "NOT NULL"
+        self.type = types[0]
+
+        if self.type_override:
+            return
+
+        if origin := get_origin(field_type):
+            if origin is Union:
+                field_type = get_args(field_type)[0]
+                self.origin = field_type
+            else:
+                field_type = origin
+                self.origin = field_type
+
+        if inspect.isclass(self.type) and issubclass(field_type, set):
+            raise TypeError(
+                "Sets cannot be encoded as text. Either override the type to BLOB, "
+                "override the codec, or change the datatype."
+            )
+
+        for mapped_type in self._type_map:
+            if issubclass(field_type, mapped_type):
+                # if isinstance(field_type, mapped_type):
+                # if issubclass(field_type, mapped_type):
+                self.sql_type = self._type_map[mapped_type]
+                break
+            # if inspect.isclass(field_type):
+            # else:
+            #     if isinstance(field_type, mapped_type):
+            #         self.sql_type = self._type_map[mapped_type]
+            #         break
+
+    def parse_default(self, dc_field):
+        if dc_field.default_factory is not MISSING:
+            default = dc_field.default_factory()
+        else:
+            default = dc_field.default
+
+        if "DEFAULT" in self.constraints or default is None or default is MISSING:
+            return
+
+        encoded = self.codec.encode(default)
+        if self.sql_type == "TEXT":
+            self.constraints["DEFAULT"] = f"DEFAULT '{encoded}'"
+        elif self.sql_type == "BLOB":
+            self.constraints["DEFAULT"] = f"DEFAULT X'{encoded.hex()}'"
+        else:
+            self.constraints["DEFAULT"] = f"DEFAULT {encoded}"
+
+    def sql_col_def(self, skip_primary=False, skip_unique=False):
+        output = [
+            self.name,
+            self.sql_create_type if self.sql_create_type else self.sql_type,
+        ]
+        for key, val in self.constraints.items():
+            if key == "PRIMARY" and skip_primary:
+                pass
+            elif key == "UNIQUE" and skip_unique:
+                pass
+            else:
+                output.append(val)
+        return " ".join(output)

dataclassdb/dataclass_sqlite_table.py

@@ -0,0 +1,152 @@
+__all__ = [
+    "DataclassTableCodec",
+    "decode_dict",
+    "get_class_codec",
+    "get_field_codec",
+    "encode_field",
+    "decode_field",
+]
+
+import functools
+import logging
+from dataclasses import asdict, fields, is_dataclass
+from typing import Any, Literal, TypeVar, get_type_hints, overload
+
+from dacite import from_dict
+
+from dataclassdb.dataclass_field_codec import identity_codec
+from dataclassdb.dataclass_sqlite_field import DataclassSqliteField
+from dataclassdb.dataclass_types import IsDataclass
+
+logger = logging.getLogger(__name__)
+
+
+def decode_dict(
+    data: dict[str, Any],
+    data_class=None,
+    prefix_class_map: dict = None,
+    remove_prefix=False,
+) -> dict[str, Any]:
+    if data is None:
+        return None
+
+    if data_class is None and prefix_class_map is None:
+        raise ValueError("Either a data class or prefix map must be provided")
+
+    if isinstance(data, list):
+        output = [
+            decode_dict(row, data_class, prefix_class_map, remove_prefix)
+            for row in data
+        ]
+        return output
+
+    output = {}
+    for key, val in data.items():
+        if "." in key:
+            prefix, split_key = key.split(".")
+            if mapped_dataclass := prefix_class_map.get(prefix, None):
+                decoded = decode_field(mapped_dataclass, split_key, val)
+            else:
+                decoded = val
+            output_key = split_key if remove_prefix else key
+            output[output_key] = decoded
+        elif data_class:
+            field_codec = get_field_codec(data_class, key)
+            output[key] = field_codec.decode(val)
+        else:
+            output[key] = val
+    return output
+
+
+T = TypeVar("T", bound=IsDataclass)
+class DataclassTableCodec:
+    """This class takes a dataclass and provides encoder and decoder (codec) from and to sqlite."""
+
+    def __contains__(self, item):
+        return item in self.class_fields
+
+    def __init__(self, data_class: T) -> None:
+        if not data_class or not is_dataclass(data_class):
+            raise ValueError(f"Provided class {data_class} is not a dataclass")
+
+        self.data_class: T = data_class
+        self.primary_keys = []
+        self.class_fields: dict[str, DataclassSqliteField] = {}
+
+        types_map: dict = get_type_hints(data_class, include_extras=True)
+        for _field in fields(data_class):
+            name = _field.name
+            field_type = types_map.get(name, None)
+            dc_field = DataclassSqliteField(field_type, _field)
+            self.class_fields[name] = dc_field
+            if dc_field.primary_key:
+                self.primary_keys.append(name)
+        self.key_match_str = " AND ".join([f"{key} = ?" for key in self.primary_keys])
+
+    @overload
+    def encode(self, obj: T, *cols, as_tuple: Literal[True]) -> tuple: ...
+    @overload
+    def encode(self, obj: T, *cols, as_tuple: Literal[False]) -> dict: ...
+    def encode(
+        self, obj: T, *cols, as_tuple: bool = True, ignore_none=True
+    ) -> tuple | dict:
+        if obj is None:
+            return None
+        output = {}
+        obj_dict = asdict(obj)
+        if not cols:
+            cols = obj_dict.keys()
+
+        for col in cols:
+            v = obj_dict.get(col, None)
+
+            if col in self.class_fields:
+                out = self.class_fields[col].encode(v)
+            else:
+                out = v
+            if out is not None or not ignore_none:
+                output[col] = out
+        return tuple(output.values()) if as_tuple else output
+
+    @overload
+    def decode(self, row_dict, as_obj: Literal[True]) -> T: ...
+    @overload
+    def decode(self, row_dict, as_obj: Literal[False]) -> dict: ...
+    def decode(self, row_dict, as_obj: bool = False) -> T | dict:
+        if row_dict is None:
+            return None
+        output = {}
+        for k, v in row_dict.items():
+            if k in self.class_fields:
+                output[k] = self.class_fields[k].decode(v)
+            else:
+                output[k] = v
+        if as_obj:
+            return from_dict(self.data_class, output)
+        else:
+            return output
+
+
+@functools.lru_cache
+def get_class_codec(data_class: IsDataclass) -> DataclassTableCodec:
+    if not data_class or not is_dataclass(data_class):
+        return None
+    else:
+        return DataclassTableCodec(data_class)
+
+
+def get_field_codec(data_class, field_name) -> DataclassSqliteField:
+    class_codec = get_class_codec(data_class)
+    if not class_codec:
+        return identity_codec
+    return class_codec.class_fields.get(field_name, identity_codec)
+
+
+def encode_field(data_class, field_name, data):
+    codec = get_field_codec(data_class, field_name)
+    return codec.encode(data)
+
+
+def decode_field(data_class, field_name, data):
+    codec = get_field_codec(data_class, field_name)
+    return codec.decode(data)

dataclassdb/dataclass_types.py

@@ -0,0 +1,22 @@
+from dataclasses import Field
+from typing import Mapping, Protocol, TypeVar, runtime_checkable
+
+T = TypeVar("T")
+S = TypeVar("S")
+
+@runtime_checkable
+class Codec(Protocol):
+    def __init__(self, *args, **kwargs): pass
+    @staticmethod
+    def encode(data: T) -> S:...
+    @staticmethod
+    def decode(data: S) -> T:...
+
+class CustomCodec(Codec):
+    def __init__(self, encode, decode):
+        self.encode = encode
+        self.decode = decode
+
+@runtime_checkable
+class IsDataclass(Protocol):
+    __dataclass_fields__: Mapping[str, Field]

dataclassdb/db_engine.py

@@ -0,0 +1,126 @@
+import logging
+import sqlite3
+
+from dataclassdb.builders.string_builder import StringBuilder
+from dataclassdb.utils import dict_factory
+
+logger = logging.getLogger()
+
+
+class DbEngine(StringBuilder):
+    def __init__(
+        self,
+        connection: sqlite3.Connection | str = "",  # type: ignore
+        **kwargs,
+    ):
+        self.connection: sqlite3.Connection
+        self.connection_owner = False
+        self.has_connection = False
+        self.connection_name = ""
+
+        if connection:
+            self.has_connection = True
+            if isinstance(connection, sqlite3.Connection):
+                self.connection = connection
+            else:
+                self.connection = sqlite3.Connection(connection)
+                self.connection_owner = True
+                self.connection_name = connection
+                self.enable_wal()
+                logger.info("Connection for %s created.", self.connection_name)
+
+        super().__init__(connection=connection, **kwargs)
+        # super(DbExecutor, self).__init__()
+
+    def __enter__(self):
+        if not self.has_connection:
+            raise ValueError("Connection must be set to use context manager")
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.connection_owner and self.has_connection:
+            logger.info("Committing DB.")
+            self.connection.commit()
+            logger.info("Closing DB.")
+            self.connection.close()
+        return False
+
+    def enable_wal(self) -> None:
+        self.execute_script(
+            "PRAGMA journal_mode=WAL;\n"
+            "PRAGMA synchronous=NORMAL;\n"
+            "PRAGMA busy_timeout = 5000;\n"
+            "PRAGMA cache_size = -20000;\n"  # 20MB
+            "PRAGMA foreign_keys = true;\n"
+            "PRAGMA temp_store = memory;"
+        )
+
+    def execute_script(self, sql_script=None, as_dict=False, commit=False):
+        if not self.has_connection:
+            raise ValueError("Connection must be set to use sql commands.")
+
+        try:
+            if sql_script is None:
+                sql_script = str(self)
+            if not sql_script:
+                raise ValueError("Query is empty")
+            cur = self.connection.cursor()  # type: ignore
+            if as_dict:
+                cur.row_factory = dict_factory
+            output = cur.executescript(sql_script)
+            if commit:
+                self.connection.commit()  # type: ignore
+            return output
+        finally:
+            self.clear
+
+    def execute_one(self, *args, sql_str=None, commit=False, as_dict=False):
+        return self.execute(
+            *args, sql_str=sql_str, commit=commit, single_row=True, as_dict=as_dict
+        )
+
+    def execute(
+        self, *args, sql_str=None, commit=False, single_row=False, as_dict=False
+    ):
+        if not self.has_connection:
+            raise ValueError("Connection must be set to use sql commands.")
+        try:
+            params = tuple(args) if args else ()
+            if sql_str is None:
+                query_str = str(self)
+            else:
+                query_str = str(sql_str)
+
+            if not query_str:
+                raise ValueError("Query is empty")
+
+            if not query_str.endswith(";"):
+                query_str = f"{query_str};"
+
+            cur = self.connection.cursor()
+            if as_dict:
+                cur.row_factory = dict_factory
+
+            if "\n" in query_str:
+                logger.debug(
+                    "Executing Query: \n    %s", query_str.replace("\n", "\n    ")
+                )
+                logger.debug(
+                    "Using Params   :\n    %s", "\n    ".join(map(str, params))
+                )
+            else:
+                logger.debug("Executing Query: %s", query_str)
+                logger.debug("Using Params   : %s", params)
+
+            cur.execute(query_str, params)
+
+            if single_row:
+                rows = cur.fetchone()
+                return rows
+            else:
+                rows = cur.fetchall()
+            if commit:
+                self.connection.commit()
+            return rows
+        finally:
+            self.clear