datamarket 0.7.41__py3-none-any.whl → 0.7.125__py3-none-any.whl

datamarket/__init__.py

@@ -1 +0,0 @@
-# init file

datamarket/exceptions/__init__.py

@@ -0,0 +1 @@
+from .main import *  # noqa: F403

datamarket/exceptions/main.py

@@ -0,0 +1,118 @@
+########################################################################################################################
+# CLASSES
+
+
+from typing import Optional
+
+from requests import Request, Response
+from requests.exceptions import HTTPError
+
+
+class ManagedHTTPError(HTTPError):
+    """Signal that this HTTP status was handled and should not be retried."""
+
+    def __init__(
+        self,
+        message: Optional[str] = None,
+        response: Optional[Response] = None,
+        request: Optional[Request] = None,
+        *args,
+        **kwargs,
+    ):
+        self.response = response
+        self.request = request or getattr(response, "request", None)
+
+        # Build a safe default message
+        if not message:
+            status = getattr(self.response, "status_code", "unknown")
+            url = getattr(self.request, "url", "unknown")
+            message = f"HTTP {status} for {url}"
+
+        self.message = message
+
+        super().__init__(message, *args, response=response, **kwargs)
+
+
+class IgnoredHTTPError(ManagedHTTPError):
+    """Exception type that signals the error should be ignored by retry logic."""
+
+    pass
+
+
+class NotFoundError(ManagedHTTPError):
+    def __init__(
+        self,
+        message: Optional[str] = None,
+        response: Optional[Response] = None,
+        request: Optional[Request] = None,
+        *args,
+        **kwargs,
+    ):
+        if not message:
+            status = getattr(response, "status_code", 404)
+            req = request or getattr(response, "request", None)
+            url = getattr(req, "url", "unknown")
+            message = f"HTTP {status} for {url}"
+        super().__init__(message, response, request, *args, **kwargs)
+
+
+class BadRequestError(ManagedHTTPError):
+    def __init__(
+        self,
+        message: Optional[str] = None,
+        response: Optional[Response] = None,
+        request: Optional[Request] = None,
+        *args,
+        **kwargs,
+    ):
+        if not message:
+            status = getattr(response, "status_code", 400)
+            req = request or getattr(response, "request", None)
+            url = getattr(req, "url", "unknown")
+            message = f"HTTP {status} for {url}"
+        super().__init__(message, response, request, *args, **kwargs)
+
+
+class EmptyResponseError(ManagedHTTPError):
+    def __init__(
+        self,
+        message: Optional[str] = None,
+        response: Optional[Response] = None,
+        request: Optional[Request] = None,
+        *args,
+        **kwargs,
+    ):
+        if not message:
+            req = request or getattr(response, "request", None)
+            url = getattr(req, "url", "unknown")
+            message = f"Empty response for {url}"
+        super().__init__(message, response, request, *args, **kwargs)
+
+
+class RedirectionDetectedError(ManagedHTTPError):
+    def __init__(
+        self,
+        message: Optional[str] = None,
+        response: Optional[Response] = None,
+        request: Optional[Request] = None,
+        *args,
+        **kwargs,
+    ):
+        if not message:
+            status = getattr(response, "status_code", 300)
+            req = request or getattr(response, "request", None)
+            url = getattr(req, "url", "unknown")
+            message = f"HTTP {status} for {url}"
+        super().__init__(message, response, request, *args, **kwargs)
+
+
+class NoWorkingProxiesError(Exception):
+    def __init__(self, message="No working proxies available"):
+        self.message = message
+        super().__init__(self.message)
+
+
+class EnsureNewIPTimeoutError(Exception):
+    def __init__(self, message="Timed out waiting for new IP"):
+        self.message = message
+        super().__init__(self.message)

datamarket/interfaces/alchemy.py

@@ -3,14 +3,16 @@
 
 import logging
 from collections.abc import MutableMapping
-from typing import Any, Iterator, List, Type, TypeVar
+from enum import Enum, auto
+from typing import Any, Iterator, List, Optional, Type, TypeVar, Union
 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 sessionmaker
+from sqlalchemy.orm import Session, sessionmaker
+from sqlalchemy.sql.expression import ClauseElement
 
 ########################################################################################################################
 # CLASSES
@@ -20,6 +22,11 @@ logger = logging.getLogger(__name__)
 ModelType = TypeVar("ModelType", bound=DeclarativeMeta)
 
 
+class CommitStrategy(Enum):
+    COMMIT_ON_SUCCESS = auto()
+    FORCE_COMMIT = auto()
+
+
 class MockContext:
     def __init__(self, column: SQLColumnExpression) -> None:
         self.current_parameters = {}
@@ -29,16 +36,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 +159,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,109 +183,192 @@ 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}...")
+    @staticmethod
+    def _log_integrity_error(ex: IntegrityError, alchemy_obj, action="insert"):
+        """
+        Compact, readable IntegrityError logger using SQLSTATE codes.
+        Consult https://www.postgresql.org/docs/current/errcodes-appendix.html for details.
+        """
 
-            self.session.add(alchemy_obj)
-            self.session.commit()
+        PG_ERROR_LABELS = {
+            "23000": "Integrity constraint violation",
+            "23001": "Restrict violation",
+            "23502": "NOT NULL violation",
+            "23503": "Foreign key violation",
+            "23505": "Unique violation",
+            "23514": "Check constraint violation",
+            "23P01": "Exclusion constraint violation",
+        }
+        code = getattr(ex.orig, "pgcode", None)
+        label = PG_ERROR_LABELS.get(code, "Integrity error (unspecified)")
 
-        except IntegrityError:
+        # Log one clean message with trace + the raw DB message separately
+        if code == "23505":  # A simple info log for unique violations
+            logger.info(f"{label} trying to {action} {alchemy_obj}")
+        else:
+            logger.error(f"{label} trying to {action} {alchemy_obj}\nPostgreSQL message: {ex.orig}")
+
+    def insert_alchemy_obj(self, alchemy_obj: ModelType, silent: bool = False) -> bool:
+        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 as ex:
+            # Rollback is handled automatically by begin_nested() context manager on error
             if not silent:
-                logger.info(f"{alchemy_obj} already in db")
+                self._log_integrity_error(ex, alchemy_obj, action="insert")
+            # Do not re-raise, allow outer transaction/loop to continue
+            return False
+
+        return True
 
-            self.session.rollback()
+    def upsert_alchemy_obj(self, alchemy_obj: ModelType, index_elements: List[str], silent: bool = False) -> bool:
+        if self.session is None:
+            raise RuntimeError("Session not active. Use 'with AlchemyInterface(...):' or call start()")
 
-    def upsert_alchemy_obj(self, alchemy_obj: ModelType, index_elements: List[str], silent: bool = False) -> None:
         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:
-            self.session.execute(statement)
-            self.session.commit()
-        except IntegrityError:
+            # Use a savepoint (nested transaction)
+            with self.session.begin_nested():
+                self.session.execute(statement)
+        except IntegrityError as ex:
+            # Rollback is handled automatically by begin_nested() context manager on error
             if not silent:
-                logger.info(f"could not upsert {alchemy_obj}")
+                self._log_integrity_error(ex, alchemy_obj, action="upsert")
+            # Do not re-raise, allow outer transaction/loop to continue
+            return False
 
-            self.session.rollback()
+        return True
 
     def windowed_query(
         self,
         stmt: Select[Any],
         order_by: List[SQLColumnExpression[Any]],
         windowsize: int,
+        commit_strategy: Union[CommitStrategy, str] = CommitStrategy.COMMIT_ON_SUCCESS,
     ) -> 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.
             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.
+            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")
+        # Parameter mapping
+        if isinstance(commit_strategy, str):
+            commit_strategy = CommitStrategy[commit_strategy.upper()]
 
-        # 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)
-        subq = inner_stmt.subquery()
-
-        # Create an outer query that selects from the subquery
-        cols = [c for c in subq.c]
-        outer_query = select(*cols)
-
-        last_row_number = 0
+        # Find id column in stmt
+        if not any(column.get("entity").id for column in stmt.column_descriptions):
+            raise Exception("Column 'id' not found in any entity of the query.")
+        id_column = stmt.column_descriptions[0]["entity"].id
 
+        last_id = 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)
-
-            # Create a FrozenResult to allow peeking at the data without consuming
-            frozen_result: FrozenResult = result.freeze()
-            chunk = frozen_result().all()
-
-            if not chunk:
-                break
-
-            # Update for next iteration
-            last_row_number = chunk[-1].row_number - 1
-
-            # Create a new Result object from the FrozenResult
-            yield_result = frozen_result()
-
-            # Remove row_number from result
-            yield_result = yield_result.columns(*range(len(cols) - 1))
-
-            yield yield_result
+            session_active = False
+            commit_needed = False
+            try:
+                self.start()
+                session_active = True
+
+                # Filter on row_number in the outer query
+                current_query = stmt.where(id_column > last_id).order_by(order_by[0], *order_by[1:]).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()
+
+                if not chunk:
+                    break
+
+                # Update for next iteration
+                last_id = chunk[-1].id
+
+                # Create a new Result object from the FrozenResult
+                yield_result = frozen_result()
+
+                yield yield_result
+                commit_needed = True
+
+            finally:
+                if session_active and self.session:
+                    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)

datamarket/interfaces/aws.py

@@ -3,6 +3,8 @@
 
 import io
 import logging
+from typing import Any, Dict, List, Optional
+
 import boto3
 
 ########################################################################################################################
@@ -12,34 +14,44 @@ logger = logging.getLogger(__name__)
 
 
 class AWSInterface:
-    def __init__(self, config):
-        self.profiles = []
+    def __init__(self, config) -> None:
+        self.profiles: List[Dict[str, Any]] = []
         self.config = config
 
-        for section in self.config.sections():
+        for section in getattr(self.config, "sections", lambda: [])():
             if section.startswith("aws:"):
                 profile_name = section.split(":", 1)[1]
+                bucket_value = self.config[section].get("buckets", "")
+                buckets = [b.strip() for b in bucket_value.split(",") if b.strip()]
+                session = boto3.Session(profile_name=profile_name)
+
                 self.profiles.append(
                     {
                         "profile": profile_name,
-                        "bucket": self.config[section]["bucket"],
-                        "session": boto3.Session(profile_name=profile_name),
+                        "buckets": buckets,
+                        "session": session,
                     }
                 )
 
         if not self.profiles:
             logger.warning("No AWS profiles found in config file")
 
-        self.current_profile = self.profiles[0] if self.profiles else None
+        self.current_profile: Optional[Dict[str, Any]] = self.profiles[0] if self.profiles else None
         self._update_resources()
 
-    def _update_resources(self):
+    def _update_resources(self) -> None:
+        """Refresh S3 resources for the current profile and set default bucket (first in list)"""
         if self.current_profile:
             self.s3 = self.current_profile["session"].resource("s3")
             self.s3_client = self.s3.meta.client
-            self.bucket = self.current_profile["bucket"]
+            buckets = self.current_profile.get("buckets", [])
+            self.bucket = buckets[0] if buckets else None
+        else:
+            self.s3 = None
+            self.s3_client = None
+            self.bucket = None
 
-    def switch_profile(self, profile_name):
+    def switch_profile(self, profile_name: str) -> None:
         for profile in self.profiles:
             if profile["profile"] == profile_name:
                 self.current_profile = profile
@@ -47,14 +59,69 @@ class AWSInterface:
                 return
         logger.warning(f"Profile {profile_name} not found")
 
-    def get_file(self, s3_path):
+    def switch_bucket(self, bucket: str) -> None:
+        if not self.current_profile:
+            logger.warning("No current AWS profile to switch bucket on")
+            return
+
+        buckets = self.current_profile.get("buckets") or []
+        if bucket not in buckets:
+            logger.warning(f"Bucket {bucket} not found in profile {self.current_profile.get('profile')}")
+            return
+
+        self.bucket = bucket
+
+    def switch_bucket_for_profile(self, profile_name: str, bucket: str) -> None:
+        """
+        Select a profile and then switch its active bucket.
+        """
+        for profile in self.profiles:
+            if profile["profile"] == profile_name:
+                self.current_profile = profile
+                self._update_resources()  # sets default bucket & s3 clients
+                self.switch_bucket(bucket)  # only sets self.bucket if valid
+                return
+        logger.warning(f"Profile {profile_name} not found")
+
+    def get_bucket_url(self) -> Optional[str]:
+        """Return active bucket URL."""
+        if not self.bucket:
+            logger.warning("No active bucket selected")
+            return None
+        region = self.s3_client.meta.region_name
+        return f"https://{self.bucket}.s3.{region}.amazonaws.com"
+
+    def get_file(self, s3_path: str):
+        if not self.bucket:
+            logger.warning("No active bucket selected")
+            return None
         try:
             return self.s3.Object(self.bucket, s3_path).get()
         except self.s3_client.exceptions.NoSuchKey:
             logger.info(f"{s3_path} does not exist")
+            return None
+
+    def file_exists(self, s3_path: str) -> bool:
+        if not self.bucket:
+            logger.warning("No active bucket selected")
+            return False
+        try:
+            self.s3_client.head_object(Bucket=self.bucket, Key=s3_path)
+            return True
+        except self.s3_client.exceptions.NoSuchKey:
+            return False
+        except Exception as e:
+            logger.error(f"Error checking existence of {s3_path}: {e}")
+            raise
 
-    def read_file_as_bytes(self, s3_path):
-        return io.BytesIO(self.get_file(s3_path)["Body"].read())
+    def read_file_as_bytes(self, s3_path: str) -> Optional[io.BytesIO]:
+        obj = self.get_file(s3_path)
+        if not obj:
+            return None
+        return io.BytesIO(obj["Body"].read())
 
-    def upload_file(self, local_path, s3_path):
-        self.s3.Bucket(self.bucket).upload_file(local_path, s3_path)
+    def upload_file(self, local_path: str, s3_path: str, **kwargs) -> None:
+        if not self.bucket:
+            logger.warning("No active bucket selected")
+            return
+        self.s3.Bucket(self.bucket).upload_file(local_path, s3_path, **kwargs)