diff --git a/src/orm_loader/backends/__init__.py b/src/orm_loader/backends/__init__.py index 3fa6888..35a99cf 100644 --- a/src/orm_loader/backends/__init__.py +++ b/src/orm_loader/backends/__init__.py @@ -1,11 +1,12 @@ from .postgres import PostgresBackend from .resolve import resolve_backend from .sqlite import SQLiteBackend -from .base import BackendCapabilities, DatabaseBackend, Dialect +from .base import BackendCapabilities, DatabaseBackend, STAGING_SCHEMA, Dialect __all__ = [ "BackendCapabilities", "DatabaseBackend", + "STAGING_SCHEMA", "Dialect", "PostgresBackend", "SQLiteBackend", diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 07028cd..4fc7089 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -10,6 +10,8 @@ import sqlalchemy.orm as so from sqlalchemy.engine import Connection, Engine +from ..helpers.sql import qualify_identifier + if TYPE_CHECKING: from ..loaders.data_classes import LoaderContext from ..tables.typing import CSVTableProtocol @@ -37,6 +39,9 @@ class Dialect(str, Enum): POSTGRESQL = "postgresql" +STAGING_SCHEMA: str = "staging" + + class DatabaseBackend(ABC): """ Abstract base class for database-specific loader behavior. @@ -45,6 +50,54 @@ class DatabaseBackend(ABC): without changing existing loader orchestration yet. """ + def __init__(self, staging_schema: str | None = None) -> None: + """ + Parameters + ---------- + staging_schema + Schema in which staging tables are created. ``None`` means no + schema qualification — staging tables land in whatever schema the + connection's search_path resolves to. Backends that support + schema-isolated staging (e.g. PostgreSQL, DuckDB) should declare + their own ``__init__`` that defaults this to ``STAGING_SCHEMA`` + rather than relying on this base default. SQLite has no schema + concept and always passes ``None``. + """ + self.staging_schema = staging_schema + + @staticmethod + @abstractmethod + def staging_name_for_table(tablename: str) -> str: + """ + Return the unqualified staging table name for a given target tablename. + + Parameters + ---------- + tablename + The target table's ``__tablename__``. + + Returns + ------- + str + The staging table name (no schema prefix). + """ + + def qualified_staging_name(self, tablename: str) -> str: + """ + Return the fully schema-qualified staging identifier. + + Parameters + ---------- + tablename + The target table's ``__tablename__``. + + Returns + ------- + str + e.g. '"staging"."_staging_concept"' or '"_staging_concept"'. + """ + return qualify_identifier(self.staging_name_for_table(tablename), self.staging_schema) + @property @abstractmethod def name(self) -> str: @@ -123,22 +176,20 @@ def create_staging_table( self, table_cls: Type["CSVTableProtocol"], session: so.Session, - staging_name: str, ) -> None: """Create a staging table for the supplied ORM table class.""" @abstractmethod def drop_staging_table( self, + table_cls: Type["CSVTableProtocol"], session: so.Session, - staging_name: str, ) -> None: """Drop a staging table if it exists.""" def load_staging_fast( self, loader_context: "LoaderContext", - staging_name: str, ) -> int | None: """ Attempt a backend-native fast-path load. @@ -179,7 +230,6 @@ def merge_replace( table_cls: Type["CSVTableProtocol"], session: so.Session, target_name: str, - staging_name: str, pk_cols: list[str], *, merge_batch_size: int | None = None, @@ -192,7 +242,6 @@ def merge_upsert( table_cls: Type["CSVTableProtocol"], session: so.Session, target_name: str, - staging_name: str, pk_cols: list[str], *, merge_batch_size: int | None = None, @@ -205,7 +254,6 @@ def merge_insert( table_cls: Type["CSVTableProtocol"], session: so.Session, target_name: str, - staging_name: str, *, merge_batch_size: int | None = None, ) -> None: diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index 6118ed6..4ebca35 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -8,7 +8,7 @@ import sqlalchemy.orm as so from ..loaders.loading_helpers import quick_load_pg -from .base import BackendCapabilities, DatabaseBackend, Dialect +from .base import BackendCapabilities, DatabaseBackend, Dialect, STAGING_SCHEMA if TYPE_CHECKING: from sqlalchemy.engine import Connection, Engine @@ -20,6 +20,13 @@ class PostgresBackend(DatabaseBackend): + def __init__(self, *, staging_schema: str | None = STAGING_SCHEMA) -> None: + super().__init__(staging_schema=staging_schema) + + @staticmethod + def staging_name_for_table(tablename: str) -> str: + return f"_staging_{tablename}" + @property def name(self) -> str: return "postgres" @@ -41,14 +48,14 @@ def create_staging_table( self, table_cls: type["CSVTableProtocol"], session: so.Session, - staging_name: str, ) -> None: table = table_cls.__table__ - session.execute(sa.text(f'DROP TABLE IF EXISTS "{staging_name}";')) + staging_ref = self.qualified_staging_name(table_cls.__tablename__) + session.execute(sa.text(f'DROP TABLE IF EXISTS {staging_ref};')) session.execute( sa.text( f''' - CREATE UNLOGGED TABLE "{staging_name}" + CREATE UNLOGGED TABLE {staging_ref} (LIKE "{table.name}" INCLUDING DEFAULTS INCLUDING CONSTRAINTS); ''' ) @@ -56,12 +63,12 @@ def create_staging_table( computed_cols = [c.name for c in table.columns if c.computed is not None] for col in computed_cols: - session.execute(sa.text(f'ALTER TABLE "{staging_name}" DROP COLUMN "{col}";')) + session.execute(sa.text(f'ALTER TABLE {staging_ref} DROP COLUMN "{col}";')) - # allows pagniation in O(N log N) time for large tables in merge_insert without needing to add an index on every staging table + # allows pagination in O(N log N) time for large tables in merge_insert without needing to add an index on every staging table session.execute( sa.text( - f'ALTER TABLE "{staging_name}" ADD COLUMN _rownum BIGINT' + f'ALTER TABLE {staging_ref} ADD COLUMN _rownum BIGINT' f" GENERATED ALWAYS AS IDENTITY (CACHE 1000);" ) ) @@ -70,20 +77,21 @@ def create_staging_table( def drop_staging_table( self, + table_cls: type["CSVTableProtocol"], session: so.Session, - staging_name: str, ) -> None: - session.execute(sa.text(f'DROP TABLE IF EXISTS "{staging_name}"')) + session.execute(sa.text(f'DROP TABLE IF EXISTS {self.qualified_staging_name(table_cls.__tablename__)}')) def load_staging_fast( self, loader_context: "LoaderContext", - staging_name: str, ) -> int | None: + tablename = loader_context.tableclass.__tablename__ return quick_load_pg( path=loader_context.path, session=loader_context.session, - tablename=staging_name, + tablename=self.staging_name_for_table(tablename), + schema=self.staging_schema, quote_mode=loader_context.quote_mode, ) @@ -131,27 +139,28 @@ def merge_replace( table_cls: type["CSVTableProtocol"], session: so.Session, target_name: str, - staging_name: str, pk_cols: list[str], *, merge_batch_size: int | None = None, ) -> None: + staging_ref = self.qualified_staging_name(table_cls.__tablename__) pk_join = " AND ".join(f't."{c}" = s."{c}"' for c in pk_cols) non_paginated_replace = sa.text( - f'DELETE FROM "{target_name}" t USING "{staging_name}" s WHERE {pk_join}' + f'DELETE FROM "{target_name}" t USING {staging_ref} s WHERE {pk_join}' ) if merge_batch_size is None: session.execute(non_paginated_replace) return - total = session.execute(sa.text(f'SELECT COUNT(*) FROM "{staging_name}"')).scalar_one() + total = session.execute(sa.text(f'SELECT COUNT(*) FROM {staging_ref}')).scalar_one() if total <= merge_batch_size: session.execute(non_paginated_replace) return - session.execute(sa.text(f'CREATE INDEX IF NOT EXISTS "{staging_name}_rownum_idx" ON "{staging_name}" (_rownum)')) + staging_name = self.staging_name_for_table(table_cls.__tablename__) + session.execute(sa.text(f'CREATE INDEX IF NOT EXISTS "{staging_name}_rownum_idx" ON {staging_ref} (_rownum)')) session.commit() start = 0 @@ -159,7 +168,7 @@ def merge_replace( end = start + merge_batch_size session.execute( sa.text( - f'DELETE FROM "{target_name}" t USING "{staging_name}" s' + f'DELETE FROM "{target_name}" t USING {staging_ref} s' f' WHERE {pk_join} AND s._rownum > :start AND s._rownum <= :end' ), {"start": start, "end": end}, @@ -172,18 +181,18 @@ def merge_upsert( table_cls: type["CSVTableProtocol"], session: so.Session, target_name: str, - staging_name: str, pk_cols: list[str], *, merge_batch_size: int | None = None, ) -> None: + staging_ref = self.qualified_staging_name(table_cls.__tablename__) insertable_cols = self._insertable_column_names(table_cls) cols_str = ", ".join(f'"{c}"' for c in insertable_cols) conflict_cols = ", ".join(f'"{c}"' for c in pk_cols) non_paginated_upsert = sa.text( f'INSERT INTO "{target_name}" ({cols_str})' - f' SELECT {cols_str} FROM "{staging_name}"' + f' SELECT {cols_str} FROM {staging_ref}' f' ON CONFLICT ({conflict_cols}) DO NOTHING' ) @@ -191,12 +200,13 @@ def merge_upsert( session.execute(non_paginated_upsert) return - total = session.execute(sa.text(f'SELECT COUNT(*) FROM "{staging_name}"')).scalar_one() + total = session.execute(sa.text(f'SELECT COUNT(*) FROM {staging_ref}')).scalar_one() if total <= merge_batch_size: session.execute(non_paginated_upsert) return - session.execute(sa.text(f'CREATE INDEX IF NOT EXISTS "{staging_name}_rownum_idx" ON "{staging_name}" (_rownum)')) + staging_name = self.staging_name_for_table(table_cls.__tablename__) + session.execute(sa.text(f'CREATE INDEX IF NOT EXISTS "{staging_name}_rownum_idx" ON {staging_ref} (_rownum)')) session.commit() start = 0 @@ -205,7 +215,7 @@ def merge_upsert( session.execute( sa.text( f'INSERT INTO "{target_name}" ({cols_str})' - f' SELECT {cols_str} FROM "{staging_name}"' + f' SELECT {cols_str} FROM {staging_ref}' f' WHERE _rownum > :start AND _rownum <= :end' f' ON CONFLICT ({conflict_cols}) DO NOTHING' ), @@ -219,23 +229,23 @@ def merge_insert( table_cls: type["CSVTableProtocol"], session: so.Session, target_name: str, - staging_name: str, *, merge_batch_size: int | None = None, ) -> None: + staging_ref = self.qualified_staging_name(table_cls.__tablename__) insertable_cols = self._insertable_column_names(table_cls) cols_str = ", ".join(f'"{c}"' for c in insertable_cols) non_paginated_insert = sa.text( f'INSERT INTO "{target_name}" ({cols_str})' - f' SELECT {cols_str} FROM "{staging_name}"' + f' SELECT {cols_str} FROM {staging_ref}' ) if merge_batch_size is None: session.execute(non_paginated_insert) return - total = session.execute(sa.text(f'SELECT COUNT(*) FROM "{staging_name}"')).scalar_one() + total = session.execute(sa.text(f'SELECT COUNT(*) FROM {staging_ref}')).scalar_one() if total <= merge_batch_size: session.execute(non_paginated_insert) return @@ -244,7 +254,8 @@ def merge_insert( # INSERT in batch-sized transactions to bound WAL per commit. # session_replication_role='replica' is session-level and persists # across commits, so FK checks stay disabled for all batches. - session.execute(sa.text(f'CREATE INDEX IF NOT EXISTS "{staging_name}_rownum_idx" ON "{staging_name}" (_rownum)')) + staging_name = self.staging_name_for_table(table_cls.__tablename__) + session.execute(sa.text(f'CREATE INDEX IF NOT EXISTS "{staging_name}_rownum_idx" ON {staging_ref} (_rownum)')) session.commit() start = 0 @@ -253,7 +264,7 @@ def merge_insert( session.execute( sa.text( f'INSERT INTO "{target_name}" ({cols_str})' - f' SELECT {cols_str} FROM "{staging_name}"' + f' SELECT {cols_str} FROM {staging_ref}' f' WHERE _rownum > :start AND _rownum <= :end' ), {"start": start, "end": end}, diff --git a/src/orm_loader/backends/sqlite.py b/src/orm_loader/backends/sqlite.py index 06b58a9..b6174fe 100644 --- a/src/orm_loader/backends/sqlite.py +++ b/src/orm_loader/backends/sqlite.py @@ -26,6 +26,10 @@ class SQLiteBackend(DatabaseBackend): + @staticmethod + def staging_name_for_table(tablename: str) -> str: + return f"_staging_{tablename}" + def __init__( self, *, @@ -33,6 +37,7 @@ def __init__( journal_mode: str = "WAL", defer_foreign_keys: bool = True, ) -> None: + super().__init__(staging_schema=None) # SQLite does not support schema-qualified staging tables self.busy_timeout_ms = busy_timeout_ms self.journal_mode = self._validate_journal_mode(journal_mode) self.defer_foreign_keys = defer_foreign_keys @@ -90,8 +95,8 @@ def create_staging_table( self, table_cls: type["CSVTableProtocol"], session: so.Session, - staging_name: str, ) -> None: + staging_name = self.staging_name_for_table(table_cls.__tablename__) session.execute(sa.text(f'DROP TABLE IF EXISTS "{staging_name}";')) metadata = sa.MetaData() @@ -105,10 +110,10 @@ def create_staging_table( def drop_staging_table( self, + table_cls: type["CSVTableProtocol"], session: so.Session, - staging_name: str, ) -> None: - session.execute(sa.text(f'DROP TABLE IF EXISTS "{staging_name}"')) + session.execute(sa.text(f'DROP TABLE IF EXISTS "{self.staging_name_for_table(table_cls.__tablename__)}"')) def disable_fk_check(self, session: so.Session) -> str | int: previous_state = session.execute(text("PRAGMA foreign_keys")).scalar() @@ -137,11 +142,11 @@ def merge_replace( table_cls: type["CSVTableProtocol"], session: so.Session, target_name: str, - staging_name: str, pk_cols: list[str], *, merge_batch_size: int | None = None, ) -> None: + staging_name = self.staging_name_for_table(table_cls.__tablename__) if len(pk_cols) == 1: pk = pk_cols[0] session.execute( @@ -176,11 +181,11 @@ def merge_upsert( table_cls: type["CSVTableProtocol"], session: so.Session, target_name: str, - staging_name: str, pk_cols: list[str], *, merge_batch_size: int | None = None, ) -> None: + staging_name = self.staging_name_for_table(table_cls.__tablename__) insertable_cols = self._insertable_column_names(table_cls) cols_str = ", ".join(f'"{c}"' for c in insertable_cols) session.execute( @@ -197,10 +202,10 @@ def merge_insert( table_cls: type["CSVTableProtocol"], session: so.Session, target_name: str, - staging_name: str, *, merge_batch_size: int | None = None, ) -> None: + staging_name = self.staging_name_for_table(table_cls.__tablename__) insertable_cols = self._insertable_column_names(table_cls) cols_str = ", ".join(f'"{c}"' for c in insertable_cols) session.execute( diff --git a/src/orm_loader/helpers/__init__.py b/src/orm_loader/helpers/__init__.py index 651ab2f..6ad5a70 100644 --- a/src/orm_loader/helpers/__init__.py +++ b/src/orm_loader/helpers/__init__.py @@ -9,6 +9,7 @@ from .metadata import Base from .discovery import get_model_by_tablename from .null_handlers import normalise_null +from .sql import qualify_identifier __all__ = [ "IngestError", @@ -23,4 +24,5 @@ "Base", "get_model_by_tablename", "normalise_null", + "qualify_identifier", ] diff --git a/src/orm_loader/helpers/bulk.py b/src/orm_loader/helpers/bulk.py index e0c44e6..a74824b 100644 --- a/src/orm_loader/helpers/bulk.py +++ b/src/orm_loader/helpers/bulk.py @@ -3,12 +3,12 @@ from sqlalchemy import Engine from sqlalchemy.orm import Session from typing import Iterator -from ..backends.resolve import resolve_backend logger = logging.getLogger(__name__) def disable_fk_check(session: Session) -> str | int: """Disable foreign-key checks for the current session and return the previous state.""" + from ..backends.resolve import resolve_backend previous_state = resolve_backend(session).disable_fk_check(session) logger.info("Disabled foreign key checks for bulk load.") if not isinstance(previous_state, (str, int)): @@ -18,6 +18,7 @@ def disable_fk_check(session: Session) -> str | int: def enable_fk_check(session: Session) -> str | int: """Enable foreign-key checks for the current session and return the previous state.""" + from ..backends.resolve import resolve_backend previous_state = resolve_backend(session).enable_fk_check(session) logger.info("Explicitly re-enabled foreign key checks.") if not isinstance(previous_state, (str, int)): @@ -27,6 +28,7 @@ def enable_fk_check(session: Session) -> str | int: def restore_fk_check(session: Session, previous_state: str | int): """Restore foreign-key checks to a previously captured backend-specific state.""" + from ..backends.resolve import resolve_backend resolve_backend(session).restore_fk_check(session, previous_state) logger.info(f"Restored foreign key checks to state: {previous_state}") @@ -43,6 +45,7 @@ def bulk_load_context( This is a thin helper over ``DatabaseBackend.bulk_load_context()``. It exists so older call sites can keep using the helper import path. """ + from ..backends.resolve import resolve_backend backend = resolve_backend(session) with backend.bulk_load_context( session, @@ -60,7 +63,7 @@ def engine_with_replica_role(engine: Engine) -> Iterator[Engine]: This is engine-scoped rather than session-scoped. It is only available on backends that explicitly implement the behaviour. """ - + from ..backends.resolve import resolve_backend backend = resolve_backend(engine) method = getattr(backend, "engine_with_replica_role", None) if method is None: diff --git a/src/orm_loader/helpers/sql.py b/src/orm_loader/helpers/sql.py new file mode 100644 index 0000000..a4bf17f --- /dev/null +++ b/src/orm_loader/helpers/sql.py @@ -0,0 +1,21 @@ +from __future__ import annotations + + +def qualify_identifier(name: str, schema: str | None) -> str: + """ + Return a double-quoted, optionally schema-qualified SQL identifier. + + Parameters + ---------- + name + The SQL identifier to qualify (e.g. a table name). + schema + Schema name to prefix. If None, returns only the quoted identifier. + Useful for backends that do not support schema-qualified identifiers (e.g. SQLite). + + Returns + ------- + str + e.g. '"staging"."_staging_foo"' or '"_staging_foo"'. + """ + return f'"{schema}"."{name}"' if schema else f'"{name}"' diff --git a/src/orm_loader/loaders/loading_helpers.py b/src/orm_loader/loaders/loading_helpers.py index 7df8571..e02e91d 100644 --- a/src/orm_loader/loaders/loading_helpers.py +++ b/src/orm_loader/loaders/loading_helpers.py @@ -11,6 +11,8 @@ import pyarrow.csv as pv import io +from ..helpers.sql import qualify_identifier + _SAFE_ENCODING = re.compile(r'^[A-Za-z][A-Za-z0-9_-]*$') logger = logging.getLogger(__name__) @@ -265,12 +267,14 @@ def quick_load_pg( path: Path, session: so.Session, tablename: str, + schema: str | None = None, quote_mode: str = "auto", ) -> int: raw_conn = session.connection().connection if not hasattr(raw_conn, "cursor"): raise RuntimeError("Expected DB-API connection for COPY") + table_ref = qualify_identifier(tablename, schema) encoding = infer_encoding(path)['encoding'] or 'utf-8' if not _SAFE_ENCODING.match(encoding): @@ -308,7 +312,7 @@ def quick_load_pg( _csv_cols = [c.strip().lower().replace('_hash', '') for c in _raw_hdr.rstrip(_nl).split(delimiter)] _cols_sql = ", ".join(f'"{c}"' for c in _csv_cols) - logger.info(f"Bulk loading {tablename} via COPY (encoding={encoding}, delimiter={delimiter})") + logger.info(f"Bulk loading {table_ref} via COPY (encoding={encoding}, delimiter={delimiter})") cur = raw_conn.cursor() try: @@ -316,7 +320,7 @@ def quick_load_pg( stream = NormalisedCSVStream(f, encoding=encoding, delimiter=delimiter) with cur.copy( f''' - COPY "{tablename}" ({_cols_sql}) + COPY {table_ref} ({_cols_sql}) FROM STDIN WITH ( {copy_options} @@ -326,7 +330,7 @@ def quick_load_pg( while data := stream.read(COPY_BLOCK_SIZE): copy.write(data) session.flush() - total = session.execute(sa.text(f'SELECT COUNT(*) FROM "{tablename}"')).scalar_one() + total = session.execute(sa.text(f'SELECT COUNT(*) FROM {table_ref}')).scalar_one() return total except Exception as e: logger.error(f"Error during bulk load via COPY: {e}") diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index f00dbbf..35641d0 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -4,7 +4,7 @@ import logging from sqlalchemy.exc import InvalidRequestError, UnboundExecutionError -from typing import Type, ClassVar, Optional, Any, Iterator +from typing import Type, Any, Iterator from pathlib import Path from contextlib import contextmanager from time import perf_counter @@ -65,30 +65,23 @@ class CSVLoadableTableInterface(ORMTableBase): """ __abstract__ = True - _staging_tablename: ClassVar[Optional[str]] = None @classmethod def staging_tablename(cls: Type[CSVTableProtocol]) -> str: """ Return the name of the staging table for this model. - If a custom staging table name has been set on the class, it is - used; otherwise a default name derived from ``__tablename__`` - is returned. - Returns ------- str - The staging table name. + The staging table name, derived from ``__tablename__``. """ - if cls._staging_tablename: # type: ignore - return cls._staging_tablename return f"_staging_{cls.__tablename__}" @classmethod def create_staging_table( - cls: Type[CSVTableProtocol], - session: so.Session + cls: Type[CSVTableProtocol], + session: so.Session, ): """ Create a fresh staging table for ingestion. @@ -110,7 +103,7 @@ def create_staging_table( """ _require_bind(session) backend = resolve_backend(session) - backend.create_staging_table(cls, session, cls.staging_tablename()) + backend.create_staging_table(cls, session) @classmethod @contextmanager @@ -227,7 +220,6 @@ def get_staging_table( cls: Type[CSVTableProtocol], session: so.Session, ) -> sa.Table: - """ Return the reflected staging table, creating it if necessary. @@ -242,24 +234,26 @@ def get_staging_table( The reflected staging table. """ engine = _require_bind(session) + backend = resolve_backend(session) inspector = sa.inspect(engine) - staging_name = cls.staging_tablename() + staging_name = backend.staging_name_for_table(cls.__tablename__) - if not inspector.has_table(staging_name): + if not inspector.has_table(staging_name, schema=backend.staging_schema): logger.debug(f"Staging table {staging_name} does not exist; recreating",) cls.create_staging_table(session) return sa.Table( staging_name, - cls.metadata, + sa.MetaData(), # throwaway — keeps staging table out of Base.metadata autoload_with=engine, + schema=backend.staging_schema, ) - @classmethod + @classmethod def load_staging( cls: Type[CSVTableProtocol], loader: LoaderInterface, - loader_context: LoaderContext + loader_context: LoaderContext, ) -> int: """ Load data into the staging table. @@ -284,27 +278,21 @@ def load_staging( backend = resolve_backend(loader_context.session) total = 0 + cls.create_staging_table(loader_context.session) + try: - cls.create_staging_table(loader_context.session) + total = backend.load_staging_fast(loader_context=loader_context) + if total is not None: + return total + except Exception as e: + loader_context.session.rollback() + logger.warning(f"Fast-path load failed for {cls.__tablename__}: {e}") + logger.info('Falling back to ORM-based load functionality') - try: - total = backend.load_staging_fast( - loader_context=loader_context, - staging_name=cls.staging_tablename(), - ) - if total is not None: - return total - except Exception as e: - loader_context.session.rollback() - logger.warning(f"Fast-path load failed for {cls.staging_tablename()}: {e}") - logger.info('Falling back to ORM-based load functionality') - - total = cls.orm_staging_load( - loader=loader, - loader_context=loader_context - ) - finally: - cls._staging_tablename = None + total = cls.orm_staging_load( + loader=loader, + loader_context=loader_context + ) return total @classmethod @@ -448,7 +436,7 @@ def load_csv( logger.info(f"Table `{cls.__tablename__}`: Merging staging data into target table") with cls.manage_indices(session, index_strategy=index_strategy): cls.merge_from_staging(session, merge_strategy=merge_strategy, merge_batch_size=merge_batch_size) - + cls.drop_staging_table(session) logger.info(f"Table `{cls.__tablename__}`: Successfully finished ingestion. Total rows: {total}") @@ -497,7 +485,6 @@ def merge_from_staging( ``upsert``, or ``insert_if_empty``). """ target = cls.__tablename__ - staging = cls.staging_tablename() pk_cols = cls.pk_names() _require_bind(session) @@ -527,14 +514,14 @@ def merge_from_staging( if merge_strategy == "replace": logger.info(f"Table `{target}`: Merge replace delete phase starting.") delete_started = perf_counter() - backend.merge_replace(cls, session, target, staging, pk_cols, merge_batch_size=merge_batch_size) + backend.merge_replace(cls, session, target, pk_cols, merge_batch_size=merge_batch_size) logger.info( f"Table `{target}`: Merge replace delete phase completed in " f"{_format_elapsed(perf_counter() - delete_started)}." ) logger.info(f"Table `{target}`: Merge insert phase starting.") insert_started = perf_counter() - backend.merge_insert(cls, session, target, staging, merge_batch_size=merge_batch_size) + backend.merge_insert(cls, session, target, merge_batch_size=merge_batch_size) logger.info( f"Table `{target}`: Merge insert phase completed in " f"{_format_elapsed(perf_counter() - insert_started)}." @@ -542,7 +529,7 @@ def merge_from_staging( elif merge_strategy == "upsert": logger.info(f"Table `{target}`: Merge upsert phase starting.") upsert_started = perf_counter() - backend.merge_upsert(cls, session, target, staging, pk_cols, merge_batch_size=merge_batch_size) + backend.merge_upsert(cls, session, target, pk_cols, merge_batch_size=merge_batch_size) logger.info( f"Table `{target}`: Merge upsert phase completed in " f"{_format_elapsed(perf_counter() - upsert_started)}." @@ -568,7 +555,7 @@ def merge_from_staging( logger.info(f"Table `{target}`: Merge insert-if-empty phase starting.") insert_started = perf_counter() - backend.merge_insert(cls, session, target, staging, merge_batch_size=merge_batch_size) + backend.merge_insert(cls, session, target, merge_batch_size=merge_batch_size) logger.info( f"Table `{target}`: Merge insert-if-empty phase completed in " f"{_format_elapsed(perf_counter() - insert_started)}." @@ -577,12 +564,20 @@ def merge_from_staging( raise ValueError(f"Unknown merge strategy '{merge_strategy}'") @classmethod - def drop_staging_table(cls: Type[CSVTableProtocol], session: so.Session): + def drop_staging_table( + cls: Type[CSVTableProtocol], + session: so.Session, + ): """ Drop the staging table if it exists. + + Parameters + ---------- + session + An active SQLAlchemy session bound to an engine. """ backend = resolve_backend(session) - backend.drop_staging_table(session, cls.staging_tablename()) + backend.drop_staging_table(cls, session) @classmethod def csv_columns(cls) -> dict[str, sa.ColumnElement[Any]]: diff --git a/src/orm_loader/tables/typing.py b/src/orm_loader/tables/typing.py index 6b88e2b..14b7b35 100644 --- a/src/orm_loader/tables/typing.py +++ b/src/orm_loader/tables/typing.py @@ -60,11 +60,6 @@ class CSVTableProtocol(ORMTableProtocol, Protocol): No assumptions are made about the underlying database or schema. """ - _staging_tablename: ClassVar[Optional[str]] = None - - @classmethod - def staging_tablename(cls) -> str: ... - @classmethod def _select_loader(cls, path: Path) -> "LoaderInterface": ... diff --git a/tests/backends/test_base_backend.py b/tests/backends/test_base_backend.py index 212f1f6..d23f691 100644 --- a/tests/backends/test_base_backend.py +++ b/tests/backends/test_base_backend.py @@ -55,46 +55,49 @@ def capabilities(self) -> BackendCapabilities: supports_fk_toggle=True, ) + @staticmethod + def staging_name_for_table(tablename: str) -> str: + return f"_staging_{tablename}" + def create_staging_table( - self, table_cls: Type[CSVTableProtocol], session: so.Session, staging_name: str + self, _table_cls: Type[CSVTableProtocol], _session: so.Session ) -> None: return None - def drop_staging_table(self, session: so.Session, staging_name: str) -> None: + def drop_staging_table( + self, _table_cls: Type[CSVTableProtocol], _session: so.Session + ) -> None: return None def merge_replace( self, - table_cls: Type[CSVTableProtocol], - session: so.Session, - target_name: str, - staging_name: str, - pk_cols: list[str], + _table_cls: Type[CSVTableProtocol], + _session: so.Session, + _target_name: str, + _pk_cols: list[str], *, - merge_batch_size: int | None = None, + _merge_batch_size: int | None = None, ) -> None: return None def merge_upsert( self, - table_cls: Type[CSVTableProtocol], - session: so.Session, - target_name: str, - staging_name: str, - pk_cols: list[str], + _table_cls: Type[CSVTableProtocol], + _session: so.Session, + _target_name: str, + _pk_cols: list[str], *, - merge_batch_size: int | None = None, + _merge_batch_size: int | None = None, ) -> None: return None def merge_insert( self, - table_cls: Type[CSVTableProtocol], - session: so.Session, - target_name: str, - staging_name: str, + _table_cls: Type[CSVTableProtocol], + _session: so.Session, + _target_name: str, *, - merge_batch_size: int | None = None, + _merge_batch_size: int | None = None, ) -> None: return None @@ -160,7 +163,7 @@ def test_fake_backend_can_implement_contract(): assert backend.supports_dialect(Dialect.POSTGRESQL) is False assert backend.resolve_index_strategy("auto") == "drop_rebuild" assert backend.resolve_index_strategy("keep") == "keep" - assert backend.load_staging_fast(cast("LoaderContext", None), "staging") is None + assert backend.load_staging_fast(cast("LoaderContext", None)) is None with backend.merge_context(cast("Type[CSVTableProtocol]", None), cast(so.Session, None)): pass diff --git a/tests/backends/test_postgres_backend.py b/tests/backends/test_postgres_backend.py index 360b331..874aba1 100644 --- a/tests/backends/test_postgres_backend.py +++ b/tests/backends/test_postgres_backend.py @@ -8,15 +8,22 @@ from sqlalchemy.dialects import postgresql from sqlalchemy.engine import Connection, Engine -from orm_loader.backends import Dialect, PostgresBackend +from orm_loader.backends import STAGING_SCHEMA, Dialect, PostgresBackend +from orm_loader.helpers.sql import qualify_identifier + +_TARGET_TABLE = "target_table" +_STAGING_TABLE = f"_staging_{_TARGET_TABLE}" +_STAGING_TABLE_WITH_SCHEMA: str = qualify_identifier(_STAGING_TABLE, STAGING_SCHEMA) + if TYPE_CHECKING: from orm_loader.tables.typing import CSVTableProtocol class _ComputedTable: + __tablename__ = _TARGET_TABLE __table__ = sa.Table( - "target_table", + _TARGET_TABLE, sa.MetaData(), sa.Column("id", sa.Integer, primary_key=True), sa.Column("name", sa.String), @@ -80,11 +87,11 @@ def test_postgres_backend_create_staging_table_drops_computed_columns(): backend = PostgresBackend() session = _FakeSession() - backend.create_staging_table(_ComputedTableCls, _sess(session), "_staging_target_table") + backend.create_staging_table(_ComputedTableCls, _sess(session)) - assert any('DROP TABLE IF EXISTS "_staging_target_table"' in sql for sql in session.statements) - assert any('CREATE UNLOGGED TABLE "_staging_target_table"' in sql for sql in session.statements) - assert any('ALTER TABLE "_staging_target_table" DROP COLUMN "slug"' in sql for sql in session.statements) + assert any(f'DROP TABLE IF EXISTS {_STAGING_TABLE_WITH_SCHEMA}' in sql for sql in session.statements) + assert any(f'CREATE UNLOGGED TABLE {_STAGING_TABLE_WITH_SCHEMA}' in sql for sql in session.statements) + assert any(f'ALTER TABLE {_STAGING_TABLE_WITH_SCHEMA} DROP COLUMN "slug"' in sql for sql in session.statements) assert session.commits == 1 @@ -92,9 +99,9 @@ def test_postgres_backend_drop_staging_table(): backend = PostgresBackend() session = _FakeSession() - backend.drop_staging_table(_sess(session), "_staging_target_table") + backend.drop_staging_table(_ComputedTableCls, _sess(session)) - assert session.statements == ['DROP TABLE IF EXISTS "_staging_target_table"'] + assert session.statements == [f'DROP TABLE IF EXISTS {_STAGING_TABLE_WITH_SCHEMA}'] def test_postgres_backend_fk_methods_emit_expected_sql(): @@ -120,36 +127,34 @@ def test_postgres_backend_merge_replace_uses_using_delete(): backend = PostgresBackend() session = _FakeSession(scalar_result=0) - backend.merge_replace(_ComputedTableCls, _sess(session), "target_table", "_staging_target_table", ["id", "name"]) + backend.merge_replace(_ComputedTableCls, _sess(session), _TARGET_TABLE, ["id", "name"]) - # No merge_batch_size → non-paginated path → single DELETE statement at index 0 sql = session.statements[0] - assert 'DELETE FROM "target_table" t' in sql - assert 'USING "_staging_target_table" s' in sql + assert f'DELETE FROM "{_TARGET_TABLE}" t' in sql + assert f'USING {_STAGING_TABLE_WITH_SCHEMA} s' in sql assert 't."id" = s."id" AND t."name" = s."name"' in sql + assert f'USING {qualify_identifier(_TARGET_TABLE, STAGING_SCHEMA)}' not in sql def test_postgres_backend_merge_insert_excludes_computed_columns(): backend = PostgresBackend() session = _FakeSession(scalar_result=0) - backend.merge_insert(_ComputedTableCls, _sess(session), "target_table", "_staging_target_table") + backend.merge_insert(_ComputedTableCls, _sess(session), _TARGET_TABLE) - # No merge_batch_size → non-paginated path → single INSERT statement at index 0 sql = session.statements[0] - assert 'INSERT INTO "target_table" ("id", "name")' in sql - assert 'SELECT "id", "name" FROM "_staging_target_table"' in sql + assert f'INSERT INTO "{_TARGET_TABLE}" ("id", "name")' in sql + assert f'SELECT "id", "name" FROM {_STAGING_TABLE_WITH_SCHEMA}' in sql def test_postgres_backend_merge_upsert_excludes_computed_columns(): backend = PostgresBackend() session = _FakeSession(scalar_result=0) - backend.merge_upsert(_ComputedTableCls, _sess(session), "target_table", "_staging_target_table", ["id"]) + backend.merge_upsert(_ComputedTableCls, _sess(session), _TARGET_TABLE, ["id"]) - # No merge_batch_size → non-paginated path → single INSERT statement at index 0 sql = session.statements[0] - assert 'INSERT INTO "target_table" ("id", "name")' in sql + assert f'INSERT INTO "{_TARGET_TABLE}" ("id", "name")' in sql assert 'ON CONFLICT ("id") DO NOTHING' in sql @@ -158,7 +163,7 @@ def test_postgres_backend_merge_replace_paginated_path(): session = _FakeSession(scalar_result=10) backend.merge_replace( - _ComputedTableCls, _sess(session), "target_table", "_staging_target_table", + _ComputedTableCls, _sess(session), _TARGET_TABLE, ["id", "name"], merge_batch_size=3, ) @@ -173,7 +178,7 @@ def test_postgres_backend_merge_insert_paginated_path(): session = _FakeSession(scalar_result=10) backend.merge_insert( - _ComputedTableCls, _sess(session), "target_table", "_staging_target_table", + _ComputedTableCls, _sess(session), _TARGET_TABLE, merge_batch_size=3, ) @@ -188,7 +193,7 @@ def test_postgres_backend_merge_upsert_paginated_path(): session = _FakeSession(scalar_result=10) backend.merge_upsert( - _ComputedTableCls, _sess(session), "target_table", "_staging_target_table", + _ComputedTableCls, _sess(session), _TARGET_TABLE, ["id"], merge_batch_size=3, ) @@ -257,6 +262,7 @@ def test_postgres_backend_enable_fk_raises_when_show_returns_non_string(): raise AssertionError("Expected RuntimeError when SHOW returns a non-string") + def test_postgres_backend_engine_with_replica_role_unregisters_listener(monkeypatch): backend = PostgresBackend() events: list[tuple[str, object, str]] = [] diff --git a/tests/backends/test_sqlite_backend.py b/tests/backends/test_sqlite_backend.py index 5b1f060..d93e5e6 100644 --- a/tests/backends/test_sqlite_backend.py +++ b/tests/backends/test_sqlite_backend.py @@ -13,10 +13,14 @@ if TYPE_CHECKING: from orm_loader.tables.typing import CSVTableProtocol +_TARGET_TABLE = "target_table" +_STAGING_TABLE = f"_staging_{_TARGET_TABLE}" + class _ComputedTable: + __tablename__ = _TARGET_TABLE __table__ = sa.Table( - "target_table", + _TARGET_TABLE, sa.MetaData(), sa.Column("id", sa.Integer, primary_key=True), sa.Column("name", sa.String), @@ -66,10 +70,10 @@ def test_sqlite_backend_identity_and_capabilities(): def test_sqlite_backend_create_staging_table(session, engine): backend = SQLiteBackend() - backend.create_staging_table(_ComputedTableCls, session, "_staging_target_table") + backend.create_staging_table(_ComputedTableCls, session) inspector = sa.inspect(engine) - assert inspector.has_table("_staging_target_table") is True - cols = inspector.get_columns("_staging_target_table") + assert inspector.has_table(_STAGING_TABLE) is True + cols = inspector.get_columns(_STAGING_TABLE) assert [c["name"] for c in cols] == ["id", "name", "slug"] assert all(c["nullable"] is True for c in cols) @@ -78,9 +82,10 @@ def test_sqlite_backend_drop_staging_table(): backend = SQLiteBackend() session = _FakeSession() - backend.drop_staging_table(_sess(session), "_staging_target_table") + backend.drop_staging_table(_ComputedTableCls, _sess(session)) + + assert session.statements == [f'DROP TABLE IF EXISTS "{_STAGING_TABLE}"'] - assert session.statements == ['DROP TABLE IF EXISTS "_staging_target_table"'] def test_sqlite_backend_disable_fk_reads_then_sets(): @@ -154,12 +159,12 @@ def test_sqlite_backend_merge_replace_single_pk(): session = _FakeSession() backend.merge_replace( - _ComputedTableCls, _sess(session), "target_table", "_staging_target_table", ["id"] + _ComputedTableCls, _sess(session), _TARGET_TABLE, ["id"] ) sql = session.statements[0] - assert 'DELETE FROM "target_table"' in sql - assert 'SELECT "id" FROM "_staging_target_table"' in sql + assert f'DELETE FROM "{_TARGET_TABLE}"' in sql + assert f'SELECT "id" FROM "{_STAGING_TABLE}"' in sql def test_sqlite_backend_merge_replace_composite_pk(): @@ -167,24 +172,24 @@ def test_sqlite_backend_merge_replace_composite_pk(): session = _FakeSession() backend.merge_replace( - _ComputedTableCls, _sess(session), "target_table", "_staging_target_table", ["id", "name"] + _ComputedTableCls, _sess(session), _TARGET_TABLE, ["id", "name"] ) sql = session.statements[0] assert "WHERE EXISTS (" in sql - assert '"target_table"."id" = "_staging_target_table"."id"' in sql - assert '"target_table"."name" = "_staging_target_table"."name"' in sql + assert f'"{_TARGET_TABLE}"."id" = "{_STAGING_TABLE}"."id"' in sql + assert f'"{_TARGET_TABLE}"."name" = "{_STAGING_TABLE}"."name"' in sql def test_sqlite_backend_merge_insert_excludes_computed_columns(): backend = SQLiteBackend() session = _FakeSession() - backend.merge_insert(_ComputedTableCls, _sess(session), "target_table", "_staging_target_table") + backend.merge_insert(_ComputedTableCls, _sess(session), _TARGET_TABLE) sql = session.statements[0] - assert 'INSERT INTO "target_table" ("id", "name")' in sql - assert 'SELECT "id", "name" FROM "_staging_target_table"' in sql + assert f'INSERT INTO "{_TARGET_TABLE}" ("id", "name")' in sql + assert f'SELECT "id", "name" FROM "{_STAGING_TABLE}"' in sql def test_sqlite_backend_merge_upsert_excludes_computed_columns(): @@ -192,11 +197,11 @@ def test_sqlite_backend_merge_upsert_excludes_computed_columns(): session = _FakeSession() backend.merge_upsert( - _ComputedTableCls, _sess(session), "target_table", "_staging_target_table", ["id"] + _ComputedTableCls, _sess(session), _TARGET_TABLE, ["id"] ) sql = session.statements[0] - assert 'INSERT OR IGNORE INTO "target_table" ("id", "name")' in sql + assert f'INSERT OR IGNORE INTO "{_TARGET_TABLE}" ("id", "name")' in sql def test_sqlite_backend_materialized_view_methods_raise(engine): diff --git a/tests/conftest.py b/tests/conftest.py index a5a9abe..d3f0970 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -6,6 +6,7 @@ import sqlalchemy.orm as so from dotenv import load_dotenv +from orm_loader.backends import STAGING_SCHEMA from tests.models import Base load_dotenv(Path(__file__).parent.parent / ".env") @@ -65,6 +66,8 @@ def pg_engine(): def pg_session(pg_engine): Session = so.sessionmaker(pg_engine, future=True) with pg_engine.begin() as conn: + conn.execute(sa.text(f"DROP SCHEMA IF EXISTS {STAGING_SCHEMA} CASCADE")) + conn.execute(sa.text(f"CREATE SCHEMA {STAGING_SCHEMA}")) Base.metadata.drop_all(conn) Base.metadata.create_all(conn) diff --git a/tests/loaders/test_pg_loader.py b/tests/loaders/test_pg_loader.py index 9877e5e..9a44cb0 100644 --- a/tests/loaders/test_pg_loader.py +++ b/tests/loaders/test_pg_loader.py @@ -1,6 +1,7 @@ import sqlalchemy as sa import pandas as pd import pytest +from orm_loader.backends import STAGING_SCHEMA from orm_loader.loaders.loading_helpers import infer_encoding, infer_delim, check_line_ending, quick_load_pg from orm_loader.config import OrmLoaderConfig @@ -17,15 +18,16 @@ def test_copy_into_staging_with_extra_identity_column(pg_session, tmp_path): staging_name = SimpleTable.staging_tablename() cols = pg_session.execute(sa.text( - "SELECT column_name FROM information_schema.columns WHERE table_name = :t" - ), {"t": staging_name}).scalars().all() + "SELECT column_name FROM information_schema.columns" + " WHERE table_schema = :s AND table_name = :t" + ), {"s": STAGING_SCHEMA, "t": staging_name}).scalars().all() assert "_rownum" in cols, "staging table should have _rownum before COPY" - total = quick_load_pg(path=csv, session=pg_session, tablename=staging_name) + total = quick_load_pg(path=csv, session=pg_session, tablename=staging_name, schema=STAGING_SCHEMA) assert total == 2 rownums = pg_session.execute( - sa.text(f'SELECT _rownum FROM "{staging_name}" ORDER BY _rownum') + sa.text(f'SELECT _rownum FROM "{STAGING_SCHEMA}"."{staging_name}" ORDER BY _rownum') ).scalars().all() assert rownums == [1, 2], "_rownum must be auto-populated by IDENTITY sequence" @@ -191,9 +193,9 @@ def test_staging_schema_matches_target(pg_session, tmp_path): cols = pg_session.execute(sa.text(""" SELECT column_name, data_type FROM information_schema.columns - WHERE table_name = :table + WHERE table_schema = :s AND table_name = :table ORDER BY ordinal_position - """), {"table": SimpleTable.staging_tablename()}).all() + """), {"s": STAGING_SCHEMA, "table": SimpleTable.staging_tablename()}).all() data_cols = [(name, dtype) for name, dtype in cols if name != "_rownum"] assert data_cols == [