datamarket 0.9.23__tar.gz → 0.9.25__tar.gz

{datamarket-0.9.23 → datamarket-0.9.25}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: datamarket
-Version: 0.9.23
+Version: 0.9.25
 Summary: Utilities that integrate advanced scraping knowledge into just one library.
 License: GPL-3.0-or-later
 Author: DataMarket

{datamarket-0.9.23 → datamarket-0.9.25}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "datamarket"
-version = "0.9.23"
+version = "0.9.25"
 description = "Utilities that integrate advanced scraping knowledge into just one library."
 authors = ["DataMarket <techsupport@datamarket.es>"]
 license = "GPL-3.0-or-later"

{datamarket-0.9.23 → datamarket-0.9.25}/src/datamarket/interfaces/alchemy.py

@@ -6,11 +6,13 @@ from collections.abc import MutableMapping
 from typing import Any, Iterator, List, Optional, Type, TypeVar
 from urllib.parse import quote_plus
 
-from sqlalchemy import DDL, FrozenResult, Result, Select, SQLColumnExpression, create_engine, func, select, text
+from sqlalchemy import DDL, FrozenResult, Result, Select, SQLColumnExpression, create_engine, text
 from sqlalchemy.dialects.postgresql import insert
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.ext.declarative import DeclarativeMeta
 from sqlalchemy.orm import Session, sessionmaker
+from sqlalchemy.sql.expression import ClauseElement
+from enum import Enum
 
 ########################################################################################################################
 # CLASSES
@@ -20,6 +22,11 @@ logger = logging.getLogger(__name__)
 ModelType = TypeVar("ModelType", bound=DeclarativeMeta)
 
 
+class CommitStrategy(Enum):
+    COMMIT_ON_SUCCESS = "commit_on_success"
+    FORCE_COMMIT = "force_commit"
+
+
 class MockContext:
     def __init__(self, column: SQLColumnExpression) -> None:
         self.current_parameters = {}
@@ -214,17 +221,43 @@ class AlchemyInterface:
         if not silent:
             logger.info(f"upserting {alchemy_obj}")
 
-        primary_keys = list(col.name for col in alchemy_obj.__table__.primary_key.columns.values())
-        obj_dict = {
+        table = alchemy_obj.__table__
+        primary_keys = list(col.name for col in table.primary_key.columns.values())
+
+        # Build the dictionary for the INSERT values
+        insert_values = {
+            col.name: getattr(alchemy_obj, col.name)
+            for col in table.columns
+            if getattr(alchemy_obj, col.name) is not None  # Include all non-None values for insert
+        }
+
+        # Build the dictionary for the UPDATE set clause
+        # Start with values from the object, excluding primary keys
+        update_set_values = {
             col.name: val
-            for col in alchemy_obj.__table__.columns
+            for col in table.columns
             if col.name not in primary_keys and (val := getattr(alchemy_obj, col.name)) is not None
         }
 
+        # Add columns with SQL-based onupdate values explicitly to the set clause
+        for column in table.columns:
+            actual_sql_expression = None
+            if column.onupdate is not None:
+                if hasattr(column.onupdate, "arg") and isinstance(column.onupdate.arg, ClauseElement):
+                    # This handles wrappers like ColumnElementColumnDefault,
+                    # where the actual SQL expression is in the .arg attribute.
+                    actual_sql_expression = column.onupdate.arg
+                elif isinstance(column.onupdate, ClauseElement):
+                    # This handles cases where onupdate might be a direct SQL expression.
+                    actual_sql_expression = column.onupdate
+
+            if actual_sql_expression is not None:
+                update_set_values[column.name] = actual_sql_expression
+
         statement = (
-            insert(alchemy_obj.__table__)
-            .values(obj_dict)
-            .on_conflict_do_update(index_elements=index_elements, set_=obj_dict)
+            insert(table)
+            .values(insert_values)
+            .on_conflict_do_update(index_elements=index_elements, set_=update_set_values)
         )
 
         try:
@@ -245,6 +278,7 @@ class AlchemyInterface:
         stmt: Select[Any],
         order_by: List[SQLColumnExpression[Any]],
         windowsize: int,
+        commit_strategy: CommitStrategy = CommitStrategy.COMMIT_ON_SUCCESS,
     ) -> Iterator[Result[Any]]:
         """
         Executes a windowed query, fetching each window in a separate, short-lived session.
@@ -253,6 +287,8 @@ class AlchemyInterface:
             stmt: The SQL select statement to execute.
             order_by: The columns to use for ordering.
             windowsize: The number of rows to fetch in each window.
+            commit_strategy: The strategy to use for committing the session after each window.
+                             Defaults to CommitStrategy.COMMIT_ON_SUCCESS.
 
         Returns:
             An iterator of Result objects, each containing a window of data.
@@ -296,4 +332,15 @@ class AlchemyInterface:
 
             finally:
                 if session_active and self.session:
-                    self.stop(commit=commit_needed)
+                    if commit_strategy == CommitStrategy.FORCE_COMMIT:
+                        # For forced commit, always attempt to commit.
+                        # The self.stop() method already handles potential exceptions during commit/rollback.
+                        self.stop(commit=True)
+                    elif commit_strategy == CommitStrategy.COMMIT_ON_SUCCESS:
+                        # Commit only if no exception occurred before yielding the result.
+                        self.stop(commit=commit_needed)
+                    else:
+                        # Fallback or error for unknown strategy, though type hinting should prevent this.
+                        # For safety, default to rollback.
+                        logger.warning(f"Unknown commit strategy: {commit_strategy}. Defaulting to rollback.")
+                        self.stop(commit=False)