Skip to content

Overhauled staging tables with schemas and centralised specifiers #22

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
nicoloesch wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

Overhauled staging tables with schemas and centralised specifiers #22

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
48631ef
Overhauled staging tables with schemas and centralised specifiers
nicoloesch Jul 8, 2026
b09f060
Fix outstanding ruff error
nicoloesch Jul 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion src/orm_loader/backends/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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",
Expand Down
60 changes: 54 additions & 6 deletions src/orm_loader/backends/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand All @@ -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:
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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:
Expand Down
63 changes: 37 additions & 26 deletions src/orm_loader/backends/postgres.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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"
Expand All @@ -41,27 +48,27 @@ 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);
'''
)
)

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);"
)
)
Expand All @@ -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,
)

Expand Down Expand Up @@ -131,35 +139,36 @@ 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
while start < total:
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},
Expand All @@ -172,31 +181,32 @@ 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'
)

if merge_batch_size is None:
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
Expand All @@ -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'
),
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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},
Expand Down
Loading
Loading