datamarket 0.9.20__tar.gz → 0.9.21__tar.gz

{datamarket-0.9.20 → datamarket-0.9.21}/PKG-INFO

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

{datamarket-0.9.20 → datamarket-0.9.21}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "datamarket"
-version = "0.9.20"
+version = "0.9.21"
 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.20 → datamarket-0.9.21}/src/datamarket/interfaces/alchemy.py

@@ -3,14 +3,14 @@
 
 import logging
 from collections.abc import MutableMapping
-from typing import Any, Iterator, List, Type, TypeVar
+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.dialects.postgresql import insert
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.ext.declarative import DeclarativeMeta
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.orm import Session, sessionmaker
 
 ########################################################################################################################
 # CLASSES
@@ -29,16 +29,58 @@ class MockContext:
 
 class AlchemyInterface:
     def __init__(self, config: MutableMapping) -> None:
+        self.session: Optional[Session] = None
         if "db" in config:
             self.config = config["db"]
-
             self.engine = create_engine(self.get_conn_str())
-            self.session = sessionmaker(bind=self.engine)()
-            self.cursor = self.session.connection().connection.cursor()
-
+            self.Session = sessionmaker(bind=self.engine)
         else:
             logger.warning("no db section in config")
 
+    def __enter__(self) -> "AlchemyInterface":
+        """Enter the runtime context related to this object (starts session)."""
+        self.start()
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """Exit the runtime context related to this object (stops session)."""
+        should_commit = exc_type is None
+        self.stop(commit=should_commit)
+
+    def start(self) -> None:
+        """Starts a new SQLAlchemy session manually."""
+        if not hasattr(self, "Session"):
+            raise AttributeError("Database configuration not initialized. Cannot create session.")
+        if self.session is not None:
+            raise RuntimeError("Session already active.")
+        self.session = self.Session()
+        logger.debug("SQLAlchemy session started manually.")
+
+    def stop(self, commit: bool = True) -> None:
+        """Stops the manually started SQLAlchemy session."""
+        if self.session is None:
+            logger.warning("No active session to stop.")
+            return
+
+        try:
+            if commit:
+                logger.debug("Committing SQLAlchemy session before stopping.")
+                self.session.commit()
+            else:
+                logger.debug("Rolling back SQLAlchemy session before stopping.")
+                self.session.rollback()
+        except Exception as e:
+            logger.error(f"Exception during session commit/rollback on stop: {e}", exc_info=True)
+            try:
+                self.session.rollback()
+            except Exception as rb_exc:
+                logger.error(f"Exception during secondary rollback attempt on stop: {rb_exc}", exc_info=True)
+            raise
+        finally:
+            logger.debug("Closing SQLAlchemy session.")
+            self.session.close()
+            self.session = None
+
     def get_conn_str(self):
         return (
             f"{self.config['engine']}://"
@@ -110,6 +152,16 @@ class AlchemyInterface:
         self.create_tables(tables)
 
     def reset_column(self, query_results: List[Result[Any]], column_name: str) -> None:
+        """
+        Reset a column to its default value for a list of query results.
+
+        Args:
+            query_results: List of query results to update
+            column_name: Name of the column to reset
+        """
+        if self.session is None:
+            raise RuntimeError("Session not active. Use 'with AlchemyInterface(...):' or call start()")
+
         if not query_results:
             logger.warning("No objects to reset column for.")
             return
@@ -124,33 +176,38 @@ class AlchemyInterface:
 
         column = table.columns[column_name]
 
+        # Determine the default value to use
         if column.server_default is not None:
-            query_results.update({column_name: text("DEFAULT")}, synchronize_session=False)
+            default_value = text("DEFAULT")
         elif column.default is not None:
             default_value = column.default.arg
             if callable(default_value):
                 default_value = default_value(MockContext(column))
-            query_results.update({column_name: default_value}, synchronize_session=False)
         else:
             raise ValueError(f"Column '{column_name}' doesn't have a default value defined.")
 
-        self.session.commit()
+        query_results.update({column_name: default_value}, synchronize_session=False)
 
     def insert_alchemy_obj(self, alchemy_obj: ModelType, silent: bool = False) -> None:
-        try:
-            if not silent:
-                logger.info(f"adding {alchemy_obj}...")
-
-            self.session.add(alchemy_obj)
-            self.session.commit()
+        if self.session is None:
+            raise RuntimeError("Session not active. Use 'with AlchemyInterface(...):' or call start()")
 
+        try:
+            # Use a savepoint (nested transaction)
+            with self.session.begin_nested():
+                if not silent:
+                    logger.info(f"adding {alchemy_obj}...")
+                self.session.add(alchemy_obj)
         except IntegrityError:
+            # Rollback is handled automatically by begin_nested() context manager on error
             if not silent:
-                logger.info(f"{alchemy_obj} already in db")
-
-            self.session.rollback()
+                logger.info(f"{alchemy_obj} already in db (savepoint rolled back)")
+            # Do not re-raise, allow outer transaction/loop to continue
 
     def upsert_alchemy_obj(self, alchemy_obj: ModelType, index_elements: List[str], silent: bool = False) -> None:
+        if self.session is None:
+            raise RuntimeError("Session not active. Use 'with AlchemyInterface(...):' or call start()")
+
         if not silent:
             logger.info(f"upserting {alchemy_obj}")
 
@@ -168,13 +225,14 @@ class AlchemyInterface:
         )
 
         try:
-            self.session.execute(statement)
-            self.session.commit()
+            # Use a savepoint (nested transaction)
+            with self.session.begin_nested():
+                self.session.execute(statement)
         except IntegrityError:
+            # Rollback is handled automatically by begin_nested() context manager on error
             if not silent:
-                logger.info(f"could not upsert {alchemy_obj}")
-
-            self.session.rollback()
+                logger.info(f"could not upsert {alchemy_obj} (savepoint rolled back)")
+            # Do not re-raise, allow outer transaction/loop to continue
 
     def windowed_query(
         self,
@@ -183,7 +241,7 @@ class AlchemyInterface:
         windowsize: int,
     ) -> Iterator[Result[Any]]:
         """
-        Executes a windowed query on the given statement.
+        Executes a windowed query, fetching each window in a separate, short-lived session.
 
         Args:
             stmt: The SQL select statement to execute.
@@ -192,14 +250,15 @@ class AlchemyInterface:
 
         Returns:
             An iterator of Result objects, each containing a window of data.
+            The session used to fetch the Result is closed immediately after yielding.
 
         More info: https://github.com/sqlalchemy/sqlalchemy/wiki/RangeQuery-and-WindowedRangeQuery
         """
         # Add row_number over the specified order
         row_number = func.row_number().over(order_by=order_by).label("row_number")
 
-        # Add the windowing column to the statement and sort by it for deterministic results
-        inner_stmt = stmt.add_columns(row_number).order_by(row_number)
+        # Add the windowing column to the statement
+        inner_stmt = stmt.with_session(None).add_columns(row_number).order_by(row_number)
         subq = inner_stmt.subquery()
 
         # Create an outer query that selects from the subquery
@@ -209,24 +268,37 @@ class AlchemyInterface:
         last_row_number = 0
 
         while True:
-            # Filter on row_number in the outer query
-            current_query = outer_query.where(subq.c.row_number > last_row_number).limit(windowsize)
-            result = self.session.execute(current_query)
+            session_active = False
+            commit_needed = False
+            try:
+                self.start()
+                session_active = True
+
+                # Filter on row_number in the outer query
+                current_query = outer_query.where(subq.c.row_number > last_row_number).limit(windowsize)
+                result = self.session.execute(current_query)
+
+                # Create a FrozenResult to allow peeking at the data without consuming
+                frozen_result: FrozenResult = result.freeze()
+                chunk = frozen_result().all()
 
-            # Create a FrozenResult to allow peeking at the data without consuming
-            frozen_result: FrozenResult = result.freeze()
-            chunk = frozen_result().all()
+                if not chunk:
+                    break
 
-            if not chunk:
-                break
+                # Update for next iteration
+                last_row_number = chunk[-1].row_number
 
-            # Update for next iteration
-            last_row_number = chunk[-1].row_number - 1
+                # Create a new Result object from the FrozenResult
+                yield_result = frozen_result()
 
-            # Create a new Result object from the FrozenResult
-            yield_result = frozen_result()
+                # Remove row_number from result before yielding
+                # Ensure we don't yield the row_number column itself
+                original_col_count = len(cols) - 1
+                yield_result = yield_result.columns(*range(original_col_count))
 
-            # Remove row_number from result
-            yield_result = yield_result.columns(*range(len(cols) - 1))
+                yield yield_result
+                commit_needed = True
 
-            yield yield_result
+            finally:
+                if session_active and self.session:
+                    self.stop(commit=commit_needed)