Improved relationship support & customizable results set classes

Author: maximebf

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "sqlorm-py"
-version = "0.3.0"
+version = "0.3.1"
 description = "A new kind or ORM that do not abstract away your database or SQL queries."
 authors = [
     {"name" = "Maxime Bouroumeau-Fuseau", email = "[email protected]"}

src/sqlorm/engine.py

@@ -465,19 +465,19 @@ def executemany(self, stmt, seq_of_parameters):
         self.after_execute.send(self, cursor=cur, stmt=stmt, params=seq_of_parameters, many=True)
         cur.close()
 
-    def fetch(self, stmt, params=None, model=None, obj=None, loader=None):
+    def fetch(self, stmt, params=None, model=None, obj=None, loader=None, resultset_class=None):
         cursor = self.cursor(stmt, params)
 
         if obj and not model:
             model = obj.__class__
         if model:
-            rs = HydratedResultSet(cursor, model)
+            rs = (resultset_class or HydratedResultSet)(cursor, model)
             if obj:
                 rs.mapper.hydrate(obj, rs.first(with_loader=False))
                 return obj
             return rs
 
-        return ResultSet(cursor, loader)
+        return (resultset_class or ResultSet)(cursor, loader)
 
     def fetchall(self, stmt, params=None, **fetch_kwargs):
         return self.fetch(stmt, params, **fetch_kwargs).all()
@@ -501,6 +501,7 @@ def fetchcomposite(
         obj=None,
         map=None,
         separator=None,
+        resultset_class=CompositeResultSet,
     ):
         if obj and not model:
             model = obj.__class__
@@ -514,7 +515,7 @@ def fetchcomposite(
         elif not map:
             map = CompositionMap.create([loader, nested])
 
-        rs = CompositeResultSet(
+        rs = resultset_class(
             self.cursor(stmt, params), map, separator or self.default_composite_separator
         )
         if obj:

src/sqlorm/mapper.py

@@ -165,13 +165,17 @@ def hydrate_new(self, data):
     def hydrate(self, obj, data, with_unknown=None):
         """Populates a model object with data from the db"""
         attrs = set()
+        relnames = {r.attribute: r for r in self.relationships}
         for key in (
             data.keys()
         ):  # avoid using .items() as some DBAPI returned objects only provide keys() (eg: sqlite3)
             if key in self.columns.names:
                 col = self.columns[key]
                 col.load(obj, data)
                 key = col.attribute
+            elif key in relnames:
+                rel = relnames[key]
+                rel.load(obj, data)
             elif with_unknown or with_unknown is None and self.allow_unknown_columns:
                 obj.__dict__[key] = data[key]  # ensure that no custom setter are used
             else:
@@ -414,8 +418,8 @@ def __init__(
         self,
         target_mapper,
         target_col=None,
-        target_attr=None,
         source_col=None,
+        target_attr=None,
         source_attr=None,
         join_type="LEFT JOIN",
         join_condition=None,
@@ -424,7 +428,7 @@ def __init__(
         lazy=True,
     ):
         self._target_mapper = target_mapper
-        self.target_col = target_col
+        self._target_col = target_col
         self._target_attr = target_attr
         self._source_col = source_col
         self._source_attr = source_attr
@@ -443,6 +447,16 @@ def target_mapper(self) -> Mapper:
     def target_table(self):
         return self.target_mapper.table
 
+    @property
+    def target_col(self):
+        if self._target_col:
+            return self._target_col
+        if not self.single:
+            raise MapperError(f"Missing target_col on relationship '{self.attribute}'")
+        if self.target_mapper.primary_key:
+            return self.target_mapper.primary_key.name
+        return "id"
+
     @property
     def target_attr(self):
         if self._target_attr:
@@ -493,8 +507,6 @@ def join_condition(self, target_alias=None, source_alias=None) -> SQL:
             )
         elif self._join_condition:
             return self._join_condition
-        if not self.target_col:
-            raise MapperError(f"Missing target_col on relationship '{self.attribute}'")
         return SQL.Col(self.target_col, table=target_alias) == SQL.Col(
             self.source_col, table=source_alias
         )
@@ -538,6 +550,23 @@ def delete_from_target(self, source_obj):
         return SQL.delete_from(self.target_table).where(
             SQL.Col(self.target_col) == SQL.Param(getattr(source_obj, source_attr))
         )
+    
+    def update_related_objs(self, source_obj, target_obj):
+        target_attr = self.target_attr
+        if not target_attr:
+            raise MapperError(
+                f"Missing target_attr on relationship '{self.attribute}'"
+            )
+        if self.single:
+            setattr(source_obj, self.source_attr, getattr(target_obj, target_attr) if target_obj else None)
+        else:
+            setattr(target_obj, target_attr, None if source_obj is None else getattr(source_obj, self.source_attr, None))
+
+    def load(self, obj, values):
+        """Sets the object attribute from the database row, using the provided load function if needed
+        (used by Mapper.hydrate())
+        """
+        obj.__dict__[self.attribute] = values[self.attribute]
 
     def __repr__(self):
         return f"<Relationship({self.attribute})>"
@@ -552,14 +581,18 @@ def __init__(self, cursor, mapper, auto_close_cursor=True):
 
 
 class HydrationMap(CompositionMap):
-    def __init__(self, mapper, nested=None, single=False):
+    def __init__(self, mapper, nested=None, single=False, _already_visited=None):
         if not isinstance(mapper, Mapper):
             mapper = Mapper.from_class(mapper)
+        if _already_visited is None:
+            _already_visited = set()
         rowid = mapper.primary_key.name if mapper.primary_key else None
         _nested = {
-            r.attribute: HydrationMap(r.target_mapper, single=r.single)
+            r.attribute: HydrationMap(r.target_mapper, single=r.single, _already_visited={mapper} | _already_visited)
             for r in mapper.relationships
+            if r.target_mapper not in _already_visited
         }
+        _already_visited.add(mapper)
         if nested:
             _nested.update({k: HydrationMap.create(v) for k, v in nested.items()})
         super().__init__(mapper.hydrate_new, _nested, rowid, single)

src/sqlorm/model.py

@@ -256,7 +256,7 @@ def target_mapper(self):
 
     @property
     def __isabstractmethod__(self):
-        # compatibility between this description __getattr__ usage and abc.ABC
+        # compatibility between this descriptor __getattr__ usage and abc.ABC
         return False
 
     def __get__(self, obj, owner=None):
@@ -267,18 +267,29 @@ def __get__(self, obj, owner=None):
         return obj.__dict__[self.attribute]
 
     def __set__(self, obj, value):
-        obj.__dict__[self.attribute] = value if self.single else self.list_class(obj, self, value)
+        if self.single:
+            obj.__dict__[self.attribute] = value
+            self.update_related_objs(obj, value)
+            flag_dirty_attr(obj, self.source_attr)
+        else:
+            obj.__dict__[self.attribute] = self.list_class(obj, self, [])
+            for item in value:
+                obj.__dict__[self.attribute].append(item) # ensure target_attr is set
 
     def is_loaded(self, obj):
         return self.attribute in obj.__dict__
 
     def fetch(self, obj):
         """Fetches the list of related objects from the database and loads it in the object"""
         r = self.target.query(self.select_from_target(obj))
-        self.__set__(obj, r.first() if self.single else r.all())
+        obj.__dict__[self.attribute] = r.first() if self.single else self.list_class(obj, self, r.all())
+
+    def load(self, obj, values):
+        value = values[self.attribute]
+        obj.__dict__[self.attribute] = value if self.single else self.list_class(obj, self, value)
 
 
-class RelatedObjectsList(object):
+class RelatedObjectsList:
     def __init__(self, obj, relationship, items):
         self.obj = obj
         self.relationship = relationship
@@ -297,20 +308,12 @@ def __contains__(self, item):
         return item in self.items
 
     def append(self, item):
-        target_attr = self.relationship.target_attr
-        if not target_attr:
-            raise MapperError(
-                f"Missing target_attr on relationship '{self.relationship.attribute}'"
-            )
-        setattr(item, target_attr, getattr(self.obj, self.relationship.source_attr))
+        self.relationship.update_related_objs(self.obj, item)
+        flag_dirty_attr(item, self.relationship.target_attr)
 
     def remove(self, item):
-        target_attr = self.relationship.target_attr
-        if not target_attr:
-            raise MapperError(
-                f"Missing target_attr on relationship '{self.relationship.attribute}'"
-            )
-        setattr(item, target_attr, None)
+        self.relationship.update_related_objs(None, item)
+        flag_dirty_attr(item, self.relationship.target_attr)
 
 
 def flag_dirty_attr(obj, attr):
@@ -367,6 +370,7 @@ def __sql__(self):
 
 class Model(BaseModel, abc.ABC):
     """Our standard model class with CRUD methods"""
+    __resultset_class__ = CompositeResultSet
 
     class Meta:
         insert_update_dirty_only: bool = (
@@ -397,12 +401,12 @@ def query(cls, stmt, params=None) -> CompositeResultSet:
         with ensure_transaction(cls.__engine__) as tx:
             rv = _signal_rv(cls.before_query.send(cls, stmt=stmt, params=params))
             if rv is False:
-                return ResultSet(None)
+                return cls.__resultset_class__(None)
             if isinstance(rv, ResultSet):
                 return rv
             if isinstance(rv, tuple):
                 stmt, params = rv
-            return tx.fetchhydrated(cls, stmt, params)
+            return tx.fetchhydrated(cls, stmt, params, resultset_class=cls.__resultset_class__)
 
     @classmethod
     def find_all(
@@ -473,8 +477,11 @@ def __init__(self, **values):
             setattr(self, k, v)
 
     def __setattr__(self, name, value):
-        self.__dict__[name] = value
-        flag_dirty_attr(self, name)
+        if isinstance(getattr(self.__class__, name, None), (ModelColumnMixin, Relationship)):
+            super().__setattr__(name, value)
+        else:
+            self.__dict__[name] = value
+            flag_dirty_attr(self, name)
 
     def refresh(self, **select_kwargs):
         stmt = self.__mapper__.select_by_pk(self.__mapper__.get_primary_key(self), **select_kwargs)

src/sqlorm/sqlfunc.py

@@ -42,6 +42,8 @@ def query_builder(*args, **kwargs):
 
 def is_sqlfunc(func):
     """Checks if func is an empty function with a python doc"""
+    if not inspect.isfunction(func):
+        return False
     doc = inspect.getdoc(func)
     src = inspect.getsource(func).strip(' "\n\r')
     return doc and src.endswith(func.__doc__.strip(' \n\r')) and not getattr(func, "sqlfunc", False)

tests/models.py

@@ -7,6 +7,7 @@ class Task(Model):
     title: str
     done: bool = Column("completed")
     user_id: Integer = Column(references="users(id)")
+    user = Relationship("User", "id", "user_id", single=True)
 
     @classmethod
     def find_todos(cls):

tests/test_model.py

@@ -175,20 +175,6 @@ def on_after_refresh(sender, obj):
         assert listener_called == 2
 
 
-def test_relationships(engine):
-    with engine:
-        user = User.get(1)
-        assert not User.tasks.is_loaded(user)
-        assert len(user.tasks) == 1
-        assert User.tasks.is_loaded(user)
-        assert isinstance(user.tasks[0], Task)
-        assert user.tasks[0].id == 1
-
-        user = User.get(1, with_rels=True)
-        assert User.tasks.is_loaded(user)
-        assert len(user.tasks) == 1
-
-
 def test_sql_methods(engine):
     assert Task.toggle.query_decorator == "update"
     assert Task.find_todos.query_decorator == "fetchall"
@@ -386,3 +372,37 @@ def on_after_delete(sender, obj):
 
         user = User.get(4)
         assert not user
+
+
+def test_relationships_many(engine):
+    with engine:
+        user = User.get(1)
+        assert not User.tasks.is_loaded(user)
+        assert len(user.tasks) == 1
+        assert User.tasks.is_loaded(user)
+        assert isinstance(user.tasks[0], Task)
+        assert user.tasks[0].id == 1
+
+        user = User.get(1, with_rels=True)
+        assert User.tasks.is_loaded(user)
+        assert len(user.tasks) == 1
+
+        user = User.get(1)
+        task = Task()
+        user.tasks.append(task)
+        assert task.user_id == user.id
+        user.tasks.remove(task)
+        assert task.user_id is None
+
+
+def test_relationships_single(engine):
+    with engine:
+        task = Task.get(1)
+        assert not Task.user.is_loaded(task)
+        assert task.user.id == 1
+        assert Task.user.is_loaded(task)
+
+        user = User.get(1)
+        task = Task()
+        task.user = user
+        assert task.user_id == user.id