Fixes and improvements after integrating in a real codebase

Author: rundef

README.md

@@ -5,16 +5,27 @@
 [![PyPI - Downloads](https://img.shields.io/pypi/dm/sqlalchemy-memory)](https://pypistats.org/packages/sqlalchemy-memory)
 
 
-**In‑memory SQLAlchemy 2.0 dialect for blazing‑fast prototyping**  
+**In‑memory SQLAlchemy 2.0 dialect for blazing‑fast prototyping**
 
-A pure‑Python, zero‑configuration SQLAlchemy 2.0 dialect that lives entirely in RAM.
-Ideal for rapid prototyping, backtesting, and demos-no external database required.
+A pure‑Python SQLAlchemy 2.0 dialect that runs entirely in RAM.
+It avoids typical database I/O and ORM overhead while maintaining full compatibility with the SQLAlchemy 2.0 Core and ORM APIs.
+Ideal for rapid prototyping, backtesting engines, simulations.
 
 ## Why ?
 
-This project was inspired by the idea of creating a fast, introspectable, no-dependency backend for SQLAlchemy. Useful for prototyping, education, and testing ORM logic without spinning up a real database engine.
+This project was inspired by the idea of building a **fast, introspectable, no-dependency backend** for SQLAlchemy.
 
-It's also perfect for apps that need a fast, in-memory store compatible with SQLAlchemy, such as backtesting engines, simulators, or tools where you don't want to maintain a separate memory layer alongside your database models.
+It is useful for:
+
+- Prototyping new applications
+
+- Educational purposes
+
+- Testing ORM logic without spinning up a real database engine
+
+Unlike traditional in-memory solutions like SQLite, `sqlalchemy-memory` fully avoids serialization, connection pooling, and driver overhead, leading to much faster in-memory performance while keeping the familiar SQLAlchemy API.
+
+It is also perfect for **applications that need a lightweight, high-performance store** compatible with SQLAlchemy, such as backtesting engines, simulators, or other tools where you don't want to maintain a separate in-memory layer alongside your database models.
 
 ## Features
 

sqlalchemy_memory/base/connection.py

@@ -1,8 +1,5 @@
-from .store import InMemoryStore
-
 class MemoryDBAPIConnection:
-    def __init__(self):
-        self.store = InMemoryStore()
+    store = None
 
     def commit(self):
         self.store.commit()
@@ -12,3 +9,9 @@ def rollback(self):
 
     def close(self):
         pass
+
+    @classmethod
+    def connect(cls, store, *args, **kwargs):
+        connection = MemoryDBAPIConnection()
+        connection.store = store
+        return connection

sqlalchemy_memory/base/dialect.py

@@ -1,9 +1,8 @@
 from sqlalchemy.engine import URL, default
 import types
-import copy
 
 from .connection import MemoryDBAPIConnection
-from ..util import _raw_dbapi_connection
+from .store import InMemoryStore
 
 class MemoryDialect(default.DefaultDialect):
     name = "memory"
@@ -12,14 +11,21 @@ class MemoryDialect(default.DefaultDialect):
     supports_native_boolean = True
     supports_statement_cache = False
 
+    _store = None
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._store = InMemoryStore()
+
     def initialize(self, connection):
         super().initialize(connection)
 
         # Turn off pool reset to preserve in-memory data
         connection.engine.pool._reset_on_return = None
 
     def create_connect_args(self, url: URL):
-        return [], {}
+        db_name = url.database or "_default"
+        return [db_name], {}
 
     @classmethod
     def import_dbapi(cls, **kwargs):
@@ -28,23 +34,21 @@ def import_dbapi(cls, **kwargs):
             paramstyle="named",
             apilevel="2.0",
             threadsafety=1,
-            Error=Exception
+            Error=Exception,
+            connect=lambda *a, **k: None
         )
-        module.connect = lambda *a, **k: MemoryDBAPIConnection()
         return module
 
-    def do_begin(self, dbapi_conn):
-        connection = _raw_dbapi_connection(dbapi_conn)
-        connection.store._snapshot = copy.deepcopy(connection.store.data)
+    def connect(self, *args, **kwargs):
+        connection = MemoryDBAPIConnection()
+        connection.store = self._store
+        return connection
 
     def do_commit(self, dbapi_conn):
-        connection = _raw_dbapi_connection(dbapi_conn)
-        connection.store.commit()
+        self._store.commit()
 
     def do_rollback(self, dbapi_conn):
-        connection = _raw_dbapi_connection(dbapi_conn)
-        connection.store.data = connection.store._snapshot
-        connection.store.rollback()
+        self._store.rollback()
 
     def has_table(self, *args, **kwargs):
         # Patch to make Base.metadata.create_all(engine) not throw any exception

sqlalchemy_memory/base/query.py

@@ -5,10 +5,28 @@
 from sqlalchemy.sql.functions import FunctionElement
 from sqlalchemy.sql import operators
 from sqlalchemy.sql.annotation import AnnotatedTable
+from sqlalchemy.orm.query import Query
 from functools import cached_property
 import fnmatch
 
-from sqlalchemy.orm.query import Query
+from ..logger import logger
+from .resolvers import DateResolver, JsonExtractResolver
+
+OPERATOR_ADAPTERS = {
+    operators.is_: lambda value: lambda x, _: x is value,
+    operators.isnot: lambda value: lambda x, _: x is not value,
+    operators.like_op: lambda value: lambda x, _: fnmatch.fnmatchcase(x or '', value.replace('%', '*').replace('_', '?')),
+    operators.not_like_op: lambda value: lambda x, _: not fnmatch.fnmatchcase(x or '', value.replace('%', '*').replace('_', '?')),
+    operators.between_op: lambda bounds: lambda x, _: bounds[0] <= x <= bounds[1],
+    operators.not_between_op: lambda bounds: lambda x, _: not (bounds[0] <= x <= bounds[1]),
+    operators.in_op: lambda values: lambda x, _: x in values,
+    operators.not_in_op: lambda values: lambda x, _: x not in values,
+}
+
+FUNCTION_RESOLVERS = {
+    "date": DateResolver,
+    "json_extract": JsonExtractResolver,
+}
 
 class MemoryQuery(Query):
     def __init__(self, entities, element):
@@ -52,17 +70,6 @@ def order_by(self, clause):
         self._order_by.append(clause)
         return self
 
-    def _extract_json_value(self, data_dict, path):
-        # Traverse nested keys for a JSON path like 'ref.abc.xyz'
-        current = data_dict or {}
-        for key in path.split('.'):
-            if not isinstance(current, dict):
-                return None
-            current = current.get(key)
-            if current is None:
-                return None
-        return current
-
     def _apply_condition(self, cond, collection):
         if not isinstance(cond, BinaryExpression):
             raise NotImplementedError(f"Unsupported condition type: {type(cond)}")
@@ -85,34 +92,21 @@ def _apply_condition(self, cond, collection):
         else:
             raise NotImplementedError(f"Unsupported RHS: {type(rhs)}")
 
-        # Handle JSON extraction: func.json_extract(column, path)
-        if isinstance(cond.left, FunctionElement) and cond.left.name.lower() == 'json_extract':
-            args = list(cond.left.clauses)
-            column_expr, path_expr = args[0], args[1]
-            attr_name = column_expr.name
-            # Determine raw path string
-            raw = path_expr.value if hasattr(path_expr, 'value') else str(path_expr).strip('"')
-
-            # Strip leading '$.' or '$'
-            if raw.startswith('$.'):
-                raw_path = raw[2:]
-            elif raw.startswith('$'):
-                raw_path = raw[1:]
-            else:
-                raw_path = raw
-
-            # Compare nested value
-            op = cond.operator
-            return [
-                item for item in collection
-                if op(
-                    self._extract_json_value(getattr(item, attr_name), raw_path),
-                    value
-                )
-            ]
+        col = cond.left
+        accessor = lambda obj, attr_name: getattr(obj, attr_name)
+
+        if isinstance(cond.left, FunctionElement):
+            fn_name = cond.left.name.lower()
+            if fn_name not in FUNCTION_RESOLVERS:
+                raise NotImplementedError(f"Unsupported LHS function: {fn_name}")
+
+            clauses = list(cond.left.clauses)
+            col = clauses[0]
+            _class = FUNCTION_RESOLVERS[fn_name]
+            _resolver = _class(clauses[1:])
+            accessor = _resolver.accessor
 
         # Extract column name (LHS) and operator
-        col = cond.left
         if not hasattr(col, "name"):
             raise NotImplementedError(f"Unsupported LHS: {col}")
         attr_name = col.name
@@ -122,43 +116,18 @@ def _apply_condition(self, cond, collection):
 
         op = cond.operator
 
-        # special‑case SQL "IS NULL" and "IS NOT NULL"
-        if value is None:
-            if op is operators.is_:
-                op = lambda x, y: x is None
-            elif op is operators.isnot:
-                op = lambda x, y: x is not None
-
-        elif op is operators.like_op:
-            fnmatch_pattern = value.replace('%', '*').replace('_', '?')
-            op = lambda x, y: fnmatch.fnmatchcase(x or '', fnmatch_pattern)
-
-        elif op is operators.not_like_op:
-            fnmatch_pattern = value.replace('%', '*').replace('_', '?')
-            op = lambda x, y: not fnmatch.fnmatchcase(x or '', fnmatch_pattern)
-
-        elif op is operators.between_op:
-            low, high = value
-            op = lambda x, _: low <= x <= high
-
-        elif op is operators.not_between_op:
-            low, high = value
-            op = lambda x, _: not (low <= x <= high)
-
-        elif op is operators.in_op:
-            op = lambda x, y: x in y
-
-        elif op is operators.not_in_op:
-            op = lambda x, y: x not in y
+        if op in OPERATOR_ADAPTERS:
+            op = OPERATOR_ADAPTERS[op](value)
 
         return [
             item for item in collection
-            if op(getattr(item, attr_name), value)
+            if op(accessor(item, attr_name), value)
         ]
 
     def _execute_query(self):
         collection = self.session.store.data.get(self.tablename, [])
         if not collection:
+            logger.debug(f"Table '{self.tablename}' is empty")
             return collection
 
         # Apply conditions
@@ -178,8 +147,7 @@ def _execute_query(self):
             else:
                 col = clause
 
-            attr = col.name
-            collection.sort(key=lambda x: getattr(x, attr), reverse=reverse)
+            collection = sorted(collection, key=lambda x: getattr(x, col.name), reverse=reverse)
 
         # Apply offset
         if self._offset is not None:

sqlalchemy_memory/base/resolvers/__init__.py

@@ -0,0 +1,7 @@
+from .date import DateResolver
+from .json_extract import JsonExtractResolver
+
+__all__ = [
+    "DateResolver",
+    "JsonExtractResolver",
+]

sqlalchemy_memory/base/resolvers/abstract.py

@@ -0,0 +1,6 @@
+class FunctionResolver:
+    def __init__(self, clauses):
+        self.clauses = clauses
+
+    def accessor(self, item, attr_name):
+        raise NotImplementedError

sqlalchemy_memory/base/resolvers/date.py

@@ -0,0 +1,7 @@
+from .abstract import FunctionResolver
+
+
+class DateResolver(FunctionResolver):
+    def accessor(self, item, attr_name):
+        value = getattr(item, attr_name)
+        return value.date() if value else None

sqlalchemy_memory/base/resolvers/json_extract.py

@@ -0,0 +1,29 @@
+from .abstract import FunctionResolver
+
+
+class JsonExtractResolver(FunctionResolver):
+    def _extract_json_value(self, data_dict, path):
+        # Traverse nested keys for a JSON path like 'ref.abc.xyz'
+        current = data_dict or {}
+        for key in path.split('.'):
+            if not isinstance(current, dict):
+                return None
+            current = current.get(key)
+            if current is None:
+                return None
+        return current
+
+    def accessor(self, item, attr_name):
+        path_expr = self.clauses[0]
+
+        raw = path_expr.value if hasattr(path_expr, 'value') else str(path_expr).strip('"')
+
+        # Strip leading '$.' or '$'
+        if raw.startswith('$.'):
+            raw_path = raw[2:]
+        elif raw.startswith('$'):
+            raw_path = raw[1:]
+        else:
+            raw_path = raw
+
+        return self._extract_json_value(getattr(item, attr_name), raw_path)

sqlalchemy_memory/base/session.py

@@ -5,6 +5,7 @@
 from sqlalchemy.engine.cursor import SimpleResultMetaData
 
 from .query import MemoryQuery
+from ..logger import logger
 
 class MemorySession(Session):
     def __init__(self, *args, **kwargs):
@@ -163,6 +164,8 @@ def _handle_update(self, statement: Update, **kwargs):
         return result
 
     def execute(self, statement, params=None, **kwargs):
+        #logger.debug(f"Executing query: {statement}")
+
         if isinstance(statement, Select):
             return self._handle_select(statement, **kwargs)
 
@@ -183,17 +186,19 @@ def merge(self, instance, **kwargs):
         """
 
         pk_name = self.store._get_primary_key_name(instance)
-        id = getattr(instance, pk_name)
-        existing = self.store.get_by_primary_key(instance, id)
+        pk_value = getattr(instance, pk_name)
+        existing = self.store.get_by_primary_key(instance, pk_value)
 
         if existing:
-            data = {
-                col.name: getattr(instance, col.name)
-                for col in instance.__table__.columns
-                if col.name != pk_name
-            }
+            self.store.mark_as_fetched(existing)
+
+            for column in instance.__table__.columns:
+                field = column.name
+                if field == pk_name:
+                    continue
+                value = getattr(instance, field)
+                setattr(existing, field, value)
 
-            self.store.update(instance.__tablename__, id, data)
             return existing
 
         else:
@@ -207,7 +212,9 @@ def flush(self, objects=None):
         pass
 
     def rollback(self, **kwargs):
+        logger.debug("Rolling back ...")
         self.store.rollback()
 
     def commit(self):
+        logger.debug("Committing ...")
         self.store.commit()