lecrapaud 0.18.7__py3-none-any.whl → 0.22.6__py3-none-any.whl

lecrapaud/db/models/base.py

@@ -10,19 +10,26 @@ from sqlalchemy.orm.attributes import InstrumentedAttribute
 from lecrapaud.db.session import get_db
 from sqlalchemy.ext.declarative import declared_attr
 from sqlalchemy.dialects.mysql import insert as mysql_insert
+from sqlalchemy import UniqueConstraint
+from sqlalchemy.inspection import inspect as sqlalchemy_inspect
 from lecrapaud.config import LECRAPAUD_TABLE_PREFIX
 
 
 def with_db(func):
-    """Decorator to allow passing an optional db session"""
+    """Decorator to provide a database session to the wrapped function.
+
+    If a db parameter is already provided, it will be used. Otherwise,
+    a new session will be created and automatically managed.
+    """
 
     @wraps(func)
     def wrapper(*args, **kwargs):
-        db = kwargs.pop("db", None)
-        if db:
-            return func(*args, db=db, **kwargs)
+        if "db" in kwargs and kwargs["db"] is not None:
+            return func(*args, **kwargs)
+
         with get_db() as db:
-            return func(*args, db=db, **kwargs)
+            kwargs["db"] = db
+            return func(*args, **kwargs)
 
     return wrapper
 
@@ -102,51 +109,6 @@ class Base(DeclarativeBase):
             ]
         return results
 
-    @classmethod
-    @with_db
-    def upsert_bulk(cls, db=None, match_fields: list[str] = None, **kwargs):
-        """
-        Performs a bulk upsert into the database using ON DUPLICATE KEY UPDATE.
-
-        Args:
-            db (Session): SQLAlchemy DB session
-            match_fields (list[str]): Fields to match on for deduplication
-            **kwargs: Column-wise keyword arguments (field_name=[...])
-        """
-        # Ensure all provided fields have values of equal length
-        value_lengths = [len(v) for v in kwargs.values()]
-        if not value_lengths or len(set(value_lengths)) != 1:
-            raise ValueError(
-                "All field values must be non-empty lists of the same length."
-            )
-
-        # Convert column-wise kwargs to row-wise list of dicts
-        items = [dict(zip(kwargs.keys(), row)) for row in zip(*kwargs.values())]
-        if not items:
-            return
-
-        stmt = mysql_insert(cls.__table__).values(items)
-
-        # Default to primary keys if match_fields not provided
-        if not match_fields:
-            match_fields = [col.name for col in cls.__table__.primary_key.columns]
-
-        # Ensure all columns to be updated are in the insert
-        update_dict = {
-            c.name: stmt.inserted[c.name]
-            for c in cls.__table__.columns
-            if c.name not in match_fields and c.name in items[0]
-        }
-
-        if not update_dict:
-            # Avoid triggering ON DUPLICATE KEY UPDATE with empty dict
-            db.execute(stmt.prefix_with("IGNORE"))
-        else:
-            upsert_stmt = stmt.on_duplicate_key_update(**update_dict)
-            db.execute(upsert_stmt)
-
-        db.commit()
-
     @classmethod
     @with_db
     def filter(cls, db=None, **kwargs):
@@ -194,33 +156,126 @@ class Base(DeclarativeBase):
 
     @classmethod
     @with_db
-    def upsert(cls, match_fields: list[str], db=None, **kwargs):
+    def upsert(cls, db=None, **kwargs):
         """
-        Upsert an instance of the model: update if found, else create.
+        Upsert an instance of the model using MySQL's ON DUPLICATE KEY UPDATE.
 
-        :param match_fields: list of field names to use for matching
         :param kwargs: all fields for creation or update
         """
-        filters = [
-            getattr(cls, field) == kwargs[field]
-            for field in match_fields
-            if field in kwargs
-        ]
+        # If an ID is provided and row exists, fall back to a standard update
+        instance_id = kwargs.get("id")
+        if instance_id is not None:
+            instance = db.get(cls, instance_id)
+            if instance:
+                for key, value in kwargs.items():
+                    if key == "id":
+                        continue
+                    setattr(instance, key, value)
+                db.commit()
+                db.refresh(instance)
+                return instance
+
+        # Use INSERT ... ON DUPLICATE KEY UPDATE
+        stmt = mysql_insert(cls.__table__).values(**kwargs)
+        stmt = stmt.on_duplicate_key_update(
+            **{k: v for k, v in kwargs.items() if k != "id"}
+        )
 
-        instance = db.query(cls).filter(*filters).first()
+        result = db.execute(stmt)
+        db.commit()
 
-        if instance:
-            for key, value in kwargs.items():
-                if key != "id":
-                    setattr(instance, key, value)
+        # Get the instance - either the newly inserted or updated one
+        # If updated, lastrowid is 0, so we need to query
+        if result.lastrowid and result.lastrowid > 0:
+            # New insert
+            instance = db.get(cls, result.lastrowid)
         else:
-            instance = cls(**kwargs)
-            db.add(instance)
+            # Updated - need to find it using unique constraint fields
+            mapper = sqlalchemy_inspect(cls)
+            instance = None
+
+            for constraint in mapper.mapped_table.constraints:
+                if isinstance(constraint, UniqueConstraint):
+                    col_names = [col.name for col in constraint.columns]
+                    if all(name in kwargs for name in col_names):
+                        filters = [
+                            getattr(cls, col_name) == kwargs[col_name]
+                            for col_name in col_names
+                        ]
+                        instance = db.query(cls).filter(*filters).first()
+                        if instance:
+                            break
+
+            # Check for single column unique constraints
+            if not instance:
+                for col in mapper.mapped_table.columns:
+                    if col.unique and col.name in kwargs:
+                        instance = (
+                            db.query(cls)
+                            .filter(getattr(cls, col.name) == kwargs[col.name])
+                            .first()
+                        )
+                        if instance:
+                            break
+
+            # If still not found, try to find by all kwargs (excluding None values)
+            if not instance:
+                instance = (
+                    db.query(cls)
+                    .filter_by(
+                        **{
+                            k: v
+                            for k, v in kwargs.items()
+                            if v is not None and k != "id"
+                        }
+                    )
+                    .first()
+                )
+
+        if instance:
+            db.refresh(instance)
 
-        db.commit()
-        db.refresh(instance)
         return instance
 
+    @classmethod
+    @with_db
+    def bulk_upsert(cls, rows: list[dict] = None, db=None, **kwargs):
+        """
+        Performs a bulk upsert into the database using ON DUPLICATE KEY UPDATE.
+
+        Args:
+            rows (list[dict]): List of dictionaries representing rows to upsert
+            db (Session): SQLAlchemy DB session
+            **kwargs: Column-wise keyword arguments (field_name=[...]) for backwards compatibility
+        """
+        # Handle both new format (rows) and legacy format (kwargs)
+        if rows is None and kwargs:
+            # Legacy format: convert column-wise kwargs to row-wise list of dicts
+            value_lengths = [len(v) for v in kwargs.values()]
+            if not value_lengths or len(set(value_lengths)) != 1:
+                raise ValueError(
+                    "All field values must be non-empty lists of the same length."
+                )
+            rows = [dict(zip(kwargs.keys(), row)) for row in zip(*kwargs.values())]
+
+        if not rows:
+            return 0
+
+        BATCH_SIZE = 200
+        total_affected = 0
+
+        for i in range(0, len(rows), BATCH_SIZE):
+            batch = rows[i : i + BATCH_SIZE]
+            stmt = mysql_insert(cls.__table__).values(batch)
+            stmt = stmt.on_duplicate_key_update(
+                **{key: stmt.inserted[key] for key in batch[0] if key != "id"}
+            )
+            result = db.execute(stmt)
+            total_affected += result.rowcount
+
+        db.commit()
+        return total_affected
+
     @classmethod
     @with_db
     def delete(cls, id: int, db=None):