From bcd334735f978190a39fbfdc38a2d38ab115d2e0 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 28 Oct 2026 04:21:34 +0000 Subject: [PATCH 01/21] Add null-handling regression tests --- tests/loaders/test_null_handling.py | 216 ++++++++++++++++++++++++++++ 1 file changed, 216 insertions(+) create mode 100644 tests/loaders/test_null_handling.py diff --git a/tests/loaders/test_null_handling.py b/tests/loaders/test_null_handling.py new file mode 100644 index 0000000..c12b68f --- /dev/null +++ b/tests/loaders/test_null_handling.py @@ -0,0 +1,216 @@ +"""Null handling: nonempty marker-like text is preserved verbatim. + +Typed columns accept a valid value or a recognised missing word and raise +otherwise. Empty and whitespace-only values are missing. +""" + +from datetime import date +from pathlib import Path +from typing import Type, cast + +import pyarrow as pa +import pytest +import sqlalchemy as sa +from oa_configurator.testing import DIALECT_PARAMS + +from orm_loader.backends import STAGING_SCHEMA, Dialect +from orm_loader.helpers import IngestError +from orm_loader.loaders.data.converters import cast_arrow_column, cast_scalar +from orm_loader.tables.typing import CSVTableProtocol +from tests.models import DemoRecord, DemoRequiredRecord + +_Demo = cast(Type[CSVTableProtocol], DemoRecord) +_DemoRequired = cast(Type[CSVTableProtocol], DemoRequiredRecord) + +MARKER_WORDS = [ + "Null", "NULL", "null", "None", "none", "NA", "N/A", "NaN", "nan", + "#N/A", "", "Contains NULL here", '"NULL"', +] + +SEVEN_ROWS = 'demo_id,demo_label\n1,NULL\n2,"NULL"\n3,"""NULL"""\n4,\n5,""\n6,"alpha,beta"\n7,"alpha\nbeta"\n' +SEVEN_ROWS_QUOTED = ( + '"demo_id","demo_label"\n"1","NULL"\n"2","NULL"\n"3","""NULL"""\n"4",""\n"5",""\n' + '"6","alpha,beta"\n"7","alpha\nbeta"\n' +) +SEVEN_EXPECTED = { + 1: "NULL", 2: "NULL", 3: '"NULL"', 4: None, 5: None, 6: "alpha,beta", 7: "alpha\nbeta", +} + + +_SESSION_FIXTURE_BY_DIALECT = { + Dialect.SQLITE: "sqlite_session", + Dialect.POSTGRESQL: "pg_session", + Dialect.DUCKDB: "duckdb_session", +} + + +@pytest.fixture(params=DIALECT_PARAMS) +def session(request): + return request.getfixturevalue(_SESSION_FIXTURE_BY_DIALECT[request.param]) + + +def _write(tmp_path: Path, table: type, text: str) -> Path: + path = tmp_path / f"{table.__tablename__}.csv" + path.write_text(text, encoding="utf-8", newline="") + return path + + +def _load(session, tmp_path: Path, table, text: str, **kwargs) -> int: + tag = None if session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA + n = table.load_csv(session, _write(tmp_path, table, text), staging_schema_tag=tag, **kwargs) + session.commit() + session.expire_all() + return n + + +def _labels(session, model) -> dict[int, str | None]: + rows = session.execute(sa.select(model.demo_id, model.demo_label).order_by(model.demo_id)).all() + return {r[0]: r[1] for r in rows} + + +# Marker-like text through each entry point + +@pytest.mark.parametrize("word", MARKER_WORDS) +def test_scalar_text_keeps_marker_words(word): + assert cast_scalar(word, sa.String()) == word + + +@pytest.mark.parametrize("word", MARKER_WORDS) +def test_arrow_text_keeps_marker_words(word): + out = cast_arrow_column(pa.array([word, "plain"]), DemoRecord.__table__.c.demo_label) + assert out.to_pylist() == [word, "plain"] + + +@pytest.mark.parametrize("word", MARKER_WORDS) +def test_dict_path_keeps_marker_words(word): + obj = DemoRecord(**DemoRecord.clean_kwargs({"demo_id": 1, "demo_label": word})) + assert obj.demo_label == word + + +# Actual missing objects + +@pytest.mark.parametrize("value", [None, float("nan"), "", " "]) +def test_scalar_missing_objects_are_null(value): + assert cast_scalar(value, sa.String()) is None + + +def test_arrow_null_bitmap_and_nan_are_null(): + col = DemoRecord.__table__.c.demo_number + assert cast_arrow_column(pa.array([1, None, 0], pa.int64()), col).to_pylist() == [1, None, 0] + assert cast_arrow_column(pa.array([1.0, float("nan")], pa.float64()), col).to_pylist() == [1, None] + + +# Whitespace and identifiers + +def test_text_is_stripped_by_default(): + assert cast_scalar(" x ", sa.String()) == "x" + + +@pytest.mark.parametrize("code", ["007", "0123456", "1.50", "+5", "12.0"]) +def test_text_is_never_canonicalised_as_a_number(code): + assert cast_scalar(code, sa.String()) == code + assert cast_arrow_column(pa.array([code]), DemoRecord.__table__.c.demo_label).to_pylist() == [code] + + +# Typed columns: valid value, missing word, or error + +@pytest.mark.parametrize("word", ["null", "None", "NA", "N/A", "NaN", "#N/A"]) +def test_typed_missing_words_are_null(word): + assert cast_scalar(word, sa.Integer()) is None + assert cast_scalar(word, sa.Date()) is None + + +@pytest.mark.parametrize("word", ["null", "None", "NA", "N/A", "NaN"]) +def test_typed_float_missing_words_are_null(word): + assert cast_scalar(word, sa.Float()) is None + + +# Persisted round trips + +def test_seven_row_fixture_preserves_text(session, tmp_path): + _load(session, tmp_path, _Demo, SEVEN_ROWS) + assert _labels(session, DemoRecord) == SEVEN_EXPECTED + + +def test_seven_row_fixture_every_field_quoted(session, tmp_path): + _load(session, tmp_path, _Demo, SEVEN_ROWS_QUOTED) + assert _labels(session, DemoRecord) == SEVEN_EXPECTED + + +def test_marker_words_survive_tab_delimited_load(session, tmp_path): + words = [w for w in MARKER_WORDS if w != '"NULL"'] + text = "demo_id\tdemo_label\n" + "".join(f"{i}\t{w}\n" for i, w in enumerate(words, 1)) + _load(session, tmp_path, _Demo, text) + assert _labels(session, DemoRecord) == dict(enumerate(words, 1)) + + +def test_literal_quotes_survive_tab_delimited_load(session, tmp_path): + values = ['"NULL"', '"word" label', "Null", "NA"] + text = "demo_id\tdemo_label\n" + "".join(f"{i}\t{v}\n" for i, v in enumerate(values, 1)) + "5\t\n" + _load(session, tmp_path, _Demo, text, quote_mode="literal") + assert _labels(session, DemoRecord) == {**dict(enumerate(values, 1)), 5: None} + + +def test_typed_values_round_trip(session, tmp_path): + text = ( + "demo_id,demo_label,demo_number,demo_flag,demo_date\n" + "1,a,0,false,20200229\n" + "2,b,NA,,\n" + "3,c,7,true,2020-01-02\n" + ) + _load(session, tmp_path, _Demo, text) + rows = session.execute(sa.select( + DemoRecord.demo_id, DemoRecord.demo_number, DemoRecord.demo_flag, DemoRecord.demo_date + ).order_by(DemoRecord.demo_id)).all() + assert [tuple(r) for r in rows] == [ + (1, 0, False, date(2020, 2, 29)), + (2, None, None, None), + (3, 7, True, date(2020, 1, 2)), + ] + + +@pytest.mark.parametrize( + "column, bad", + [ + ("demo_date", "20200231"), + ("demo_date", "Jan 2 2020"), + ("demo_date", "02/01/2020"), + ("demo_number", "12abc"), + ("demo_number", "0x1F"), + ("demo_number", "1.5"), + ], +) +def test_unreadable_typed_value_raises_and_loads_nothing(session, tmp_path, column, bad): + cols = {"demo_number": "", "demo_date": ""} + cols[column] = bad + text = f"demo_id,demo_label,demo_number,demo_date\n1,a,,\n2,b,{cols['demo_number']},{cols['demo_date']}\n" + with pytest.raises(IngestError, match=column): + _load(session, tmp_path, _Demo, text) + session.rollback() + assert session.execute(sa.select(sa.func.count()).select_from(DemoRecord)).scalar_one() == 0 + + +# Required columns + +def test_required_text_accepts_marker_words(session, tmp_path): + words = ["Null", "None", "NA", "N/A", "NaN"] + text = "demo_id,demo_label\n" + "".join(f"{i},{w}\n" for i, w in enumerate(words, 1)) + _load(session, tmp_path, _DemoRequired, text) + assert _labels(session, DemoRequiredRecord) == dict(enumerate(words, 1)) + + +@pytest.mark.parametrize("missing", ["", '""', " "]) +def test_required_text_empty_raises_and_loads_nothing(session, tmp_path, missing): + text = f"demo_id,demo_label\n1,ok\n2,{missing}\n" + with pytest.raises(IngestError, match="demo_label"): + _load(session, tmp_path, _DemoRequired, text) + session.rollback() + assert session.execute(sa.select(sa.func.count()).select_from(DemoRequiredRecord)).scalar_one() == 0 + + +def test_required_empty_raises_atomically_with_batching(session, tmp_path): + text = "demo_id,demo_label\n1,a\n2,b\n3,\n" + with pytest.raises(IngestError, match="demo_label"): + _load(session, tmp_path, _DemoRequired, text, merge_batch_size=1) + session.rollback() + assert session.execute(sa.select(sa.func.count()).select_from(DemoRequiredRecord)).scalar_one() == 0 From bbbb1fb32a86a20bc63fbcc71e30b899f91424f2 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 28 Oct 2026 04:22:14 +0000 Subject: [PATCH 02/21] Add demo models for null-handling tests and ignore local benchmarks --- .gitignore | 4 +++- tests/models.py | 24 ++++++++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 8475748..6b40e38 100644 --- a/.gitignore +++ b/.gitignore @@ -212,4 +212,6 @@ OMOP_CDM*.csv .vscode/ .DS_Store _temp/ -notebooks/ \ No newline at end of file +notebooks/ +# Local benchmarks +benchmarks/ diff --git a/tests/models.py b/tests/models.py index 7ad0adb..3d6be90 100644 --- a/tests/models.py +++ b/tests/models.py @@ -1,4 +1,5 @@ +from datetime import date from enum import Enum import sqlalchemy as sa @@ -201,3 +202,26 @@ class SelfReferentialTable(Base, CSVLoadableTableInterface): sa.Integer, sa.ForeignKey(f"{SchemaRole.PRIMARY.value}.self_referential_table.id", use_alter=True), ) + + +class DemoRecord(Base, CSVLoadableTableInterface): + """Optional text label plus one column of each typed kind, for null-handling tests.""" + + __tablename__ = "demo_record" + __table_args__ = {"schema": SchemaRole.PRIMARY.value} + + demo_id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) + demo_label: so.Mapped[str | None] = so.mapped_column(sa.String, nullable=True) + demo_number: so.Mapped[int | None] = so.mapped_column(sa.Integer, nullable=True) + demo_flag: so.Mapped[bool | None] = so.mapped_column(sa.Boolean, nullable=True) + demo_date: so.Mapped[date | None] = so.mapped_column(sa.Date, nullable=True) + + +class DemoRequiredRecord(Base, CSVLoadableTableInterface): + """Same shape as DemoRecord with a required text label.""" + + __tablename__ = "demo_required_record" + __table_args__ = {"schema": SchemaRole.PRIMARY.value} + + demo_id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) + demo_label: so.Mapped[str] = so.mapped_column(sa.String, nullable=False) From ca15ae1d9f5e0d2664500c6cac9fe4f0bf80293a Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Thu, 29 Oct 2026 01:20:56 +0000 Subject: [PATCH 03/21] Single file-to-sink pipeline with conversion policy and column handlers --- README.md | 15 +- docs/loaders/context.md | 49 +- docs/loaders/conversion.md | 100 ++++ docs/loaders/helpers.md | 70 --- docs/loaders/index.md | 117 ++--- docs/loaders/loaders.md | 92 ---- docs/loaders/readers.md | 57 +++ docs/loaders/staging.md | 36 ++ docs/tables/index.md | 2 +- docs/tables/loadable_table.md | 8 +- docs/tables/typing.md | 7 +- mkdocs.yml | 9 +- pyproject.toml | 10 +- src/orm_loader/backends/__init__.py | 2 + src/orm_loader/backends/base.py | 162 +++++-- src/orm_loader/backends/duckdb.py | 71 ++- src/orm_loader/backends/postgres.py | 68 ++- src/orm_loader/backends/sqlite.py | 57 ++- src/orm_loader/conversion/__init__.py | 37 ++ .../conversion/handlers/__init__.py | 21 + src/orm_loader/conversion/handlers/base.py | 137 ++++++ .../conversion/handlers/booleans.py | 40 ++ src/orm_loader/conversion/handlers/dates.py | 182 +++++++ src/orm_loader/conversion/handlers/numbers.py | 70 +++ src/orm_loader/conversion/handlers/text.py | 52 ++ src/orm_loader/conversion/policy.py | 220 +++++++++ src/orm_loader/conversion/values.py | 94 ++++ src/orm_loader/helpers/__init__.py | 15 +- src/orm_loader/helpers/bulk.py | 32 +- src/orm_loader/helpers/errors.py | 84 ++++ src/orm_loader/helpers/null_handlers.py | 24 - src/orm_loader/helpers/sqlite.py | 47 -- src/orm_loader/loaders/__init__.py | 22 +- src/orm_loader/loaders/data/__init__.py | 6 - src/orm_loader/loaders/data/converters.py | 452 ------------------ src/orm_loader/loaders/data_classes.py | 243 ---------- src/orm_loader/loaders/loader.py | 212 ++++++++ src/orm_loader/loaders/loader_interface.py | 299 ------------ src/orm_loader/loaders/loading_helpers.py | 348 -------------- src/orm_loader/loaders/processor.py | 259 ++++++++++ src/orm_loader/loaders/readers/__init__.py | 51 ++ src/orm_loader/loaders/readers/base.py | 75 +++ src/orm_loader/loaders/readers/delimited.py | 141 ++++++ src/orm_loader/loaders/readers/detection.py | 119 +++++ src/orm_loader/loaders/readers/parquet.py | 53 ++ src/orm_loader/tables/__init__.py | 8 +- src/orm_loader/tables/loadable_table.py | 209 ++++---- src/orm_loader/tables/orm_table.py | 125 +---- src/orm_loader/tables/typing.py | 29 +- tests/backends/test_base_backend.py | 33 +- tests/backends/test_duckdb_backend.py | 19 +- tests/backends/test_postgres_backend.py | 76 ++- tests/backends/test_shared_backend.py | 19 +- tests/backends/test_sinks.py | 87 ++++ tests/backends/test_sqlite_backend.py | 10 +- tests/conftest.py | 33 +- tests/conversion/test_boolean_handler.py | 17 + tests/conversion/test_date_handlers.py | 142 ++++++ tests/conversion/test_missing_values.py | 28 ++ tests/conversion/test_number_handlers.py | 40 ++ tests/conversion/test_policy.py | 109 +++++ tests/conversion/test_text_handler.py | 60 +++ tests/loaders/test_casting.py | 312 ------------ tests/loaders/test_dedupe.py | 32 +- tests/loaders/test_helpers.py | 39 +- tests/loaders/test_load_policy.py | 283 +++++++++++ tests/loaders/test_loader_e2e.py | 241 +++------- tests/loaders/test_null_handling.py | 132 ++--- tests/loaders/test_parquet_loader.py | 91 +--- tests/loaders/test_pg_loader.py | 427 ----------------- tests/loaders/test_pipeline.py | 144 ++++++ tests/loaders/test_processor.py | 140 ++++++ tests/loaders/test_readers.py | 115 +++++ tests/loaders/test_schema_translate_map.py | 20 +- tests/loaders/test_split_connection.py | 12 +- ...st_pandas_loader.py => test_stage_file.py} | 92 ++-- tests/models.py | 62 +-- tests/support.py | 38 ++ tests/tables/test_typing_protocols.py | 12 +- uv.lock | 75 ++- 80 files changed, 4154 insertions(+), 3424 deletions(-) create mode 100644 docs/loaders/conversion.md delete mode 100644 docs/loaders/helpers.md delete mode 100644 docs/loaders/loaders.md create mode 100644 docs/loaders/readers.md create mode 100644 docs/loaders/staging.md create mode 100644 src/orm_loader/conversion/__init__.py create mode 100644 src/orm_loader/conversion/handlers/__init__.py create mode 100644 src/orm_loader/conversion/handlers/base.py create mode 100644 src/orm_loader/conversion/handlers/booleans.py create mode 100644 src/orm_loader/conversion/handlers/dates.py create mode 100644 src/orm_loader/conversion/handlers/numbers.py create mode 100644 src/orm_loader/conversion/handlers/text.py create mode 100644 src/orm_loader/conversion/policy.py create mode 100644 src/orm_loader/conversion/values.py delete mode 100644 src/orm_loader/helpers/null_handlers.py delete mode 100644 src/orm_loader/helpers/sqlite.py delete mode 100644 src/orm_loader/loaders/data/__init__.py delete mode 100644 src/orm_loader/loaders/data/converters.py delete mode 100644 src/orm_loader/loaders/data_classes.py create mode 100644 src/orm_loader/loaders/loader.py delete mode 100644 src/orm_loader/loaders/loader_interface.py delete mode 100644 src/orm_loader/loaders/loading_helpers.py create mode 100644 src/orm_loader/loaders/processor.py create mode 100644 src/orm_loader/loaders/readers/__init__.py create mode 100644 src/orm_loader/loaders/readers/base.py create mode 100644 src/orm_loader/loaders/readers/delimited.py create mode 100644 src/orm_loader/loaders/readers/detection.py create mode 100644 src/orm_loader/loaders/readers/parquet.py create mode 100644 tests/backends/test_sinks.py create mode 100644 tests/conversion/test_boolean_handler.py create mode 100644 tests/conversion/test_date_handlers.py create mode 100644 tests/conversion/test_missing_values.py create mode 100644 tests/conversion/test_number_handlers.py create mode 100644 tests/conversion/test_policy.py create mode 100644 tests/conversion/test_text_handler.py delete mode 100644 tests/loaders/test_casting.py create mode 100644 tests/loaders/test_load_policy.py delete mode 100644 tests/loaders/test_pg_loader.py create mode 100644 tests/loaders/test_pipeline.py create mode 100644 tests/loaders/test_processor.py create mode 100644 tests/loaders/test_readers.py rename tests/loaders/{test_pandas_loader.py => test_stage_file.py} (53%) create mode 100644 tests/support.py diff --git a/README.md b/README.md index f6c3953..18739b3 100644 --- a/README.md +++ b/README.md @@ -40,25 +40,26 @@ class MyTable(ORMTableBase, Base): ``` You can inherit from it directly or pick it up through one of the higher-level mixins. -2. CSV-based ingestion mixins +2. File ingestion mixins -`CSVLoadableTableInterface` adds staged file loading to ORM tables. It can use pandas or PyArrow loaders, and on PostgreSQL it can use a fast `COPY` path when the input is clean enough. +`FileLoadableTableInterface` adds staged file loading to ORM tables. Files are read with PyArrow on every backend, interpreted under a `ConversionPolicy`, and written to staging through each backend's bulk path. Features include: * staging table creation and cleanup -* chunked loading for large files -* optional casting and deduplication before insert +* files streamed in bounded batches, so memory does not grow with the row count +* optional conversion to column types and deduplication of the staging table * backend-specific merge behaviour -* PostgreSQL fast-path loading with ORM fallback +* Arrow-based loading with a configurable conversion policy on PostgreSQL, DuckDB and SQLite +* delimited text (tab or comma delimiter, detected encoding) and Parquet input * backend-aware index handling during merge ```python -class MyTable(CSVLoadableTableInterface, ORMTableBase, Base): +class MyTable(FileLoadableTableInterface, ORMTableBase, Base): __tablename__ = "my_table" ``` -The main extension points here are loader choice, column mapping, and the normal SQLAlchemy model definitions themselves. Most downstream projects do not need to override much beyond `csv_columns()` and the model schema. +The main extension points here are reader choice and the normal SQLAlchemy model definitions themselves. Most downstream projects do not need to override much beyond the model schema. 3. Structured serialisation and hashing diff --git a/docs/loaders/context.md b/docs/loaders/context.md index 59a4abe..00f728d 100644 --- a/docs/loaders/context.md +++ b/docs/loaders/context.md @@ -1,4 +1,4 @@ -# Loader Context and Diagnostics +# Loader Context and Report This page documents the shared data structures used by the loaders subsystem. @@ -22,47 +22,28 @@ on globals or implicit configuration. | `session` | Active SQLAlchemy session | | `path` | Path to the input file | | `staging_table` | SQLAlchemy Table used for staging | -| `chunksize` | Optional chunk size | -| `normalise` | Whether to cast values to ORM types | -| `dedupe` | Whether to deduplicate incoming data | -| `quote_mode` | CSV quoting mode, resolved to the same concrete mode on both the PostgreSQL COPY fast-path and the pandas ORM fallback | +| `chunksize` | Rows per batch for the default reader of formats that batch by rows | +| `normalise` | Whether to interpret values under the policy; when false, values are written as read | +| `dedupe` | Whether the staging table keeps only the first row for each primary key | +| `quote_mode` | Quoting mode for delimited text: `csv`, `literal`, `by_delimiter` or `auto` | +| `policy` | `ConversionPolicy`: how values are interpreted and what happens to unreadable ones | +| `rejects_dir` | Directory that rejected rows are written to | +| `report` | `LoadReport` row accounting | -::: orm_loader.loaders.data_classes.LoaderContext +::: orm_loader.loaders.loader.LoaderContext --- -## Casting statistics +## LoadReport -Loaders track **casting failures** during normalisation to support -debugging and auditability. +`LoadReport` accounts for every row of a load: rows read, loaded, rejected and removed as duplicates. `DedupeResult` holds the key columns and the rows removed. A load that cannot account for every row raises. -### ColumnCastingStats - -Tracks casting failures for a single column: - -- total failure count -- representative example values - -### TableCastingStats - -Aggregates column-level statistics for a table. - -This enables loaders to emit warnings such as: - -- unexpected nulls in required columns -- values that could not be cast to target types - -Statistics are logged, not raised as exceptions. - -::: orm_loader.loaders.data_classes.ColumnCastingStats -::: orm_loader.loaders.data_classes.TableCastingStats +::: orm_loader.loaders.loader.LoadReport +::: orm_loader.backends.base.DedupeResult --- ## Design notes -- Casting failures do **not** abort loads -- Rows violating required constraints are dropped explicitly -- Examples are capped to avoid log flooding - -This design favours **observability over strict enforcement**. +- Unreadable values and missing required values raise by default, so no row disappears without a trace +- With `ErrorPolicy(on_error="reject_row")` each rejected row is written to a rejects file with its raw text and the reason diff --git a/docs/loaders/conversion.md b/docs/loaders/conversion.md new file mode 100644 index 0000000..b279256 --- /dev/null +++ b/docs/loaders/conversion.md @@ -0,0 +1,100 @@ +# Conversion + +`orm_loader.conversion` interprets raw values for every file load, whatever the file format. + +```mermaid +flowchart LR + A[Arrow column] --> B[Policy picks the column's handler] + B --> C[clean] + C --> D[detect missing] + D --> E[parse] + E --> F[check] + F --> G[values and invalid mask] +``` + +--- + +## ConversionPolicy + +`ConversionPolicy` is a frozen settings object passed to `load_file(policy=...)`. It holds the settings shared by every column and one handler for each kind of column type. + +| Field | Meaning | +|------|------| +| `nulls` | `NullPolicy`: what counts as a missing value | +| `errors` | `ErrorPolicy`: what happens to unreadable and missing required values | +| `text`, `integer`, `floating`, `numeric`, `boolean`, `date`, `datetime` | The handler for that kind of column, each with its own settings | +| `handlers` | Handlers for further column types, or replacements, by SQLAlchemy type class | +| `columns` | Handlers for single columns by `(table_name, column_name)` | + +```python +policy = ConversionPolicy( + nulls=NullPolicy(text_markers=frozenset({"NULL"})), + errors=ErrorPolicy(on_error="reject_row", max_rejects=0.01), + date=DateHandler(formats=("dd/mm/yyyy",)), + columns={("concept", "concept_name"): TextHandler(markers=frozenset({"N/A"}))}, +) +``` + +::: orm_loader.conversion.ConversionPolicy + +--- + +## NullPolicy + +- Text columns keep every nonempty value verbatim. Only empty and whitespace-only text is null. +- `text_markers` adds exact, case-sensitive words that are null in text columns. `TextHandler(markers=...)` replaces them for the columns it serves. +- `marker_quoting="unquoted_only"` matches a marker only when the source did not quote it. +- Typed columns accept a valid value, a missing word (`null`, `none`, `na`, `n/a`, `nan` and similar) or a real null, and fail otherwise. + +::: orm_loader.conversion.NullPolicy + +--- + +## ErrorPolicy + +`on_error` is `raise` (default) or `reject_row`. `max_rejects` stops a `reject_row` load once rejected rows pass a count or a fraction of the rows read. + +::: orm_loader.conversion.ErrorPolicy + +--- + +## Handlers + +The handler for a column is the first match of: the entry in `ConversionPolicy.columns` and the handler for the most specific class of the column type. A column type with no handler is passed through unchanged. + +| Handler | Column types | Accepted text | Settings | +|------|------|------|------| +| `TextHandler` | `String`, `Text`, `Enum` | Any nonempty text; longer than the column is an error | `strip`, `markers` | +| `IntegerHandler` | `Integer` | Optional sign and digits | none | +| `FloatHandler` | `Float` | Decimal or exponent form, no `Infinity` | none | +| `NumericHandler` | `Numeric` | Decimal form, converted to the declared precision and scale without rounding | none | +| `BooleanHandler` | `Boolean` | `true`, `t`, `yes`, `y`, `1`, `on` and their opposites | `true_words`, `false_words` | +| `DateHandler` | `Date` | One of `formats`, validated against the calendar | `formats`, `parser` | +| `DateTimeHandler` | `DateTime` | ISO date and time, or a date in one of `formats` | `formats`, `parser` | + +Dates default to `yyyymmdd`, `dd-mon-yyyy` and `yyyy-mm-dd`. Ambiguous forms such as `02/01/2020` are never guessed; declare `DateHandler(formats=("dd/mm/yyyy",))` or `("mm/dd/yyyy",)`. A timestamp column reads date-only values with its own `formats`. + +`DateHandler(parser=...)` replaces the date parser with a function that has the signature of `parse_dates`. + +### Adding a column type + +Subclass `ConvertingHandler`, implement `arrow_type` and `parse`, and pass the handler in `ConversionPolicy.handlers` for the type class, or in `ConversionPolicy.columns` for a single column. `clean`, `missing` and `check` can be overridden. + +```python +class HexHandler(ConvertingHandler): + def arrow_type(self, sa_type): + return pa.int64() + + def parse(self, cleaned, sa_type): + ... # return Interpretation(values, invalid) + +policy = ConversionPolicy(handlers={MyHexType: HexHandler()}) +``` + +### ConvertingHandler + +::: orm_loader.conversion.ConvertingHandler + +### ColumnHandler + +::: orm_loader.conversion.ColumnHandler diff --git a/docs/loaders/helpers.md b/docs/loaders/helpers.md deleted file mode 100644 index 8842dd3..0000000 --- a/docs/loaders/helpers.md +++ /dev/null @@ -1,70 +0,0 @@ -# Loader Helper Utilities - -This page covers the low-level functions that support the loader implementations. - ---- - -## Encoding and delimiter detection - -### `infer_encoding(path)` - -Detects file encoding using `chardet`. - -- ASCII is normalised to UTF-8 -- Only a small prefix of the file is inspected - -### `infer_delim(path)` - -Detects CSV delimiter by inspecting the first line. - -- Prefers tabs over commas if more frequent - ---- - -## Duplicate detection (PyArrow) - -### `arrow_drop_duplicates(table, pk_names)` - -Removes duplicate rows from a PyArrow table based on primary key columns. - -- Sorts by PKs -- Keeps the first occurrence -- Handles chunked arrays safely - -Used by `ParquetLoader` for internal deduplication. - ---- - -## Batch-oriented CSV parsing - -### `conservative_load_parquet(...)` - -Despite the name, this helper reads delimited text with PyArrow and yields batches: - -- strict column inclusion -- malformed row skipping -- chunked batch iteration - -This is used by the PyArrow-based loader path. - ---- - -## PostgreSQL fast-path loading - -### `quick_load_pg(...)` - -Loads CSV files into a PostgreSQL staging table using `COPY`. - -### Characteristics - -- Fast -- Bypasses ORM row construction -- Works best on clean input - -### Failure handling - -- Errors trigger rollback -- `CSVLoadableTableInterface` falls back to ORM-based loading -- Failures are noisy on purpose - -This helper is only used when explicitly supported by the database. diff --git a/docs/loaders/index.md b/docs/loaders/index.md index 3d41b0f..53d3ac4 100644 --- a/docs/loaders/index.md +++ b/docs/loaders/index.md @@ -1,100 +1,43 @@ # Loaders -The `orm_loader.loaders` module provides conservative, schema-aware file -loading into ORM-backed staging tables. - -This subsystem is designed to handle: - -- untrusted or messy source files -- large datasets requiring chunked processing -- repeatable staged loads -- dialect-specific optimisations (e.g. PostgreSQL COPY) -- explicit, inspectable failure modes - -Loaders are intentionally **infrastructure-only**: -they do not embed domain rules or business semantics. - ---- - -## Core concepts - -### LoaderContext - -[`LoaderContext`](context.md) - -A `LoaderContext` object carries the state required to load one file: - -- target ORM table -- database session -- staging table -- file path -- operational flags (chunking, deduplication, normalisation) - -This makes loader behaviour explicit, testable, and side-effect free. - ---- - -### LoaderInterface - -[`LoaderInterface`](loaders.md) - -All loaders implement a common interface: - -- `orm_file_load(ctx)` — orchestrates file ingestion -- `dedupe(data, ctx)` — defines deduplication semantics - -Concrete implementations mainly differ in how they read and transform incoming data. - ---- - -### Staging tables - -Loaders always write to **staging tables**, never directly to production -tables. - -This gives you: - -- safe rollback -- repeatable merges -- bulk loading optimisations - -Final merge semantics are handled by the table mixins, not by loaders. - ---- - -## Provided loaders - -| Loader | Use case | -|------|----------| -| `PandasLoader` | Flexible CSV and TSV ingestion | -| `ParquetLoader` | Columnar or batch-oriented ingestion | - -Both loaders share the same lifecycle and guarantees. +`orm_loader.loaders` loads files into staging tables. A file takes the same path on every backend and in every format: + +```mermaid +flowchart LR + A[File] --> B[FileReader] + B --> C[Arrow record batches] + C --> D[BatchProcessor with the ConversionPolicy] + D --> E[Backend sink] + E --> F[Staging table] +``` + +Loading is infrastructure only. It carries no domain rules, streams files in bounded batches, and fails visibly instead of changing data. + +| Page | Covers | +|------|------| +| [Readers](readers.md) | `FileReader`, the provided readers, format detection and how to add a format | +| [Conversion](conversion.md) | The policy that turns raw values into column types, and the column handlers | +| [Staging and sinks](staging.md) | `stage_file`, the backend sinks, deduplication and `normalise` | +| [Context and report](context.md) | `LoaderContext`, `LoadReport` and row accounting | --- -## Loader lifecycle +## Lifecycle -1. Detect file format and encoding -2. Read data in chunks or batches -3. Optionally normalise to ORM column types -4. Optionally deduplicate within the incoming data -5. Insert into staging table -6. Return row count +1. A reader is chosen from the file suffix, or passed explicitly, and yields the file as batches. +2. Each batch is converted to the ORM column types under the conversion policy. With `normalise=False` values are written as read. +3. The backend's bulk sink writes the batch into the staging table. +4. The staging table is optionally deduplicated. +5. The number of rows left in staging is returned, and the table layer merges them into the target. -Final merge behaviour belongs to the table mixins and backend layer, not to the loader itself. +Staging tables isolate a load from the production table, which makes rollback safe and merges repeatable. Merge behaviour belongs to the table mixins and the backend layer, not to the loader. --- ## Guarantees -The loaders subsystem guarantees: - -- deterministic ingestion behaviour -- no silent data loss -- explicit logging of dropped or malformed rows -- isolation from domain-specific rules - -It does **not** guarantee correctness of the source data. +- Ingestion is deterministic and no value is silently changed or lost. +- Every row read is loaded, rejected to a rejects file, or dropped as a duplicate. A load that cannot account for every row raises. +- Domain rules stay outside the loader. -[`Helpers`](helpers.md) +The loader does not guarantee the correctness of the source data. diff --git a/docs/loaders/loaders.md b/docs/loaders/loaders.md deleted file mode 100644 index 1fa8728..0000000 --- a/docs/loaders/loaders.md +++ /dev/null @@ -1,92 +0,0 @@ -# Loader Implementations - -This page documents the concrete loader implementations provided by -`orm_loader`. - -All loaders implement the same interface. The difference is in how they read data and how much work they do before rows reach the staging table. - ---- - -## LoaderInterface - -`LoaderInterface` defines the contract for all loaders. - -### Required methods - -- `orm_file_load(ctx)` -- `dedupe(data, ctx)` - -### Shared behaviour - -All loaders: - -- load into staging tables only -- respect `LoaderContext` flags -- return row counts -- leave final merge behaviour to the table layer - ---- - -## PandasLoader - -`PandasLoader` uses pandas to read and process files. - -### Characteristics - -- Works well with CSV and TSV inputs -- Easy to debug and inspect -- Supports chunked loading -- Flexible transformation pipeline - -### Trade-offs - -- Slower for very large datasets -- Higher memory overhead than columnar approaches - -### Best suited for - -- initial data exploration -- messy or inconsistent files -- pipelines requiring heavy cleaning or inspection - ---- - -## ParquetLoader - -`ParquetLoader` uses PyArrow for columnar ingestion. - -### Characteristics - -- Efficient for very large datasets -- Supports Parquet and CSV inputs -- Batch-oriented processing -- Lower memory overhead - -### Trade-offs - -- More complex pipeline -- Less flexible row-wise transformations - -### Best suited for - -- high-volume ingestion -- repeated production loads -- columnar data sources - ---- - -## Deduplication behaviour - -Deduplication here means deduplicating within the incoming data before it is inserted into staging. The merge step is what decides what happens when incoming rows overlap with existing target rows. - ---- - -## Normalisation behaviour - -When enabled, loaders: - -- cast values to ORM column types -- drop rows violating required constraints -- log casting failures with examples - -No schema changes are performed at the loader layer. diff --git a/docs/loaders/readers.md b/docs/loaders/readers.md new file mode 100644 index 0000000..ab41898 --- /dev/null +++ b/docs/loaders/readers.md @@ -0,0 +1,57 @@ +# Readers + +A reader yields a file as Arrow record batches holding the model columns present in the file. Text formats yield raw strings, so the conversion policy sees exactly what the file contains. Typed formats yield their stored types. + +`open_reader` picks the reader from the file suffix: `.parquet` is read as Parquet, any other suffix as delimited text. A table chooses its reader in `_select_reader`, and `load_file(reader=...)` passes one explicitly. + +| Reader | Input | +|------|------| +| `DelimitedTextReader` | Delimited text with a tab or comma delimiter and detected encoding | +| `ParquetReader` | Parquet files, in batches of `chunksize` rows | + +--- + +## FileReader + +A new format is a new `FileReader` subclass that implements `batches()` and sets `columns` and `empty`. + +::: orm_loader.loaders.readers.base.FileReader + +--- + +## DelimitedTextReader + +Reads raw strings in 4 MB blocks, so memory does not grow with the row count. Header names are lowercased and stripped of quotes and a `_hash` suffix, and file columns that are not model columns are ignored. An unquoted empty field is null and a quoted empty field is an empty string. + +`quote_mode` selects how double quotes are read: + +| Mode | Behaviour | +|------|------| +| `csv` | Double quotes quote fields | +| `literal` | Double quotes are data | +| `by_delimiter` | `literal` for tab-delimited files, `csv` otherwise | +| `auto` | Samples the file and picks the mode that gives consistent column counts | + +::: orm_loader.loaders.readers.delimited.DelimitedTextReader + +--- + +## ParquetReader + +::: orm_loader.loaders.readers.parquet.ParquetReader + +--- + +## open_reader + +::: orm_loader.loaders.readers.open_reader + +--- + +## Detection + +The delimited reader detects the encoding with `chardet` (ASCII is normalised to UTF-8, and only a small prefix of the file is inspected), the delimiter from the first line (tabs win over commas when more frequent), and the quote mode when `quote_mode="auto"`. + +::: orm_loader.loaders.readers.detection.infer_encoding +::: orm_loader.loaders.readers.detection.infer_delim +::: orm_loader.loaders.readers.detection.resolve_quote_mode diff --git a/docs/loaders/staging.md b/docs/loaders/staging.md new file mode 100644 index 0000000..645d27a --- /dev/null +++ b/docs/loaders/staging.md @@ -0,0 +1,36 @@ +# Staging and Sinks + +`stage_file(ctx, reader)` connects a reader, the batch processor and the sink of the session's backend. It loads into the staging table, respects the `LoaderContext` flags, returns the number of rows left in staging, and leaves the merge to the table layer. + +::: orm_loader.loaders.loader.stage_file + +--- + +## Interpretation + +- Values are interpreted by column type, as described in [Conversion](conversion.md). +- With the default policy an unreadable value or a missing required value raises and nothing is loaded. +- With `ErrorPolicy(on_error="reject_row")` those rows are written to `/.rejects.csv` and the rest load. `rejects_dir` defaults to a `rejects` directory inside the input file's directory. `LoadReport` accounts for every row read. +- With `normalise=False` raw text is passed to the staging table unchanged. + +--- + +## Sinks + +Each backend uses the fastest bulk interface its engine offers. + +| Backend | Sink | +|------|------| +| PostgreSQL | One `COPY` from Arrow-written CSV | +| DuckDB | Insert from a temporary view over the Arrow batch, committed every five batches | +| SQLite | Native `executemany` per batch | + +A failure in the sink is raised. There is no fallback to another loading path. + +--- + +## Deduplication + +With `dedupe` set, the staging table keeps the first row for each primary key, in source order. The backend removes the others after the load and reports a `DedupeResult` on `LoadReport` with the rows before and the rows removed. The merge step decides what happens when incoming rows overlap with existing target rows. + +No schema changes are performed at the loader layer. diff --git a/docs/tables/index.md b/docs/tables/index.md index 2602257..25aa971 100644 --- a/docs/tables/index.md +++ b/docs/tables/index.md @@ -14,7 +14,7 @@ All components are: - [`ORMTableBase`](orm_table.md) - [`SerialisableTableInterface`](serialisable_table.md) -- [`CSVLoadableTableInterface`](loadable_table.md) +- [`FileLoadableTableInterface`](loadable_table.md) --- diff --git a/docs/tables/loadable_table.md b/docs/tables/loadable_table.md index 51ebfe4..bc92091 100644 --- a/docs/tables/loadable_table.md +++ b/docs/tables/loadable_table.md @@ -4,13 +4,13 @@ Infrastructure for staged, file-based ingestion into ORM tables. Supports: - staged file loading into backend-specific staging tables -- PostgreSQL fast-path `COPY` with ORM fallback +- bulk staging sinks per backend: PostgreSQL `COPY`, DuckDB insert from Arrow, SQLite `executemany` - backend-aware merge strategies -- pandas and PyArrow-based loader paths +- Arrow-based file readers (delimited text and Parquet) - index handling during merge --- -## CSVLoadableTableInterface +## FileLoadableTableInterface -::: orm_loader.tables.loadable_table.CSVLoadableTableInterface +::: orm_loader.tables.loadable_table.FileLoadableTableInterface diff --git a/docs/tables/typing.md b/docs/tables/typing.md index 583717e..3df8595 100644 --- a/docs/tables/typing.md +++ b/docs/tables/typing.md @@ -2,8 +2,7 @@ Structural protocols defining the expected interfaces for: - ORM table classes -- CSV-loadable tables -- Parquet-loadable tables +- file-loadable tables - serialisable row instances Protocols allow infrastructure to remain decoupled from @@ -17,9 +16,9 @@ concrete base classes. --- -## CSVTableProtocol +## FileTableProtocol -::: orm_loader.tables.typing.CSVTableProtocol +::: orm_loader.tables.typing.FileTableProtocol --- diff --git a/mkdocs.yml b/mkdocs.yml index 5af8883..6c46877 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -40,6 +40,8 @@ plugins: - src options: show_source: false + show_root_heading: false + show_root_toc_entry: false docstring_style: numpy nav: @@ -61,6 +63,7 @@ nav: - Model Descriptors: registry/descriptors.md - Loaders: - Overview: loaders/index.md - - Loader Context & Diagnostics: loaders/context.md - - Loader Implementations: loaders/loaders.md - - Loader Helpers: loaders/helpers.md + - Readers: loaders/readers.md + - Conversion: loaders/conversion.md + - Staging and Sinks: loaders/staging.md + - Context and Report: loaders/context.md diff --git a/pyproject.toml b/pyproject.toml index 01616fb..3ed6f36 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,7 +11,6 @@ requires-python = ">=3.12" dependencies = [ "chardet>=5.2.0", "oa-configurator>=2.1.0,<3", - "pandas>=2.3.3", "pyarrow>=23.0.1", "sqlalchemy>=2.0.45", ] @@ -54,6 +53,7 @@ duckdb = [ ] dev = [ "oa-configurator[postgres,duckdb]>=2.0.0,<3", + "pandas>=2.3.3", "pytest>=9.0.3", "ty>=0.0.59", "ruff>=0.14.11", @@ -89,3 +89,11 @@ addopts = "-ra -m 'not db_dialect'" [tool.pyright] reportMissingTypeStubs = false + +# Temporary ty override until #43 picks it up in a follow-up +# with pyarrow-stubs +[[tool.ty.overrides]] +include = ["src/orm_loader/conversion/**", "src/orm_loader/loaders/processor.py", "src/orm_loader/loaders/readers/**"] + +[tool.ty.overrides.rules] +unresolved-attribute = "ignore" diff --git a/src/orm_loader/backends/__init__.py b/src/orm_loader/backends/__init__.py index 56cc81c..4aae24a 100644 --- a/src/orm_loader/backends/__init__.py +++ b/src/orm_loader/backends/__init__.py @@ -5,6 +5,7 @@ from .base import ( BackendCapabilities, DatabaseBackend, + DedupeResult, Dialect, STAGING_SCHEMA, staging_schema_claim, @@ -20,6 +21,7 @@ "BackendCapabilities", "ConcurrentRefreshNotEligibleError", "DatabaseBackend", + "DedupeResult", "Dialect", "DuckDBBackend", "MaterializationError", diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 74baa60..f151d5a 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -1,12 +1,13 @@ from __future__ import annotations from abc import ABC, abstractmethod -from collections.abc import Generator, Iterator, Mapping +from collections.abc import Generator, Iterable, Iterator, Mapping, Sequence from contextlib import AbstractContextManager, contextmanager from dataclasses import dataclass from functools import wraps from typing import TYPE_CHECKING, Any, Callable, ParamSpec, Type, TypeVar, cast +import pyarrow as pa import sqlalchemy as sa import sqlalchemy.orm as so from oa_configurator import Dialect, ResolvedDatabase, SchemaClaim, open_connection, qualified @@ -14,9 +15,33 @@ from sqlalchemy.sql.compiler import IdentifierPreparer if TYPE_CHECKING: - from ..loaders.data_classes import LoaderContext + from ..loaders.loader import LoaderContext from ..mappers.materialised_view_contracts import MaterializedViewIndex - from ..tables.typing import CSVTableProtocol + from ..tables.typing import FileTableProtocol + + +@dataclass(frozen=True) +class DedupeResult: + """What deduplicating a staging table removed. + + Attributes + ---------- + key_columns + Primary key columns that identify a duplicate. + rows_before + Rows in the staging table before deduplication. + rows_removed + Rows removed because an earlier row had the same key. + """ + + key_columns: tuple[str, ...] + rows_before: int + rows_removed: int + + @property + def rows_after(self) -> int: + """Rows left in the staging table.""" + return self.rows_before - self.rows_removed @dataclass(frozen=True) @@ -28,7 +53,6 @@ class BackendCapabilities: opt into capabilities explicitly. """ - supports_fast_load: bool = False supports_unlogged_staging: bool = False supports_fk_toggle: bool = False supports_materialized_views: bool = False @@ -292,7 +316,7 @@ def _as_connection( def _insertable_column_names( self, - table_cls: Type["CSVTableProtocol"], + table_cls: Type["FileTableProtocol"], ) -> list[str]: """ Return column names safe to include in generic insert statements. @@ -304,7 +328,7 @@ def _insertable_column_names( def deferred_fk_constraints( self, - table_cls: Type["CSVTableProtocol"], + table_cls: Type["FileTableProtocol"], ) -> tuple[sa.ForeignKeyConstraint, ...]: """FK constraints on *table_cls* that this backend never enforces at the DB level. @@ -315,7 +339,7 @@ def deferred_fk_constraints( def validate_deferred_fks( self, - table_cls: Type["CSVTableProtocol"], + table_cls: Type["FileTableProtocol"], session: so.Session, ) -> None: """Raise if any row in *table_cls* violates one of its deferred_fk_constraints(). @@ -375,7 +399,7 @@ def _validate_one_deferred_fk(fk: sa.ForeignKeyConstraint, session: so.Session) @abstractmethod def create_staging_table( self, - table_cls: Type["CSVTableProtocol"], + table_cls: Type["FileTableProtocol"], session: so.Session, ) -> None: """Create a staging table for the supplied ORM table class.""" @@ -383,22 +407,106 @@ def create_staging_table( @abstractmethod def drop_staging_table( self, - table_cls: Type["CSVTableProtocol"], + table_cls: Type["FileTableProtocol"], session: so.Session, ) -> None: """Drop a staging table if it exists.""" - def load_staging_fast( + @abstractmethod + def insert_batches( self, - loader_context: "LoaderContext", - ) -> int | None: + ctx: "LoaderContext", + batches: Iterable[pa.Table], + columns: Sequence[str], + ) -> int: """ - Attempt a backend-native fast-path load. + Write typed batches into the load's staging table. + + Each backend uses the fastest bulk interface its engine offers. - Return the inserted row count when handled, or ``None`` when the - backend has no fast-path loader for the given context. + Parameters + ---------- + ctx + Context of the load; its staging table receives the rows. + batches + Typed tables, one per batch. + columns + Names of the columns the batches carry. + + Returns + ------- + int + Number of rows written. + + Notes + ----- + A backend may commit between batches to bound memory, which also ends + the caller's open transaction. """ - return None + + @staticmethod + def _staging_reference(ctx: "LoaderContext", connection: Connection) -> str: + """ + Return the quoted, schema-qualified name of the load's staging table. + + The schema is resolved through the connection's ``schema_translate_map``, + for raw SQL that SQLAlchemy does not translate itself. + + Parameters + ---------- + ctx + Context of the load. + connection + Connection the SQL runs on. + + Returns + ------- + str + The staging table name as the connection resolves it. + """ + staging = ctx.staging_table + translate = connection.get_execution_options().get("schema_translate_map") or {} + schema = translate.get(staging.schema, staging.schema) + return qualified(connection, staging.name, physical_schema=schema) + + def deduplicate_staging( + self, + ctx: "LoaderContext", + key_columns: Sequence[str], + ) -> "DedupeResult": + """ + Keep the first row for each primary key in the load's staging table. + + The first row is the one with the lowest row key, which follows the + order of the source file. + + Parameters + ---------- + ctx + Context of the load; its staging table is deduplicated. + key_columns + Primary key columns that identify a duplicate. + + Returns + ------- + DedupeResult + How many rows the staging table held and how many were removed. + """ + staging = ctx.staging_table + arguments = {"mapper": ctx.tableclass} + row_key = self._pagination_column(ctx.tableclass, staging) + ranked = sa.select( + row_key.label("row_key"), + sa.func.row_number() + .over(partition_by=[staging.c[name] for name in key_columns], order_by=row_key) + .label("position"), + ).subquery() + duplicates = sa.select(ranked.c.row_key).where(ranked.c.position > 1) + count = sa.select(sa.func.count()).select_from(staging) + before = ctx.session.execute(count, bind_arguments=arguments).scalar_one() + ctx.session.execute(sa.delete(staging).where(row_key.in_(duplicates)), bind_arguments=arguments) + after = ctx.session.execute(count, bind_arguments=arguments).scalar_one() + return DedupeResult(tuple(key_columns), before, before - after) @staticmethod @abstractmethod @@ -429,7 +537,7 @@ def restore_fk_check( def check_replace_feasible( self, - table_cls: Type["CSVTableProtocol"], + table_cls: Type["FileTableProtocol"], pk_cols: list[str], session: so.Session, ) -> None: @@ -445,7 +553,7 @@ def check_replace_feasible( @abstractmethod def merge_replace( self, - table_cls: Type["CSVTableProtocol"], + table_cls: Type["FileTableProtocol"], session: so.Session, pk_cols: list[str], *, @@ -463,7 +571,7 @@ def merge_replace( @abstractmethod def merge_upsert( self, - table_cls: Type["CSVTableProtocol"], + table_cls: Type["FileTableProtocol"], session: so.Session, pk_cols: list[str], *, @@ -479,7 +587,7 @@ def merge_upsert( @abstractmethod def merge_insert( self, - table_cls: Type["CSVTableProtocol"], + table_cls: Type["FileTableProtocol"], session: so.Session, *, merge_batch_size: int | None = None, @@ -493,7 +601,7 @@ def merge_insert( def _pagination_column( self, - table_cls: Type["CSVTableProtocol"], + table_cls: Type["FileTableProtocol"], staging: sa.Table, ) -> sa.ColumnElement[int]: """Column this backend orders a staging-table scan by, for pagination. @@ -506,7 +614,7 @@ def _pagination_column( def _ensure_pagination_index( self, - table_cls: Type["CSVTableProtocol"], + table_cls: Type["FileTableProtocol"], staging: sa.Table, session: so.Session, ) -> None: @@ -527,7 +635,7 @@ def _commits_between_pagination_batches(self) -> bool: def _paginate_batches( self, - table_cls: Type["CSVTableProtocol"], + table_cls: Type["FileTableProtocol"], staging: sa.Table, session: so.Session, merge_batch_size: int | None, @@ -560,7 +668,7 @@ def _paginate_batches( def _commit_batch_if_needed( self, - table_cls: Type["CSVTableProtocol"], + table_cls: Type["FileTableProtocol"], session: so.Session, batch_range: sa.ColumnElement[bool] | None, *, @@ -572,7 +680,7 @@ def _commit_batch_if_needed( Called from every merge_*() method after each batch of rows is processed. Instead of committing directly, a batch that violates a deferred FK is caught before the next batch starts. Without this, an earlier batch's commit cannot - be undone once a later batch's load_csv()-level validation fails. + be undone once a later batch's load_file()-level validation fails. No-op when there is nothing to validate to not add cost. This is true for every backend but DuckDB, and for DuckDB whenever table_cls has no deferred/ @@ -585,7 +693,7 @@ def _commit_batch_if_needed( session.commit() def _self_referential_deferred_fks( - self, table_cls: Type["CSVTableProtocol"] + self, table_cls: Type["FileTableProtocol"] ) -> tuple[sa.ForeignKeyConstraint, ...]: """Deferred FKs on table_cls that point back at table_cls itself. @@ -600,7 +708,7 @@ def _self_referential_deferred_fks( def merge_context( self, - table_cls: Type["CSVTableProtocol"], + table_cls: Type["FileTableProtocol"], session: so.Session, ) -> AbstractContextManager[None]: """Disable backend FK checks for a table's merge and restore them afterward. diff --git a/src/orm_loader/backends/duckdb.py b/src/orm_loader/backends/duckdb.py index 2bfaaf0..02f494e 100644 --- a/src/orm_loader/backends/duckdb.py +++ b/src/orm_loader/backends/duckdb.py @@ -1,8 +1,9 @@ from __future__ import annotations -from collections.abc import Iterator -from typing import TYPE_CHECKING, Any, Type, cast +from collections.abc import Iterable, Iterator, Sequence +from typing import TYPE_CHECKING, Any, ClassVar, Type, cast +import pyarrow as pa import sqlalchemy as sa import sqlalchemy.orm as so from oa_configurator import ( @@ -21,7 +22,8 @@ from .base import BackendCapabilities, DatabaseBackend, DeferredForeignKeyViolationError, Dialect if TYPE_CHECKING: - from ..tables.typing import CSVTableProtocol + from ..loaders.loader import LoaderContext + from ..tables.typing import FileTableProtocol class DuckDBBackend(DatabaseBackend): @@ -43,10 +45,13 @@ def identifier_preparer(self) -> IdentifierPreparer: return duckdb_sqlalchemy.Dialect().identifier_preparer + # Batches per commit of the staging table. DuckDB keeps uncommitted rows in memory, so this bounds + # the memory of a load. + STAGING_COMMIT_EVERY: ClassVar[int] = 5 + @property def capabilities(self) -> BackendCapabilities: return BackendCapabilities( - supports_fast_load=False, supports_unlogged_staging=False, supports_fk_toggle=False, supports_materialized_views=False, @@ -71,7 +76,7 @@ def _unenforced_constraints( if self.resolved is None: raise ValueError( "DuckDBBackend needs resolved= to classify foreign-key " - "enforcement. Pass it to load_csv(...) or resolve_backend(...)." + "enforcement. Pass it to load_file(...) or resolve_backend(...)." ) return unenforced_foreign_keys( tables, resolved=self.resolved, physical_schemas=self.physical_schemas @@ -79,7 +84,7 @@ def _unenforced_constraints( def deferred_fk_constraints( self, - table_cls: Type["CSVTableProtocol"], + table_cls: Type["FileTableProtocol"], ) -> tuple[sa.ForeignKeyConstraint, ...]: """FK constraints on *table_cls* this backend's database will not enforce: - omitted (cross-database/cross-schema), or @@ -93,7 +98,7 @@ def deferred_fk_constraints( unenforced = {entry.constraint for entry in self._unenforced_constraints(tuple(metadata.tables.values()))} return tuple(fk for fk in own_fks if fk in unenforced) - def validate_deferred_fks(self, table_cls: Type["CSVTableProtocol"], session: so.Session) -> None: + def validate_deferred_fks(self, table_cls: Type["FileTableProtocol"], session: so.Session) -> None: """Validate every deferred FK on table_cls, including a cross-database one. A cross-schema or cyclic deferred FK points to a table in the same @@ -218,7 +223,7 @@ def _staging_rownum_sequence_ref(self, tablename: str) -> str: def create_staging_table( self, - table_cls: type["CSVTableProtocol"], + table_cls: type["FileTableProtocol"], session: so.Session, ) -> None: bind_arguments = {"mapper": table_cls} @@ -253,7 +258,7 @@ def create_staging_table( def drop_staging_table( self, - table_cls: type["CSVTableProtocol"], + table_cls: type["FileTableProtocol"], session: so.Session, ) -> None: bind_arguments = {"mapper": table_cls} @@ -298,7 +303,7 @@ def restore_fk_check( def _paginate_batches( self, - table_cls: type["CSVTableProtocol"], + table_cls: type["FileTableProtocol"], staging: sa.Table, session: so.Session, merge_batch_size: int | None, @@ -318,7 +323,7 @@ def _paginate_batches( def _self_referential_batches( self, - table_cls: type["CSVTableProtocol"], + table_cls: type["FileTableProtocol"], staging: sa.Table, session: so.Session, merge_batch_size: int | None, @@ -383,9 +388,45 @@ def _commits_between_pagination_batches(self) -> bool: """ return True + def insert_batches( + self, + ctx: "LoaderContext", + batches: Iterable[pa.Table], + columns: Sequence[str], + ) -> int: + """Insert the batches into the staging table straight from Arrow. + + Notes + ----- + - The fastest bulk insertion method is DuckDB's ``INSERT ... SELECT`` from a + temporary view over the Arrow memory. + - The staging table is committed every ``STAGING_COMMIT_EVERY`` batches, + because DuckDB keeps uncommitted rows in memory and one transaction would + hold the whole file. + + Notes + ----- + Committing the staging table also ends the caller's open transaction. + """ + batch = sa.table("_orm_loader_batch", *(sa.column(c) for c in columns)) + insert = ctx.staging_table.insert().from_select(list(columns), sa.select(*batch.c)) + bind_arguments = {"mapper": ctx.tableclass} + total = 0 + for index, table in enumerate(batches, start=1): + native = ctx.session.connection(bind_arguments=bind_arguments).connection + native.register("_orm_loader_batch", table) + try: + ctx.session.execute(insert, bind_arguments=bind_arguments) + finally: + native.unregister("_orm_loader_batch") + total += table.num_rows + if index % self.STAGING_COMMIT_EVERY == 0: + ctx.session.commit() + return total + def check_replace_feasible( self, - table_cls: type["CSVTableProtocol"], + table_cls: type["FileTableProtocol"], pk_cols: list[str], session: so.Session, ) -> None: @@ -423,7 +464,7 @@ def check_replace_feasible( def merge_replace( self, - table_cls: type["CSVTableProtocol"], + table_cls: type["FileTableProtocol"], session: so.Session, pk_cols: list[str], *, @@ -524,7 +565,7 @@ def _process_batch(batch_range: sa.ColumnElement[bool] | None) -> None: def merge_upsert( self, - table_cls: type["CSVTableProtocol"], + table_cls: type["FileTableProtocol"], session: so.Session, pk_cols: list[str], *, @@ -561,7 +602,7 @@ def _upsert(select_: Any) -> sa.Insert: def merge_insert( self, - table_cls: type["CSVTableProtocol"], + table_cls: type["FileTableProtocol"], session: so.Session, *, merge_batch_size: int | None = None, diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index 43c3e23..25efbd6 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -1,9 +1,12 @@ from __future__ import annotations -from collections.abc import Mapping +import io +from collections.abc import Iterable, Mapping, Sequence from contextlib import contextmanager from typing import TYPE_CHECKING, Any +import pyarrow as pa +import pyarrow.csv as pacsv import sqlalchemy as sa import sqlalchemy.event as sae import sqlalchemy.orm as so @@ -30,9 +33,9 @@ from oa_configurator import ResolvedDatabase from sqlalchemy.engine import Connection, Engine - from ..loaders.data_classes import LoaderContext + from ..loaders.loader import LoaderContext from ..mappers.materialised_view_contracts import MaterializedViewIndex - from ..tables.typing import CSVTableProtocol + from ..tables.typing import FileTableProtocol _VALID_PG_REPLICATION_ROLES = frozenset({"origin", "local", "replica"}) @@ -72,7 +75,6 @@ def identifier_preparer(self) -> IdentifierPreparer: @property def capabilities(self) -> BackendCapabilities: return BackendCapabilities( - supports_fast_load=True, supports_unlogged_staging=True, supports_fk_toggle=True, supports_materialized_views=True, @@ -81,7 +83,7 @@ def capabilities(self) -> BackendCapabilities: def create_staging_table( self, - table_cls: type["CSVTableProtocol"], + table_cls: type["FileTableProtocol"], session: so.Session, ) -> None: table = table_cls.__table__ @@ -127,7 +129,7 @@ def create_staging_table( def drop_staging_table( self, - table_cls: type["CSVTableProtocol"], + table_cls: type["FileTableProtocol"], session: so.Session, ) -> None: staging_name = self.staging_name_for_table(table_cls.__tablename__) @@ -136,21 +138,41 @@ def drop_staging_table( bind_arguments={"mapper": table_cls}, ) - def load_staging_fast( + def insert_batches( self, - loader_context: "LoaderContext", - ) -> int | None: - from ..loaders.loading_helpers import quick_load_pg - - tablename = loader_context.tableclass.__tablename__ - return quick_load_pg( - path=loader_context.path, - session=loader_context.session, - tablename=self.staging_name_for_table(tablename), - schema=self.staging_schema, - quote_mode=loader_context.quote_mode, - mapper=loader_context.tableclass, - ) + ctx: "LoaderContext", + batches: Iterable[pa.Table], + columns: Sequence[str], + ) -> int: + """Stream the batches into the staging table with one COPY, as CSV written by Arrow. + + An unquoted empty field is NULL and a quoted one an empty string, so + the null policy's decisions reach the table unchanged. A failure in + any batch aborts the COPY. + """ + connection = ctx.session.connection(bind_arguments={"mapper": ctx.tableclass}) + raw_conn = connection.connection + if not hasattr(raw_conn, "cursor"): + raise RuntimeError("Expected DB-API connection for COPY") + table_ref = self._staging_reference(ctx, connection) + column_list = ", ".join(self.identifier_preparer.quote(c) for c in columns) + options = pacsv.WriteOptions(include_header=False, quoting_style="needed") + total = 0 + cursor = raw_conn.cursor() + try: + with cursor.copy(f"COPY {table_ref} ({column_list}) FROM STDIN WITH (FORMAT csv)") as copy: + for table in batches: + buffer = io.BytesIO() + pacsv.write_csv(table, buffer, options) + copy.write(buffer.getvalue()) + total += table.num_rows + ctx.session.flush() + return total + except Exception: + ctx.session.rollback() + raise + finally: + cursor.close() @staticmethod def _normalize_fk_check_state(previous_state: str | int) -> str: @@ -217,7 +239,7 @@ def _commits_between_pagination_batches(self) -> bool: def merge_replace( self, - table_cls: type["CSVTableProtocol"], + table_cls: type["FileTableProtocol"], session: so.Session, pk_cols: list[str], *, @@ -239,7 +261,7 @@ def merge_replace( def merge_upsert( self, - table_cls: type["CSVTableProtocol"], + table_cls: type["FileTableProtocol"], session: so.Session, pk_cols: list[str], *, @@ -270,7 +292,7 @@ def _upsert(select_: Any) -> sa.Insert: def merge_insert( self, - table_cls: type["CSVTableProtocol"], + table_cls: type["FileTableProtocol"], session: so.Session, *, merge_batch_size: int | None = None, diff --git a/src/orm_loader/backends/sqlite.py b/src/orm_loader/backends/sqlite.py index 14be15f..a0fae64 100644 --- a/src/orm_loader/backends/sqlite.py +++ b/src/orm_loader/backends/sqlite.py @@ -1,11 +1,13 @@ from __future__ import annotations -from collections.abc import Mapping +from collections.abc import Iterable, Mapping, Sequence import logging import sqlite3 from pathlib import Path from typing import TYPE_CHECKING, Any +import pyarrow as pa +import pyarrow.compute as pc import sqlalchemy as sa import sqlalchemy.orm as so from sqlalchemy import event, text @@ -20,7 +22,8 @@ from oa_configurator import ResolvedDatabase from sqlalchemy.engine import Connection, Engine - from ..tables.typing import CSVTableProtocol + from ..loaders.loader import LoaderContext + from ..tables.typing import FileTableProtocol logger = logging.getLogger(__name__) @@ -103,7 +106,6 @@ def identifier_preparer(self) -> IdentifierPreparer: @property def capabilities(self) -> BackendCapabilities: return BackendCapabilities( - supports_fast_load=False, supports_unlogged_staging=False, supports_fk_toggle=True, supports_materialized_views=False, @@ -114,9 +116,44 @@ def capabilities(self) -> BackendCapabilities: def default_index_strategy(self) -> str: return "keep" + def insert_batches( + self, + ctx: "LoaderContext", + batches: Iterable[pa.Table], + columns: Sequence[str], + ) -> int: + """Insert the batches with one native ``executemany`` per batch. + + The statement goes straight to the DB-API without per-value bind + processing, so the connection's schema translation is applied here. + Dates, timestamps and decimals are passed as the text SQLAlchemy + itself stores for them. + """ + connection = ctx.session.connection(bind_arguments={"mapper": ctx.tableclass}) + preparer = connection.dialect.identifier_preparer + table_ref = self._staging_reference(ctx, connection) + sql = ( + f"INSERT INTO {table_ref} ({', '.join(preparer.quote(c) for c in columns)}) " + f"VALUES ({', '.join('?' for _ in columns)})" + ) + total = 0 + for table in batches: + values = [self._storable(table[name]) for name in columns] + connection.exec_driver_sql(sql, list(zip(*values))) + total += table.num_rows + return total + + @staticmethod + def _storable(column: pa.ChunkedArray) -> list[Any]: + """Return a column as Python values SQLite can bind.""" + kind = column.type + if pa.types.is_date(kind) or pa.types.is_timestamp(kind) or pa.types.is_decimal(kind): + column = pc.cast(column, pa.string()) + return column.to_pylist() + def create_staging_table( self, - table_cls: type["CSVTableProtocol"], + table_cls: type["FileTableProtocol"], session: so.Session, ) -> None: staging_name = self.staging_name_for_table(table_cls.__tablename__) @@ -140,7 +177,7 @@ def create_staging_table( def drop_staging_table( self, - table_cls: type["CSVTableProtocol"], + table_cls: type["FileTableProtocol"], session: so.Session, ) -> None: staging_name = self.staging_name_for_table(table_cls.__tablename__) @@ -177,13 +214,13 @@ def restore_fk_check( session.execute(text(f"PRAGMA foreign_keys = {safe_state}"), **execute_options) def _pagination_column( - self, table_cls: type["CSVTableProtocol"], staging: sa.Table + self, table_cls: type["FileTableProtocol"], staging: sa.Table ) -> sa.ColumnElement[int]: """SQLite's implicit rowid is already gapless and indexed.""" return sa.literal_column("rowid") def _ensure_pagination_index( - self, table_cls: type["CSVTableProtocol"], staging: sa.Table, session: so.Session + self, table_cls: type["FileTableProtocol"], staging: sa.Table, session: so.Session ) -> None: """No-op: rowid needs no index, and the staging table has no _rownum column.""" return @@ -198,7 +235,7 @@ def _commits_between_pagination_batches(self) -> bool: def merge_replace( self, - table_cls: type["CSVTableProtocol"], + table_cls: type["FileTableProtocol"], session: so.Session, pk_cols: list[str], *, @@ -235,7 +272,7 @@ def _delete(extra: sa.ColumnElement[bool] | None = None) -> sa.Delete: def merge_upsert( self, - table_cls: type["CSVTableProtocol"], + table_cls: type["FileTableProtocol"], session: so.Session, pk_cols: list[str], *, @@ -263,7 +300,7 @@ def _upsert(select_: Any) -> sa.Insert: def merge_insert( self, - table_cls: type["CSVTableProtocol"], + table_cls: type["FileTableProtocol"], session: so.Session, *, merge_batch_size: int | None = None, diff --git a/src/orm_loader/conversion/__init__.py b/src/orm_loader/conversion/__init__.py new file mode 100644 index 0000000..69da4fe --- /dev/null +++ b/src/orm_loader/conversion/__init__.py @@ -0,0 +1,37 @@ +from .handlers import ( + BooleanHandler, + ColumnHandler, + ConvertingHandler, + DateFormats, + DateHandler, + DateTimeHandler, + FloatHandler, + IntegerHandler, + NumberHandler, + NumericHandler, + PassthroughHandler, + TextHandler, + parse_dates, +) +from .policy import ConversionPolicy, ErrorPolicy, NullPolicy +from .values import Interpretation + +__all__ = [ + "BooleanHandler", + "ColumnHandler", + "ConversionPolicy", + "ConvertingHandler", + "DateFormats", + "DateHandler", + "DateTimeHandler", + "ErrorPolicy", + "FloatHandler", + "IntegerHandler", + "Interpretation", + "NullPolicy", + "NumberHandler", + "NumericHandler", + "PassthroughHandler", + "TextHandler", + "parse_dates", +] diff --git a/src/orm_loader/conversion/handlers/__init__.py b/src/orm_loader/conversion/handlers/__init__.py new file mode 100644 index 0000000..2230205 --- /dev/null +++ b/src/orm_loader/conversion/handlers/__init__.py @@ -0,0 +1,21 @@ +from .base import ColumnHandler, ConvertingHandler, PassthroughHandler +from .booleans import BooleanHandler +from .dates import DateFormats, DateHandler, DateTimeHandler, parse_dates +from .numbers import FloatHandler, IntegerHandler, NumberHandler, NumericHandler +from .text import TextHandler + +__all__ = [ + "BooleanHandler", + "ColumnHandler", + "ConvertingHandler", + "DateFormats", + "DateHandler", + "DateTimeHandler", + "FloatHandler", + "IntegerHandler", + "NumberHandler", + "NumericHandler", + "PassthroughHandler", + "TextHandler", + "parse_dates", +] diff --git a/src/orm_loader/conversion/handlers/base.py b/src/orm_loader/conversion/handlers/base.py new file mode 100644 index 0000000..d8dc910 --- /dev/null +++ b/src/orm_loader/conversion/handlers/base.py @@ -0,0 +1,137 @@ +"""The contract every column handler fulfils, and the staged template most handlers share.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any, ClassVar + +import pyarrow as pa +import pyarrow.compute as pc +from sqlalchemy.types import TypeEngine + +from ..values import Interpretation, as_array, is_text_like, no_failures, NULL_STRING, safe_cast + +if TYPE_CHECKING: + from ..policy import NullPolicy + +class ColumnHandler(ABC): + """Interprets the raw values of the columns of one kind of column type. + + A handler is a frozen settings object. It is selected for a column by + ``ConversionPolicy.for_column`` and holds no state about the column. + """ + + @abstractmethod + def interpret( + self, arr: pa.Array | pa.ChunkedArray, sa_type: TypeEngine[Any], nulls: NullPolicy + ) -> Interpretation: + """Interpret a raw column without raising on unreadable values. + + Parameters + ---------- + arr + Raw column. + sa_type + Column type the values are converted to. + nulls + What counts as a missing value. + + Returns + ------- + Interpretation + Converted values with nulls where a value is missing or + unreadable, and a mask that is true where a value was unreadable. + """ + + def expected(self, sa_type: TypeEngine[Any]) -> str: + """Return the name reported in an error for an unreadable value.""" + return type(sa_type).__name__ + + +class PassthroughHandler(ColumnHandler): + """Returns a column unchanged, for column types that are not converted.""" + + def interpret( + self, arr: pa.Array | pa.ChunkedArray, sa_type: TypeEngine[Any], nulls: NullPolicy + ) -> Interpretation: + arr = as_array(arr) + return Interpretation(arr, no_failures(len(arr))) + + +class ConvertingHandler(ColumnHandler): + """Converts a column in four stages: clean, detect missing, parse, check. + + Text input follows the stages. Input that already has a type, such as a + Parquet column, is cast directly with floating NaN read as missing. + Subclasses supply ``arrow_type`` and ``parse`` and may override ``clean``, + ``missing`` and ``check``. + + Attributes + ---------- + distinct + Run the stages once per distinct value and gather the results back, + for types whose parse step is expensive per value. + """ + + distinct: ClassVar[bool] = False + + @abstractmethod + def arrow_type(self, sa_type: TypeEngine[Any]) -> pa.DataType: + """Return the Arrow type of the converted values.""" + + @abstractmethod + def parse(self, cleaned: pa.Array, sa_type: TypeEngine[Any]) -> Interpretation: + """Parse cleaned text in which missing values are null. + + Parameters + ---------- + cleaned + String array with missing values set to null. + sa_type + Column type the values are converted to. + + Returns + ------- + Interpretation + Parsed values and a mask of values that did not fit. + """ + + def clean(self, raw: pa.Array, trimmed: pa.Array) -> pa.Array: + """Return the text that is parsed, the trimmed text by default.""" + return trimmed + + def missing(self, trimmed: pa.Array, nulls: NullPolicy) -> pa.Array: + """Return a mask of values that count as missing, by the typed rules by default.""" + return nulls.typed_missing(trimmed) + + def check(self, values: pa.Array, sa_type: TypeEngine[Any]) -> pa.Array | None: + """Return a mask of parsed values the column cannot hold, or None when it holds any.""" + return None + + def interpret( + self, arr: pa.Array | pa.ChunkedArray, sa_type: TypeEngine[Any], nulls: NullPolicy + ) -> Interpretation: + arr = as_array(arr) + if not is_text_like(arr.type): + return self._interpret_typed(arr, sa_type) + if self.distinct: + encoded = arr if pa.types.is_dictionary(arr.type) else pc.dictionary_encode(arr) + values, invalid = self._interpret_text(encoded.dictionary, sa_type, nulls) + invalid = pc.fill_null(pc.take(invalid, encoded.indices), False) + return Interpretation(pc.take(values, encoded.indices), invalid) + if pa.types.is_dictionary(arr.type): + arr = arr.dictionary_decode() + return self._interpret_text(arr, sa_type, nulls) + + def _interpret_text(self, raw: pa.Array, sa_type: TypeEngine[Any], nulls: NullPolicy) -> Interpretation: + trimmed = pc.utf8_trim_whitespace(raw) + missing = self.missing(trimmed, nulls) + cleaned = pc.if_else(missing, NULL_STRING, self.clean(raw, trimmed)) + values, invalid = self.parse(cleaned, sa_type) + rejected = self.check(values, sa_type) + return Interpretation(values, invalid if rejected is None else pc.or_(invalid, rejected)) + + def _interpret_typed(self, arr: pa.Array, sa_type: TypeEngine[Any]) -> Interpretation: + if pa.types.is_floating(arr.type): + arr = pc.if_else(pc.fill_null(pc.is_nan(arr), False), pa.scalar(None, arr.type), arr) + return Interpretation(*safe_cast(arr, self.arrow_type(sa_type))) diff --git a/src/orm_loader/conversion/handlers/booleans.py b/src/orm_loader/conversion/handlers/booleans.py new file mode 100644 index 0000000..4807150 --- /dev/null +++ b/src/orm_loader/conversion/handlers/booleans.py @@ -0,0 +1,40 @@ +"""Handler for boolean columns.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import pyarrow as pa +import pyarrow.compute as pc +from sqlalchemy.types import TypeEngine + +from ..values import Interpretation +from .base import ConvertingHandler + + +@dataclass(frozen=True) +class BooleanHandler(ConvertingHandler): + """Reads case-insensitive words as true or false. + + Attributes + ---------- + true_words + Lowercase words read as true. + false_words + Lowercase words read as false. + """ + + true_words: tuple[str, ...] = ("true", "t", "yes", "y", "1", "on") + false_words: tuple[str, ...] = ("false", "f", "no", "n", "0", "off") + + def arrow_type(self, sa_type: TypeEngine[Any]) -> pa.DataType: + return pa.bool_() + + def parse(self, cleaned: pa.Array, sa_type: TypeEngine[Any]) -> Interpretation: + lowered = pc.utf8_lower(cleaned) + is_true = pc.fill_null(pc.is_in(lowered, value_set=pa.array(self.true_words, pa.string())), False) + is_false = pc.fill_null(pc.is_in(lowered, value_set=pa.array(self.false_words, pa.string())), False) + known = pc.or_(is_true, is_false) + values = pc.if_else(known, is_true, pa.scalar(None, pa.bool_())) + return Interpretation(values, pc.and_(pc.is_valid(cleaned), pc.invert(known))) diff --git a/src/orm_loader/conversion/handlers/dates.py b/src/orm_loader/conversion/handlers/dates.py new file mode 100644 index 0000000..cf19fc1 --- /dev/null +++ b/src/orm_loader/conversion/handlers/dates.py @@ -0,0 +1,182 @@ +"""Handlers for date and timestamp columns, with the date formats and the default date parser.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any, ClassVar + +import pyarrow as pa +import pyarrow.compute as pc +from sqlalchemy.types import TypeEngine + +from ..values import NULL_STRING, Interpretation, cast_each, safe_cast +from .base import ConvertingHandler + +DateShape = Callable[[pa.Array], tuple[pa.Array, pa.Array]] +DateParser = Callable[[pa.Array, tuple[str, ...]], tuple[pa.Array, pa.Array]] + + +class DateFormats: + """Supported date format names, each a shape check plus ISO reassembly. + + A shape function takes trimmed strings and returns a boolean array that + is true where the string fits the format, and the same strings + reassembled as ISO dates. + + Attributes + ---------- + DEFAULT + Mutually distinguishable formats tried in order when none is declared. + MONTHS + Lowercase month abbreviations used by ``dd-mon-yyyy``. + SHAPES + Shape function by format name. + """ + + DEFAULT: ClassVar[tuple[str, ...]] = ("yyyymmdd", "dd-mon-yyyy", "yyyy-mm-dd") + MONTHS: ClassVar[tuple[str, ...]] = ( + "jan", "feb", "mar", "apr", "may", "jun", "jul", "aug", "sep", "oct", "nov", "dec", + ) + + @staticmethod + def yyyymmdd(values: pa.Array) -> tuple[pa.Array, pa.Array]: + ok = pc.and_(pc.equal(pc.utf8_length(values), 8), pc.utf8_is_digit(values)) + cut = pc.utf8_slice_codeunits + return ok, pc.binary_join_element_wise(cut(values, 0, 4), cut(values, 4, 6), cut(values, 6, 8), "-") + + @staticmethod + def yyyy_mm_dd(values: pa.Array) -> tuple[pa.Array, pa.Array]: + return pc.match_substring_regex(values, r"^\d{4}-\d{2}-\d{2}$"), values + + @staticmethod + def dd_mon_yyyy(values: pa.Array) -> tuple[pa.Array, pa.Array]: + ok = pc.match_substring_regex(values, r"^\d{2}-[A-Za-z]{3}-\d{4}$") + cut = pc.utf8_slice_codeunits + month = pc.index_in(pc.utf8_lower(cut(values, 3, 6)), value_set=pa.array(DateFormats.MONTHS)) + month = pc.utf8_lpad(pc.cast(pc.add(month, 1), pa.string()), 2, "0") + return ok, pc.binary_join_element_wise(cut(values, 7, 11), month, cut(values, 0, 2), "-") + + @staticmethod + def dd_mm_yyyy(values: pa.Array) -> tuple[pa.Array, pa.Array]: + ok = pc.match_substring_regex(values, r"^\d{2}/\d{2}/\d{4}$") + cut = pc.utf8_slice_codeunits + return ok, pc.binary_join_element_wise(cut(values, 6, 10), cut(values, 3, 5), cut(values, 0, 2), "-") + + @staticmethod + def mm_dd_yyyy(values: pa.Array) -> tuple[pa.Array, pa.Array]: + ok = pc.match_substring_regex(values, r"^\d{2}/\d{2}/\d{4}$") + cut = pc.utf8_slice_codeunits + return ok, pc.binary_join_element_wise(cut(values, 6, 10), cut(values, 0, 2), cut(values, 3, 5), "-") + + @classmethod + def shape(cls, name: str) -> DateShape: + """Return the shape function for a format name.""" + return { + "yyyymmdd": cls.yyyymmdd, + "yyyy-mm-dd": cls.yyyy_mm_dd, + "dd-mon-yyyy": cls.dd_mon_yyyy, + "dd/mm/yyyy": cls.dd_mm_yyyy, + "mm/dd/yyyy": cls.mm_dd_yyyy, + }[name] + + @classmethod + def names(cls) -> tuple[str, ...]: + """Return every supported format name.""" + return ("yyyymmdd", "yyyy-mm-dd", "dd-mon-yyyy", "dd/mm/yyyy", "mm/dd/yyyy") + + +def parse_dates(values: pa.Array, formats: tuple[str, ...]) -> tuple[pa.Array, pa.Array]: + """Parse trimmed distinct date strings against the declared formats. + + A value is matched to the first format whose shape it fits, reassembled + as an ISO string and validated by Arrow's string-to-date cast, so no + calendar logic lives here. Null inputs stay null and are not invalid. + + Parameters + ---------- + values + String array with missing values already set to null. + formats + Ordered names from ``DateFormats.names()``. + + Returns + ------- + tuple[pa.Array, pa.Array] + A date32 array with nulls for missing and invalid values, and a + boolean array that is true where a value is invalid. + """ + iso = pa.nulls(len(values), pa.string()) + matched = pa.array([False] * len(values), pa.bool_()) + for name in formats: + ok, candidate = DateFormats.shape(name)(values) + ok = pc.and_(pc.fill_null(ok, False), pc.is_valid(candidate)) + ok = pc.and_(ok, pc.invert(matched)) + iso = pc.if_else(ok, candidate, iso) + matched = pc.or_(matched, ok) + invalid = pc.and_(pc.is_valid(values), pc.invert(matched)) + try: + return pc.cast(iso, pa.date32()), invalid + except pa.ArrowInvalid: + dates, calendar_failed = cast_each(iso, pa.date32()) + return dates, pc.or_(invalid, calendar_failed) + + +@dataclass(frozen=True) +class DateHandler(ConvertingHandler): + """Parses dates against an ordered list of declared formats. + + Values are parsed once per distinct value. A value that fits none of the + formats, or is not a calendar date, is unreadable. Forms such as + ``02/01/2020`` are never guessed and need a declared format. + + Attributes + ---------- + formats + Ordered names from ``DateFormats.names()``. + parser + Replaceable function with the signature of ``parse_dates``. + """ + + distinct: ClassVar[bool] = True + + formats: tuple[str, ...] = DateFormats.DEFAULT + parser: DateParser = parse_dates + + def __post_init__(self) -> None: + unknown = [name for name in self.formats if name not in DateFormats.names()] + if unknown: + raise ValueError(f"Unknown date formats {unknown}; choose from {sorted(DateFormats.names())}") + + def arrow_type(self, sa_type: TypeEngine[Any]) -> pa.DataType: + return pa.date32() + + def parse(self, cleaned: pa.Array, sa_type: TypeEngine[Any]) -> Interpretation: + return Interpretation(*self.parser(cleaned, self.formats)) + + +@dataclass(frozen=True) +class DateTimeHandler(DateHandler): + """Reads ISO timestamps, and date-only values in the declared date formats. + + A timestamp has date and time separated by ``T`` or a space, with + optional seconds and up to six fractional digits. A date-only value + becomes midnight. + """ + + distinct: ClassVar[bool] = False + + ISO: ClassVar[str] = r"^\d{4}-\d{2}-\d{2}[T ]\d{2}:\d{2}(:\d{2}(\.\d{1,6})?)?$" + + def arrow_type(self, sa_type: TypeEngine[Any]) -> pa.DataType: + return pa.timestamp("us") + + def parse(self, cleaned: pa.Array, sa_type: TypeEngine[Any]) -> Interpretation: + is_iso = pc.fill_null(pc.match_substring_regex(cleaned, self.ISO), False) + stamps, stamp_failed = safe_cast(pc.if_else(is_iso, cleaned, NULL_STRING), pa.timestamp("us")) + rest = pc.dictionary_encode(pc.if_else(is_iso, NULL_STRING, cleaned)) + dates, invalid = self.parser(rest.dictionary, self.formats) + dates = pc.take(dates, rest.indices) + invalid = pc.fill_null(pc.take(invalid, rest.indices), False) + values = pc.coalesce(stamps, pc.cast(dates, pa.timestamp("us"))) + return Interpretation(values, pc.or_(invalid, stamp_failed)) diff --git a/src/orm_loader/conversion/handlers/numbers.py b/src/orm_loader/conversion/handlers/numbers.py new file mode 100644 index 0000000..315335f --- /dev/null +++ b/src/orm_loader/conversion/handlers/numbers.py @@ -0,0 +1,70 @@ +"""Handlers for integer, floating point and fixed precision columns.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, ClassVar + +import pyarrow as pa +import pyarrow.compute as pc +from sqlalchemy.types import TypeEngine + +from ..values import NULL_STRING, Interpretation, safe_cast +from .base import ConvertingHandler + + +@dataclass(frozen=True) +class NumberHandler(ConvertingHandler): + """Accepts text matching ``PATTERN`` and casts it to the Arrow type. + + Attributes + ---------- + PATTERN + Regular expression a trimmed value must match. + """ + + PATTERN: ClassVar[str] + + def parse(self, cleaned: pa.Array, sa_type: TypeEngine[Any]) -> Interpretation: + fits = pc.fill_null(pc.match_substring_regex(cleaned, self.PATTERN), False) + candidate = pc.replace_substring_regex(pc.if_else(fits, cleaned, NULL_STRING), r"^\+", "") + values, cast_failed = safe_cast(candidate, self.arrow_type(sa_type)) + invalid = pc.and_(pc.is_valid(cleaned), pc.invert(fits)) + return Interpretation(values, pc.or_(invalid, cast_failed)) + + +@dataclass(frozen=True) +class IntegerHandler(NumberHandler): + """Accepts an optionally signed run of digits.""" + + PATTERN: ClassVar[str] = r"^[+-]?\d+$" + + def arrow_type(self, sa_type: TypeEngine[Any]) -> pa.DataType: + return pa.int64() + + +@dataclass(frozen=True) +class FloatHandler(NumberHandler): + """Accepts decimal and exponent forms; infinity is not accepted.""" + + PATTERN: ClassVar[str] = r"^[+-]?(\d+\.?\d*|\.\d+)([eE][+-]?\d+)?$" + + def arrow_type(self, sa_type: TypeEngine[Any]) -> pa.DataType: + return pa.float64() + + +@dataclass(frozen=True) +class NumericHandler(NumberHandler): + """Accepts decimal forms and keeps the column's declared precision and scale. + + A value with more digits than the column holds is unreadable. A column + without a declared precision and scale is read as a float. + """ + + PATTERN: ClassVar[str] = r"^[+-]?(\d+\.?\d*|\.\d+)$" + + def arrow_type(self, sa_type: TypeEngine[Any]) -> pa.DataType: + precision, scale = getattr(sa_type, "precision", None), getattr(sa_type, "scale", None) + if precision is not None and scale is not None: + return pa.decimal128(precision, scale) + return pa.float64() diff --git a/src/orm_loader/conversion/handlers/text.py b/src/orm_loader/conversion/handlers/text.py new file mode 100644 index 0000000..e2758db --- /dev/null +++ b/src/orm_loader/conversion/handlers/text.py @@ -0,0 +1,52 @@ +"""Handler for string and text columns.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import pyarrow as pa +import pyarrow.compute as pc +from sqlalchemy.types import TypeEngine + +from ..values import Interpretation, no_failures +from .base import ConvertingHandler + +if TYPE_CHECKING: + from ..policy import NullPolicy + + +@dataclass(frozen=True) +class TextHandler(ConvertingHandler): + """Keeps every nonempty value verbatim and rejects values longer than the column. + + Attributes + ---------- + strip + Surrounding whitespace is removed from stored text. + markers + Exact, case-sensitive values treated as missing, ignoring surrounding + whitespace. Replaces ``NullPolicy.text_markers`` for the columns this + handler serves; None uses the policy markers. + """ + + strip: bool = True + markers: frozenset[str] | None = None + + def arrow_type(self, sa_type: TypeEngine[Any]) -> pa.DataType: + return pa.string() + + def clean(self, raw: pa.Array, trimmed: pa.Array) -> pa.Array: + return trimmed if self.strip else raw + + def missing(self, trimmed: pa.Array, nulls: NullPolicy) -> pa.Array: + return nulls.text_missing(trimmed, self.markers) + + def parse(self, cleaned: pa.Array, sa_type: TypeEngine[Any]) -> Interpretation: + return Interpretation(cleaned, no_failures(len(cleaned))) + + def check(self, values: pa.Array, sa_type: TypeEngine[Any]) -> pa.Array | None: + length = getattr(sa_type, "length", None) + if not length: + return None + return pc.fill_null(pc.greater(pc.utf8_length(values), length), False) diff --git a/src/orm_loader/conversion/policy.py b/src/orm_loader/conversion/policy.py new file mode 100644 index 0000000..bd8f343 --- /dev/null +++ b/src/orm_loader/conversion/policy.py @@ -0,0 +1,220 @@ +"""The settings of a load: what is missing, what happens on error, and how each type is read.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any, ClassVar, Literal + +import pyarrow as pa +import pyarrow.compute as pc +from sqlalchemy.types import ( + Boolean, + Date, + DateTime, + Float, + Integer, + Numeric, + String, + TypeEngine, +) + +from .handlers import ( + BooleanHandler, + ColumnHandler, + DateHandler, + DateTimeHandler, + FloatHandler, + IntegerHandler, + NumericHandler, + PassthroughHandler, + TextHandler, +) +from .values import no_failures + + +@dataclass(frozen=True) +class NullPolicy: + """What counts as a missing value, shared by every column handler. + + Text columns keep every nonempty value verbatim apart from the markers. + Typed columns treat empty text and the typed missing words as missing. + + Attributes + ---------- + empty_is_null + Empty and whitespace-only text becomes null in text columns. + text_markers + Exact, case-sensitive text values treated as missing in text columns. + Matching ignores surrounding whitespace. + marker_quoting + ``"ignore"`` matches markers whether or not the source quoted them. + ``"unquoted_only"`` matches only markers the source did not quote, + which needs a source that records quoting and import-level markers. + typed_missing_words + Lowercase words treated as missing in non-text columns. + """ + + DEFAULT_TYPED_MISSING: ClassVar[frozenset[str]] = frozenset( + {"null", "none", "na", "n/a", "nan", "#n/a", "-nan", ""} + ) + + empty_is_null: bool = True + text_markers: frozenset[str] = frozenset() + marker_quoting: Literal["ignore", "unquoted_only"] = "ignore" + typed_missing_words: frozenset[str] = DEFAULT_TYPED_MISSING + + def __post_init__(self) -> None: + if self.marker_quoting not in ("ignore", "unquoted_only"): + raise ValueError(f"marker_quoting must be 'ignore' or 'unquoted_only', not {self.marker_quoting!r}") + + def text_missing(self, trimmed: pa.Array, markers: frozenset[str] | None = None) -> pa.Array: + """Return a mask of trimmed text values that are missing in a text column. + + Parameters + ---------- + trimmed + String array without surrounding whitespace. + markers + Markers that replace ``text_markers``, or None to use them. + + Returns + ------- + pa.Array + Boolean array that is true where the value is missing. + """ + words = self.text_markers if markers is None else markers + if self.marker_quoting != "ignore": + words = frozenset() + return self._mask(trimmed, words, empty=self.empty_is_null, casefold=False) + + def typed_missing(self, trimmed: pa.Array) -> pa.Array: + """Return a mask of trimmed text values that are missing in a typed column.""" + return self._mask(trimmed, self.typed_missing_words, empty=True, casefold=True) + + @staticmethod + def _mask(trimmed: pa.Array, words: frozenset[str], *, empty: bool, casefold: bool) -> pa.Array: + mask = no_failures(len(trimmed)) + if empty: + mask = pc.or_(mask, pc.fill_null(pc.equal(pc.utf8_length(trimmed), 0), False)) + if words: + probe = pc.utf8_lower(trimmed) if casefold else trimmed + hit = pc.is_in(probe, value_set=pa.array(sorted(words), pa.string())) + mask = pc.or_(mask, pc.fill_null(hit, False)) + return mask + + +@dataclass(frozen=True) +class ErrorPolicy: + """What happens to a row with an unreadable value or a missing required value. + + Attributes + ---------- + on_error + ``"raise"`` aborts on the first batch containing such a row. + ``"reject_row"`` drops those rows and writes them to a rejects file. + max_rejects + With ``reject_row``, abort once rejected rows exceed this count, or + this fraction of rows read when it is a float between 0 and 1. + """ + + on_error: Literal["raise", "reject_row"] = "raise" + max_rejects: float | None = None + + def __post_init__(self) -> None: + if self.on_error not in ("raise", "reject_row"): + raise ValueError(f"on_error must be 'raise' or 'reject_row', not {self.on_error!r}") + + +@dataclass(frozen=True) +class ConversionPolicy: + """The settings of one load, shared by every column. + + The handler for a column is, in order, the entry in ``columns`` and the + handler for the most specific class of the column type among + ``handlers`` and the built-in fields. + + Attributes + ---------- + nulls + What counts as a missing value. + errors + What happens to unreadable and missing required values. + text + Handler for string and text columns. + integer + Handler for integer columns. + floating + Handler for floating point columns. + numeric + Handler for fixed precision columns. + boolean + Handler for boolean columns. + date + Handler for date columns. + datetime + Handler for timestamp columns. + handlers + Handlers for further column types, or replacements for the built-in + ones, by SQLAlchemy type class. + columns + Handlers for single columns by ``(table_name, column_name)``. + """ + + nulls: NullPolicy = field(default_factory=NullPolicy) + errors: ErrorPolicy = field(default_factory=ErrorPolicy) + text: TextHandler = field(default_factory=TextHandler) + integer: IntegerHandler = field(default_factory=IntegerHandler) + floating: FloatHandler = field(default_factory=FloatHandler) + numeric: NumericHandler = field(default_factory=NumericHandler) + boolean: BooleanHandler = field(default_factory=BooleanHandler) + date: DateHandler = field(default_factory=DateHandler) + datetime: DateTimeHandler = field(default_factory=DateTimeHandler) + handlers: Mapping[type, ColumnHandler] = field(default_factory=dict) + columns: Mapping[tuple[str, str], ColumnHandler] = field(default_factory=dict) + + def __post_init__(self) -> None: + if self.nulls.marker_quoting == "unquoted_only": + texts = [self.text, *self.columns.values(), *self.handlers.values()] + if any(isinstance(h, TextHandler) and h.markers is not None for h in texts): + raise ValueError("marker_quoting='unquoted_only' supports NullPolicy.text_markers only") + + def for_column( + self, sa_type: TypeEngine[Any], table_name: str | None = None, column_name: str | None = None + ) -> ColumnHandler: + """Return the handler that interprets a column. + + Parameters + ---------- + sa_type + Column type. + table_name + Table the column belongs to. + column_name + The column. + + Returns + ------- + ColumnHandler + The handler for the column; a passthrough handler for a column + type with none. + """ + if table_name and column_name: + handler = self.columns.get((table_name, column_name)) + if handler is not None: + return handler + # Dispatch default first and then overwrite with any registered handlers + by_type: dict[type, ColumnHandler] = { + String: self.text, + Integer: self.integer, + Float: self.floating, + Numeric: self.numeric, + Boolean: self.boolean, + Date: self.date, + DateTime: self.datetime, + **self.handlers, + } + for cls in type(sa_type).__mro__: + if cls in by_type: + return by_type[cls] + return PassthroughHandler() diff --git a/src/orm_loader/conversion/values.py b/src/orm_loader/conversion/values.py new file mode 100644 index 0000000..3bd2056 --- /dev/null +++ b/src/orm_loader/conversion/values.py @@ -0,0 +1,94 @@ +"""Result type and Arrow helpers shared by the column handlers.""" + +from __future__ import annotations + +from typing import Any, NamedTuple + +import pyarrow as pa +import pyarrow.compute as pc + +NULL_STRING = pa.scalar(None, pa.string()) + + +class Interpretation(NamedTuple): + """The outcome of interpreting one column. + + Attributes + ---------- + values + Converted values, null where a value is missing or unreadable. + invalid + Boolean array that is true where a value could not be read. + """ + + values: pa.Array + invalid: pa.Array + + +def as_array(arr: pa.Array | pa.ChunkedArray) -> pa.Array: + """Return a chunked array as one contiguous array.""" + if isinstance(arr, pa.ChunkedArray): + return arr.combine_chunks() if arr.num_chunks else pa.array([], arr.type) + return arr + + +def is_text_like(arr_type: pa.DataType) -> bool: + """Return True for string arrays and dictionary arrays of strings.""" + if pa.types.is_dictionary(arr_type): + arr_type = arr_type.value_type + return pa.types.is_string(arr_type) or pa.types.is_large_string(arr_type) + + +def no_failures(length: int) -> pa.Array: + """Return a boolean array of ``length`` false values.""" + return pa.array([False] * length, pa.bool_()) + + +def cast_each(values: pa.Array, target: pa.DataType) -> tuple[pa.Array, pa.Array]: + """Cast one element at a time. + + Parameters + ---------- + values + Array to cast. + target + Arrow type to cast to. + + Returns + ------- + tuple[pa.Array, pa.Array] + The cast array with nulls where a cast failed, and a boolean array + that is true where it failed. + """ + out: list[Any] = [] + failed: list[bool] = [] + for i in range(len(values)): + try: + out.append(pc.cast(values[i : i + 1], target)[0].as_py()) + failed.append(False) + except (pa.ArrowInvalid, pa.ArrowNotImplementedError): + out.append(None) + failed.append(True) + return pa.array(out, target), pa.array(failed, pa.bool_()) + + +def safe_cast(values: pa.Array, target: pa.DataType) -> tuple[pa.Array, pa.Array]: + """Cast the whole array, falling back to element-wise casts on failure. + + Parameters + ---------- + values + Array to cast. + target + Arrow type to cast to. + + Returns + ------- + tuple[pa.Array, pa.Array] + The cast array with nulls where a cast failed, and a boolean array + that is true where it failed. + """ + try: + return pc.cast(values, target), no_failures(len(values)) + except (pa.ArrowInvalid, pa.ArrowNotImplementedError): + return cast_each(values, target) diff --git a/src/orm_loader/helpers/__init__.py b/src/orm_loader/helpers/__init__.py index 2c4ed42..55533f9 100644 --- a/src/orm_loader/helpers/__init__.py +++ b/src/orm_loader/helpers/__init__.py @@ -1,29 +1,22 @@ -from .errors import IngestError, ValidationError +from .errors import ConversionError, IngestError, RejectLimitError, RequiredValueError, ValidationError from .bootstrap import bootstrap, create_db, create_tables from .engines import dispose_engines -from .sqlite import ( - attach_sqlite_bulk_load_pragmas, - explain_sqlite_fk_error, - restore_sqlite_journal_mode, -) from .bulk import bulk_load_context, engine_with_replica_role from .metadata import Base from .discovery import get_model_by_tablename -from .null_handlers import normalise_null __all__ = [ + "ConversionError", + "RejectLimitError", + "RequiredValueError", "IngestError", "ValidationError", "bootstrap", "create_db", "create_tables", "dispose_engines", - "attach_sqlite_bulk_load_pragmas", - "explain_sqlite_fk_error", - "restore_sqlite_journal_mode", "bulk_load_context", "engine_with_replica_role", "Base", "get_model_by_tablename", - "normalise_null", ] diff --git a/src/orm_loader/helpers/bulk.py b/src/orm_loader/helpers/bulk.py index 459ccf5..d425b09 100644 --- a/src/orm_loader/helpers/bulk.py +++ b/src/orm_loader/helpers/bulk.py @@ -1,36 +1,8 @@ -import logging from contextlib import contextmanager from sqlalchemy import Engine from sqlalchemy.orm import Session from typing import Generator -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)): - logger.error(f"Unexpected FK state type: {type(previous_state)}. Expected str or int.") - raise TypeError(f"Expected previous FK state to be str or int, got {type(previous_state)}") - return previous_state - -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)): - logger.error(f"Unexpected FK state type: {type(previous_state)}. Expected str or int.") - raise TypeError(f"Expected previous FK state to be str or int, got {type(previous_state)}") - return previous_state - -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}") @contextmanager def bulk_load_context( @@ -42,8 +14,8 @@ def bulk_load_context( """ Wrap a trusted bulk operation in backend-aware session settings. - This is a thin helper over ``DatabaseBackend.bulk_load_context()``. - It exists so older call sites can keep using the helper import path. + Foreign-key checks are disabled and autoflush is suspended for the + duration, then restored. See ``DatabaseBackend.bulk_load_context``. """ from ..backends.resolve import resolve_backend backend = resolve_backend(session) diff --git a/src/orm_loader/helpers/errors.py b/src/orm_loader/helpers/errors.py index 2e32a28..665f226 100644 --- a/src/orm_loader/helpers/errors.py +++ b/src/orm_loader/helpers/errors.py @@ -4,3 +4,87 @@ class IngestError(Exception): class ValidationError(Exception): """Raised when schema or specification validation fails.""" + + +class ConversionError(IngestError): + """Raised when values cannot be read as their column type. + + Parameters + ---------- + table_name + Table the column belongs to, if known. + column_name + Column the values were read for, if known. + expected + Name of the column type. + count + Number of unreadable values. + samples + First few unreadable raw values. + """ + + def __init__( + self, + *, + table_name: str | None, + column_name: str | None, + expected: str, + count: int, + samples: list, + ) -> None: + self.table_name = table_name + self.column_name = column_name + self.expected = expected + self.count = count + self.samples = samples + target = ".".join(part for part in (table_name, column_name) if part) or "column" + hint = " Declare the source format with DateHandler(formats=...)." if "Date" in expected else "" + super().__init__( + f"{target}: {count} value(s) cannot be read as {expected}. Examples: {samples}.{hint}" + ) + + +class RequiredValueError(IngestError): + """Raised when a required column has missing values. + + Parameters + ---------- + table_name + Table the column belongs to. + column_name + The required column. + count + Number of rows missing a value. + rows + First few data row numbers (starting at 1) that are missing a value. + """ + + def __init__(self, *, table_name: str, column_name: str, count: int, rows: list[int]) -> None: + self.table_name = table_name + self.column_name = column_name + self.count = count + self.rows = rows + super().__init__( + f"{table_name}.{column_name}: {count} row(s) have no value for a required column. " + f"First data rows: {rows}." + ) + + +class RejectLimitError(IngestError): + """Raised when rejected rows exceed the policy's ``max_rejects``. + + Parameters + ---------- + rejected + Number of rows rejected so far. + read + Number of rows read so far. + limit + The configured limit. + """ + + def __init__(self, *, rejected: int, read: int, limit: float) -> None: + self.rejected = rejected + self.read = read + self.limit = limit + super().__init__(f"{rejected} of {read} rows rejected, over the limit of {limit}.") diff --git a/src/orm_loader/helpers/null_handlers.py b/src/orm_loader/helpers/null_handlers.py deleted file mode 100644 index 09a5464..0000000 --- a/src/orm_loader/helpers/null_handlers.py +++ /dev/null @@ -1,24 +0,0 @@ -import pandas as pd -from typing import Any - -_NULL_STRINGS = {"", "nan", "null", "none", "na", "n/a"} - -def normalise_null(value: Any) -> Any | None: - if value is None: - return None - - # pandas / numpy NaN - try: - if pd.isna(value): - return None - except Exception: - pass - - # string garbage - if isinstance(value, str): - s = value.strip().lower() - if s in _NULL_STRINGS: - return None - return value # keep legit strings - - return value diff --git a/src/orm_loader/helpers/sqlite.py b/src/orm_loader/helpers/sqlite.py deleted file mode 100644 index f3e2f24..0000000 --- a/src/orm_loader/helpers/sqlite.py +++ /dev/null @@ -1,47 +0,0 @@ -from pathlib import Path - -from sqlalchemy.engine import Engine -from sqlalchemy.exc import IntegrityError - -from ..backends.sqlite import SQLiteBackend - - -def attach_sqlite_bulk_load_pragmas( - engine: Engine, - *, - busy_timeout_ms: int = 60000, - journal_mode: str = "WAL", - defer_foreign_keys: bool = True, -) -> None: - """ - Install SQLite connect hooks aimed at heavy local write workloads. - - The hook currently sets ``busy_timeout``, journal mode, and foreign-key - enforcement, and can also enable deferred foreign-key checking for the - connection. - - Note that this is the replacement for old ``enable_sqlite_foreign_keys()`` - workaround - this should be no longer needed. - """ - SQLiteBackend( - busy_timeout_ms=busy_timeout_ms, - journal_mode=journal_mode, - defer_foreign_keys=defer_foreign_keys, - ).install_engine_hooks(engine) - - -def explain_sqlite_fk_error( - session, exc: IntegrityError, raise_error: bool = True, *, mapper=None -): - """Log SQLite foreign-key check details before re-raising an error.""" - SQLiteBackend().explain_fk_error(session, exc, mapper=mapper, raise_error=raise_error) - - -def restore_sqlite_journal_mode(db_path: Path) -> None: - """ - Checkpoint WAL contents and switch the database back to ``DELETE`` mode. - - Call this after disposing active SQLite connections. Reconnecting - through an engine that still installs WAL hooks will enable WAL again. - """ - SQLiteBackend().restore_journal_mode(db_path) diff --git a/src/orm_loader/loaders/__init__.py b/src/orm_loader/loaders/__init__.py index 2be4aa2..d9705ca 100644 --- a/src/orm_loader/loaders/__init__.py +++ b/src/orm_loader/loaders/__init__.py @@ -1,14 +1,12 @@ -from .loader_interface import LoaderInterface, PandasLoader, ParquetLoader -from .data_classes import LoaderContext, TableCastingStats -from .loading_helpers import infer_delim, infer_encoding, quick_load_pg +from .loader import LoadReport, LoaderContext, stage_file +from .readers import DelimitedTextReader, FileReader, ParquetReader, open_reader __all__ = [ - "LoaderInterface", - "LoaderContext", - "PandasLoader", - "TableCastingStats", - "infer_delim", - "infer_encoding", - "quick_load_pg", - "ParquetLoader", -] \ No newline at end of file + "DelimitedTextReader", + "FileReader", + "LoadReport", + "LoaderContext", + "ParquetReader", + "open_reader", + "stage_file", +] diff --git a/src/orm_loader/loaders/data/__init__.py b/src/orm_loader/loaders/data/__init__.py deleted file mode 100644 index 782e1cc..0000000 --- a/src/orm_loader/loaders/data/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -from .converters import perform_cast, cast_arrow_column - -__all__ = [ - "perform_cast", - "cast_arrow_column", -] \ No newline at end of file diff --git a/src/orm_loader/loaders/data/converters.py b/src/orm_loader/loaders/data/converters.py deleted file mode 100644 index c162a0e..0000000 --- a/src/orm_loader/loaders/data/converters.py +++ /dev/null @@ -1,452 +0,0 @@ -from sqlalchemy import ColumnElement -from sqlalchemy.types import Integer, Float, Boolean, Date, DateTime, String, Text, TypeEngine -from typing import Any, Callable -from enum import Enum -import math -from dataclasses import dataclass -import pyarrow as pa -import pyarrow.compute as pc -import pandas as pd -import re -from datetime import datetime, date -from dateutil import parser - -from ..data_classes import TableCastingStats - -_NUMERIC_RE = re.compile(r"^[+-]?\d+(\.\d+)?$") - -_AVAILABLE_DATE_FORMATS = ( - "%Y%m%d", # 20170824 (athena standard) - "%d-%b-%Y", # 24-AUG-2017 (oncology-branch vocab) - "%Y-%m-%d", # 2017-08-24 (ISO) - "%d/%m/%Y", # 24/08/2017 -) - - -_ARROW_TYPE_MAP: dict[type, pa.DataType] = { - Integer: pa.int64(), - Float: pa.float64(), - Boolean: pa.bool_(), - Date: pa.date32(), - DateTime: pa.timestamp("us"), -} - -@dataclass(frozen=True) -class CastRule: - sa_type: type - scalar: Callable[[Any, Any], Any] - arrow: Callable[..., Any] | None = None # optional vectorised impl, never actually set/read today - - -_NULL_STRINGS = { - "null", - "none", - "na", - "n/a", - "nan", -} - -def _normalise_null(value: Any) -> Any | None: - if value is None: - return None - - if value is pd.NA or value is pd.NaT: - return None - - if isinstance(value, float) and math.isnan(value): - return None - - if isinstance(value, str): - s = value.strip().lower() - if s == "": - return None - if s in _NULL_STRINGS: - return None - - return value - - -_ARROW_NULL_STRINGS = pa.array(sorted(_NULL_STRINGS | {""}), type=pa.string()) - - -def _normalise_null_arrow(arr: pa.Array | pa.ChunkedArray) -> pa.Array | pa.ChunkedArray: - arrow_type = arr.type - - if pa.types.is_string(arrow_type) or pa.types.is_large_string(arrow_type): - probe = pc.utf8_lower(pc.utf8_trim_whitespace(arr)) # type: ignore - return pc.if_else( # type: ignore - pc.is_in(probe, value_set=_ARROW_NULL_STRINGS), # type: ignore - pa.nulls(len(arr), arrow_type), - arr, - ) - - if pa.types.is_floating(arrow_type): - return pc.if_else( # type: ignore - pc.is_nan(arr), # type: ignore - pa.nulls(len(arr), arrow_type), - arr, - ) - - return arr - - -def _to_numeric_string(value: str | None) -> str | None: - if value is None: - return None - - if not _NUMERIC_RE.match(value): - return value - - if "." in value: - f = float(value) - if f.is_integer(): - return str(int(f)) - return str(f) - - return str(int(value)) - -def _to_number(value: Any) -> int | float | None: - if value is None: - return None - - if isinstance(value, float): - if math.isnan(value): - return None - if value.is_integer(): - return int(value) - raise ValueError(f"Non-integer float: {value}") - - if isinstance(value, int): - return value - - if isinstance(value, str): - s = value.strip() - if s == "": - return None - - s = _to_numeric_string(s) - if s is None: - return None - - return int(s) - - -def _to_int(value: Any) -> int | None: - n = _to_number(value) - if n is None: - return None - - if isinstance(n, int): - return n - - raise ValueError(f"Non-integer numeric value: {value}") - -def _cast_string(value: Any, sa_type: TypeEngine[Any]) -> str | None: - if value is None: - return None - - if isinstance(value, float) and math.isnan(value): - return None - - s = str(value).strip() - if s == "": - return None - - s = _to_numeric_string(s) or "" - - if isinstance(sa_type, (String, Text)) and sa_type.length: - if len(s) > sa_type.length: - return s[: sa_type.length] - - return s - -def _to_float(value: Any) -> float | None: - if value is None: - return None - - if isinstance(value, float): - if math.isnan(value): - return None - return float(value) - - if isinstance(value, int): - return float(value) - - if isinstance(value, str): - s = value.strip() - if s == "": - return None - try: - return float(s) - except ValueError: - return None - - return None - - -CAST_RULES: list[CastRule] = [ - CastRule(Integer, lambda v, _: _to_int(v) if v is not None else None), - CastRule(Float, lambda v, _: _to_float(v) if v is not None else None), - CastRule(Boolean, lambda v, _: _to_bool(v)), - CastRule(Date, lambda v, _: _parse_date(v)), - CastRule(DateTime,lambda v, _: _parse_datetime(v)), - CastRule(String, _cast_string), - CastRule(Text, _cast_string), -] - -# Per-column overrides, keyed by (table_name, column_name). Checked ahead of -# CAST_RULES, since dispatch there is purely by SQLAlchemy type and can't -# target one specific column (e.g. a plain String(1) column with no -# type-level signal of the enum it's meant to hold) or safely disambiguate -# sa.Enum from its own String supertype without a caller telling it which -# column actually means "enum" -_COLUMN_CAST_RULES: dict[tuple[str, str], Callable[[Any], Any]] = {} - - -def _enum_member_scalar(enum_type: type[Enum]) -> Callable[[Any], Any]: - # Built once per registration rather than per row: every CSV load reads with - # dtype=str, so an int-valued member arrives as "1" -- which neither - # enum_type(value) nor enum_type(str(value)) can match. - _by_str_value = {str(member.value).strip(): member.name for member in enum_type} - - def _scalar(value: Any) -> Any: - if value is None: - return None - - # Match by .value (the raw display text a real data source uses), - # store by .name -- sa.Enum's own default column-storage convention. - # The staging-to-target merge is a plain SQL copy with no Python-level - # type translation, so this must already be exactly what the column - # expects to find; .name is that default for an unmodified sa.Enum. - - # Already the right member: only reachable via direct cast_scalar use, - # since file sources yield raw scalars. Handled ahead of the str() - # fallback, which would mangle it to "SomeEnum.MEMBER". - if isinstance(value, enum_type): - return value.name - - # Raw value first, so enums whose member values aren't strings - # (IntEnum and friends -- sa.Enum still persists .name for those) match - # without a lossy str() round-trip. - try: - return enum_type(value).name - except (ValueError, KeyError, TypeError): - pass - - s = str(value).strip() - if s == "": - return None - try: - return enum_type(s).name - except (ValueError, KeyError, TypeError): - pass - - try: - return _by_str_value[s] - except KeyError: - # Raise rather than return None, so an unmatched value goes through - # the on_error/TableCastingStats path like every other cast failure. - raise ValueError(f"{value!r} matches no {enum_type.__name__} member value") - return _scalar - - -def register_column_cast_rule( - table_name: str, - column_name: str, - scalar: Callable[[Any], Any] | None = None, - *, - enum_type: type[Enum] | None = None, -) -> None: - """ - Register custom cast/validation logic for one specific column. - - If `enum_type` is given, the raw value is matched against its members - by `.value` and coerced to the matching member's `.name`. Member values - need not be strings -- an `IntEnum` matches both `1` (a parquet int64 - column) and `"1"` (any CSV column, since those are read with dtype=str). - - This is `sa.Enum`'s default column-storage convention, and covers genuine - `sa.Enum(enum_type)` columns with no further setup. An unmatched value - goes through the existing on_error/TableCastingStats path and the - column is set to None, same as every other cast failure today. - - For anything that departs from that default convention (a plain - String/Char column with an implied enum and no `sa.Enum` storage - convention at all (e.g. OMOP CDM's `concept.standard_concept`: - `String(1)`, values 'S'/'C'/NULL, where the column expects `.value` - directly, not `.name`), or a genuine `sa.Enum` column configured with a - custom `values_callable` such as in HemOnc) use `scalar` instead. - Resolution logic for the scalar callable is then tailored to what - that specific column actually expects to store (e.g. - `lambda v: SomeEnum(v).value` for the plain-String case above). - - Registering the same (table_name, column_name) again replaces the - previous rule rather than stacking, so repeated registration is - safe/idempotent. - """ - if enum_type is not None: - if scalar is not None: - raise ValueError("register_column_cast_rule requires exactly one of `scalar` or `enum_type`.") - resolved_scalar = _enum_member_scalar(enum_type) - elif scalar is not None: - resolved_scalar = scalar - else: - raise ValueError("register_column_cast_rule requires exactly one of `scalar` or `enum_type`.") - - _COLUMN_CAST_RULES[(table_name, column_name)] = resolved_scalar - -def _dateutil_fallback(value: str) -> datetime | None: - try: - dt = parser.parse( - value, - dayfirst=True, - yearfirst=False, - fuzzy=False, - ) - except (ValueError, OverflowError): - return None - - normalised = dt.strftime("%Y-%m-%d") - if normalised not in value: - return None - - return dt - -def _parse_date(value: Any) -> date | None: - if isinstance(value, date) and not isinstance(value, datetime): - return value - if isinstance(value, datetime): - return value.date() - if isinstance(value, str): - for fmt in _AVAILABLE_DATE_FORMATS: - try: - return datetime.strptime(value, fmt).date() - except ValueError: - continue - return None - -def _parse_datetime(value: Any) -> datetime | None: - if isinstance(value, datetime): - return value - if isinstance(value, date): - return datetime.combine(value, datetime.min.time()) - if isinstance(value, str): - try: - return datetime.fromisoformat(value) - except ValueError: - pass - - # Fallback to date-only formats + midnight - d = _parse_date(value) - if d: - return datetime.combine(d, datetime.min.time()) - - return _dateutil_fallback(value) - -def _to_bool(value: Any) -> bool | None: - if value is None or (isinstance(value, float) and math.isnan(value)): - return None - if isinstance(value, bool): - return value - s = str(value).strip().lower() - if s in {"true", "t", "yes", "y", "1"}: - return True - if s in {"false", "f", "no", "n", "0"}: - return False - return None - -def cast_scalar( - value: Any, - sa_type: TypeEngine[Any], - *, - on_error: Callable[[Any], None] | None = None, - table_name: str | None = None, - column_name: str | None = None, -) -> Any: - value = _normalise_null(value) - if value is None: - return None - - column_rule = _COLUMN_CAST_RULES.get((table_name, column_name)) if table_name and column_name else None - if column_rule is not None: - try: - return column_rule(value) - except Exception: - if on_error: - on_error(value) - return None - - for rule in CAST_RULES: - if isinstance(sa_type, rule.sa_type): - try: - return rule.scalar(value, sa_type) - except Exception: - if on_error: - on_error(value) - return None - - return value - - -def perform_cast( - value: Any, - sa_type: TypeEngine[Any], - *, - on_error: Callable[[Any], None] | None, - table_name: str | None = None, - column_name: str | None = None, -) -> Any: - return cast_scalar(value, sa_type, on_error=on_error, table_name=table_name, column_name=column_name) - - -def cast_arrow_column(arr: pa.Array, sa_col: ColumnElement[Any], stats: TableCastingStats | None = None) -> pa.Array: - arr = _normalise_null_arrow(arr) - - column_rule = _COLUMN_CAST_RULES.get((sa_col.table.name, sa_col.name)) - if column_rule is not None: - values: list[Any] = [] - for v in arr: - raw = v.as_py() - try: - values.append(column_rule(raw) if raw is not None else None) - except Exception: - if stats: - stats.record(column=sa_col.name, value=raw) - values.append(None) - return pa.array(values, type=arr.type) - - for rule in CAST_RULES: - if isinstance(sa_col.type, rule.sa_type): - # Use Arrow native cast if available - arrow_type = _ARROW_TYPE_MAP.get(rule.sa_type) - if arrow_type: - try: - return pc.cast(arr, arrow_type) - except pa.ArrowInvalid: - validity = pc.is_valid(arr) # type: ignore - invalid_mask = pc.invert(validity) # type: ignore - invalid_count = pc.sum(invalid_mask).as_py() # type: ignore - if invalid_count == 0: - return arr - - bad_values = [ - v.as_py() - for v, bad in zip(arr, invalid_mask) - if bad - ][:3] - if stats: - stats.record( - column=sa_col.name, - value={ - "count": invalid_count, - "examples": bad_values, - "reason": f"Arrow cast to {arrow_type} failed" - }, - ) - return arr - # fallback: scalar apply - return pa.array( - [rule.scalar(v.as_py(), sa_col) for v in arr], - type=arr.type, - ) - return arr \ No newline at end of file diff --git a/src/orm_loader/loaders/data_classes.py b/src/orm_loader/loaders/data_classes.py deleted file mode 100644 index e36e9b9..0000000 --- a/src/orm_loader/loaders/data_classes.py +++ /dev/null @@ -1,243 +0,0 @@ -from typing import Any, Type, List, Dict, TYPE_CHECKING -from dataclasses import dataclass, field -import sqlalchemy as sa -import sqlalchemy.orm as so -from pathlib import Path -import pandas as pd -import pyarrow as pa -from logging import getLogger - -logger = getLogger(__name__) - -if TYPE_CHECKING: - from ..tables.typing import CSVTableProtocol - -def _clean_nulls(v): - if v is None: - return None - if v is pd.NA or v is pd.NaT: - return None - if isinstance(v, float) and pd.isna(v): - return None - return v -""" -Loader Data Structures -====================== - -This module defines shared data structures and base interfaces used by -file loaders. - -It includes: -- the LoaderContext coordination object -- the abstract LoaderInterface -- casting statistics helpers for diagnostics and logging - -No file I/O or database-specific loading logic is implemented here. -""" - - -@dataclass(frozen=True) -class LoaderContext: - - """ - Immutable context object passed through loader operations. - - The LoaderContext encapsulates all state required to load a file - into a staging table, without relying on global variables. - - Attributes - ---------- - tableclass - ORM table class being loaded. - session - Active SQLAlchemy session. - path - Path to the source data file. - staging_table - SQLAlchemy Table object representing the staging table. - chunksize - Optional chunk size for incremental loading. - normalise - Whether to apply type casting / normalisation. - dedupe - Whether to perform deduplication (pre-insertion for source issues) - quote_mode - Quoting mode, resolved to the same concrete mode on both the - PostgreSQL COPY fast-path and the pandas ORM fallback. - staging_schema_tag - schema_translate_map tag the staging table lives in, passed to - resolve_backend() so every backend resolution within this load - resolves the same schema. - """ - tableclass: Type["CSVTableProtocol"] - session: so.Session - path: Path - staging_table: sa.Table - - chunksize: int | None = None - normalise: bool = True - dedupe: bool = True - quote_mode: str = "auto" - staging_schema_tag: str | None = None - -class LoaderInterface: - - """ - Abstract interface for file loaders. - - Loader implementations are responsible for: - - reading input files - - optional chunking - - optional deduplication - - optional normalisation - - loading into a staging table - - Concrete loaders must implement ``orm_file_load`` and ``dedupe``. - """ - - @classmethod - def orm_file_load(cls, ctx: LoaderContext) -> int: - """ - Load a file into the staging table. - - Parameters - ---------- - ctx - Loader context. - - Returns - ------- - int - Number of rows loaded. - """ - raise NotImplementedError - - @classmethod - def _load_chunk( - cls, - staging_cls: sa.Table, - session: so.Session, - dataframe: pd.DataFrame, - mapper: Any | None = None, - ) -> int: - """ - Load a single DataFrame chunk into the staging table. - - Parameters - ---------- - staging_cls - SQLAlchemy Table object representing the staging table. - session - Active SQLAlchemy session. - dataframe - DataFrame containing rows to insert. - - Returns - ------- - int - Number of rows inserted. - """ - if dataframe.empty: - return 0 - - raw_records = dataframe.to_dict(orient="records") - records = [ - { - f'{k}': _clean_nulls(v) - for k, v in record.items() - } - for record in raw_records - ] - - session.execute( - staging_cls.insert(), - records, - bind_arguments={"mapper": mapper}, - ) - session.flush() - session.expunge_all() - - return len(dataframe) - - @classmethod - def dedupe(cls, data: pd.DataFrame | pa.Table, ctx: LoaderContext) -> Any: - """ - Deduplicate incoming data. - - Concrete implementations define deduplication semantics. - - Parameters - ---------- - data - Input data. - ctx - Loader context. - - Returns - ------- - Any - Deduplicated data. - """ - raise NotImplementedError - - -@dataclass -class ColumnCastingStats: - """ - Casting statistics for a single column. - """ - count: int = 0 - examples: List[Any] = field(default_factory=list) - - def record(self, value: Any, example_limit: int = 3): - self.count += 1 - if len(self.examples) < example_limit: - self.examples.append(value) - -@dataclass -class TableCastingStats: - """ - Aggregated casting statistics for a table, keyed by column. - """ - table_name: str - columns: Dict[str, ColumnCastingStats] = field(default_factory=dict) - - def record( - self, - *, - column: str, - value: Any, - example_limit: int = 3, - ): - """ - Record a casting failure for a column. - """ - if column not in self.columns: - self.columns[column] = ColumnCastingStats() - self.columns[column].record(value, example_limit=example_limit) - - @property - def total_failures(self) -> int: - """ - Total number of casting failures. - """ - return sum(stats.count for stats in self.columns.values()) - - def has_failures(self) -> bool: - """ - Whether any casting failures occurred. - """ - return self.total_failures > 0 - - def to_dict(self) -> dict[str, dict[str, Any]]: - """ - Return a dictionary representation of casting statistics. - """ - return { - col: { - "count": stats.count, - "examples": stats.examples, - } - for col, stats in self.columns.items() - } - diff --git a/src/orm_loader/loaders/loader.py b/src/orm_loader/loaders/loader.py new file mode 100644 index 0000000..6ca55bd --- /dev/null +++ b/src/orm_loader/loaders/loader.py @@ -0,0 +1,212 @@ +"""File loading +============ + +A reader yields a file as Arrow batches, the batch processor interprets them +under the load's conversion policy and the backend's sink writes them to the +staging table. Every backend and file type takes the same path. +""" + +from __future__ import annotations + +import logging +from collections.abc import Iterable +from itertools import chain +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, Type + +import sqlalchemy as sa +import sqlalchemy.orm as so + +from ..backends.base import DedupeResult +from ..backends.resolve import resolve_backend +from ..conversion import ConversionPolicy +from ..helpers import IngestError +from .processor import BatchProcessor +from .readers import FileReader + +if TYPE_CHECKING: + from ..tables.typing import FileTableProtocol + +logger = logging.getLogger(__name__) + + +@dataclass +class LoadReport: + """Row accounting for one load. + + Attributes + ---------- + rows_read + Data rows read from the file. + rows_loaded + Rows written to the staging table. + rows_rejected + Rows dropped because a value was unreadable or a required value was + missing, with ``ErrorPolicy.on_error`` set to ``"reject_row"``. + dedupe + Result of deduplicating the staging table, if it was requested. + rejects_path + File the rejected rows were written to, if any. + """ + + rows_read: int = 0 + rows_loaded: int = 0 + rows_rejected: int = 0 + dedupe: DedupeResult | None = None + rejects_path: Path | None = None + + @property + def rows_deduplicated(self) -> int: + """Rows removed as duplicates, zero when deduplication did not run.""" + return self.dedupe.rows_removed if self.dedupe else 0 + + @property + def accounted(self) -> bool: + """Whether every row read was loaded, rejected or dropped as a duplicate.""" + return self.rows_read == self.rows_loaded + self.rows_rejected + self.rows_deduplicated + + +@dataclass(frozen=True) +class LoaderContext: + """Immutable context passed through a load. + + The context holds all state required to load a file into a staging + table, without relying on global variables. + + Attributes + ---------- + tableclass + ORM table class being loaded. + session + Active SQLAlchemy session. + path + Path to the source data file. + staging_table + SQLAlchemy Table object representing the staging table. + chunksize + Rows per batch for the default reader of formats that batch by rows. + normalise + Whether to interpret values under the policy. + dedupe + Whether the staging table keeps only the first row for each primary key. + quote_mode + Quoting mode of the default reader for delimited text: ``"csv"``, + ``"literal"``, ``"by_delimiter"`` or ``"auto"``. + staging_schema_tag + schema_translate_map tag the staging table lives in, passed to + resolve_backend() so every backend resolution within this load + resolves the same schema. + policy + How values are interpreted and what happens to unreadable ones. + rejects_dir + Directory that rejected rows are written to. Defaults to a + ``rejects`` directory inside the source file's directory. + report + Row accounting, filled in as the load runs. + """ + + tableclass: Type["FileTableProtocol"] + session: so.Session + path: Path + staging_table: sa.Table + + chunksize: int | None = None + normalise: bool = True + dedupe: bool = True + quote_mode: str = "auto" + staging_schema_tag: str | None = None + policy: ConversionPolicy = field(default_factory=ConversionPolicy) + rejects_dir: Path | None = None + report: LoadReport = field(default_factory=LoadReport) + + +def _require_columns_present(tableclass: Type["FileTableProtocol"], present_columns: Iterable[str]) -> None: + """Raise before staging if a required column is entirely absent from the file. + + Partial-column loads are not supported, so every merge strategy + (including upsert) needs a complete row shape and a missing required + column is always a data problem. + + Parameters + ---------- + tableclass + Table being loaded. + present_columns + Model columns present in the file. + + Raises + ------ + IngestError + A non-nullable column without a default is not in the file. + """ + present = set(present_columns) + missing = sorted(c for c in tableclass.required_columns() if c not in present) + if missing: + raise IngestError( + f"{tableclass.__tablename__}: source data is missing required column(s) {missing} (non-nullable, no default)" + ) + + +def stage_file(ctx: LoaderContext, reader: FileReader) -> int: + """Load a file into the staging table. + + The reader streams the file in bounded batches, so memory does not grow + with the row count. A file without data rows loads nothing and is not + checked for required columns. With ``ctx.dedupe`` set, the staging table keeps the + first row for each primary key. + + Parameters + ---------- + ctx + Loader context. + reader + Reader that yields the file as Arrow batches. + + Returns + ------- + int + Number of rows left in the staging table. + + Raises + ------ + IngestError + The file cannot be read, is missing a required column, or the rows + read are not all accounted for. + ConversionError + A value cannot be read as its column type. + RequiredValueError + A required column has no value. + RejectLimitError + Rejected rows exceed ``ErrorPolicy.max_rejects``. + """ + table_name = ctx.tableclass.__tablename__ + batches = reader.batches() + first = None if reader.empty else next(batches, None) + if first is None: + logger.info("File %s has no data rows; skipping load for %s", ctx.path.name, table_name) + return 0 + if ctx.normalise: + _require_columns_present(ctx.tableclass, reader.columns) + rejects_dir = ctx.rejects_dir or ctx.path.parent / "rejects" + processor = BatchProcessor( + ctx.tableclass, + reader.columns, + ctx.policy, + ctx.report, + rejects_dir / f"{ctx.path.stem}.rejects.csv", + ctx.normalise, + ) + backend = resolve_backend(ctx.session, staging_schema_tag=ctx.staging_schema_tag, mapper=ctx.tableclass) + total = backend.insert_batches(ctx, processor.stream(chain([first], batches)), reader.columns) + keys = [c for c in ctx.tableclass.pk_names() if c in reader.columns] + if ctx.dedupe and keys: + ctx.report.dedupe = backend.deduplicate_staging(ctx, keys) + total = ctx.report.dedupe.rows_after + ctx.report.rows_loaded += total + if not ctx.report.accounted: + raise IngestError( + f"{table_name}: read {ctx.report.rows_read} rows but loaded {ctx.report.rows_loaded}, " + f"rejected {ctx.report.rows_rejected} and dropped {ctx.report.rows_deduplicated} duplicates" + ) + return total diff --git a/src/orm_loader/loaders/loader_interface.py b/src/orm_loader/loaders/loader_interface.py deleted file mode 100644 index 2b02fee..0000000 --- a/src/orm_loader/loaders/loader_interface.py +++ /dev/null @@ -1,299 +0,0 @@ -from __future__ import annotations -from typing import Any, Iterable, Type, TYPE_CHECKING -import csv as _csv -import pandas as pd -import logging -import pyarrow as pa -import pyarrow.dataset as ds -import pyarrow.compute as pc -from functools import reduce -from .data_classes import LoaderContext, TableCastingStats, LoaderInterface -from .loading_helpers import infer_delim, infer_encoding, conservative_load_parquet, arrow_drop_duplicates, resolve_quote_mode -from .data import perform_cast, cast_arrow_column -from ..helpers import IngestError - -if TYPE_CHECKING: - from ..tables.typing import CSVTableProtocol - -logger = logging.getLogger(__name__) - -""" -File Loader Implementations -=========================== - -This module provides concrete loader implementations for ingesting -CSV- and Parquet-based datasets into staging tables. - -Loaders are intentionally conservative and designed to handle: -- untrusted data sources -- incremental loads -- partial failures -- schema drift - -Where supported (PostgreSQL), fast-path COPY-based loading is available -via helper utilities. - - -These loader interfaces implement a very conservative loading strategy for handling data -from untrusted sources and accommodating updates and deletes for incremental loads. - -""" - -def _require_columns_present(tableclass: Type["CSVTableProtocol"], present_columns: Iterable[str]) -> None: - """ - Raise before staging if a genuinely required column is entirely - absent from the incoming data. - - No support for partial-column load, which means every merge strategy - (including upsert) requires a complete row shape, so a required column's - absence is always a real data problem, not a legitimate partial update. - """ - missing = sorted(c for c in tableclass.required_columns() if c not in present_columns) - if missing: - raise IngestError( - f"{tableclass.__tablename__}: source data is missing required column(s) {missing} (non-nullable, no default)" - ) - -@staticmethod -def _normalise_columns(df: pd.DataFrame) -> pd.DataFrame: - df = df.copy() - df.columns = [c.lower().replace('_hash', '').strip() for c in df.columns] - return df - -class PandasLoader(LoaderInterface): - - """ - For initial dataloads, pandasloader may not be sufficiently performant for very large files. - However, it provides a very flexible and easy-to-debug pathway for data ingestion, especially - when dealing with CSV files. It allows for chunked processing, which helps manage memory usage - effectively. Supports significant transformation or cleaning before being loaded into the database. - """ - - @classmethod - def dedupe(cls, data: pd.DataFrame | pa.Table, ctx: LoaderContext) -> Any: - if not isinstance(data, pd.DataFrame): - df = data.to_pandas() - else: - df = data - if df.empty: - return df - pk_names = ctx.tableclass.pk_names() - before = len(df) - df = df.drop_duplicates(subset=pk_names, keep='first') - dropped_internal = before - len(df) - if dropped_internal > 0: - logger.info(f"Dropped {dropped_internal} duplicate rows internally in staging for {ctx.tableclass.__tablename__}") - return df.copy() - - @classmethod - def cast_to_model(cls, data: pd.DataFrame | pa.Table, ctx: LoaderContext) -> Any: - if not isinstance(data, pd.DataFrame): - df = data.to_pandas() - else: - df = data - if df.empty: - return df - - table_name = ctx.tableclass.__tablename__ - stats = TableCastingStats(table_name=table_name) - - model_columns = ctx.tableclass.model_columns() - for col_name, sa_col in model_columns.items(): - if col_name not in df.columns: - continue - - def _on_cast_error(value, *, _col=col_name): - stats.record(column=_col, value=value) - - df[col_name] = df[col_name].map( - lambda v: perform_cast( - v, sa_col.type, on_error=_on_cast_error, table_name=table_name, column_name=col_name - ) - ) - - _require_columns_present(ctx.tableclass, df.columns) - required_cols = list(ctx.tableclass.required_columns()) - - if required_cols: - null_mask = df[required_cols].isna() - for col in required_cols: - null_count = int(null_mask[col].sum()) # type: ignore[arg-type] - if null_count > 0: - logger.warning( - "Found %d rows with unexpected nulls in %s.%s", - null_count, - table_name, - col, - ) - # Drop rows violating required constraints - df = df.loc[~null_mask.any(axis=1)] - if stats.has_failures(): - for col, col_stats in stats.columns.items(): - logger.warning(f"CAST {table_name}.{col}: {col_stats.count} row(s) failed. Examples: {col_stats.examples}") - - return df - - @classmethod - def orm_file_load(cls, ctx: LoaderContext) -> int: - """ - Load a file into a staging table, delegating chunking to pandas. - - If chunksize is None, pandas returns a single DataFrame, which we - normalise to a one-element iterator for unified processing. - """ - - delimiter = infer_delim(ctx.path) - encoding = infer_encoding(ctx.path)['encoding'] - - # Resolve to the same concrete mode the COPY fast-path would use, so a - # file that falls back here is parsed identically. Without this the - # fallback ignored quote_mode entirely and applied RFC-4180 quoting - # unconditionally — re-breaking on the stray quotes that pushed a - # literal-mode file off the fast-path. - quote_mode = resolve_quote_mode(ctx.quote_mode, ctx.path, delimiter, encoding or "utf-8") - quoting = _csv.QUOTE_NONE if quote_mode == "literal" else _csv.QUOTE_MINIMAL - - try: - reader = pd.read_csv( - ctx.path, - delimiter=delimiter, - dtype=str, - chunksize=ctx.chunksize, - encoding=encoding, - quoting=quoting, - ) - except pd.errors.EmptyDataError: - logger.info(f"File {ctx.path.name} is empty — skipping load for {ctx.tableclass.__tablename__}") - return 0 - - logger.info(f"Detected encoding {encoding} for file {ctx.path.name}") - logger.info(f"Detected delimiter '{delimiter}' for file {ctx.path.name}") - logger.info(f"Loading with chunksize '{ctx.chunksize}' for file {ctx.path.name}") - chunks = (reader,) if isinstance(reader, pd.DataFrame) else reader - - total = 0 - for i, chunk in enumerate(chunks): - logger.debug(f"Processing chunk {i} with {len(chunk)} rows for {ctx.tableclass.__tablename__}") - chunk = _normalise_columns(chunk) - if ctx.dedupe: - chunk = cls.dedupe(chunk, ctx) - if ctx.normalise: - chunk = cls.cast_to_model(chunk, ctx) - total += cls._load_chunk( - staging_cls=ctx.staging_table, - session=ctx.session, - dataframe=chunk, - mapper=ctx.tableclass, - ) - return total - -class ParquetLoader(LoaderInterface): - - """ - Overhead from this loader is worthwhile with processing and cleaning of very large files. - """ - - @classmethod - def cast_to_model(cls, data: pa.Table, ctx: LoaderContext) -> pa.Table: - if data.num_rows == 0: - return data - - table_name = ctx.tableclass.__tablename__ - stats = TableCastingStats(table_name=table_name) - model_columns = ctx.tableclass.model_columns() - arrays: dict[str, pa.Array] = {} - for col_name, sa_col in model_columns.items(): - if col_name not in data.schema.names: - continue - arr = data[col_name] - - arrays[col_name] = cast_arrow_column( - arr, - sa_col, - stats=stats, - ) - - out = pa.table(arrays) - _require_columns_present(ctx.tableclass, out.schema.names) - required_cols = list(ctx.tableclass.required_columns()) - - if required_cols: - masks = [pc.is_valid(out[c]) for c in required_cols] # type: ignore - valid_mask = reduce(pc.and_, masks) # type: ignore - - dropped = out.num_rows - pc.sum(valid_mask).as_py() # type: ignore - if dropped > 0: - logger.warning( - "Dropped %d rows with unexpected nulls in %s", - dropped, - table_name, - ) - - out = out.filter(valid_mask) - - if stats.has_failures(): - for col, col_stats in stats.columns.items(): - logger.warning(f"CAST {table_name}.{col}: {col_stats.count} failures. Examples: {col_stats.examples}") - - return out - - @classmethod - def dedupe(cls, data: pa.Table, ctx: LoaderContext) -> pa.Table: - if data.num_rows == 0: - return data - - pk_names = ctx.tableclass.pk_names() - deduped = arrow_drop_duplicates(data, pk_names) - dropped = data.num_rows - deduped.num_rows - if dropped > 0: - logger.info( - "Dropped %d duplicate rows internally for %s", - dropped, - ctx.tableclass.__tablename__, - ) - return deduped - - @classmethod - def _scan_batches(cls, ctx: LoaderContext): - suffix = ctx.path.suffix.lower() - model_columns = ctx.tableclass.model_columns() - wanted_cols = list(model_columns.keys()) - logger.info(f"Scanning batches for {ctx.tableclass.__tablename__}") - if suffix == ".parquet": - dataset = ds.dataset(ctx.path, format="parquet") - yield from dataset.to_batches(batch_size=ctx.chunksize or 64_000) - - elif suffix in {".csv", ".tsv"}: - yield from conservative_load_parquet(ctx.path, wanted_cols=wanted_cols, chunksize=ctx.chunksize) - else: - raise ValueError(f"Unsupported file type: {ctx.path}") - - - @classmethod - def orm_file_load(cls, ctx: LoaderContext) -> int: - total = 0 - for record_batch in cls._scan_batches(ctx): - if record_batch.num_rows == 0: - continue - data: pa.Table | pa.RecordBatch = record_batch - if ctx.normalise: - data = cls.cast_to_model(data, ctx=ctx) - if ctx.dedupe: - data = cls.dedupe(data, ctx) - - if isinstance(data, pa.RecordBatch) or isinstance(data, pa.Table): - df = data.to_pandas() - else: - df = data - if df.empty: - - continue - - total += cls._load_chunk( - staging_cls=ctx.staging_table, - session=ctx.session, - dataframe=df, - mapper=ctx.tableclass, - ) - - return total diff --git a/src/orm_loader/loaders/loading_helpers.py b/src/orm_loader/loaders/loading_helpers.py deleted file mode 100644 index dbf5bbf..0000000 --- a/src/orm_loader/loaders/loading_helpers.py +++ /dev/null @@ -1,348 +0,0 @@ -from __future__ import annotations - -import csv as _csv -import io -import logging -import re -from pathlib import Path -from typing import Any - -import chardet -import pyarrow as pa -import pyarrow.compute as pc -import pyarrow.csv as pv -import sqlalchemy as sa -import sqlalchemy.orm as so -from oa_configurator import qualified - -_SAFE_ENCODING = re.compile(r'^[A-Za-z][A-Za-z0-9_-]*$') - -logger = logging.getLogger(__name__) -COPY_BLOCK_SIZE = 8192 - -""" -Loader Helper Functions -======================= - -Utility functions supporting file loading and ingestion workflows. - -Includes helpers for: -- delimiter and encoding detection -- conservative CSV parsing via pyarrow -- duplicate detection in columnar data -- PostgreSQL COPY-based bulk loading - -These helpers are intentionally low-level and stateless. -""" - -class NormalisedCSVStream(io.RawIOBase): - def __init__(self, f, encoding: str, delimiter: str): - self._f = f - self._encoding = encoding - self._delimiter = delimiter - self._sent_header = False - self._buffer = b"" - self._eof = False - - def readable(self): - return True - - def read(self, size=-1): - if self._eof: - return b"" - - out = bytearray() - - # Send rewritten header once - if not self._sent_header: - header = self._f.readline().decode(self._encoding) - newline = check_line_ending(header) - cols = header.rstrip(newline).split(self._delimiter) - lowered = [c.strip().strip('"').lower().replace('_hash', '') for c in cols] - new_header = (self._delimiter.join(lowered) + "\n").encode(self._encoding) - out.extend(new_header) - self._sent_header = True - - while size < 0 or len(out) < size: - chunk = self._f.read(8192) - if not chunk: - self._eof = True - break - - # Normalize CRLF/CR to LF per chunk - chunk = chunk.replace(b"\r\n", b"\n").replace(b"\r", b"\n") - out.extend(chunk) - - if size > 0 and len(out) >= size: - break - - return bytes(out) - -def infer_encoding(file): - with open(file, 'rb') as infile: - encoding = chardet.detect(infile.read(10000)) - if encoding['encoding'] == 'ascii': - encoding['encoding'] = 'utf-8' # utf-8 valid superset of ascii, so being more conservative here just because it flakes occasionally - return encoding - -def infer_delim(file): - with open(file, 'r') as infile: - line = infile.readline() - tabs = line.count('\t') - commas = line.count(',') - if tabs > commas: - return '\t' - return ',' - - -def infer_quote_mode( - path: Path, - delimiter: str, - encoding: str = "utf-8", - sample_rows: int = 2000, -) -> str: - """Return 'csv' or 'literal' by comparing column-count consistency under both - quoting interpretations across a sample of rows. - - - 'csv' → standard RFC-4180 quoting; surrounding double-quotes are stripped - and embedded delimiters/newlines inside quotes are preserved. - - 'literal' → double-quote has no special meaning; every byte is stored as-is. - - Defaults to 'csv' when both modes produce identical output (no quoting in play) - or when the evidence is tied. Callers can always override by passing an - explicit value instead of relying on auto-detection. - """ - with open(path, encoding=encoding, errors="replace", newline="") as f: - lines = [f.readline() for _ in range(sample_rows + 1)] - - raw = "".join(ln for ln in lines if ln) - if not raw: - return "csv" - - try: - rows_csv = list(_csv.reader(io.StringIO(raw), delimiter=delimiter)) - except _csv.Error: - return "literal" - - try: - rows_lit = list( - _csv.reader(io.StringIO(raw), delimiter=delimiter, quoting=_csv.QUOTE_NONE) - ) - except _csv.Error: - return "csv" - - if not rows_csv: - return "csv" - - ncols = len(rows_csv[0]) - if ncols <= 1: - return "csv" - - # No difference between modes → no quoting is active, csv is the safe default - if rows_csv == rows_lit: - return "csv" - - data_csv = rows_csv[1:] - data_lit = rows_lit[1:] if len(rows_lit) > 1 else [] - - if not data_csv: - return "csv" - - csv_ok = sum(1 for r in data_csv if len(r) == ncols) - lit_ok = sum(1 for r in data_lit if len(r) == ncols) - - # Prefer csv on a tie; only choose literal when it is strictly more consistent - return "literal" if lit_ok > csv_ok else "csv" - - -def resolve_quote_mode( - quote_mode: str, - path: Path, - delimiter: str, - encoding: str = "utf-8", -) -> str: - """Resolve a requested quote_mode to a concrete ``"csv"`` or ``"literal"``. - - Recognised inputs: - - - ``"csv"`` / ``"literal"`` — returned unchanged (explicit override). - - ``"auto"`` — legacy content sniff via :func:`infer_quote_mode`. Compares - column-count consistency under both interpretations across a sample of - rows. Supports RFC-4180 quoting (embedded delimiters/newlines) but can - guess wrong when the discriminating rows fall outside the sample. - - ``"by_delimiter"`` — derive the mode from the delimiter alone, no content - scan: tab-delimited input never embeds the delimiter in a field so any - double-quote is literal data (``"literal"``); for non-tab delimiters this - resolves to ``"csv"`` where RFC-4180 quoting can be meaningful. Deterministic - and immune to the sampling gamble; use it when the source guarantees tab - fields never contain embedded tabs/newlines. - """ - if quote_mode in ("csv", "literal"): - return quote_mode - if quote_mode == "by_delimiter": - return "literal" if delimiter == "\t" else "csv" - if quote_mode == "auto": - return infer_quote_mode(path, delimiter=delimiter, encoding=encoding) - raise ValueError(f"Unknown quote_mode: {quote_mode!r}") - - -def arrow_drop_duplicates( - table: pa.Table, - pk_names: list[str], -) -> pa.Table: - if table.num_rows == 0: - return table - - sort_keys = [(name, "ascending") for name in pk_names] - sorted_idx = pc.sort_indices(table, sort_keys=sort_keys) # type: ignore - sorted_table = table.take(sorted_idx) - diffs = [] - for name in pk_names: - col = sorted_table[name] - previous_arr = col[:-1] - this_arr = col[1:] - diffs.append( - pc.not_equal(previous_arr, this_arr) # type: ignore - ) - keep_tail = diffs[0] - for d in diffs[1:]: - keep_tail = pc.or_(keep_tail, d) # type: ignore - keep = pc.fill_null(keep_tail, True) - if isinstance(keep, pa.ChunkedArray): - keep = keep.combine_chunks() - keep = pa.concat_arrays([ - pa.array([True], type=pa.bool_()), - keep, - ]) - deduped = sorted_table.filter(keep) - - return deduped - - -def conservative_load_parquet(path: Path, wanted_cols: list[str], chunksize: int | None = None) -> pa.Table: - delimiter = infer_delim(path) - encoding = infer_encoding(path)["encoding"] - convert_opts = pv.ConvertOptions( - strings_can_be_null=True, - include_columns=wanted_cols, - ) - - def _invalid_row_handler(row): - logger.warning("Skipping malformed CSV row: %r", row[:200]) - return "skip" - - parse_opts = pv.ParseOptions( - delimiter=delimiter, - ignore_empty_lines=True, - quote_char=False, - invalid_row_handler=_invalid_row_handler - ) - read_opts = pv.ReadOptions( - block_size=chunksize or 64_000, - encoding=encoding, - use_threads=True, - ) - if chunksize: - read_opts.block_size = chunksize - with pv.open_csv( - path, - read_options=read_opts, - parse_options=parse_opts, - convert_options=convert_opts, - ) as reader: - for batch in reader: - yield batch - -def check_line_ending(raw_header: str) -> str: - if raw_header.endswith("\r\n"): - return "\r\n" - elif raw_header.endswith("\n"): - return "\n" - elif raw_header.endswith("\r"): - return "\r" - else: - logger.warning("Unable to detect line ending from header: %r. Defaulting to '\\n'", raw_header) - return "\n" - -def quick_load_pg( - *, - path: Path, - session: so.Session, - tablename: str, - schema: str | None = None, - quote_mode: str = "auto", - mapper: Any | None = None, -) -> int: - connection_arguments = {"bind_arguments": {"mapper": mapper}} if mapper is not None else {} - connection = session.connection(**connection_arguments) - raw_conn = connection.connection - if not hasattr(raw_conn, "cursor"): - raise RuntimeError("Expected DB-API connection for COPY") - - table_ref = qualified(connection, tablename, physical_schema=schema) - - encoding = infer_encoding(path)['encoding'] or 'utf-8' - if not _SAFE_ENCODING.match(encoding): - raise ValueError(f"Unsafe encoding value from chardet: {encoding!r}") - delimiter = infer_delim(path) - quote_mode = resolve_quote_mode(quote_mode, path, delimiter, encoding) - logger.info(f"Using quote_mode={quote_mode!r} for {path.name} (delimiter={delimiter!r})") - if quote_mode == "csv": - copy_options = f""" - FORMAT csv, - HEADER true, - DELIMITER E'{delimiter}', - ENCODING '{encoding}' - """ - elif quote_mode == "literal": - copy_options = f""" - FORMAT csv, - HEADER true, - DELIMITER E'{delimiter}', - QUOTE E'\\x01', - ESCAPE E'\\x01', - ENCODING '{encoding}' - """ - else: - raise ValueError(f"Unknown quote_mode: {quote_mode}") - - # Peek at the CSV header to build an explicit column list for COPY. - # Without this, PostgreSQL expects ALL table columns including internal staging - # columns like _rownum (GENERATED ALWAYS AS IDENTITY), which the CSV doesn't have. - with open(path, "rb") as _f_peek: - _raw_hdr = _f_peek.readline().decode(encoding) - _nl = check_line_ending(_raw_hdr) - # _hash is an internal convention for encrypted/hashed columns; strip it so - # CSV headers map to the base column names that PostgreSQL COPY expects. - # also strip quotes and lowercase - _csv_cols = [c.strip().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 {table_ref} via COPY (encoding={encoding}, delimiter={delimiter})") - - cur = raw_conn.cursor() - try: - with open(path, "rb") as f: - stream = NormalisedCSVStream(f, encoding=encoding, delimiter=delimiter) - with cur.copy( - f''' - COPY {table_ref} ({_cols_sql}) - FROM STDIN - WITH ( - {copy_options} - ) - ''' - ) as copy: - while data := stream.read(COPY_BLOCK_SIZE): - copy.write(data) - session.flush() - total = session.execute( - sa.text(f'SELECT COUNT(*) FROM {table_ref}'), **connection_arguments - ).scalar_one() - return total - except Exception as e: - logger.error(f"Error during bulk load via COPY: {e}") - session.rollback() - raise - finally: - cur.close() diff --git a/src/orm_loader/loaders/processor.py b/src/orm_loader/loaders/processor.py new file mode 100644 index 0000000..d6272fc --- /dev/null +++ b/src/orm_loader/loaders/processor.py @@ -0,0 +1,259 @@ +"""Interpretation of raw batches into typed rows ready for a staging sink.""" + +from __future__ import annotations + +import logging +import os +from collections.abc import Iterable, Iterator +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import pyarrow as pa +import pyarrow.compute as pc +import pyarrow.csv as pv + +from ..conversion import ConversionPolicy +from ..conversion.values import no_failures +from ..helpers import ConversionError, IngestError, RejectLimitError, RequiredValueError + +if TYPE_CHECKING: + from ..tables.typing import FileTableProtocol + from .loader import LoadReport + +logger = logging.getLogger(__name__) + + +class RejectsWriter: + """Writes rejected rows, with their raw text and why they were rejected, to a CSV file. + + The file has the columns ``row`` (data row number starting at 1), + ``column`` (first column that failed), ``reason`` and then every column + read from the source as raw text. + + Parameters + ---------- + path + File to write. + """ + + def __init__(self, path: Path) -> None: + self.path = path + self._writer: pv.CSVWriter | None = None + + def check_writable(self) -> None: + """Raise IngestError unless the file can be created. + + Raises + ------ + IngestError + The directory, or the nearest existing parent of a directory that + does not exist yet, is not writable. + """ + directory = self.path.parent + while not directory.exists() and directory != directory.parent: + directory = directory.parent + if not directory.is_dir() or not os.access(directory, os.W_OK): + raise IngestError(f"Cannot write rejected rows to {self.path.parent}: directory is not writable") + + def write(self, rows: pa.Table) -> None: + """Append rejected rows, creating the directory and file on first use.""" + if self._writer is None: + self.path.parent.mkdir(parents=True, exist_ok=True) + self._writer = pv.CSVWriter(str(self.path), rows.schema) + self._writer.write_table(rows) + + def close(self) -> None: + """Close the file if it was opened.""" + if self._writer is not None: + self._writer.close() + self._writer = None + + +class BatchProcessor: + """Turn raw batches into typed batches under the load's conversion policy. + + A batch with an unreadable value or a missing required value either + raises or has those rows written to the rejects file, depending on + ``ErrorPolicy.on_error``. + + Parameters + ---------- + tableclass + Table the rows are loaded into. + columns + Model columns present in the reader. + policy + How values are interpreted and what happens to unreadable ones. + report + Row accounting, updated as batches are processed. + rejects_path + File that rejected rows are written to. + normalise + Whether to interpret values. When false, batches pass through. + """ + + def __init__( + self, + tableclass: type[FileTableProtocol], + columns: list[str], + policy: ConversionPolicy, + report: LoadReport, + rejects_path: Path, + normalise: bool = True, + ) -> None: + self.columns = columns + self.policy = policy + self.report = report + self.normalise = normalise + self.table_name = tableclass.__tablename__ + self.model_columns = tableclass.model_columns() + self.handlers = { + name: policy.for_column(self.model_columns[name].type, self.table_name, name) for name in columns + } + self.required = [c for c in tableclass.required_columns() if c in columns] + self.rejects = RejectsWriter(rejects_path) + if policy.errors.on_error == "reject_row": + self.rejects.check_writable() + + def stream(self, batches: Iterable[pa.RecordBatch]) -> Iterator[pa.Table]: + """Yield typed tables for the raw batches, closing the rejects file afterwards.""" + try: + for raw in batches: + table = self.process(raw) + if table.num_rows: + yield table + finally: + self.rejects.close() + + def process(self, raw: pa.RecordBatch) -> pa.Table: + """Interpret one raw batch. + + Parameters + ---------- + raw + Batch of raw text columns. + + Returns + ------- + pa.Table + Typed rows that passed every check. + + Raises + ------ + ConversionError + A value cannot be read as its column type and ``on_error`` is + ``"raise"``. + RequiredValueError + A required column has no value and ``on_error`` is ``"raise"``. + RejectLimitError + Rejected rows exceed ``max_rejects``. + """ + first_row = self.report.rows_read + 1 + self.report.rows_read += raw.num_rows + if not self.normalise: + return pa.Table.from_batches([raw]) + results = { + name: self.handlers[name].interpret(raw.column(name), self.model_columns[name].type, self.policy.nulls) + for name in self.columns + } + failures = self._failures(results) + if failures: + if self.policy.errors.on_error == "raise": + self._raise(raw, results, first_row) + return self._reject(raw, results, failures, first_row) + return pa.table({name: results[name].values for name in self.columns}) + + def _failures(self, results: dict[str, Any]) -> dict[str, tuple[pa.Array, str]]: + found: dict[str, tuple[pa.Array, str]] = {} + for name in self.columns: + invalid = results[name].invalid + if pc.any(invalid).as_py(): + found[name] = (invalid, f"cannot be read as {self.handlers[name].expected(self.model_columns[name].type)}") + for name in self.required: + missing = pc.and_(pc.is_null(results[name].values), pc.invert(results[name].invalid)) + if pc.any(missing).as_py(): + found[name] = (pc.or_(found[name][0], missing) if name in found else missing, "required value is missing") + return found + + def _conversion_error(self, name: str, raw: pa.Array | pa.ChunkedArray, invalid: pa.Array) -> ConversionError | None: + """Describe the unreadable values of a column, or return None when there are none. + + Parameters + ---------- + name + The column. + raw + The raw column the mask refers to. + invalid + Boolean mask that is true where a value could not be read. + + Returns + ------- + ConversionError | None + An error with the count and first three raw values, or None. + """ + count = pc.sum(pc.cast(invalid, pa.int64())).as_py() or 0 + if not count: + return None + if isinstance(raw, pa.ChunkedArray): + raw = raw.combine_chunks() + if pa.types.is_dictionary(raw.type): + raw = raw.dictionary_decode() + return ConversionError( + table_name=self.table_name, + column_name=name, + expected=self.handlers[name].expected(self.model_columns[name].type), + count=count, + samples=raw.filter(invalid).slice(0, 3).to_pylist(), + ) + + def _raise(self, raw: pa.RecordBatch, results: dict[str, Any], first_row: int) -> None: + for name in self.columns: + error = self._conversion_error(name, raw.column(name), results[name].invalid) + if error is not None: + raise error + for name in self.required: + missing = pc.is_null(results[name].values) + count = pc.sum(pc.cast(missing, pa.int64())).as_py() or 0 + if count: + rows = [first_row + index for index in pc.indices_nonzero(missing).slice(0, 3).to_pylist()] + raise RequiredValueError( + table_name=self.table_name, column_name=name, count=count, rows=rows + ) + + def _reject( + self, + raw: pa.RecordBatch, + results: dict[str, Any], + failures: dict[str, tuple[pa.Array, str]], + first_row: int, + ) -> pa.Table: + rejected = no_failures(raw.num_rows) + column = pa.nulls(raw.num_rows, pa.string()) + reason = pa.nulls(raw.num_rows, pa.string()) + for name, (mask, why) in failures.items(): + first_hit = pc.and_(mask, pc.invert(rejected)) + column = pc.if_else(first_hit, pa.scalar(name), column) + reason = pc.if_else(first_hit, pa.scalar(why), reason) + rejected = pc.or_(rejected, mask) + row_numbers = pc.add(pa.array(range(raw.num_rows), pa.int64()), first_row) + plain = [c.dictionary_decode() if pa.types.is_dictionary(c.type) else c for c in raw.columns] + listing = pa.table( + [row_numbers, column, reason, *plain], names=["row", "column", "reason", *raw.schema.names] + ) + dropped = listing.filter(rejected) + self.rejects.write(dropped) + self.report.rows_rejected += dropped.num_rows + self.report.rejects_path = self.rejects.path + self._check_limit() + keep = pc.invert(rejected) + return pa.table({n: results[n].values for n in self.columns}).filter(keep) + + def _check_limit(self) -> None: + limit = self.policy.errors.max_rejects + if limit is None: + return + rejected, read = self.report.rows_rejected, self.report.rows_read + over = rejected / read > limit if isinstance(limit, float) and limit < 1 else rejected > limit + if over: + raise RejectLimitError(rejected=rejected, read=read, limit=limit) diff --git a/src/orm_loader/loaders/readers/__init__.py b/src/orm_loader/loaders/readers/__init__.py new file mode 100644 index 0000000..8eeeb22 --- /dev/null +++ b/src/orm_loader/loaders/readers/__init__.py @@ -0,0 +1,51 @@ +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import sqlalchemy as sa + +from ...conversion import NullPolicy +from .base import FileReader +from .delimited import DelimitedTextReader +from .parquet import ParquetReader + + +def open_reader( + path: Path, + model_columns: Mapping[str, sa.ColumnElement[Any]], + *, + nulls: NullPolicy | None = None, + quote_mode: str = "auto", + chunksize: int | None = None, +) -> FileReader: + """Open a file with the reader for its format. + + A ``.parquet`` suffix selects Parquet; every other suffix is read as + delimited text. + + Parameters + ---------- + path + File to read. + model_columns + Model columns by name; only these are read. + nulls + Null policy of the load. + quote_mode + Quote handling for delimited text. + chunksize + Rows per batch for formats that batch by rows. + + Returns + ------- + FileReader + The opened reader. + """ + if path.suffix.lower() == ".parquet": + return ParquetReader(path, model_columns, nulls, chunksize) + return DelimitedTextReader(path, model_columns, nulls, quote_mode) + + + + +__all__ = ["DelimitedTextReader", "FileReader", "ParquetReader", "open_reader"] diff --git a/src/orm_loader/loaders/readers/base.py b/src/orm_loader/loaders/readers/base.py new file mode 100644 index 0000000..050d4a3 --- /dev/null +++ b/src/orm_loader/loaders/readers/base.py @@ -0,0 +1,75 @@ +"""The contract every input format fulfils to feed a load.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Iterator, Mapping +from pathlib import Path +from typing import Any + +import pyarrow as pa +import pyarrow.compute as pc +import sqlalchemy as sa + +from ...conversion import NullPolicy + + +class FileReader(ABC): + """A file read as Arrow record batches of the model's columns. + + A reader yields batches holding only the model columns present in the + file, named as the model names them. Text formats yield raw strings, so + the null policy sees exactly what the file contains. Typed formats yield + typed columns. Markers that are quote-sensitive are applied while + reading, because quoting is not visible afterwards. + + Parameters + ---------- + path + File to read. + model_columns + Model columns by name; only these are read. + nulls + Null policy whose quote-sensitive markers are applied while reading. + + Attributes + ---------- + columns + Model columns present in the file, in file order. + empty + True when the file has no content at all. + """ + + def __init__( + self, + path: Path, + model_columns: Mapping[str, sa.ColumnElement[Any]], + nulls: NullPolicy | None = None, + ) -> None: + self.path = path + self.model_columns = model_columns + self.nulls = nulls or NullPolicy() + self.columns: list[str] = [] + self.empty = False + + @abstractmethod + def batches(self) -> Iterator[pa.RecordBatch]: + """Yield the file's rows as non-empty record batches.""" + + def reader_markers(self) -> frozenset[str]: + """Return the markers that must be applied while reading.""" + if self.nulls.marker_quoting == "unquoted_only": + return self.nulls.text_markers + return frozenset() + + @staticmethod + def null_markers(batch: pa.RecordBatch, markers: frozenset[str]) -> pa.RecordBatch: + """Set string values equal to a marker, ignoring padding, to null.""" + values = pa.array(sorted(markers), pa.string()) + arrays = [] + for column in batch.columns: + if pa.types.is_string(column.type): + hit = pc.fill_null(pc.is_in(pc.utf8_trim_whitespace(column), value_set=values), False) + column = pc.if_else(hit, pa.scalar(None, pa.string()), column) + arrays.append(column) + return pa.RecordBatch.from_arrays(arrays, names=batch.schema.names) diff --git a/src/orm_loader/loaders/readers/delimited.py b/src/orm_loader/loaders/readers/delimited.py new file mode 100644 index 0000000..836634e --- /dev/null +++ b/src/orm_loader/loaders/readers/delimited.py @@ -0,0 +1,141 @@ +"""CSV and other delimited text files as record batches of raw text.""" + +from __future__ import annotations + +import codecs +import logging +from collections.abc import Iterator, Mapping +from pathlib import Path +from typing import Any + +import pyarrow as pa +import pyarrow.csv as pv +import sqlalchemy as sa + +from ...conversion import NullPolicy +from ...helpers import IngestError +from .detection import infer_delim, infer_encoding, resolve_quote_mode +from .base import FileReader + +logger = logging.getLogger(__name__) + +QUOTE_SAMPLE_LINES = 50 + + +class DelimitedTextReader(FileReader): + """A delimited text file read as raw strings in bounded blocks. + + Date columns are dictionary-encoded. An unquoted empty field is null and + a quoted empty field is an empty string. Header names are lowercased and + stripped of quotes and a ``_hash`` suffix. File columns that are not + model columns are ignored. + + Parameters + ---------- + path + File to read. + model_columns + Model columns by name; only these are read. + nulls + Null policy whose quote-sensitive markers are applied while reading. + quote_mode + ``"csv"`` reads double quotes as field quoting, ``"literal"`` keeps + them as data, ``"by_delimiter"`` chooses ``"literal"`` for tab + delimited files and ``"csv"`` otherwise, ``"auto"`` samples the file. + block_size + Bytes per block; memory use follows this, not the row count. + """ + + def __init__( + self, + path: Path, + model_columns: Mapping[str, sa.ColumnElement[Any]], + nulls: NullPolicy | None = None, + quote_mode: str = "auto", + block_size: int = 4 << 20, + ) -> None: + super().__init__(path, model_columns, nulls) + self.block_size = block_size + if path.stat().st_size == 0: + self.empty = True + return + self.delimiter = infer_delim(path) + self.encoding = infer_encoding(path)["encoding"] or "utf-8" + try: + codecs.lookup(self.encoding) + except LookupError as exc: + raise IngestError(f"{path.name}: unsupported encoding {self.encoding!r}") from exc + self.quote = resolve_quote_mode(quote_mode, path, self.delimiter, self.encoding) + self.file_columns = self._header_names() + self.columns = [name for name in self.file_columns if name in model_columns] + if self.nulls.marker_quoting == "unquoted_only" and self.nulls.text_markers: + self._warn_if_everything_is_quoted() + + def _header_names(self) -> list[str]: + with open(self.path, "rb") as handle: + header = handle.readline().decode(self.encoding, errors="replace") + header = header.lstrip("").rstrip("\r\n") + names = [ + name.strip().strip('"').strip().lower().replace("_hash", "") for name in header.split(self.delimiter) + ] + if len(set(names)) != len(names): + raise IngestError(f"{self.path.name}: duplicate column names in header {names}") + return names + + def _warn_if_everything_is_quoted(self) -> None: + with open(self.path, "rb") as handle: + handle.readline() + lines = [handle.readline().decode(self.encoding, errors="replace") for _ in range(QUOTE_SAMPLE_LINES)] + fields = [f for line in lines if line.strip() for f in line.rstrip("\r\n").split(self.delimiter)] + if fields and all(f.startswith('"') for f in fields): + logger.warning( + "%s quotes every field, so quote-sensitive text markers %s can never match", + self.path.name, + sorted(self.nulls.text_markers), + ) + + def batches(self) -> Iterator[pa.RecordBatch]: + """Yield the file's rows in blocks of ``block_size`` bytes.""" + if self.empty: + return + string, encoded = pa.string(), pa.dictionary(pa.int32(), pa.string()) + column_types = { + name: encoded if isinstance(self.model_columns[name].type, sa.Date) else string + for name in self.columns + } + try: + reader = pv.open_csv( + self.path, + read_options=pv.ReadOptions( + column_names=self.file_columns, + skip_rows=1, + block_size=self.block_size, + encoding=self.encoding, + ), + parse_options=pv.ParseOptions( + delimiter=self.delimiter, + quote_char='"' if self.quote == "csv" else False, + escape_char=False, + newlines_in_values=self.quote == "csv", + ), + convert_options=pv.ConvertOptions( + column_types=column_types, + include_columns=self.columns, + null_values=[""] + sorted(self.reader_markers()), + strings_can_be_null=True, + quoted_strings_can_be_null=False, + ), + ) + with reader: + for batch in reader: + if batch.num_rows: + yield batch + except pa.ArrowInvalid as exc: + if "Empty CSV file" in str(exc): + return + hint = ( + " The file may contain literal double quotes; try quote_mode='literal' or 'by_delimiter'." + if self.quote == "csv" + else "" + ) + raise IngestError(f"{self.path.name}: {exc}.{hint}") from exc diff --git a/src/orm_loader/loaders/readers/detection.py b/src/orm_loader/loaders/readers/detection.py new file mode 100644 index 0000000..2814cdf --- /dev/null +++ b/src/orm_loader/loaders/readers/detection.py @@ -0,0 +1,119 @@ +"""Detection of the encoding, delimiter and quote mode of a delimited text file.""" + +from __future__ import annotations + +import csv as _csv +import io +import logging +from pathlib import Path + +import chardet + +logger = logging.getLogger(__name__) + +def infer_encoding(file): + with open(file, 'rb') as infile: + encoding = chardet.detect(infile.read(10000)) + if encoding['encoding'] == 'ascii': + encoding['encoding'] = 'utf-8' # utf-8 valid superset of ascii, so being more conservative here just because it flakes occasionally + return encoding + +def infer_delim(file): + with open(file, 'r') as infile: + line = infile.readline() + tabs = line.count('\t') + commas = line.count(',') + if tabs > commas: + return '\t' + return ',' + + +def infer_quote_mode( + path: Path, + delimiter: str, + encoding: str = "utf-8", + sample_rows: int = 2000, +) -> str: + """Return 'csv' or 'literal' by comparing column-count consistency under both + quoting interpretations across a sample of rows. + + - 'csv' → standard RFC-4180 quoting; surrounding double-quotes are stripped + and embedded delimiters/newlines inside quotes are preserved. + - 'literal' → double-quote has no special meaning; every byte is stored as-is. + + Defaults to 'csv' when both modes produce identical output (no quoting in play) + or when the evidence is tied. Callers can always override by passing an + explicit value instead of relying on auto-detection. + """ + with open(path, encoding=encoding, errors="replace", newline="") as f: + lines = [f.readline() for _ in range(sample_rows + 1)] + + raw = "".join(ln for ln in lines if ln) + if not raw: + return "csv" + + try: + rows_csv = list(_csv.reader(io.StringIO(raw), delimiter=delimiter)) + except _csv.Error: + return "literal" + + try: + rows_lit = list( + _csv.reader(io.StringIO(raw), delimiter=delimiter, quoting=_csv.QUOTE_NONE) + ) + except _csv.Error: + return "csv" + + if not rows_csv: + return "csv" + + ncols = len(rows_csv[0]) + if ncols <= 1: + return "csv" + + # No difference between modes → no quoting is active, csv is the safe default + if rows_csv == rows_lit: + return "csv" + + data_csv = rows_csv[1:] + data_lit = rows_lit[1:] if len(rows_lit) > 1 else [] + + if not data_csv: + return "csv" + + csv_ok = sum(1 for r in data_csv if len(r) == ncols) + lit_ok = sum(1 for r in data_lit if len(r) == ncols) + + # Prefer csv on a tie; only choose literal when it is strictly more consistent + return "literal" if lit_ok > csv_ok else "csv" + + +def resolve_quote_mode( + quote_mode: str, + path: Path, + delimiter: str, + encoding: str = "utf-8", +) -> str: + """Resolve a requested quote_mode to a concrete ``"csv"`` or ``"literal"``. + + Recognised inputs: + + - ``"csv"`` / ``"literal"`` — returned unchanged (explicit override). + - ``"auto"`` — legacy content sniff via :func:`infer_quote_mode`. Compares + column-count consistency under both interpretations across a sample of + rows. Supports RFC-4180 quoting (embedded delimiters/newlines) but can + guess wrong when the discriminating rows fall outside the sample. + - ``"by_delimiter"`` — derive the mode from the delimiter alone, no content + scan: tab-delimited input never embeds the delimiter in a field so any + double-quote is literal data (``"literal"``); for non-tab delimiters this + resolves to ``"csv"`` where RFC-4180 quoting can be meaningful. Deterministic + and immune to the sampling gamble; use it when the source guarantees tab + fields never contain embedded tabs/newlines. + """ + if quote_mode in ("csv", "literal"): + return quote_mode + if quote_mode == "by_delimiter": + return "literal" if delimiter == "\t" else "csv" + if quote_mode == "auto": + return infer_quote_mode(path, delimiter=delimiter, encoding=encoding) + raise ValueError(f"Unknown quote_mode: {quote_mode!r}") diff --git a/src/orm_loader/loaders/readers/parquet.py b/src/orm_loader/loaders/readers/parquet.py new file mode 100644 index 0000000..8fe2563 --- /dev/null +++ b/src/orm_loader/loaders/readers/parquet.py @@ -0,0 +1,53 @@ +"""Parquet files as record batches.""" + +from __future__ import annotations + +from collections.abc import Iterator, Mapping +from pathlib import Path +from typing import Any + +import pyarrow as pa +import pyarrow.dataset as ds +import sqlalchemy as sa + +from ...conversion import NullPolicy +from .base import FileReader + + +class ParquetReader(FileReader): + """A Parquet file read with its stored types. + + Parameters + ---------- + path + File to read. + model_columns + Model columns by name; only these are read. + nulls + Null policy whose quote-sensitive markers are applied while reading. A + Parquet string was never quoted, so every marker match counts as + unquoted. + chunksize + Rows per batch. + """ + + def __init__( + self, + path: Path, + model_columns: Mapping[str, sa.ColumnElement[Any]], + nulls: NullPolicy | None = None, + chunksize: int | None = None, + ) -> None: + super().__init__(path, model_columns, nulls) + self.chunksize = chunksize or 64_000 + self._dataset = ds.dataset(path, format="parquet") + self.columns = [name for name in self._dataset.schema.names if name in model_columns] + + def batches(self) -> Iterator[pa.RecordBatch]: + """Yield the stored rows in batches of ``chunksize`` rows.""" + markers = self.reader_markers() + for batch in self._dataset.to_batches(columns=self.columns, batch_size=self.chunksize): + if markers: + batch = self.null_markers(batch, markers) + if batch.num_rows: + yield batch diff --git a/src/orm_loader/tables/__init__.py b/src/orm_loader/tables/__init__.py index de740ac..c9edaea 100644 --- a/src/orm_loader/tables/__init__.py +++ b/src/orm_loader/tables/__init__.py @@ -1,14 +1,14 @@ from .allocators import IdAllocator -from .loadable_table import CSVLoadableTableInterface +from .loadable_table import FileLoadableTableInterface from .orm_table import ORMTableBase from .serialisable_table import SerialisableTableInterface -from .typing import ORMTableProtocol, CSVTableProtocol +from .typing import ORMTableProtocol, FileTableProtocol __all__ = [ "ORMTableBase", - "CSVLoadableTableInterface", + "FileLoadableTableInterface", "SerialisableTableInterface", "IdAllocator", "ORMTableProtocol", - "CSVTableProtocol", + "FileTableProtocol", ] \ No newline at end of file diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index 65a64a0..c6261d9 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -1,4 +1,11 @@ # pyright: reportPrivateUsage=false +"""Staged, file-based loading for ORM tables, whatever the file format. + +A table loads a file into a staging table through a reader, the conversion +policy and the backend's bulk sink, then merges the staging table into the +target table. Readers choose the file formats; backends choose the merge +and index handling. +""" import logging from collections.abc import Mapping from contextlib import contextmanager @@ -12,9 +19,10 @@ from sqlalchemy.exc import InvalidRequestError, UnboundExecutionError from ..backends.resolve import resolve_backend -from ..loaders.loader_interface import LoaderContext, LoaderInterface, PandasLoader, ParquetLoader +from ..conversion import ConversionPolicy +from ..loaders import FileReader, LoadReport, LoaderContext, open_reader, stage_file from .orm_table import ORMTableBase -from .typing import CSVTableProtocol +from .typing import FileTableProtocol logger = logging.getLogger(__name__) @@ -32,33 +40,15 @@ def _require_bind(session: so.Session, mapper: Any | None = None) -> sa.Engine | raise RuntimeError("Session is not bound to an engine") from exc -""" -CSV Loadable Table Mixins -================================== - -This module provides mixins that add staged, file-based ingestion -capabilities to SQLAlchemy ORM-mapped tables. - -The functionality here is intentionally infrastructure-focused and -model-agnostic, supporting: -- CSV-based bulk ingestion via staging tables -- optional fast-path database COPY operations -- database-portable merge strategies -- pluggable loader implementations - -No schema semantics or domain rules are imposed. -""" - - -class CSVLoadableTableInterface(ORMTableBase): +class FileLoadableTableInterface(ORMTableBase): """ - Mixin for ORM tables that support staged CSV-based ingestion. + Mixin for ORM tables that load from files through a staging table. - This interface implements a database-portable ingestion workflow - based on temporary staging tables. It supports: + A file is read in batches by a reader, converted under a conversion + policy, written to a staging table by the backend's bulk sink, and + merged into the target table. It supports: - dialect-aware staging table creation - - fast-path COPY-based loading where available - - ORM-based fallback loading + - delimited text and Parquet files, and other formats through readers - configurable merge strategies - explicit staging table lifecycle management @@ -70,7 +60,7 @@ class CSVLoadableTableInterface(ORMTableBase): @classmethod def create_staging_table( - cls: Type[CSVTableProtocol], + cls: Type[FileTableProtocol], session: so.Session, *, staging_schema_tag: str | None = None, @@ -104,7 +94,7 @@ def create_staging_table( @classmethod @contextmanager def manage_indices( - cls: Type['CSVTableProtocol'], + cls: Type['FileTableProtocol'], session: so.Session, index_strategy: str = "auto", *, @@ -218,7 +208,7 @@ def manage_indices( @classmethod def get_staging_table( - cls: Type[CSVTableProtocol], + cls: Type[FileTableProtocol], session: so.Session, *, staging_schema_tag: str | None = None, @@ -275,22 +265,19 @@ def get_staging_table( @classmethod def load_staging( - cls: Type[CSVTableProtocol], - loader: LoaderInterface, + cls: Type[FileTableProtocol], loader_context: LoaderContext, + reader: FileReader, ) -> int: """ - Load data into the staging table. - - This method attempts a fast-path database-native load where - supported, falling back to an ORM-based loader if necessary. + Create the staging table and load the file into it. Parameters ---------- - loader - Loader implementation used for ORM-based loading. loader_context Context object containing session, path, and load options. + reader + Reader that yields the file as Arrow batches. Returns ------- @@ -298,74 +285,42 @@ def load_staging( Number of rows loaded into the staging table. """ _require_bind(loader_context.session, cls) - - backend = resolve_backend( - loader_context.session, staging_schema_tag=loader_context.staging_schema_tag, mapper=cls - ) - total = 0 - cls.create_staging_table(loader_context.session, staging_schema_tag=loader_context.staging_schema_tag) - - try: - 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') - - total = cls.orm_staging_load( - loader=loader, - loader_context=loader_context - ) - return total + return stage_file(loader_context, reader) @classmethod - def orm_staging_load( - cls: Type[CSVTableProtocol], - loader: LoaderInterface, - loader_context: LoaderContext - ) -> int: + def _select_reader(cls: Type[FileTableProtocol], ctx: LoaderContext) -> FileReader: """ - Load data into the staging table using an ORM-based loader. + Select the reader for a file. Override to read other formats. - Returns - ------- - int - Number of rows loaded. - """ - return loader.orm_file_load(ctx=loader_context) - - @classmethod - def _select_loader(cls: Type[CSVTableProtocol], path: Path) -> LoaderInterface: - """ - Select an appropriate loader based on file type. + The default picks the reader from the file suffix: Parquet for + ``.parquet`` and delimited text for every other suffix. Parameters ---------- - path - Path to the input file. + ctx + Context of the load. Returns ------- - LoaderInterface - A loader instance suitable for the file type. + FileReader + Reader for the file. """ - suffix = path.suffix.lower() - if suffix == ".parquet": - return ParquetLoader() - else: - return PandasLoader() - + return open_reader( + ctx.path, + ctx.tableclass.model_columns(), + nulls=ctx.policy.nulls, + quote_mode=ctx.quote_mode, + chunksize=ctx.chunksize, + ) @classmethod - def load_csv( - cls: Type[CSVTableProtocol], + def load_file( + cls: Type[FileTableProtocol], session: so.Session, path: Path, *, - loader: LoaderInterface | None = None, + reader: FileReader | None = None, normalise: bool = True, dedupe: bool = False, chunksize: int | None = None, @@ -377,12 +332,16 @@ def load_csv( validate_deferred_fks: bool = True, resolved: ResolvedDatabase | None = None, physical_schemas: Mapping[str, str | None] | None = None, + policy: ConversionPolicy | None = None, + rejects_dir: Path | None = None, + report: LoadReport | None = None, ) -> int: """ - Load a CSV (or CSV-like) file into the target table. + Load a file into the target table. - This method orchestrates the full staged ingestion lifecycle: + The file name without its suffix must equal the table name. This + method orchestrates the full staged ingestion lifecycle: - staging table creation - file loading - merge into the target table @@ -393,22 +352,25 @@ def load_csv( session An active SQLAlchemy session. path - Path to the input CSV or Parquet file. - loader - Optional explicit loader instance. + Path to the input file. + reader + Reader for the file. The table's ``_select_reader`` chooses one + when omitted, and ``chunksize`` and ``quote_mode`` then configure it. normalise - Whether to apply table-level normalisation. + Whether to convert values to the column types under the policy. + When false, values are written as read. dedupe Whether to deduplicate incoming rows. chunksize - Optional chunk size for incremental loading. + Rows per batch for the default reader of formats that batch by + rows. merge_strategy Merge strategy to apply (e.g. ``replace``, ``upsert``, or ``insert_if_empty``). quote_mode - Quoting mode. Governs parsing on both the PostgreSQL COPY - fast-path and the pandas ORM fallback, so a file that falls back - to the ORM path is still parsed identically. + How double quotes in delimited text are read: ``"csv"`` (field + quoting), ``"literal"`` (kept as data), ``"by_delimiter"`` + (literal for tab-delimited files) or ``"auto"`` (sample the file). index_strategy Index handling strategy during merge. Use ``"auto"`` to let the backend choose a sensible default. @@ -434,6 +396,18 @@ def load_csv( physical_schemas A schema tag's physical name beyond resolved's own configured ones. Forwarded alongside resolved. + policy + How values are interpreted and what happens to unreadable ones. + The default policy preserves nonempty text and treats only empty + values as missing in text columns. + rejects_dir + Directory that rejected rows are written to when the policy sets + ``on_error="reject_row"``, as ``.rejects.csv``. + Defaults to a ``rejects`` directory inside the source file's + directory. + report + Receives the load's row accounting. A new report is used when + omitted. Returns ------- @@ -441,11 +415,11 @@ def load_csv( Number of rows loaded into staging before merge. """ - logger.debug(f"Table `{cls.__tablename__}`: Loading CSV from {path}") + logger.debug(f"Table `{cls.__tablename__}`: Loading file {path}") if path.stem.lower() != cls.__tablename__: raise ValueError( - f"CSV filename '{path.name}' does not match table '{cls.__tablename__}'" + f"File name '{path.name}' does not match table '{cls.__tablename__}'" ) if merge_strategy == "insert_if_empty": @@ -483,14 +457,14 @@ def load_csv( dedupe=dedupe, quote_mode=quote_mode, staging_schema_tag=staging_schema_tag, + policy=policy or ConversionPolicy(), + rejects_dir=rejects_dir, + report=report or LoadReport(), ) - if loader is None: - loader = cls._select_loader(path) - # Load to staging (Indices are already excluded via updated create_staging_table) logger.info(f"Table `{cls.__tablename__}`: Loading data into staging table") - total = cls.load_staging(loader=loader, loader_context=loader_context) + total = cls.load_staging(loader_context, reader or cls._select_reader(loader_context)) # Merge staging to target (Wrapped in our index dropper!) logger.info(f"Table `{cls.__tablename__}`: Merging staging data into target table") @@ -524,7 +498,7 @@ def load_csv( @classmethod def _target_has_rows( - cls: Type[CSVTableProtocol], + cls: Type[FileTableProtocol], session: so.Session, ) -> bool: """ @@ -539,7 +513,7 @@ def _target_has_rows( @classmethod def merge_from_staging( - cls: Type[CSVTableProtocol], + cls: Type[FileTableProtocol], session: so.Session, merge_strategy: str = "replace", *, @@ -565,9 +539,9 @@ def merge_from_staging( validate_deferred_fks Forwarded to the backend's merge methods, which validate deferred FKs before every intermediate commit a batched merge - makes. See ``load_csv()``. + makes. See ``load_file()``. resolved, physical_schemas - Forwarded to ``resolve_backend()``. See ``load_csv()``. + Forwarded to ``resolve_backend()``. See ``load_file()``. """ target = cls.__tablename__ pk_cols = cls.pk_names() @@ -658,7 +632,7 @@ def merge_from_staging( @classmethod def drop_staging_table( - cls: Type[CSVTableProtocol], + cls: Type[FileTableProtocol], session: so.Session, *, staging_schema_tag: str | None = None, @@ -676,20 +650,3 @@ def drop_staging_table( """ backend = resolve_backend(session, staging_schema_tag=staging_schema_tag, mapper=cls) backend.drop_staging_table(cls, session) - - @classmethod - def csv_columns(cls) -> dict[str, sa.ColumnElement[Any]]: - """ - Return a mapping of CSV column names to model columns. - - By default this is equivalent to :meth:`model_columns`. - Override this method to implement custom column mappings. - - Returns - ------- - dict[str, sqlalchemy.ColumnElement] - Mapping of input column names to SQLAlchemy columns. - """ - cols = cls.model_columns() - computed_names = {c.name for c in cls.__table__.columns if c.computed is not None} # type: ignore - return {k: v for k, v in cols.items() if k not in computed_names} diff --git a/src/orm_loader/tables/orm_table.py b/src/orm_loader/tables/orm_table.py index 29fe02d..e015eb4 100644 --- a/src/orm_loader/tables/orm_table.py +++ b/src/orm_loader/tables/orm_table.py @@ -1,10 +1,9 @@ import sqlalchemy as sa import sqlalchemy.orm as so -from sqlalchemy.exc import NoInspectionAvailable, StatementError +from sqlalchemy.exc import NoInspectionAvailable from typing import Any import logging from .allocators import IdAllocator -from ..helpers import normalise_null logger = logging.getLogger(__name__) @@ -122,46 +121,6 @@ def pk_names(cls) -> list[str]: """ return [c.key for c in cls.pk_columns() if c.key is not None] - @classmethod - def pk_values(cls, obj: Any) -> dict[str, Any]: - """ - Extract primary key values from an ORM instance. - - Parameters - ---------- - obj - An ORM-mapped instance of this class. - - Returns - ------- - dict[str, Any] - A dictionary mapping primary key column names to values. - """ - return {c.key: getattr(obj, c.key) for c in cls.pk_columns() if c.key is not None} - - @classmethod - def pk_tuple(cls, obj: Any) -> tuple[Any, ...]: - """ - Extract primary key values from an ORM instance as a tuple. - - Values are returned in mapper-defined primary key order. - - Parameters - ---------- - obj - An ORM-mapped instance of this class. - - Returns - ------- - tuple - A tuple of primary key values. - """ - return tuple( - getattr(obj, c.key) - for c in cls.pk_columns() - if c.key is not None - ) - @classmethod def model_columns(cls) -> dict[str, sa.ColumnElement[Any]]: """ @@ -260,85 +219,3 @@ def allocator(cls, session: so.Session) -> IdAllocator: return IdAllocator(cls.max_id(session)) - @classmethod - def clean_kwargs( - cls, - data: dict[str, Any], - *, - drop_nulls: bool = True, - strict: bool = False, - ) -> dict[str, Any]: - """ - Normalise inbound constructor kwargs: - - nan / NaT / 'nan' / '' -> None - - optionally drop nulls - - optionally validate required columns - """ - cols = cls.model_columns() - - cleaned: dict[str, Any] = {} - for k, v in data.items(): - if k not in cols: - continue # ignore unknown keys safely - - v2 = normalise_null(v) - - if v2 is None and drop_nulls: - continue - - cleaned[k] = v2 - - if strict: - missing = cls.required_columns() - cleaned.keys() - if missing: - raise ValueError( - f"Missing required fields for {cls.__name__}: {sorted(missing)}" - ) - - return cleaned - - - @classmethod - def from_dict( - cls, - data: dict[str, Any], - *, - drop_nulls: bool = True, - strict: bool = False, - debug: bool = False, - ): - cleaned = cls.clean_kwargs( - data, - drop_nulls=drop_nulls, - strict=strict, - ) - - try: - return cls(**cleaned) - - except TypeError as e: - # bad keyword, missing arg, etc - unknown_keys = set(data.keys()) - set(cls.model_columns().keys()) - msg = ( - f"Failed to construct {cls.__name__} from dict.\n" - f"Error: {e}\n" - f"Known columns: {sorted(cls.model_columns().keys())}\n" - f"Unknown keys: {sorted(unknown_keys)}\n" - f"Cleaned kwargs: {cleaned}" - ) - raise TypeError(msg) from e - - except StatementError as e: - # SQLAlchemy type coercion / bind errors - msg = ( - f"SQLAlchemy failed to bind values for {cls.__name__}.\n" - f"Cleaned kwargs: {cleaned}\n" - ) - raise ValueError(msg) from e - - except Exception as e: - msg = ( - f"Unexpected error constructing {cls.__name__}.\n" - f"Cleaned kwargs: {cleaned}\n" - ) - raise RuntimeError(msg) from e \ No newline at end of file diff --git a/src/orm_loader/tables/typing.py b/src/orm_loader/tables/typing.py index 0d74efd..c03e21d 100644 --- a/src/orm_loader/tables/typing.py +++ b/src/orm_loader/tables/typing.py @@ -7,7 +7,8 @@ if TYPE_CHECKING: from oa_configurator import ResolvedDatabase - from ..loaders import LoaderContext, LoaderInterface + from ..conversion import ConversionPolicy + from ..loaders import FileReader, LoadReport, LoaderContext class ToDictKwargs(TypedDict, total=False): include_nulls: bool @@ -51,53 +52,53 @@ def model_columns(cls) -> dict[str, sa.ColumnElement[Any]]: ... def required_columns(cls) -> set[str]: ... @runtime_checkable -class CSVTableProtocol(ORMTableProtocol, Protocol): +class FileTableProtocol(ORMTableProtocol, Protocol): """ - Structural protocol for ORM tables that support CSV-based ingestion. + Structural protocol for ORM tables that load from files. This protocol captures the expected interface for tables that can - participate in staged, file-based loading workflows using CSV input. + participate in staged, file-based loading workflows. It defines the contract for: - staging table management - - loader selection - - CSV ingestion and merging semantics + - reader selection + - file ingestion and merging semantics No assumptions are made about the underlying database or schema. """ @classmethod - def _select_loader(cls, path: Path) -> "LoaderInterface": ... + def _select_reader(cls, ctx: "LoaderContext") -> "FileReader": ... @classmethod def create_staging_table(cls, session: so.Session, *, staging_schema_tag: str | None = None) -> None: ... @classmethod - def load_staging(cls: Type["CSVTableProtocol"], loader: "LoaderInterface", loader_context: "LoaderContext") -> int: ... + def load_staging(cls: Type["FileTableProtocol"], loader_context: "LoaderContext", reader: "FileReader") -> int: ... @classmethod - def load_csv( + def load_file( cls, session: so.Session, path: Path, *, - loader: Optional["LoaderInterface"] = None, + reader: Optional["FileReader"] = None, normalise: bool = True, dedupe: bool = False, chunksize: int | None = None, merge_strategy: str = "replace", - quote_mode: str = "csv", + quote_mode: str = "auto", index_strategy: str = "auto", merge_batch_size: int | None = None, staging_schema_tag: str | None = None, validate_deferred_fks: bool = True, resolved: Optional["ResolvedDatabase"] = None, physical_schemas: Optional[Mapping[str, str | None]] = None, + policy: Optional["ConversionPolicy"] = None, + rejects_dir: Optional[Path] = None, + report: Optional["LoadReport"] = None, ) -> int: ... - @classmethod - def orm_staging_load(cls, loader: "LoaderInterface", loader_context: "LoaderContext") -> int: ... - @classmethod def get_staging_table(cls, session: so.Session, *, staging_schema_tag: str | None = None) -> sa.Table: ... diff --git a/tests/backends/test_base_backend.py b/tests/backends/test_base_backend.py index 1a2db64..3a6537e 100644 --- a/tests/backends/test_base_backend.py +++ b/tests/backends/test_base_backend.py @@ -20,8 +20,7 @@ ) if TYPE_CHECKING: - from orm_loader.loaders.data_classes import LoaderContext - from orm_loader.tables.typing import CSVTableProtocol + from orm_loader.tables.typing import FileTableProtocol class _BlockPsycopg(importlib.abc.MetaPathFinder): @@ -40,6 +39,13 @@ class FakeBackend(DatabaseBackend): def __init__(self) -> None: self.calls: list[tuple[str, object]] = [] + def insert_batches(self, ctx, batches, columns) -> int: + total = 0 + for table in batches: + ctx.session.execute(ctx.staging_table.insert(), table.to_pylist()) + total += table.num_rows + return total + @property def name(self) -> str: return "fake" @@ -57,7 +63,6 @@ def identifier_preparer(self): @property def capabilities(self) -> BackendCapabilities: return BackendCapabilities( - supports_fast_load=True, supports_fk_toggle=True, ) @@ -66,18 +71,18 @@ def staging_name_for_table(tablename: str) -> str: return f"_staging_{tablename}" def create_staging_table( - self, _table_cls: Type[CSVTableProtocol], _session: so.Session + self, _table_cls: Type[FileTableProtocol], _session: so.Session ) -> None: return None def drop_staging_table( - self, _table_cls: Type[CSVTableProtocol], _session: so.Session + self, _table_cls: Type[FileTableProtocol], _session: so.Session ) -> None: return None def merge_replace( self, - _table_cls: Type[CSVTableProtocol], + _table_cls: Type[FileTableProtocol], _session: so.Session, _target_name: str, _pk_cols: list[str], @@ -88,7 +93,7 @@ def merge_replace( def merge_upsert( self, - _table_cls: Type[CSVTableProtocol], + _table_cls: Type[FileTableProtocol], _session: so.Session, _target_name: str, _pk_cols: list[str], @@ -99,7 +104,7 @@ def merge_upsert( def merge_insert( self, - _table_cls: Type[CSVTableProtocol], + _table_cls: Type[FileTableProtocol], _session: so.Session, _target_name: str, *, @@ -108,7 +113,7 @@ def merge_insert( return None def _pagination_column( - self, _table_cls: Type[CSVTableProtocol], _staging: sa.Table + self, _table_cls: Type[FileTableProtocol], _staging: sa.Table ) -> sa.ColumnElement[int]: return sa.literal_column("rowid") @@ -165,13 +170,12 @@ class _ComputedTable: ) -_ComputedTableCls = cast("Type[CSVTableProtocol]", _ComputedTable) +_ComputedTableCls = cast("Type[FileTableProtocol]", _ComputedTable) def test_backend_capabilities_defaults(): caps = BackendCapabilities() - assert caps.supports_fast_load is False assert caps.supports_unlogged_staging is False assert caps.supports_fk_toggle is False assert caps.supports_materialized_views is False @@ -187,22 +191,20 @@ def test_fake_backend_can_implement_contract(): assert backend.name == "fake" assert backend.dialect == Dialect.SQLITE - assert backend.capabilities.supports_fast_load is True assert backend.capabilities.supports_fk_toggle is True assert backend.supports_dialect(Dialect.SQLITE) is True 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)) is None - with backend.merge_context(cast("Type[CSVTableProtocol]", None), cast(so.Session, None)): + with backend.merge_context(cast("Type[FileTableProtocol]", None), cast(so.Session, None)): pass def test_require_capability_passes_for_supported_feature(): backend = FakeBackend() - backend._require_capability("supports_fast_load", "fast loading") + backend._require_capability("supports_fk_toggle", "foreign key toggling") def test_require_capability_raises_for_unsupported_feature(): @@ -329,7 +331,6 @@ class CapableBackend(FakeBackend): @property def capabilities(self) -> BackendCapabilities: return BackendCapabilities( - supports_fast_load=True, supports_fk_toggle=True, supports_materialized_views=True, ) diff --git a/tests/backends/test_duckdb_backend.py b/tests/backends/test_duckdb_backend.py index 7043220..519232c 100644 --- a/tests/backends/test_duckdb_backend.py +++ b/tests/backends/test_duckdb_backend.py @@ -26,15 +26,15 @@ ) if TYPE_CHECKING: - from orm_loader.tables.typing import CSVTableProtocol - -_ComputedTableCls = cast("Type[CSVTableProtocol]", ComputedColumnTable) -_CycleATableCls = cast("Type[CSVTableProtocol]", CycleATable) -_CycleBTableCls = cast("Type[CSVTableProtocol]", CycleBTable) -_ReferencedParentTableCls = cast("Type[CSVTableProtocol]", ReferencedParentTable) -_SelfReferentialTableCls = cast("Type[CSVTableProtocol]", SelfReferentialTable) -_VocabReferencingTableCls = cast("Type[CSVTableProtocol]", VocabReferencingTable) -_VocabSchemaTableCls = cast("Type[CSVTableProtocol]", VocabSchemaTable) + from orm_loader.tables.typing import FileTableProtocol + +_ComputedTableCls = cast("Type[FileTableProtocol]", ComputedColumnTable) +_CycleATableCls = cast("Type[FileTableProtocol]", CycleATable) +_CycleBTableCls = cast("Type[FileTableProtocol]", CycleBTable) +_ReferencedParentTableCls = cast("Type[FileTableProtocol]", ReferencedParentTable) +_SelfReferentialTableCls = cast("Type[FileTableProtocol]", SelfReferentialTable) +_VocabReferencingTableCls = cast("Type[FileTableProtocol]", VocabReferencingTable) +_VocabSchemaTableCls = cast("Type[FileTableProtocol]", VocabSchemaTable) _TARGET_TABLE = ComputedColumnTable.__tablename__ _STAGING_TABLE = f"_staging_{_TARGET_TABLE}" @@ -160,7 +160,6 @@ def test_duckdb_backend_identity_and_capabilities(): assert backend.name == "duckdb" assert backend.dialect == Dialect.DUCKDB assert backend.supports_dialect(Dialect.DUCKDB) is True - assert backend.capabilities.supports_fast_load is False assert backend.capabilities.supports_unlogged_staging is False assert backend.capabilities.supports_fk_toggle is False assert backend.capabilities.supports_materialized_views is False diff --git a/tests/backends/test_postgres_backend.py b/tests/backends/test_postgres_backend.py index a720fc0..8bb8843 100644 --- a/tests/backends/test_postgres_backend.py +++ b/tests/backends/test_postgres_backend.py @@ -13,18 +13,18 @@ from oa_configurator.testing import isolated_test_schema from orm_loader.backends import STAGING_SCHEMA, Dialect, PostgresBackend from orm_loader.backends.resolve import resolve_backend -from tests.models import ComputedColumnTable +from tests.models import ComputedColumnTable, EnumTable, Role as AuthorRole, SimpleTable _TARGET_TABLE = ComputedColumnTable.__tablename__ _STAGING_TABLE = f"_staging_{_TARGET_TABLE}" _PREPARER = postgresql.dialect().identifier_preparer _STAGING_TABLE_WITH_SCHEMA: str = qualified(_PREPARER, _STAGING_TABLE, physical_schema=STAGING_SCHEMA) -_ComputedTableCls = cast("Type[CSVTableProtocol]", ComputedColumnTable) +_ComputedTableCls = cast("Type[FileTableProtocol]", ComputedColumnTable) if TYPE_CHECKING: - from orm_loader.tables.typing import CSVTableProtocol + from orm_loader.tables.typing import FileTableProtocol class _FakeSession: @@ -85,7 +85,6 @@ def test_postgres_backend_identity_and_capabilities(): assert backend.name == "postgres" assert backend.dialect == Dialect.POSTGRESQL assert backend.supports_dialect(Dialect.POSTGRESQL) is True - assert backend.capabilities.supports_fast_load is True assert backend.capabilities.supports_unlogged_staging is True assert backend.capabilities.supports_fk_toggle is True assert backend.capabilities.supports_materialized_views is True @@ -640,3 +639,72 @@ def _remove(target, name, *_) -> None: "SET session_replication_role = DEFAULT", "SHOW session_replication_role", ] + + +# --------------------------------------------------------------------------- +# Loading files into real Postgres +# --------------------------------------------------------------------------- + + +def _write_simple(tmp_path, rows: list[tuple[int, str]]): + path = tmp_path / "test_table.csv" + path.write_text("id,name\n" + "".join(f"{i},{n}\n" for i, n in rows)) + return path + + +def _simple_rows(session) -> list[tuple[int, str]]: + return [(r.id, r.name) for r in session.execute(sa.select(SimpleTable).order_by(SimpleTable.id)).scalars()] + + +def test_staging_table_has_the_target_columns_and_an_identity_row_number(pg_session): + SimpleTable.create_staging_table(pg_session, staging_schema_tag=STAGING_SCHEMA) + staging_name = resolve_backend(pg_session).staging_name_for_table(SimpleTable.__tablename__) + cols = pg_session.execute(sa.text( + "SELECT column_name, data_type FROM information_schema.columns" + " WHERE table_schema = :s AND table_name = :t ORDER BY ordinal_position" + ), {"s": STAGING_SCHEMA, "t": staging_name}).all() + assert [(n, d) for n, d in cols if n != "_rownum"] == [("id", "integer"), ("name", "character varying")] + assert ("_rownum", "bigint") in [(n, d) for n, d in cols] + + +def test_load_file_merges_into_the_target(pg_session, tmp_path): + inserted = SimpleTable.load_file(pg_session, _write_simple(tmp_path, [(1, "alpha"), (2, "beta")]), staging_schema_tag=STAGING_SCHEMA) + pg_session.commit() + assert inserted == 2 + assert _simple_rows(pg_session) == [(1, "alpha"), (2, "beta")] + + +def test_a_large_file_loads_in_full(pg_session, tmp_path): + rows = [(i, f"name_{i}") for i in range(1, 10_000)] + inserted = SimpleTable.load_file(pg_session, _write_simple(tmp_path, rows), staging_schema_tag=STAGING_SCHEMA) + pg_session.commit() + assert inserted == 9999 + assert pg_session.execute(sa.select(sa.func.count()).select_from(SimpleTable)).scalar_one() == 9999 + + +def test_upsert_does_not_update_existing_rows(pg_session, tmp_path): + SimpleTable.load_file(pg_session, _write_simple(tmp_path, [(1, "alpha")]), staging_schema_tag=STAGING_SCHEMA) + pg_session.commit() + SimpleTable.load_file( + pg_session, _write_simple(tmp_path, [(1, "alpha_updated")]), merge_strategy="upsert", staging_schema_tag=STAGING_SCHEMA + ) + pg_session.commit() + assert _simple_rows(pg_session) == [(1, "alpha")] + + +def test_insert_if_empty_loads_an_empty_target_and_refuses_a_non_empty_one(pg_session, tmp_path): + path = _write_simple(tmp_path, [(1, "alpha"), (2, "beta")]) + assert SimpleTable.load_file(pg_session, path, merge_strategy="insert_if_empty", staging_schema_tag=STAGING_SCHEMA) == 2 + pg_session.commit() + assert _simple_rows(pg_session) == [(1, "alpha"), (2, "beta")] + with pytest.raises(ValueError, match="is not empty; cannot use merge strategy 'insert_if_empty'"): + SimpleTable.load_file(pg_session, path, merge_strategy="insert_if_empty", staging_schema_tag=STAGING_SCHEMA) + + +def test_a_native_enum_column_loads_member_names(pg_session, tmp_path): + path = tmp_path / "enum_table.csv" + path.write_text("id,role\n1,FIRST_AUTHOR\n2,LAST_AUTHOR\n3,\n") + assert EnumTable.load_file(pg_session, path, quote_mode="csv") == 3 + pg_session.commit() + rows = dict(pg_session.execute(sa.select(EnumTable.id, EnumTable.role)).all()) + assert rows == {1: AuthorRole.FIRST_AUTHOR, 2: AuthorRole.LAST_AUTHOR, 3: None} diff --git a/tests/backends/test_shared_backend.py b/tests/backends/test_shared_backend.py index b85da3f..f86fd9e 100644 --- a/tests/backends/test_shared_backend.py +++ b/tests/backends/test_shared_backend.py @@ -9,24 +9,16 @@ from oa_configurator.testing import DIALECT_PARAMS from orm_loader.backends import STAGING_SCHEMA, DatabaseBackend from orm_loader.backends.resolve import resolve_backend +from tests.conftest import session_for from tests.models import ComputedColumnTable, CompositeTable if TYPE_CHECKING: import sqlalchemy.orm as so - from orm_loader.tables.typing import CSVTableProtocol + from orm_loader.tables.typing import FileTableProtocol -_ComputedTableCls = cast("Type[CSVTableProtocol]", ComputedColumnTable) -_CompositeTableCls = cast("Type[CSVTableProtocol]", CompositeTable) - -_SESSION_FIXTURE_BY_DIALECT = { - Dialect.POSTGRESQL: "pg_session", - Dialect.SQLITE: "sqlite_session", - Dialect.DUCKDB: "duckdb_session", -} -assert set(_SESSION_FIXTURE_BY_DIALECT) == {param.values[0] for param in DIALECT_PARAMS}, ( - "_SESSION_FIXTURE_BY_DIALECT is missing an entry for a dialect DIALECT_PARAMS now covers." -) +_ComputedTableCls = cast("Type[FileTableProtocol]", ComputedColumnTable) +_CompositeTableCls = cast("Type[FileTableProtocol]", CompositeTable) # merge_replace deletes on sqlite/postgres but updates in place on DuckDB # (see DuckDBBackend.merge_replace); these tests check the delete-only behavior. @@ -43,8 +35,7 @@ def merge_backend(request: pytest.FixtureRequest) -> tuple[DatabaseBackend, "so. pytest's collection-time fixturenames computation (the usual pg_db-in-fixturenames auto-detection can't see it). """ - session_fixture = _SESSION_FIXTURE_BY_DIALECT[request.param] - session = request.getfixturevalue(session_fixture) + session = session_for(request, request.param) staging_schema_tag = STAGING_SCHEMA if request.param != Dialect.SQLITE else None return resolve_backend(session, staging_schema_tag=staging_schema_tag), session diff --git a/tests/backends/test_sinks.py b/tests/backends/test_sinks.py new file mode 100644 index 0000000..32e074c --- /dev/null +++ b/tests/backends/test_sinks.py @@ -0,0 +1,87 @@ +"""The sink contract: every backend writes the same typed batches to staging the same way.""" + +from datetime import date +from typing import Type, cast + +import pyarrow as pa +import pytest +import sqlalchemy as sa + +from orm_loader.backends import STAGING_SCHEMA, resolve_backend +from orm_loader.loaders import LoaderContext +from orm_loader.tables.typing import FileTableProtocol +from tests.models import DemoRecord + +_Demo = cast(Type[FileTableProtocol], DemoRecord) +COLUMNS = ["demo_id", "demo_label", "demo_number", "demo_flag", "demo_date"] + + +def _sink(session, tmp_path): + tag = None if session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA + _Demo.create_staging_table(session, staging_schema_tag=tag) + staging = _Demo.get_staging_table(session, staging_schema_tag=tag) + ctx = LoaderContext( + tableclass=_Demo, session=session, path=tmp_path / "demo_record.csv", staging_table=staging, staging_schema_tag=tag + ) + return resolve_backend(session, staging_schema_tag=tag, mapper=_Demo), ctx, staging + + +def _batch(i: int, label: str | None) -> pa.Table: + return pa.table({ + "demo_id": [i], + "demo_label": pa.array([label], pa.string()), + "demo_number": [i * 10], + "demo_flag": [i % 2 == 0], + "demo_date": pa.array([date(2020, 1, i)], pa.date32()), + }) + + +def _stored(session, staging): + """Return the staged rows in insertion order (SQLite numbers rows with its implicit rowid).""" + order = staging.c._rownum if "_rownum" in staging.c else sa.literal_column("rowid") + return session.execute(sa.select(staging).order_by(order)).all() + + +def test_typed_values_and_nulls_are_stored_across_more_batches_than_a_commit_interval(dialect_session, tmp_path): + backend, ctx, staging = _sink(dialect_session, tmp_path) + batches = [_batch(i, f"row {i}" if i % 2 else None) for i in range(1, 13)] + assert backend.insert_batches(ctx, batches, COLUMNS) == 12 + stored = _stored(dialect_session, staging) + assert [r.demo_id for r in stored] == list(range(1, 13)) + assert stored[0].demo_label == "row 1" and stored[1].demo_label is None + assert stored[1].demo_flag is True and stored[1].demo_date == date(2020, 1, 2) + assert stored[2].demo_number == 30 + + +def test_rows_keep_the_order_they_were_written_in(dialect_session, tmp_path): + backend, ctx, staging = _sink(dialect_session, tmp_path) + ids = [5, 3, 9, 1, 7] + backend.insert_batches(ctx, [_batch(i, None) for i in ids], COLUMNS) + assert [r.demo_id for r in _stored(dialect_session, staging)] == ids + + +def test_no_batches_write_nothing(dialect_session, tmp_path): + backend, ctx, staging = _sink(dialect_session, tmp_path) + assert backend.insert_batches(ctx, [], COLUMNS) == 0 + assert _stored(dialect_session, staging) == [] + + +@pytest.mark.parametrize( + "label", ["", "a,b", 'say "hi"', "line\nbreak", "tab\there", "ünïcode ✓", "NULL", "\\N"] +) +def test_text_is_stored_verbatim_and_an_empty_string_stays_distinct_from_null(dialect_session, tmp_path, label): + backend, ctx, staging = _sink(dialect_session, tmp_path) + backend.insert_batches(ctx, [_batch(1, label), _batch(2, None)], COLUMNS) + stored = _stored(dialect_session, staging) + assert (stored[0].demo_label, stored[1].demo_label) == (label, None) + + +def test_a_failing_batch_source_is_raised_not_swallowed(dialect_session, tmp_path): + backend, ctx, _ = _sink(dialect_session, tmp_path) + + def batches(): + yield _batch(1, "a") + raise RuntimeError("boom") + + with pytest.raises(RuntimeError, match="boom"): + backend.insert_batches(ctx, batches(), COLUMNS) diff --git a/tests/backends/test_sqlite_backend.py b/tests/backends/test_sqlite_backend.py index 8d5301c..45af16d 100644 --- a/tests/backends/test_sqlite_backend.py +++ b/tests/backends/test_sqlite_backend.py @@ -10,11 +10,10 @@ from sqlalchemy.dialects import sqlite as sqlite_dialect from orm_loader.backends import Dialect, SQLiteBackend -from orm_loader.helpers.sqlite import attach_sqlite_bulk_load_pragmas from tests.models import ComputedColumnTable if TYPE_CHECKING: - from orm_loader.tables.typing import CSVTableProtocol + from orm_loader.tables.typing import FileTableProtocol _TARGET_TABLE = ComputedColumnTable.__tablename__ _STAGING_TABLE = f"_staging_{_TARGET_TABLE}" @@ -42,7 +41,7 @@ def scalar(self): return _Result(self.scalar_result) -_ComputedTableCls = cast("Type[CSVTableProtocol]", ComputedColumnTable) +_ComputedTableCls = cast("Type[FileTableProtocol]", ComputedColumnTable) def _sess(s: _FakeSession) -> so.Session: @@ -55,7 +54,6 @@ def test_sqlite_backend_identity_and_capabilities(): assert backend.name == "sqlite" assert backend.dialect == Dialect.SQLITE assert backend.supports_dialect(Dialect.SQLITE) is True - assert backend.capabilities.supports_fast_load is False assert backend.capabilities.supports_unlogged_staging is False assert backend.capabilities.supports_fk_toggle is True assert backend.capabilities.supports_materialized_views is False @@ -235,11 +233,11 @@ def test_sqlite_backend_restore_journal_mode(tmp_path: Path): assert str(journal_mode).lower() == "delete" -def test_attach_sqlite_bulk_load_pragmas_installs_backend_hook(tmp_path: Path): +def test_install_engine_hooks_sets_bulk_load_pragmas(tmp_path: Path): db_path = tmp_path / "attached.db" sqlite_engine = sa.create_engine(f"sqlite:///{db_path}", future=True) - attach_sqlite_bulk_load_pragmas(sqlite_engine, busy_timeout_ms=45000) + SQLiteBackend(busy_timeout_ms=45000).install_engine_hooks(sqlite_engine) with sqlite_engine.connect() as conn: busy_timeout = conn.execute(sa.text("PRAGMA busy_timeout")).scalar_one() diff --git a/tests/conftest.py b/tests/conftest.py index 489b672..e43f259 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -7,8 +7,8 @@ import sqlalchemy.orm as so from dotenv import load_dotenv -from oa_configurator import Role -from oa_configurator.testing import ScopedTestSchema, isolated_test_database, scoped_test_schema +from oa_configurator import Dialect, Role +from oa_configurator.testing import DIALECT_PARAMS, ScopedTestSchema, isolated_test_database, scoped_test_schema from orm_loader.backends import staging_schema_claim from orm_loader.config import OrmLoaderConfig from tests.models import Base @@ -96,3 +96,32 @@ def schema_scoped_session( table.create(conn, checkfirst=True) with so.Session(scoped.engine) as sqlite_session: yield scoped, sqlite_session + + +# --------------------------------------------------------------------------- +# One session per supported dialect +# --------------------------------------------------------------------------- + +_SESSION_FIXTURE_BY_DIALECT = { + Dialect.SQLITE: "sqlite_session", + Dialect.POSTGRESQL: "pg_session", + Dialect.DUCKDB: "duckdb_session", +} +assert set(_SESSION_FIXTURE_BY_DIALECT) == {param.values[0] for param in DIALECT_PARAMS}, ( + "_SESSION_FIXTURE_BY_DIALECT is missing an entry for a dialect DIALECT_PARAMS now covers." +) + + +def session_for(request: pytest.FixtureRequest, dialect: Dialect) -> so.Session: + """Return the session fixture of one dialect. + + The lookup is dynamic, so the dialect's own mark and ``forked`` come from + ``DIALECT_PARAMS`` on the parametrised value rather than from fixture names. + """ + return request.getfixturevalue(_SESSION_FIXTURE_BY_DIALECT[dialect]) + + +@pytest.fixture(params=DIALECT_PARAMS) +def dialect_session(request: pytest.FixtureRequest) -> so.Session: + """A session on each supported dialect in turn.""" + return session_for(request, request.param) diff --git a/tests/conversion/test_boolean_handler.py b/tests/conversion/test_boolean_handler.py new file mode 100644 index 0000000..4187f7e --- /dev/null +++ b/tests/conversion/test_boolean_handler.py @@ -0,0 +1,17 @@ +import pytest +import sqlalchemy as sa + +from tests.support import read, unreadable + + +@pytest.mark.parametrize( + "raw, expected", + [("true", True), ("T", True), ("Yes", True), ("on", True), ("1", True), + ("false", False), ("f", False), ("NO", False), ("off", False), ("0", False)], +) +def test_boolean_words(raw, expected): + assert read(raw, sa.Boolean()) is expected + + +def test_boolean_rejects_unknown_words(): + assert unreadable("maybe", sa.Boolean()) diff --git a/tests/conversion/test_date_handlers.py b/tests/conversion/test_date_handlers.py new file mode 100644 index 0000000..bd457d0 --- /dev/null +++ b/tests/conversion/test_date_handlers.py @@ -0,0 +1,142 @@ +from datetime import date, datetime + +import numpy as np +import pyarrow as pa +import pyarrow.compute as pc +import pytest +import sqlalchemy as sa + +from orm_loader.conversion import ( + ConversionPolicy, + DateFormats, + DateHandler, + parse_dates, +) +from tests.support import interpret, read, unreadable + + +@pytest.mark.parametrize( + "raw, expected", + [ + ("20200229", date(2020, 2, 29)), + ("24-AUG-2017", date(2017, 8, 24)), + ("24-aug-2017", date(2017, 8, 24)), + ("2017-08-24", date(2017, 8, 24)), + (" 20200101 ", date(2020, 1, 1)), + ], +) +def test_both_athena_formats_and_iso_parse_by_default(raw, expected): + assert read(raw, sa.Date()) == expected + + +@pytest.mark.parametrize( + "raw", + ["20200231", "19000229", "2020022", "+2020229", "202002299", "abcdefgh", + "31-FEB-2020", "01-XXX-2020", "2019-02-29", "Jan 2 2020", "2020-13-45"], +) +def test_impossible_or_misshapen_dates_are_unreadable(raw): + assert unreadable(raw, sa.Date()) + + +def test_ambiguous_date_needs_a_declared_format(): + assert unreadable("02/01/2020", sa.Date()) + assert read("02/01/2020", sa.Date(), ConversionPolicy(date=DateHandler(formats=("dd/mm/yyyy",)))) == date(2020, 1, 2) + assert read("02/01/2020", sa.Date(), ConversionPolicy(date=DateHandler(formats=("mm/dd/yyyy",)))) == date(2020, 2, 1) + + +def test_unknown_date_format_name_is_rejected(): + with pytest.raises(ValueError, match="Unknown date formats"): + DateHandler(formats=("yyyy/dd/mm",)) + + +def test_date_parser_is_replaceable(): + calls = [] + + def parser(values, formats): + calls.append(len(values)) + return pa.array([date(2000, 1, 1)] * len(values), pa.date32()), pa.array([False] * len(values)) + + values, _ = interpret(["a", "a", "b"], sa.Date(), ConversionPolicy(date=DateHandler(parser=parser))) + assert values == [date(2000, 1, 1)] * 3 + assert calls == [2] + + +def test_dictionary_encoded_dates_parse_the_same(): + column = pc.dictionary_encode(pa.array(["20200229", None, "", "20200301", "20200229"])) + assert interpret(column, sa.Date())[0] == [ + date(2020, 2, 29), None, None, date(2020, 3, 1), date(2020, 2, 29), + ] + + +@pytest.mark.parametrize( + "raw, expected", + [ + ("2017-08-24T12:34:56", datetime(2017, 8, 24, 12, 34, 56)), + ("2017-08-24 12:34", datetime(2017, 8, 24, 12, 34)), + ("2017-08-24 12:34:56.123456", datetime(2017, 8, 24, 12, 34, 56, 123456)), + ("20170824", datetime(2017, 8, 24)), + ("24-AUG-2017", datetime(2017, 8, 24)), + ], +) +def test_datetime_accepts_iso_and_date_only_forms(raw, expected): + assert read(raw, sa.DateTime()) == expected + + +@pytest.mark.parametrize("raw", ["2017-02-30 01:00:00", "2017-08-24 25:00:00", "2017-08-24 something", "garbage"]) +def test_datetime_rejects_invalid_values(raw): + assert unreadable(raw, sa.DateTime()) + + +def _yyyymmdd(days: np.ndarray) -> pa.Array: + years = days.astype("datetime64[Y]").astype(int) + 1970 + months = days.astype("datetime64[M]").astype(int) % 12 + 1 + dom = (days - days.astype("datetime64[M]")).astype(int) + 1 + + def pad(values: np.ndarray, width: int) -> pa.Array: + return pc.utf8_lpad(pa.array(values).cast(pa.string()), width, "0") + + return pc.binary_join_element_wise(pad(years, 4), pad(months, 2), pad(dom, 2), "") + + +def test_every_calendar_date_from_year_1_to_9999_parses_to_itself(): + days = np.arange(np.datetime64("0001-01-01"), np.datetime64("10000-01-01"), dtype="datetime64[D]") + dates, invalid = parse_dates(_yyyymmdd(days), ("yyyymmdd",)) + assert not pc.any(invalid).as_py() + assert dates.cast(pa.int32()).to_numpy().tolist() == days.astype(int).tolist() + + +@pytest.mark.parametrize("year", [1900, 2000, 2019, 2020, 2100]) +def test_month_day_grid_accepts_exactly_what_the_calendar_accepts(year): + grid = [(m, d) for m in range(100) for d in range(100)] + values = pa.array([f"{year:04d}{m:02d}{d:02d}" for m, d in grid]) + _, invalid = parse_dates(values, ("yyyymmdd",)) + expected = [] + for m, d in grid: + try: + date(year, m, d) + expected.append(False) + except ValueError: + expected.append(True) + assert invalid.to_pylist() == expected + + +def test_arrow_strptime_still_rolls_impossible_dates_over(): + rolled = pc.strptime(pa.array(["20200231"]), format="%Y%m%d", unit="s") + assert rolled.cast(pa.date32()).to_pylist() != [None], ( + "Arrow strptime now rejects impossible dates (apache/arrow#54643): " + "replace the ISO reassembly in parse_dates with strptime" + ) + + +def test_default_formats_are_mutually_distinguishable(): + probes = ["20200229", "29-FEB-2020", "2020-02-29"] + for name in DateFormats.DEFAULT: + matches = [bool(DateFormats.shape(name)(pa.array([p]))[0][0].as_py()) for p in probes] + assert sum(matches) == 1 + + +def test_typed_input_is_cast_without_parsing(): + assert read(date(2020, 5, 17), sa.Date()) == date(2020, 5, 17) + assert read(datetime(2020, 5, 17, 14, 30), sa.Date()) == date(2020, 5, 17) + assert read(datetime(2020, 5, 17, 14, 30, 5), sa.DateTime()) == datetime(2020, 5, 17, 14, 30, 5) + assert read(date(2020, 5, 17), sa.DateTime()) == datetime(2020, 5, 17) diff --git a/tests/conversion/test_missing_values.py b/tests/conversion/test_missing_values.py new file mode 100644 index 0000000..5b1530f --- /dev/null +++ b/tests/conversion/test_missing_values.py @@ -0,0 +1,28 @@ +from datetime import date + +import pyarrow as pa +import pytest +import sqlalchemy as sa + +from tests.support import interpret, read + + +def test_zero_false_and_dates_are_values(): + assert read(0, sa.Integer()) == 0 + assert read(False, sa.Boolean()) is False + assert read(date(2020, 2, 29), sa.Date()) == date(2020, 2, 29) + + +def test_arrow_nulls_and_nan_arrays(): + assert interpret(pa.nulls(3, pa.int64()), sa.Integer())[0] == [None] * 3 + assert interpret(pa.nulls(2, pa.string()), sa.String())[0] == [None, None] + for dtype in (pa.float32(), pa.float64()): + values, invalid = interpret(pa.array([1.0, float("nan"), 0.0], dtype), sa.Float()) + assert values == [1.0, None, 0.0] + assert not any(invalid) + + +@pytest.mark.parametrize("word", ["null", "None", "NA", "N/A", "NaN", "#N/A", "-nan", ""]) +def test_typed_columns_read_missing_words_as_null(word): + for sa_type in (sa.Integer(), sa.Float(), sa.Date(), sa.Boolean(), sa.DateTime()): + assert read(word, sa_type) is None diff --git a/tests/conversion/test_number_handlers.py b/tests/conversion/test_number_handlers.py new file mode 100644 index 0000000..afcdad3 --- /dev/null +++ b/tests/conversion/test_number_handlers.py @@ -0,0 +1,40 @@ +from decimal import Decimal + +import pytest +import sqlalchemy as sa + +from tests.support import read, unreadable + + +@pytest.mark.parametrize("raw", ["12abc", "0x1F", "1.0", "1e3", "--1", "Infinity"]) +def test_integer_rejects_non_integer_text(raw): + assert unreadable(raw, sa.Integer()) + + +def test_integer_accepts_sign_and_leading_zeros(): + assert read("+5", sa.Integer()) == 5 + assert read("-007", sa.Integer()) == -7 + + +@pytest.mark.parametrize("raw", ["Infinity", "inf", "-inf", "1,5", "abc"]) +def test_float_rejects_non_numbers(raw): + assert unreadable(raw, sa.Float()) + + +def test_float_accepts_decimal_and_exponent_forms(): + assert [read(v, sa.Float()) for v in ("1.80", "2", ".5", "1e3", "+2.5")] == [1.8, 2.0, 0.5, 1000.0, 2.5] + + +def test_numeric_keeps_declared_precision_and_rejects_overflow_of_scale(): + assert read("1.50", sa.Numeric(10, 2)) == Decimal("1.50") + assert unreadable("1.555", sa.Numeric(10, 2)) + + +def test_integer_trims_padding_and_reads_blank_and_missing_words_as_null(): + assert read(" 7 ", sa.Integer()) == 7 + assert [read(v, sa.Integer()) for v in ("", " NULL ", "NaN")] == [None, None, None] + + +def test_typed_input_is_cast_without_parsing(): + assert read(42.0, sa.Integer()) == 42 + assert unreadable(42.5, sa.Integer()) diff --git a/tests/conversion/test_policy.py b/tests/conversion/test_policy.py new file mode 100644 index 0000000..34ab58b --- /dev/null +++ b/tests/conversion/test_policy.py @@ -0,0 +1,109 @@ +import pyarrow as pa +import pytest +import sqlalchemy as sa + +from orm_loader.conversion import ( + BooleanHandler, + ConversionPolicy, + ConvertingHandler, + DateHandler, + DateTimeHandler, + ErrorPolicy, + FloatHandler, + IntegerHandler, + Interpretation, + NullPolicy, + NumericHandler, + PassthroughHandler, + TextHandler, +) +from tests.support import interpret, read, unreadable + + +class _Money(sa.Numeric): + pass + + +@pytest.mark.parametrize( + "sa_type, handler", + [ + (sa.String(), TextHandler), + (sa.Text(), TextHandler), + (sa.Enum("a", "b"), TextHandler), + (sa.Integer(), IntegerHandler), + (sa.BigInteger(), IntegerHandler), + (sa.Float(), FloatHandler), + (sa.Numeric(10, 2), NumericHandler), + (_Money(10, 2), NumericHandler), + (sa.Boolean(), BooleanHandler), + (sa.Date(), DateHandler), + (sa.DateTime(), DateTimeHandler), + (sa.TIMESTAMP(), DateTimeHandler), + (sa.LargeBinary(), PassthroughHandler), + ], +) +def test_the_most_specific_type_class_selects_the_handler(sa_type, handler): + assert type(ConversionPolicy().for_column(sa_type)) is handler + + +def test_a_type_without_a_handler_is_returned_unchanged(): + raw = pa.array([b"a", None]) + assert interpret(raw, sa.LargeBinary())[0] == [b"a", None] + + +class _Hex(ConvertingHandler): + def arrow_type(self, sa_type): + return pa.int64() + + def parse(self, cleaned, sa_type): + values = [int(v, 16) if v is not None else None for v in cleaned.to_pylist()] + return Interpretation(pa.array(values, pa.int64()), pa.array([False] * len(values))) + + +def test_a_new_type_is_one_handler_and_one_registration(): + policy = ConversionPolicy(handlers={sa.LargeBinary: _Hex()}) + assert interpret(["ff", " 10 ", "", "NA"], sa.LargeBinary(), policy)[0] == [255, 16, None, None] + + +def test_policy_handlers_replace_the_built_in_one_for_their_class(): + policy = ConversionPolicy(handlers={sa.Integer: _Hex()}) + assert read("ff", sa.Integer(), policy) == 255 + assert read("1.5", sa.Float(), policy) == 1.5 + + +def test_a_handler_carries_its_own_settings(): + policy = ConversionPolicy(boolean=BooleanHandler(true_words=("si",), false_words=("no",))) + assert read("SI", sa.Boolean(), policy) is True + assert read("no", sa.Boolean(), policy) is False + assert unreadable("yes", sa.Boolean(), policy) + + +def test_date_and_datetime_formats_are_declared_separately(): + policy = ConversionPolicy(date=DateHandler(formats=("dd/mm/yyyy",))) + assert read("02/01/2020", sa.Date(), policy).isoformat() == "2020-01-02" + assert unreadable("02/01/2020", sa.DateTime(), policy) + + +def test_column_handlers_win_over_type_handlers(): + policy = ConversionPolicy(columns={("t", "c"): IntegerHandler()}) + assert read("7", sa.String(), policy, table_name="t", column_name="c") == 7 + assert read("7", sa.String(), policy, table_name="t", column_name="other") == "7" + + +def test_text_markers_in_a_text_handler_replace_the_policy_markers(): + policy = ConversionPolicy( + nulls=NullPolicy(text_markers=frozenset({"NULL"})), + text=TextHandler(markers=frozenset({"-"})), + ) + assert read("-", sa.String(), policy) is None + assert read("NULL", sa.String(), policy) == "NULL" + + +def test_unknown_error_mode_is_rejected(): + with pytest.raises(ValueError, match="on_error"): + ErrorPolicy(on_error="ignore") # type: ignore[arg-type] + + +def test_unknown_marker_quoting_is_rejected(): + with pytest.raises(ValueError, match="marker_quoting"): + NullPolicy(marker_quoting="sometimes") # type: ignore[arg-type] diff --git a/tests/conversion/test_text_handler.py b/tests/conversion/test_text_handler.py new file mode 100644 index 0000000..dc923da --- /dev/null +++ b/tests/conversion/test_text_handler.py @@ -0,0 +1,60 @@ +import pytest +import sqlalchemy as sa + +from orm_loader.conversion import ConversionPolicy, NullPolicy, TextHandler +from tests.support import read, unreadable + +MARKER = ConversionPolicy(nulls=NullPolicy(text_markers=frozenset({"NULL"}))) + + +@pytest.mark.parametrize( + "raw, expected", + [ + ("NULL", None), + (" NULL ", None), + ('"NULL"', '"NULL"'), + ("Null", "Null"), + ("null", "null"), + ("NULLIFY", "NULLIFY"), + ("Contains NULL here", "Contains NULL here"), + ("NA", "NA"), + ("", None), + (" ", None), + ], +) +def test_text_marker_matches_only_the_exact_word(raw, expected): + assert read(raw, sa.String(), MARKER) == expected + + +def test_default_policy_has_no_text_markers(): + assert read("NULL", sa.String()) == "NULL" + + +def test_column_markers_replace_the_import_markers_for_their_column(): + policy = ConversionPolicy( + nulls=NullPolicy(text_markers=frozenset({"NULL"})), + columns={("t", "a"): TextHandler(markers=frozenset({"N/A"}))}, + ) + assert read("NULL", sa.String(), policy, table_name="t", column_name="a") == "NULL" + assert read("N/A", sa.String(), policy, table_name="t", column_name="a") is None + assert read("NULL", sa.String(), policy, table_name="t", column_name="b") is None + assert read("N/A", sa.String(), policy, table_name="t", column_name="b") == "N/A" + + +def test_policies_do_not_leak_between_conversions(): + assert read("NULL", sa.String(), MARKER) is None + assert read("NULL", sa.String()) == "NULL" + assert read("NULL", sa.String(), MARKER) is None + + +def test_padding_is_stripped_unless_disabled(): + assert read(" x ", sa.String()) == "x" + assert read(" x ", sa.String(), ConversionPolicy(text=TextHandler(strip=False))) == " x " + + +def test_blank_text_is_kept_as_empty_when_empty_is_not_null(): + assert read(" ", sa.String(), ConversionPolicy(nulls=NullPolicy(empty_is_null=False))) == "" + + +def test_text_longer_than_the_column_raises(): + assert unreadable("ab", sa.String(1)) diff --git a/tests/loaders/test_casting.py b/tests/loaders/test_casting.py deleted file mode 100644 index 4a0ac70..0000000 --- a/tests/loaders/test_casting.py +++ /dev/null @@ -1,312 +0,0 @@ -from enum import Enum, IntEnum - -import sqlalchemy as sa -from orm_loader.loaders.data.converters import ( - _COLUMN_CAST_RULES, - cast_scalar, - perform_cast, - register_column_cast_rule, -) -import pytest -from datetime import date, datetime - -def test_perform_cast_integer(): - assert perform_cast("123", sa.Integer(), on_error=None) == 123 - -def test_perform_cast_invalid_integer_returns_none(): - errors = [] - def on_error(v): errors.append(v) - - result = perform_cast("abc", sa.Integer(), on_error=on_error) - - assert result is None - assert errors == ["abc"] - -def test_perform_cast_boolean(): - assert perform_cast("yes", sa.Boolean(), on_error=None) is True - assert perform_cast("no", sa.Boolean(), on_error=None) is False - -def test_integer_from_numeric_string(): - assert perform_cast("42", sa.Integer(), on_error=None) == 42 - -def test_integer_from_float_string(): - assert perform_cast("42.0", sa.Integer(), on_error=None) == 42 - -def test_integer_from_actual_float(): - assert perform_cast(42.0, sa.Integer(), on_error=None) == 42 - -def test_integer_whitespace(): - assert perform_cast(" 7 ", sa.Integer(), on_error=None) == 7 - -def test_integer_invalid_string_returns_none_and_records_error(): - errors = [] - result = perform_cast("Episode", sa.Integer(), on_error=errors.append) - assert result is None - assert errors == ["Episode"] - -def test_integer_empty_string_is_null_not_error(): - errors = [] - assert perform_cast("", sa.Integer(), on_error=errors.append) is None - assert errors == [] - -def test_date_from_date(): - d = date(2020, 5, 17) - assert perform_cast(d, sa.Date(), on_error=None) == d - -def test_date_from_datetime(): - dt = datetime(2020, 5, 17, 14, 30) - assert perform_cast(dt, sa.Date(), on_error=None) == date(2020, 5, 17) - -def test_date_yyyymmdd(): - assert perform_cast("20170824", sa.Date(), on_error=None) == date(2017, 8, 24) - -def test_date_dd_mmm_yyyy(): - assert perform_cast("24-AUG-2017", sa.Date(), on_error=None) == date(2017, 8, 24) - -def test_date_iso(): - assert perform_cast("2017-08-24", sa.Date(), on_error=None) == date(2017, 8, 24) - -def test_date_dd_mm_yyyy(): - assert perform_cast("24/08/2017", sa.Date(), on_error=None) == date(2017, 8, 24) - -def test_date_rejects_iso_datetime(): - assert perform_cast("2017-08-24T12:00:00", sa.Date(), on_error=None) is None - -def test_date_rejects_invalid_date(): - assert perform_cast("2017-99-99", sa.Date(), on_error=None) is None - -def test_date_rejects_fuzzy(): - assert perform_cast("Aug 24 2017", sa.Date(), on_error=None) is None - -def test_datetime_from_datetime(): - dt = datetime(2020, 5, 17, 14, 30, 5) - assert perform_cast(dt, sa.DateTime(), on_error=None) == dt - -def test_datetime_from_date(): - d = date(2020, 5, 17) - assert perform_cast(d, sa.DateTime(), on_error=None) == datetime(2020, 5, 17, 0, 0) - -def test_datetime_iso_basic(): - assert perform_cast( - "2017-08-24T12:34:56", - sa.DateTime(), - on_error=None, - ) == datetime(2017, 8, 24, 12, 34, 56) - -def test_datetime_iso_with_seconds_only(): - assert perform_cast( - "2017-08-24T00:00:00", - sa.DateTime(), - on_error=None, - ) == datetime(2017, 8, 24, 0, 0, 0) - - -def test_datetime_from_yyyymmdd(): - assert perform_cast( - "20170824", - sa.DateTime(), - on_error=None, - ) == datetime(2017, 8, 24, 0, 0) - -def test_datetime_from_dd_mmm_yyyy(): - assert perform_cast( - "24-AUG-2017", - sa.DateTime(), - on_error=None, - ) == datetime(2017, 8, 24, 0, 0) - -def test_datetime_from_iso_date(): - assert perform_cast( - "2017-08-24", - sa.DateTime(), - on_error=None, - ) == datetime(2017, 8, 24, 0, 0) - -def test_datetime_dateutil_strict_accept(): - assert perform_cast( - "2017-08-24 something", - sa.DateTime(), - on_error=None, - ) is None - -def test_datetime_dateutil_exact_match(): - assert perform_cast( - "2017-08-24", - sa.DateTime(), - on_error=None, - ) == datetime(2017, 8, 24, 0, 0) - - -def test_datetime_rejects_invalid(): - assert perform_cast("2017-99-99", sa.DateTime(), on_error=None) is None - -def test_datetime_rejects_garbage(): - assert perform_cast("not a date", sa.DateTime(), on_error=None) is None - - -@pytest.mark.parametrize( - "raw, expected", - [ - ("NULL", None), - (" null ", None), - ("NaN", None), - ("N/A", None), - ("none", None), - ("", None), - (None, None), - ], -) -def test_string_null_normalisation(raw, expected): - assert cast_scalar(raw, sa.String()) is expected - - -@pytest.mark.parametrize( - "raw, expected", - [ - ("NULL", None), - (" NULL ", None), - ("NaN", None), - ("", None), - ("123", 123), - ], -) -def test_numeric_null_normalisation(raw, expected): - assert cast_scalar(raw, sa.Integer()) == expected - - -def test_cast_float_accepts_decimal_strings(): - assert cast_scalar("1.80", sa.Float()) == 1.8 - assert cast_scalar("2.50", sa.Float()) == 2.5 - assert cast_scalar("2", sa.Float()) == 2.0 - - -def test_cast_int_rejects_decimal_strings(): - assert cast_scalar("1.80", sa.Integer()) is None - assert cast_scalar("2.5", sa.Integer()) is None - assert cast_scalar("2", sa.Integer()) == 2 - - -class Role(str, Enum): - FIRST_AUTHOR = "first author" - LAST_AUTHOR = "last author" - - -@pytest.fixture(autouse=True) -def _clear_column_cast_rules(): - # _COLUMN_CAST_RULES is process-global state; isolate each test so - # registrations here can't leak into other test modules (or later tests - # in this file) that also call register_column_cast_rule. - _COLUMN_CAST_RULES.clear() - yield - _COLUMN_CAST_RULES.clear() - - -def test_register_column_cast_rule_requires_scalar_or_enum_type(): - with pytest.raises(ValueError, match="exactly one"): - register_column_cast_rule("t", "c") - - -def test_register_column_cast_rule_rejects_both_scalar_and_enum_type(): - with pytest.raises(ValueError, match="exactly one"): - register_column_cast_rule("t", "c", scalar=lambda v: v, enum_type=Role) - - -def test_column_cast_rule_enum_type_resolves_known_value(): - # Matches by .value (the raw display text) but returns .name -- sa.Enum's - # own default column-storage convention. - register_column_cast_rule("authors", "role", enum_type=Role) - assert cast_scalar("first author", sa.String(), table_name="authors", column_name="role") == "FIRST_AUTHOR" - - -def test_column_cast_rule_enum_type_records_unknown_value(): - errors = [] - register_column_cast_rule("authors", "role", enum_type=Role) - result = cast_scalar( - "BOGUS", sa.String(), on_error=errors.append, table_name="authors", column_name="role" - ) - assert result is None - assert errors == ["BOGUS"] - - -def test_column_cast_rule_takes_precedence_over_type_based_rules(): - # sa.Enum is itself a subclass of sa.String, so without column-specific - # precedence the existing String CastRule would match first and this - # would just pass "first author" through untouched, never resolving it - # to the Role member. - register_column_cast_rule("authors", "role", enum_type=Role) - result = cast_scalar("first author", sa.Enum(Role), table_name="authors", column_name="role") - assert result == "FIRST_AUTHOR" - - -def test_column_with_no_registered_rule_is_unaffected(): - register_column_cast_rule("authors", "role", enum_type=Role) - # A different column on the same table falls through to ordinary type-based casting - assert cast_scalar("42", sa.Integer(), table_name="authors", column_name="other") == 42 - - -def test_perform_cast_threads_table_and_column_name(): - register_column_cast_rule("authors", "role", enum_type=Role) - assert ( - perform_cast("last author", sa.String(), on_error=None, table_name="authors", column_name="role") - == "LAST_AUTHOR" - ) - - -def test_register_column_cast_rule_replaces_not_stacks(): - register_column_cast_rule("t", "c", scalar=lambda v: "first") - register_column_cast_rule("t", "c", scalar=lambda v: "second") - assert cast_scalar("anything", sa.String(), table_name="t", column_name="c") == "second" - - -def test_custom_scalar_cast_rule(): - register_column_cast_rule("t", "c", scalar=lambda v: v.upper()) - assert cast_scalar("shout", sa.String(), table_name="t", column_name="c") == "SHOUT" - -class Grade(IntEnum): - LOW = 1 - HIGH = 2 - - -def test_enum_type_matches_non_string_member_values(): - # sa.Enum(IntEnum) is legal and still persists .name, so an int-valued - # source column has to resolve without a lossy str() round-trip. - register_column_cast_rule("g", "grade", enum_type=Grade) - assert cast_scalar(1, sa.Enum(Grade), table_name="g", column_name="grade") == "LOW" - assert cast_scalar(2.0, sa.Enum(Grade), table_name="g", column_name="grade") == "HIGH" - - -def test_enum_type_still_matches_stringified_non_string_values(): - # A CSV source reads as str even for an int-valued enum, so the str() - # fallback has to stay behind the raw-value attempt. - register_column_cast_rule("g", "grade", enum_type=Grade) - assert cast_scalar("1", sa.Enum(Grade), table_name="g", column_name="grade") == "LOW" - - -def test_enum_type_accepts_an_already_resolved_member(): - register_column_cast_rule("authors", "role", enum_type=Role) - assert ( - cast_scalar(Role.FIRST_AUTHOR, sa.Enum(Role), table_name="authors", column_name="role") - == "FIRST_AUTHOR" - ) - register_column_cast_rule("g", "grade", enum_type=Grade) - assert cast_scalar(Grade.HIGH, sa.Enum(Grade), table_name="g", column_name="grade") == "HIGH" - - -def test_enum_type_unknown_value_still_fails(): - # The widened matching must not start accepting non-members. - register_column_cast_rule("g", "grade", enum_type=Grade) - errors = [] - assert cast_scalar( - 99, sa.Enum(Grade), table_name="g", column_name="grade", on_error=errors.append - ) is None - assert errors == [99] - - -@pytest.mark.parametrize("raw", [" ", "", None]) -def test_enum_type_blank_values_are_null_not_failures(raw): - register_column_cast_rule("authors", "role", enum_type=Role) - errors = [] - assert cast_scalar( - raw, sa.Enum(Role), table_name="authors", column_name="role", on_error=errors.append - ) is None - assert errors == [] diff --git a/tests/loaders/test_dedupe.py b/tests/loaders/test_dedupe.py index af46839..3887630 100644 --- a/tests/loaders/test_dedupe.py +++ b/tests/loaders/test_dedupe.py @@ -1,21 +1,20 @@ -import pyarrow as pa -from typing import cast, Type -from oa_configurator import Role -from orm_loader.loaders.loading_helpers import arrow_drop_duplicates +from typing import Type, cast + import pandas as pd import sqlalchemy as sa import sqlalchemy.orm as so +from oa_configurator import Role from sqlalchemy.orm import DeclarativeBase -from orm_loader.tables.loadable_table import CSVLoadableTableInterface -from orm_loader.tables.typing import CSVTableProtocol -from orm_loader.loaders.loader_interface import PandasLoader + +from orm_loader.tables.loadable_table import FileLoadableTableInterface +from orm_loader.tables.typing import FileTableProtocol class Base(DeclarativeBase): pass -class DedupTable(Base, CSVLoadableTableInterface): +class DedupTable(Base, FileLoadableTableInterface): __tablename__ = "dedup_table" __table_args__ = {"schema": Role.PRIMARY.value} @@ -23,19 +22,7 @@ class DedupTable(Base, CSVLoadableTableInterface): value: so.Mapped[str] = so.mapped_column(sa.String, nullable=False) -_DedupTable = cast(Type[CSVTableProtocol], DedupTable) - - -def test_arrow_drop_duplicates_simple(): - table = pa.table({ - "id": [1, 1, 2], - "x": ["a", "a", "b"], - }) - - deduped = arrow_drop_duplicates(table, ["id"]) - assert deduped.num_rows == 2 - assert deduped["id"].to_pylist() == [1, 2] - +_DedupTable = cast(Type[FileTableProtocol], DedupTable) def test_internal_deduplication(sqlite_session, sqlite_engine, tmp_path): @@ -50,10 +37,9 @@ def test_internal_deduplication(sqlite_session, sqlite_engine, tmp_path): ] ).to_csv(csv, index=False) - inserted = _DedupTable.load_csv( + inserted = _DedupTable.load_file( sqlite_session, csv, - loader=PandasLoader(), dedupe=True, ) sqlite_session.commit() diff --git a/tests/loaders/test_helpers.py b/tests/loaders/test_helpers.py index 75775ba..cd2dade 100644 --- a/tests/loaders/test_helpers.py +++ b/tests/loaders/test_helpers.py @@ -2,11 +2,7 @@ import pytest -import io - -from orm_loader.loaders.data_classes import ColumnCastingStats, TableCastingStats -from orm_loader.loaders.loading_helpers import ( - NormalisedCSVStream, +from orm_loader.loaders.readers.detection import ( infer_delim, infer_encoding, infer_quote_mode, @@ -14,33 +10,6 @@ ) -def test_normalised_csv_stream_strips_quoted_header_with_leading_whitespace(): - # A space after the delimiter is ordinary, valid CSV ('"id", "name"') -- - # stripping the quote before trimming whitespace leaves a stray leading - # quote on any token but the first. - raw = b'"id", "name"\n1,alpha\n' - stream = NormalisedCSVStream(io.BytesIO(raw), encoding="utf-8", delimiter=",") - assert stream.read() == b"id,name\n1,alpha\n" - -def test_column_casting_stats_records_examples(): - stats = ColumnCastingStats() - stats.record("bad1") - stats.record("bad2") - stats.record("bad3") - stats.record("bad4") # should not be stored - - assert stats.count == 4 - assert len(stats.examples) == 3 - -def test_table_casting_stats_aggregation(): - stats = TableCastingStats(table_name="test") - stats.record(column="a", value="x") - stats.record(column="a", value="y") - stats.record(column="b", value="z") - - assert stats.total_failures == 3 - assert stats.has_failures() is True - def test_infer_delim_csv(tmp_path): p = tmp_path / "x.csv" p.write_text("a,b,c\n1,2,3\n") @@ -51,6 +20,12 @@ def test_infer_delim_tsv(tmp_path): p.write_text("a\tb\tc\n1\t2\t3\n") assert infer_delim(p) == "\t" +def test_infer_encoding_promotes_ascii_to_utf8(tmp_path): + p = tmp_path / "ascii.csv" + p.write_bytes(b"id,name\n1,alpha\n") + assert infer_encoding(p)["encoding"] == "utf-8" + + def test_infer_encoding_utf8(tmp_path): p = tmp_path / "x.csv" p.write_text("hello") diff --git a/tests/loaders/test_load_policy.py b/tests/loaders/test_load_policy.py new file mode 100644 index 0000000..5f3f2da --- /dev/null +++ b/tests/loaders/test_load_policy.py @@ -0,0 +1,283 @@ +"""Loading under a null policy: markers, rejected rows, row accounting and sources.""" + +import logging +import os +from datetime import date +from pathlib import Path +from typing import Type, cast + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +import sqlalchemy as sa + +from orm_loader.backends import STAGING_SCHEMA +from orm_loader.conversion import ConversionPolicy, DateHandler, ErrorPolicy, NullPolicy, TextHandler +from orm_loader.helpers import ConversionError, IngestError, RejectLimitError, RequiredValueError +from orm_loader.backends import DedupeResult +from orm_loader.loaders import LoadReport +from orm_loader.tables.typing import FileTableProtocol +from tests.models import DemoRecord, DemoRequiredRecord + +_Demo = cast(Type[FileTableProtocol], DemoRecord) +_DemoRequired = cast(Type[FileTableProtocol], DemoRequiredRecord) + +SEVEN_ROWS = 'demo_id,demo_label\n1,NULL\n2,"NULL"\n3,"""NULL"""\n4,\n5,""\n6,"alpha,beta"\n7,"alpha\nbeta"\n' +MARKER = ConversionPolicy(nulls=NullPolicy(text_markers=frozenset({"NULL"}))) + + +def _load(dialect_session, tmp_path: Path, table, text: str, name: str | None = None, **kwargs) -> int: + path = tmp_path / f"{name or table.__tablename__}.csv" + path.write_text(text, encoding="utf-8", newline="") + tag = None if dialect_session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA + loaded = table.load_file(dialect_session, path, staging_schema_tag=tag, **kwargs) + dialect_session.commit() + dialect_session.expire_all() + return loaded + + +def _labels(dialect_session, model) -> dict[int, str | None]: + rows = dialect_session.execute(sa.select(model.demo_id, model.demo_label).order_by(model.demo_id)).all() + return {r[0]: r[1] for r in rows} + + +def test_text_marker_is_missing_whether_or_not_quoted(dialect_session, tmp_path): + _load(dialect_session, tmp_path, _Demo, SEVEN_ROWS, policy=MARKER) + assert _labels(dialect_session, DemoRecord) == { + 1: None, 2: None, 3: '"NULL"', 4: None, 5: None, 6: "alpha,beta", 7: "alpha\nbeta", + } + + +def test_text_marker_matches_exactly_and_ignores_padding(dialect_session, tmp_path): + text = "demo_id,demo_label\n1,NULL\n2, NULL \n3,Null\n4,null\n5,NULLIFY\n6,Contains NULL here\n" + _load(dialect_session, tmp_path, _Demo, text, policy=MARKER) + assert _labels(dialect_session, DemoRecord) == { + 1: None, 2: None, 3: "Null", 4: "null", 5: "NULLIFY", 6: "Contains NULL here", + } + + +def test_marker_in_a_required_column_raises(dialect_session, tmp_path): + text = "demo_id,demo_label\n1,ok\n2,NULL\n" + with pytest.raises(RequiredValueError, match="demo_label"): + _load(dialect_session, tmp_path, _DemoRequired, text, policy=MARKER) + + +def test_column_markers_apply_only_to_their_column(dialect_session, tmp_path): + policy = ConversionPolicy(columns={("demo_record", "demo_label"): TextHandler(markers=frozenset({"N/A"}))}) + _load(dialect_session, tmp_path, _Demo, "demo_id,demo_label\n1,N/A\n2,NULL\n", policy=policy) + assert _labels(dialect_session, DemoRecord) == {1: None, 2: "NULL"} + + +def test_quote_sensitive_marker_matches_only_unquoted(dialect_session, tmp_path): + policy = ConversionPolicy(nulls=NullPolicy(text_markers=frozenset({"NULL"}), marker_quoting="unquoted_only")) + _load(dialect_session, tmp_path, _Demo, SEVEN_ROWS, policy=policy) + assert _labels(dialect_session, DemoRecord) == { + 1: None, 2: "NULL", 3: '"NULL"', 4: None, 5: None, 6: "alpha,beta", 7: "alpha\nbeta", + } + + +def test_quote_sensitive_marker_warns_when_every_field_is_quoted(dialect_session, tmp_path, caplog): + policy = ConversionPolicy(nulls=NullPolicy(text_markers=frozenset({"NULL"}), marker_quoting="unquoted_only")) + text = '"demo_id","demo_label"\n"1","NULL"\n"2","x"\n' + with caplog.at_level(logging.WARNING): + _load(dialect_session, tmp_path, _Demo, text, policy=policy) + assert "quotes every field" in caplog.text + assert _labels(dialect_session, DemoRecord) == {1: "NULL", 2: "x"} + + +def test_quote_sensitive_markers_reject_column_markers(): + with pytest.raises(ValueError, match="unquoted_only"): + ConversionPolicy( + nulls=NullPolicy(marker_quoting="unquoted_only"), + columns={("t", "c"): TextHandler(markers=frozenset({"x"}))}, + ) + + +def test_both_athena_date_formats_and_iso_load_in_one_file(dialect_session, tmp_path): + text = "demo_id,demo_date\n1,20200229\n2,01-JAN-2020\n3,2020-01-02\n4,\n5,NA\n" + _load(dialect_session, tmp_path, _Demo, text) + rows = dialect_session.execute(sa.select(DemoRecord.demo_id, DemoRecord.demo_date).order_by(DemoRecord.demo_id)).all() + assert [tuple(r) for r in rows] == [ + (1, date(2020, 2, 29)), (2, date(2020, 1, 1)), (3, date(2020, 1, 2)), (4, None), (5, None), + ] + + +def test_ambiguous_date_needs_a_declared_format(dialect_session, tmp_path): + text = "demo_id,demo_date\n1,02/01/2020\n" + with pytest.raises(ConversionError, match="DateHandler"): + _load(dialect_session, tmp_path, _Demo, text) + _load(dialect_session, tmp_path, _Demo, text, policy=ConversionPolicy(date=DateHandler(formats=("dd/mm/yyyy",)))) + assert dialect_session.execute(sa.select(DemoRecord.demo_date)).scalar_one() == date(2020, 1, 2) + + +MESSY = ( + "demo_id,demo_label,demo_number,demo_date\n" + "1,ok,1,20200101\n" + "2,,2,20200102\n" + "3,bad number,x,20200103\n" + "4,bad date,4,20200231\n" + "5,fine,5,20200105\n" +) + + +def _rejects(path: Path) -> list[dict]: + import pyarrow.csv as pv + + header = [h.strip('"') for h in path.read_text().splitlines()[0].split(",")] + options = pv.ConvertOptions(column_types={name: pa.string() for name in header}) + return pv.read_csv(path, convert_options=options).to_pylist() + + +def test_reject_row_loads_good_rows_and_writes_the_rest(dialect_session, tmp_path): + report = LoadReport() + policy = ConversionPolicy(errors=ErrorPolicy(on_error="reject_row")) + loaded = _load(dialect_session, tmp_path, _Demo, MESSY, policy=policy, report=report) + assert loaded == 3 + assert _labels(dialect_session, DemoRecord) == {1: "ok", 2: None, 5: "fine"} + assert (report.rows_read, report.rows_loaded, report.rows_rejected) == (5, 3, 2) + assert report.accounted + assert report.rejects_path == tmp_path / "rejects" / "demo_record.rejects.csv" + rejected = _rejects(report.rejects_path) + assert [(r["row"], r["column"], r["demo_number"], r["demo_date"]) for r in rejected] == [ + ("3", "demo_number", "x", "20200103"), + ("4", "demo_date", "4", "20200231"), + ] + assert rejected[0]["reason"] == "cannot be read as Integer" + + +def test_reject_row_also_rejects_missing_required_values(dialect_session, tmp_path): + report = LoadReport() + text = "demo_id,demo_label\n1,ok\n2,\n3,NA\n" + loaded = _load( + dialect_session, tmp_path, _DemoRequired, text, policy=ConversionPolicy(errors=ErrorPolicy(on_error="reject_row")), report=report + ) + assert loaded == 2 + assert _labels(dialect_session, DemoRequiredRecord) == {1: "ok", 3: "NA"} + [rejected] = _rejects(report.rejects_path) + assert (rejected["row"], rejected["column"], rejected["reason"]) == ("2", "demo_label", "required value is missing") + + +def test_rejects_go_to_the_requested_directory(dialect_session, tmp_path): + target = tmp_path / "elsewhere" + report = LoadReport() + _load( + dialect_session, tmp_path, _Demo, MESSY, + policy=ConversionPolicy(errors=ErrorPolicy(on_error="reject_row")), rejects_dir=target, report=report, + ) + assert report.rejects_path == target / "demo_record.rejects.csv" + assert len(_rejects(report.rejects_path)) == 2 + assert not (tmp_path / "rejects").exists() + + +def test_unwritable_rejects_directory_fails_before_loading(dialect_session, tmp_path): + blocked = tmp_path / "blocked" + blocked.mkdir() + blocked.chmod(0o500) + try: + if os.access(blocked, os.W_OK): + pytest.skip("directory permissions are not enforced for this user") + with pytest.raises(IngestError, match="not writable"): + _load( + dialect_session, tmp_path, _Demo, MESSY, + policy=ConversionPolicy(errors=ErrorPolicy(on_error="reject_row")), rejects_dir=blocked / "inner", + ) + finally: + blocked.chmod(0o700) + assert dialect_session.execute(sa.select(sa.func.count()).select_from(DemoRecord)).scalar_one() == 0 + + +def test_clean_file_writes_no_rejects_file(dialect_session, tmp_path): + report = LoadReport() + _load( + dialect_session, tmp_path, _Demo, "demo_id,demo_label\n1,a\n", + policy=ConversionPolicy(errors=ErrorPolicy(on_error="reject_row")), report=report, + ) + assert report.rejects_path is None + assert not (tmp_path / "rejects").exists() + + +@pytest.mark.parametrize("limit", [1, 0.3]) +def test_max_rejects_stops_the_load(dialect_session, tmp_path, limit): + policy = ConversionPolicy(errors=ErrorPolicy(on_error="reject_row", max_rejects=limit)) + with pytest.raises(RejectLimitError): + _load(dialect_session, tmp_path, _Demo, MESSY, policy=policy) + + +@pytest.mark.parametrize("limit", [2, 0.5]) +def test_max_rejects_allows_a_load_within_the_limit(dialect_session, tmp_path, limit): + policy = ConversionPolicy(errors=ErrorPolicy(on_error="reject_row", max_rejects=limit)) + assert _load(dialect_session, tmp_path, _Demo, MESSY, policy=policy) == 3 + + +def test_default_policy_raises_on_the_first_bad_batch(dialect_session, tmp_path): + with pytest.raises(ConversionError, match="demo_number"): + _load(dialect_session, tmp_path, _Demo, MESSY) + assert dialect_session.execute(sa.select(sa.func.count()).select_from(DemoRecord)).scalar_one() == 0 + + +def test_report_counts_a_clean_load(dialect_session, tmp_path): + report = LoadReport() + _load(dialect_session, tmp_path, _Demo, "demo_id,demo_label\n1,a\n2,b\n", report=report) + assert (report.rows_read, report.rows_loaded, report.rows_rejected, report.rows_deduplicated) == (2, 2, 0, 0) + assert report.dedupe is None + + +def test_dedupe_keeps_the_first_row_and_reports_what_it_removed(dialect_session, tmp_path): + text = "demo_id,demo_label\n1,first\n2,only\n1,second\n1,third\n3,last\n2,again\n" + report = LoadReport() + loaded = _load(dialect_session, tmp_path, _Demo, text, dedupe=True, report=report) + assert loaded == 3 + assert _labels(dialect_session, DemoRecord) == {1: "first", 2: "only", 3: "last"} + assert report.dedupe == DedupeResult(("demo_id",), rows_before=6, rows_removed=3) + assert report.dedupe.rows_after == 3 + assert (report.rows_read, report.rows_loaded, report.rows_deduplicated) == (6, 3, 3) + assert report.accounted + + +def test_dedupe_reports_zero_when_nothing_is_duplicated(dialect_session, tmp_path): + report = LoadReport() + _load(dialect_session, tmp_path, _Demo, "demo_id,demo_label\n1,a\n2,b\n", dedupe=True, report=report) + assert report.dedupe == DedupeResult(("demo_id",), rows_before=2, rows_removed=0) + + +def test_dedupe_works_across_read_batches(sqlite_session, tmp_path): + rows = ["demo_id,demo_label"] + [f"{i},abcdefghijkl" for i in range(1, 450_001)] + ["1,again", "450000,again"] + report = LoadReport() + loaded = _load(sqlite_session, tmp_path, _Demo, "\n".join(rows) + "\n", dedupe=True, report=report) + assert loaded == 450_000 + assert report.dedupe.rows_removed == 2 + assert report.accounted + labels = _labels(sqlite_session, DemoRecord) + assert labels[1] == "abcdefghijkl" and labels[450_000] == "abcdefghijkl" + + +def test_dedupe_counts_rejected_rows_separately(dialect_session, tmp_path): + report = LoadReport() + text = "demo_id,demo_number\n1,1\n1,2\n2,x\n" + _load(dialect_session, tmp_path, _Demo, text, dedupe=True, report=report, policy=ConversionPolicy(errors=ErrorPolicy(on_error="reject_row"))) + assert (report.rows_read, report.rows_rejected, report.rows_deduplicated, report.rows_loaded) == (3, 1, 1, 1) + assert report.accounted + + +def test_parquet_goes_through_the_same_policy(dialect_session, tmp_path): + table = pa.table({ + "demo_id": pa.array([1, 2, 3, 4], pa.int64()), + "demo_label": pa.array(["NULL", "x", None, ""], pa.string()), + "demo_number": pa.array([0, None, 3, 4], pa.int64()), + "demo_date": pa.array([date(2020, 2, 29), None, date(2020, 1, 1), None], pa.date32()), + }) + path = tmp_path / "demo_record.parquet" + pq.write_table(table, path) + tag = None if dialect_session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA + _Demo.load_file(dialect_session, path, staging_schema_tag=tag, policy=MARKER) + dialect_session.commit() + rows = dialect_session.execute( + sa.select(DemoRecord.demo_id, DemoRecord.demo_label, DemoRecord.demo_number, DemoRecord.demo_date) + .order_by(DemoRecord.demo_id) + ).all() + assert [tuple(r) for r in rows] == [ + (1, None, 0, date(2020, 2, 29)), + (2, "x", None, None), + (3, None, 3, date(2020, 1, 1)), + (4, None, 4, None), + ] diff --git a/tests/loaders/test_loader_e2e.py b/tests/loaders/test_loader_e2e.py index 5cf7bf2..725a72d 100644 --- a/tests/loaders/test_loader_e2e.py +++ b/tests/loaders/test_loader_e2e.py @@ -1,7 +1,6 @@ from typing import Type, cast import logging -import numpy as np import pandas as pd import pytest import sqlalchemy as sa @@ -11,44 +10,30 @@ from orm_loader.backends import resolve_backend from orm_loader.backends.base import DeferredForeignKeyViolationError -from orm_loader.helpers import IngestError -from orm_loader.loaders.data.converters import _COLUMN_CAST_RULES, register_column_cast_rule -from orm_loader.loaders.data_classes import _clean_nulls -from orm_loader.loaders.loader_interface import PandasLoader -from orm_loader.tables.loadable_table import CSVLoadableTableInterface -from orm_loader.tables.typing import CSVTableProtocol +from orm_loader.helpers import IngestError, RequiredValueError +from orm_loader.tables.loadable_table import FileLoadableTableInterface +from orm_loader.tables.typing import FileTableProtocol from tests.models import ( Base, CompositeTable, CycleATable, EnumTable, - Flag, - ImpliedEnumTable, RequiredTable, - Role, SimpleTable, VocabSchemaTable, ) # Typed aliases: Pylance cannot verify SQLAlchemy metaclass-generated attrs -# satisfy CSVTableProtocol structurally, so we cast once per class here. -_SimpleTable = cast(Type[CSVTableProtocol], SimpleTable) -_RequiredTable = cast(Type[CSVTableProtocol], RequiredTable) -_CompositeTable = cast(Type[CSVTableProtocol], CompositeTable) -_EnumTable = cast(Type[CSVTableProtocol], EnumTable) -_ImpliedEnumTable = cast(Type[CSVTableProtocol], ImpliedEnumTable) -_VocabSchemaTable = cast(Type[CSVTableProtocol], VocabSchemaTable) -_CycleATable = cast(Type[CSVTableProtocol], CycleATable) +# satisfy FileTableProtocol structurally, so we cast once per class here. +_SimpleTable = cast(Type[FileTableProtocol], SimpleTable) +_RequiredTable = cast(Type[FileTableProtocol], RequiredTable) +_CompositeTable = cast(Type[FileTableProtocol], CompositeTable) +_EnumTable = cast(Type[FileTableProtocol], EnumTable) +_VocabSchemaTable = cast(Type[FileTableProtocol], VocabSchemaTable) +_CycleATable = cast(Type[FileTableProtocol], CycleATable) -@pytest.fixture(autouse=True) -def _clear_column_cast_rules(): - _COLUMN_CAST_RULES.clear() - yield - _COLUMN_CAST_RULES.clear() - - -def test_initial_csv_load(sqlite_session, tmp_path): +def test_initial_file_load(sqlite_session, tmp_path): csv_path = tmp_path / "test_table.csv" pd.DataFrame( @@ -59,13 +44,10 @@ def test_initial_csv_load(sqlite_session, tmp_path): ] ).to_csv(csv_path, index=False, sep="\t") - loader = PandasLoader() - - inserted = _SimpleTable.load_csv( + inserted = _SimpleTable.load_file( sqlite_session, csv_path, dedupe=False, - loader=loader, ) sqlite_session.commit() @@ -80,7 +62,7 @@ def test_initial_csv_load(sqlite_session, tmp_path): ] -def test_initial_csv_load_for_a_non_primary_schema_table(sqlite_session, tmp_path): +def test_initial_file_load_for_a_non_primary_schema_table(sqlite_session, tmp_path): """SQLite has no real schema concept, so every Role folds to None on this connection (see oa_configurator's SQLiteTestStrategy). A VOCAB-tagged table's load path must not error out just because @@ -93,8 +75,8 @@ def test_initial_csv_load_for_a_non_primary_schema_table(sqlite_session, tmp_pat [{"id": 1, "name": "alpha"}, {"id": 2, "name": "beta"}] ).to_csv(csv_path, index=False, sep="\t") - inserted = _VocabSchemaTable.load_csv( - sqlite_session, csv_path, dedupe=False, loader=PandasLoader() + inserted = _VocabSchemaTable.load_file( + sqlite_session, csv_path, dedupe=False ) sqlite_session.commit() @@ -119,13 +101,10 @@ def test_replace_merge_strategy(sqlite_session, tmp_path): ] ).to_csv(csv_path, index=False, sep="\t") - loader = PandasLoader() - - _SimpleTable.load_csv( + _SimpleTable.load_file( sqlite_session, csv_path, dedupe=False, - loader=loader, ) sqlite_session.commit() @@ -137,11 +116,10 @@ def test_replace_merge_strategy(sqlite_session, tmp_path): ] ).to_csv(csv_path, index=False, sep="\t") - replaced = _SimpleTable.load_csv( + replaced = _SimpleTable.load_file( sqlite_session, csv_path, dedupe=False, - loader=loader, merge_strategy="replace", ) sqlite_session.commit() @@ -157,17 +135,14 @@ def test_replace_merge_strategy(sqlite_session, tmp_path): ] -def test_empty_csv_is_noop(sqlite_session, tmp_path): +def test_empty_file_is_noop(sqlite_session, tmp_path): csv_path = tmp_path / "test_table.csv" csv_path.touch() - loader = PandasLoader() - - inserted = _SimpleTable.load_csv( + inserted = _SimpleTable.load_file( sqlite_session, csv_path, dedupe=False, - loader=loader, ) sqlite_session.commit() @@ -177,7 +152,7 @@ def test_empty_csv_is_noop(sqlite_session, tmp_path): assert rows == [] -def test_required_column_violation_drops_rows(sqlite_session, tmp_path): +def test_required_column_violation_raises(sqlite_session, tmp_path): csv = tmp_path / "required_table.csv" pd.DataFrame( [ @@ -186,19 +161,16 @@ def test_required_column_violation_drops_rows(sqlite_session, tmp_path): ] ).to_csv(csv, index=False) - inserted = _RequiredTable.load_csv( - sqlite_session, - csv, - loader=PandasLoader(), - dedupe=False, - ) - sqlite_session.commit() - - assert inserted == 1 + with pytest.raises(RequiredValueError, match="name"): + _RequiredTable.load_file( + sqlite_session, + csv, + dedupe=False, + ) def test_required_column_entirely_missing_raises(sqlite_session, tmp_path): - # Distinct from test_required_column_violation_drops_rows above: here + # Distinct from test_required_column_violation_raises above: here # the required column isn't present with a bad value, it's absent from # the source file entirely. No merge strategy supports a partial-column # load, so this must raise up front rather than silently proceed to a @@ -207,71 +179,13 @@ def test_required_column_entirely_missing_raises(sqlite_session, tmp_path): pd.DataFrame([{"id": 1}, {"id": 2}]).to_csv(csv, index=False) with pytest.raises(IngestError, match="required_table.*name"): - _RequiredTable.load_csv( + _RequiredTable.load_file( sqlite_session, csv, - loader=PandasLoader(), dedupe=False, ) -def test_enum_column_cast_rule_resolves_known_and_drops_unknown(sqlite_session, tmp_path): - # sa.Enum is itself a subclass of sa.String, so without a registered - # column-specific rule this would just pass every raw value through - # untouched via the existing String rule -- no membership check at all. - register_column_cast_rule("enum_table", "role", enum_type=Role) - - csv = tmp_path / "enum_table.csv" - pd.DataFrame( - [ - {"id": 1, "role": "first author"}, - {"id": 2, "role": "BOGUS"}, - {"id": 3, "role": "last author"}, - ] - ).to_csv(csv, index=False) - - total = _EnumTable.load_csv(sqlite_session, csv, loader=PandasLoader(), dedupe=False) - sqlite_session.commit() - - assert total == 3 - rows = dict(sqlite_session.execute(sa.select(EnumTable.id, EnumTable.role)).all()) - assert rows == {1: Role.FIRST_AUTHOR, 2: None, 3: Role.LAST_AUTHOR} - - -def test_implied_enum_on_plain_string_column(sqlite_session, tmp_path): - # The OMOP CDM concept.standard_concept/invalid_reason shape: a plain - # String(1) column with no type-level enum signal at all -- sa.Enum-keyed - # dispatch could never reach this, only a per-column rule can. This column - # expects the raw code itself ('S'), not enum_type='s default .name output - # ('STANDARD', which wouldn't even fit in String(1)) -- a genuinely custom - # case, so `scalar` is used directly rather than the `enum_type` shortcut. - def _to_flag_value(raw): - if raw is None: - return None - s = str(raw).strip() - if not s: - return None - return Flag(s).value - - register_column_cast_rule("implied_enum_table", "flag", scalar=_to_flag_value) - - csv = tmp_path / "implied_enum_table.csv" - pd.DataFrame( - [ - {"id": 1, "flag": "S"}, - {"id": 2, "flag": "BOGUS"}, - {"id": 3, "flag": "C"}, - ] - ).to_csv(csv, index=False) - - total = _ImpliedEnumTable.load_csv(sqlite_session, csv, loader=PandasLoader(), dedupe=False) - sqlite_session.commit() - - assert total == 3 - rows = dict(sqlite_session.execute(sa.select(ImpliedEnumTable.id, ImpliedEnumTable.flag)).all()) - assert rows == {1: "S", 2: None, 3: "C"} - - def test_composite_pk_dedup(sqlite_session, tmp_path): csv = tmp_path / "composite_table.csv" pd.DataFrame( @@ -282,10 +196,9 @@ def test_composite_pk_dedup(sqlite_session, tmp_path): ] ).to_csv(csv, index=False) - inserted = _CompositeTable.load_csv( + inserted = _CompositeTable.load_file( sqlite_session, csv, - loader=PandasLoader(), dedupe=True, ) sqlite_session.commit() @@ -327,8 +240,7 @@ def test_merge_strategies(sqlite_session, tmp_path, merge_strategy, expected_row ] ).to_csv(csv_path, index=False, sep="\t") - loader = PandasLoader() - SimpleTable.load_csv(sqlite_session, csv_path, dedupe=False, loader=loader) + SimpleTable.load_file(sqlite_session, csv_path, dedupe=False) sqlite_session.commit() pd.DataFrame( @@ -338,11 +250,10 @@ def test_merge_strategies(sqlite_session, tmp_path, merge_strategy, expected_row ] ).to_csv(csv_path, index=False, sep="\t") - inserted = SimpleTable.load_csv( + inserted = SimpleTable.load_file( sqlite_session, csv_path, dedupe=False, - loader=loader, merge_strategy=merge_strategy, ) sqlite_session.commit() @@ -368,12 +279,10 @@ def test_insert_if_empty_merge_strategy(sqlite_session, tmp_path): ] ).to_csv(csv_path, index=False, sep="\t") - loader = PandasLoader() - inserted = _SimpleTable.load_csv( + inserted = _SimpleTable.load_file( sqlite_session, csv_path, dedupe=False, - loader=loader, merge_strategy="insert_if_empty", ) sqlite_session.commit() @@ -393,18 +302,16 @@ def test_insert_if_empty_raises_on_non_empty_target(sqlite_session, sqlite_engin pd.DataFrame([{"id": 1, "name": "alpha"}]).to_csv(csv_path, index=False, sep="\t") - loader = PandasLoader() - _SimpleTable.load_csv(sqlite_session, csv_path, dedupe=False, loader=loader) + _SimpleTable.load_file(sqlite_session, csv_path, dedupe=False) sqlite_session.commit() pd.DataFrame([{"id": 2, "name": "beta"}]).to_csv(csv_path, index=False, sep="\t") with pytest.raises(ValueError, match="is not empty; cannot use merge strategy 'insert_if_empty'"): - _SimpleTable.load_csv( + _SimpleTable.load_file( sqlite_session, csv_path, dedupe=False, - loader=loader, merge_strategy="insert_if_empty", ) @@ -426,11 +333,10 @@ def test_empty_target_routes_merge_to_insert_if_empty(sqlite_session, tmp_path, caplog.set_level(logging.INFO, logger="orm_loader.tables.loadable_table") - inserted = _SimpleTable.load_csv( + inserted = _SimpleTable.load_file( sqlite_session, csv_path, dedupe=False, - loader=PandasLoader(), merge_strategy=merge_strategy, ) sqlite_session.commit() @@ -462,10 +368,9 @@ def test_staging_table_is_created_and_dropped(sqlite_session, sqlite_engine, tmp pd.DataFrame([{"id": 1, "name": "alpha"}]).to_csv(csv_path, index=False) - _SimpleTable.load_csv( + _SimpleTable.load_file( sqlite_session, csv_path, - loader=PandasLoader(), dedupe=False, ) sqlite_session.commit() @@ -485,7 +390,7 @@ def test_composite_pk_replace_merge(sqlite_session, tmp_path): ] ).to_csv(csv, index=False) - CompositeTable.load_csv(sqlite_session, csv, loader=PandasLoader()) + CompositeTable.load_file(sqlite_session, csv) sqlite_session.commit() pd.DataFrame( @@ -494,10 +399,9 @@ def test_composite_pk_replace_merge(sqlite_session, tmp_path): ] ).to_csv(csv, index=False) - CompositeTable.load_csv( + CompositeTable.load_file( sqlite_session, csv, - loader=PandasLoader(), merge_strategy="replace", ) sqlite_session.commit() @@ -519,29 +423,14 @@ def test_filename_must_match_tablename(sqlite_session, tmp_path): pd.DataFrame([{"id": 1, "name": "x"}]).to_csv(csv, index=False) with pytest.raises(ValueError, match="does not match table"): - SimpleTable.load_csv( + SimpleTable.load_file( sqlite_session, csv, - loader=PandasLoader(), ) -def test_clean_nulls_basic(): - assert _clean_nulls(None) is None - assert _clean_nulls(pd.NA) is None - assert _clean_nulls(float("nan")) is None - assert _clean_nulls(np.nan) is None - - -def test_clean_nulls_passthrough(): - assert _clean_nulls("") == "" - assert _clean_nulls("nan") == "nan" # string 'nan' must not be converted - assert _clean_nulls(0) == 0 - assert _clean_nulls("S") == "S" - - def test_nullable_column_with_nan_does_not_crash(sqlite_session, sqlite_engine, tmp_path): - class NullableTable(Base, CSVLoadableTableInterface): + class NullableTable(Base, FileLoadableTableInterface): __tablename__ = "nullable_table" __table_args__ = {"schema": SchemaRole.PRIMARY.value} @@ -549,7 +438,7 @@ class NullableTable(Base, CSVLoadableTableInterface): flag: so.Mapped[str | None] = so.mapped_column(sa.String, nullable=True) Base.metadata.create_all(sqlite_engine) - _NullableTable = cast(Type[CSVTableProtocol], NullableTable) + _NullableTable = cast(Type[FileTableProtocol], NullableTable) csv = tmp_path / "nullable_table.csv" pd.DataFrame( @@ -559,10 +448,9 @@ class NullableTable(Base, CSVLoadableTableInterface): ] ).to_csv(csv, index=False) - inserted = _NullableTable.load_csv( + inserted = _NullableTable.load_file( sqlite_session, csv, - loader=PandasLoader(), dedupe=False, ) sqlite_session.commit() @@ -578,7 +466,7 @@ class NullableTable(Base, CSVLoadableTableInterface): def test_embedded_newline_in_field_is_preserved(sqlite_session, sqlite_engine, tmp_path): - class TextTable(Base, CSVLoadableTableInterface): + class TextTable(Base, FileLoadableTableInterface): __tablename__ = "text_table" __table_args__ = {"schema": SchemaRole.PRIMARY.value} @@ -586,17 +474,16 @@ class TextTable(Base, CSVLoadableTableInterface): name: so.Mapped[str] = so.mapped_column(sa.String) Base.metadata.create_all(sqlite_engine) - _TextTable = cast(Type[CSVTableProtocol], TextTable) + _TextTable = cast(Type[FileTableProtocol], TextTable) csv = tmp_path / "text_table.csv" # Properly quoted CSV with embedded newline csv.write_text('id\tname\n1\t"hello\nworld"\n') - _TextTable.load_csv( + _TextTable.load_file( sqlite_session, csv, - loader=PandasLoader(), dedupe=False, ) sqlite_session.commit() @@ -606,7 +493,7 @@ class TextTable(Base, CSVLoadableTableInterface): def test_embedded_tab_in_field(sqlite_session, sqlite_engine, tmp_path): - class TextTable2(Base, CSVLoadableTableInterface): + class TextTable2(Base, FileLoadableTableInterface): __tablename__ = "tab_table" __table_args__ = {"schema": SchemaRole.PRIMARY.value} @@ -614,15 +501,14 @@ class TextTable2(Base, CSVLoadableTableInterface): name: so.Mapped[str] = so.mapped_column(sa.String) Base.metadata.create_all(sqlite_engine) - _TextTable2 = cast(Type[CSVTableProtocol], TextTable2) + _TextTable2 = cast(Type[FileTableProtocol], TextTable2) csv = tmp_path / "tab_table.csv" csv.write_text('id\tname\n1\t"foo\tbar"\n') - _TextTable2.load_csv( + _TextTable2.load_file( sqlite_session, csv, - loader=PandasLoader(), dedupe=False, ) sqlite_session.commit() @@ -656,7 +542,7 @@ def test_auto_strategy_keeps_indices_on_sqlite(sqlite_session, sqlite_engine, tm csv_path, index=False, sep="\t" ) - _SimpleTable.load_csv(sqlite_session, csv_path, loader=PandasLoader(), index_strategy="auto") + _SimpleTable.load_file(sqlite_session, csv_path, index_strategy="auto") sqlite_session.commit() assert not any("DROP INDEX" in s.upper() for s in ddl_log) @@ -672,7 +558,7 @@ def test_explicit_keep_preserves_indices(sqlite_session, sqlite_engine, tmp_path csv_path = tmp_path / "test_table.csv" pd.DataFrame([{"id": 1, "name": "alpha"}]).to_csv(csv_path, index=False, sep="\t") - _SimpleTable.load_csv(sqlite_session, csv_path, loader=PandasLoader(), index_strategy="keep") + _SimpleTable.load_file(sqlite_session, csv_path, index_strategy="keep") sqlite_session.commit() assert not any("DROP INDEX" in s.upper() for s in ddl_log) @@ -689,7 +575,7 @@ def test_explicit_drop_rebuild_on_sqlite_restores_index(sqlite_session, sqlite_e csv_path, index=False, sep="\t" ) - _SimpleTable.load_csv(sqlite_session, csv_path, loader=PandasLoader(), index_strategy="drop_rebuild") + _SimpleTable.load_file(sqlite_session, csv_path, index_strategy="drop_rebuild") sqlite_session.commit() assert any("DROP INDEX" in s.upper() for s in ddl_log) @@ -709,8 +595,7 @@ def test_drop_rebuild_logging_shows_merge_phases(sqlite_session, tmp_path, caplo ] ).to_csv(csv_path, index=False, sep="\t") - loader = PandasLoader() - _SimpleTable.load_csv(sqlite_session, csv_path, dedupe=False, loader=loader) + _SimpleTable.load_file(sqlite_session, csv_path, dedupe=False) sqlite_session.commit() pd.DataFrame( @@ -722,11 +607,10 @@ def test_drop_rebuild_logging_shows_merge_phases(sqlite_session, tmp_path, caplo caplog.set_level(logging.INFO, logger="orm_loader.tables.loadable_table") - _SimpleTable.load_csv( + _SimpleTable.load_file( sqlite_session, csv_path, dedupe=False, - loader=loader, merge_strategy="replace", index_strategy="drop_rebuild", ) @@ -767,7 +651,7 @@ def test_invalid_index_strategy_raises(sqlite_session, tmp_path): pd.DataFrame([{"id": 1, "name": "alpha"}]).to_csv(csv_path, index=False, sep="\t") with pytest.raises(ValueError, match="Unknown index_strategy"): - _SimpleTable.load_csv(sqlite_session, csv_path, loader=PandasLoader(), index_strategy="not-valid") + _SimpleTable.load_file(sqlite_session, csv_path, index_strategy="not-valid") # from hypothesis import given, strategies as st @@ -789,7 +673,7 @@ def test_invalid_index_strategy_raises(sqlite_session, tmp_path): # BaseLocal = declarative_base() -# class FuzzTable(BaseLocal, CSVLoadableTableInterface): +# class FuzzTable(BaseLocal, FileLoadableTableInterface): # __tablename__ = "fuzz_table" # id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) # txt: so.Mapped[str] = so.mapped_column(sa.String) @@ -801,7 +685,7 @@ def test_invalid_index_strategy_raises(sqlite_session, tmp_path): # csv = tmp_path / "fuzz_table.csv" # pd.DataFrame([{"id": 1, "txt": s}]).to_csv(csv, index=False, sep="\t") -# FuzzTable.load_csv(sqlite_session, csv, loader=PandasLoader(), dedupe=False) +# FuzzTable.load_file(sqlite_session, csv, dedupe=False) # sqlite_session.commit() # rows = sqlite_session.execute(sa.select(FuzzTable)).scalars().all() @@ -829,11 +713,10 @@ def test_sqlite_failed_batched_replace_keeps_original_target_rows(sqlite_session ).to_csv(csv_path, index=False, sep="\t") with pytest.raises(sa.exc.IntegrityError): - SimpleTable.load_csv( + SimpleTable.load_file( sqlite_session, csv_path, dedupe=False, - loader=PandasLoader(), merge_strategy="replace", merge_batch_size=2, ) @@ -845,7 +728,7 @@ def test_sqlite_failed_batched_replace_keeps_original_target_rows(sqlite_session @pytest.mark.duckdb @pytest.mark.db_dialect -def test_validate_deferred_fks_catches_an_orphan_introduced_by_load_csv(duckdb_session, duckdb_db, tmp_path): +def test_validate_deferred_fks_catches_an_orphan_introduced_by_load_file(duckdb_session, duckdb_db, tmp_path): """cycle_a_table.b_id is part of a genuine FK cycle with cycle_b_table, so DuckDB never got a real constraint for it (oa-configurator's _skip_deferred_fk_on_duckdb no-oped the ALTER TABLE ADD CONSTRAINT): a @@ -862,11 +745,10 @@ def test_validate_deferred_fks_catches_an_orphan_introduced_by_load_csv(duckdb_s pd.DataFrame([{"id": 1, "b_id": 999}]).to_csv(csv_path, index=False, sep="\t") with pytest.raises(DeferredForeignKeyViolationError, match="cycle_a_table"): - _CycleATable.load_csv( + _CycleATable.load_file( duckdb_session, csv_path, dedupe=False, - loader=PandasLoader(), validate_deferred_fks=True, resolved=duckdb_db.resolved, ) @@ -879,7 +761,7 @@ def test_validate_deferred_fks_catches_an_orphan_introduced_by_load_csv(duckdb_s @pytest.mark.duckdb @pytest.mark.db_dialect -def test_load_csv_rejects_replace_with_keep_index_strategy_on_duckdb_upfront(duckdb_session, tmp_path): +def test_load_file_rejects_replace_with_keep_index_strategy_on_duckdb_upfront(duckdb_session, tmp_path): """merge_strategy='replace' does an UPDATE in place on DuckDB, which breaks with a live index on an updated column. This combination must be rejected before the staging load even starts, not surface as DuckDB's @@ -888,11 +770,10 @@ def test_load_csv_rejects_replace_with_keep_index_strategy_on_duckdb_upfront(duc pd.DataFrame([{"id": 1, "name": "a"}]).to_csv(csv_path, index=False, sep="\t") with pytest.raises(NotImplementedError, match="merge_strategy='replace' with index_strategy='keep'"): - SimpleTable.load_csv( + SimpleTable.load_file( duckdb_session, csv_path, dedupe=False, - loader=PandasLoader(), merge_strategy="replace", index_strategy="keep", ) diff --git a/tests/loaders/test_null_handling.py b/tests/loaders/test_null_handling.py index c12b68f..c97d7f5 100644 --- a/tests/loaders/test_null_handling.py +++ b/tests/loaders/test_null_handling.py @@ -11,16 +11,15 @@ import pyarrow as pa import pytest import sqlalchemy as sa -from oa_configurator.testing import DIALECT_PARAMS -from orm_loader.backends import STAGING_SCHEMA, Dialect +from orm_loader.backends import STAGING_SCHEMA from orm_loader.helpers import IngestError -from orm_loader.loaders.data.converters import cast_arrow_column, cast_scalar -from orm_loader.tables.typing import CSVTableProtocol +from orm_loader.tables.typing import FileTableProtocol from tests.models import DemoRecord, DemoRequiredRecord +from tests.support import interpret, read -_Demo = cast(Type[CSVTableProtocol], DemoRecord) -_DemoRequired = cast(Type[CSVTableProtocol], DemoRequiredRecord) +_Demo = cast(Type[FileTableProtocol], DemoRecord) +_DemoRequired = cast(Type[FileTableProtocol], DemoRequiredRecord) MARKER_WORDS = [ "Null", "NULL", "null", "None", "none", "NA", "N/A", "NaN", "nan", @@ -37,129 +36,105 @@ } -_SESSION_FIXTURE_BY_DIALECT = { - Dialect.SQLITE: "sqlite_session", - Dialect.POSTGRESQL: "pg_session", - Dialect.DUCKDB: "duckdb_session", -} - - -@pytest.fixture(params=DIALECT_PARAMS) -def session(request): - return request.getfixturevalue(_SESSION_FIXTURE_BY_DIALECT[request.param]) - - def _write(tmp_path: Path, table: type, text: str) -> Path: path = tmp_path / f"{table.__tablename__}.csv" path.write_text(text, encoding="utf-8", newline="") return path -def _load(session, tmp_path: Path, table, text: str, **kwargs) -> int: - tag = None if session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA - n = table.load_csv(session, _write(tmp_path, table, text), staging_schema_tag=tag, **kwargs) - session.commit() - session.expire_all() +def _load(dialect_session, tmp_path: Path, table, text: str, **kwargs) -> int: + tag = None if dialect_session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA + n = table.load_file(dialect_session, _write(tmp_path, table, text), staging_schema_tag=tag, **kwargs) + dialect_session.commit() + dialect_session.expire_all() return n -def _labels(session, model) -> dict[int, str | None]: - rows = session.execute(sa.select(model.demo_id, model.demo_label).order_by(model.demo_id)).all() +def _labels(dialect_session, model) -> dict[int, str | None]: + rows = dialect_session.execute(sa.select(model.demo_id, model.demo_label).order_by(model.demo_id)).all() return {r[0]: r[1] for r in rows} # Marker-like text through each entry point @pytest.mark.parametrize("word", MARKER_WORDS) -def test_scalar_text_keeps_marker_words(word): - assert cast_scalar(word, sa.String()) == word - - -@pytest.mark.parametrize("word", MARKER_WORDS) -def test_arrow_text_keeps_marker_words(word): - out = cast_arrow_column(pa.array([word, "plain"]), DemoRecord.__table__.c.demo_label) - assert out.to_pylist() == [word, "plain"] - - -@pytest.mark.parametrize("word", MARKER_WORDS) -def test_dict_path_keeps_marker_words(word): - obj = DemoRecord(**DemoRecord.clean_kwargs({"demo_id": 1, "demo_label": word})) - assert obj.demo_label == word +def test_text_keeps_marker_words(word): + out, _ = interpret([word, "plain"], DemoRecord.__table__.c.demo_label.type) + assert out == [word, "plain"] # Actual missing objects -@pytest.mark.parametrize("value", [None, float("nan"), "", " "]) -def test_scalar_missing_objects_are_null(value): - assert cast_scalar(value, sa.String()) is None +@pytest.mark.parametrize("value", [None, "", " "]) +def test_missing_text_is_null(value): + assert read(value, sa.String()) is None def test_arrow_null_bitmap_and_nan_are_null(): col = DemoRecord.__table__.c.demo_number - assert cast_arrow_column(pa.array([1, None, 0], pa.int64()), col).to_pylist() == [1, None, 0] - assert cast_arrow_column(pa.array([1.0, float("nan")], pa.float64()), col).to_pylist() == [1, None] + assert interpret(pa.array([1, None, 0], pa.int64()), col.type)[0] == [1, None, 0] + assert interpret(pa.array([1.0, float("nan")], pa.float64()), col.type)[0] == [1, None] # Whitespace and identifiers def test_text_is_stripped_by_default(): - assert cast_scalar(" x ", sa.String()) == "x" + assert read(" x ", sa.String()) == "x" @pytest.mark.parametrize("code", ["007", "0123456", "1.50", "+5", "12.0"]) def test_text_is_never_canonicalised_as_a_number(code): - assert cast_scalar(code, sa.String()) == code - assert cast_arrow_column(pa.array([code]), DemoRecord.__table__.c.demo_label).to_pylist() == [code] + assert read(code, DemoRecord.__table__.c.demo_label.type) == code # Typed columns: valid value, missing word, or error @pytest.mark.parametrize("word", ["null", "None", "NA", "N/A", "NaN", "#N/A"]) def test_typed_missing_words_are_null(word): - assert cast_scalar(word, sa.Integer()) is None - assert cast_scalar(word, sa.Date()) is None + assert read(word, sa.Integer()) is None + assert read(word, sa.Date()) is None @pytest.mark.parametrize("word", ["null", "None", "NA", "N/A", "NaN"]) def test_typed_float_missing_words_are_null(word): - assert cast_scalar(word, sa.Float()) is None + assert read(word, sa.Float()) is None # Persisted round trips -def test_seven_row_fixture_preserves_text(session, tmp_path): - _load(session, tmp_path, _Demo, SEVEN_ROWS) - assert _labels(session, DemoRecord) == SEVEN_EXPECTED +def test_seven_row_fixture_preserves_text(dialect_session, tmp_path): + _load(dialect_session, tmp_path, _Demo, SEVEN_ROWS) + assert _labels(dialect_session, DemoRecord) == SEVEN_EXPECTED -def test_seven_row_fixture_every_field_quoted(session, tmp_path): - _load(session, tmp_path, _Demo, SEVEN_ROWS_QUOTED) - assert _labels(session, DemoRecord) == SEVEN_EXPECTED +def test_seven_row_fixture_every_field_quoted(dialect_session, tmp_path): + _load(dialect_session, tmp_path, _Demo, SEVEN_ROWS_QUOTED) + assert _labels(dialect_session, DemoRecord) == SEVEN_EXPECTED -def test_marker_words_survive_tab_delimited_load(session, tmp_path): +def test_marker_words_survive_tab_delimited_load(dialect_session, tmp_path): words = [w for w in MARKER_WORDS if w != '"NULL"'] text = "demo_id\tdemo_label\n" + "".join(f"{i}\t{w}\n" for i, w in enumerate(words, 1)) - _load(session, tmp_path, _Demo, text) - assert _labels(session, DemoRecord) == dict(enumerate(words, 1)) + _load(dialect_session, tmp_path, _Demo, text) + assert _labels(dialect_session, DemoRecord) == dict(enumerate(words, 1)) -def test_literal_quotes_survive_tab_delimited_load(session, tmp_path): +def test_literal_quotes_survive_tab_delimited_load(dialect_session, tmp_path): values = ['"NULL"', '"word" label', "Null", "NA"] text = "demo_id\tdemo_label\n" + "".join(f"{i}\t{v}\n" for i, v in enumerate(values, 1)) + "5\t\n" - _load(session, tmp_path, _Demo, text, quote_mode="literal") - assert _labels(session, DemoRecord) == {**dict(enumerate(values, 1)), 5: None} + _load(dialect_session, tmp_path, _Demo, text, quote_mode="literal") + assert _labels(dialect_session, DemoRecord) == {**dict(enumerate(values, 1)), 5: None} -def test_typed_values_round_trip(session, tmp_path): +def test_typed_values_round_trip(dialect_session, tmp_path): text = ( "demo_id,demo_label,demo_number,demo_flag,demo_date\n" "1,a,0,false,20200229\n" "2,b,NA,,\n" "3,c,7,true,2020-01-02\n" ) - _load(session, tmp_path, _Demo, text) - rows = session.execute(sa.select( + _load(dialect_session, tmp_path, _Demo, text) + rows = dialect_session.execute(sa.select( DemoRecord.demo_id, DemoRecord.demo_number, DemoRecord.demo_flag, DemoRecord.demo_date ).order_by(DemoRecord.demo_id)).all() assert [tuple(r) for r in rows] == [ @@ -180,37 +155,34 @@ def test_typed_values_round_trip(session, tmp_path): ("demo_number", "1.5"), ], ) -def test_unreadable_typed_value_raises_and_loads_nothing(session, tmp_path, column, bad): +def test_unreadable_typed_value_raises_and_loads_nothing(dialect_session, tmp_path, column, bad): cols = {"demo_number": "", "demo_date": ""} cols[column] = bad text = f"demo_id,demo_label,demo_number,demo_date\n1,a,,\n2,b,{cols['demo_number']},{cols['demo_date']}\n" with pytest.raises(IngestError, match=column): - _load(session, tmp_path, _Demo, text) - session.rollback() - assert session.execute(sa.select(sa.func.count()).select_from(DemoRecord)).scalar_one() == 0 + _load(dialect_session, tmp_path, _Demo, text) + assert dialect_session.execute(sa.select(sa.func.count()).select_from(DemoRecord)).scalar_one() == 0 # Required columns -def test_required_text_accepts_marker_words(session, tmp_path): +def test_required_text_accepts_marker_words(dialect_session, tmp_path): words = ["Null", "None", "NA", "N/A", "NaN"] text = "demo_id,demo_label\n" + "".join(f"{i},{w}\n" for i, w in enumerate(words, 1)) - _load(session, tmp_path, _DemoRequired, text) - assert _labels(session, DemoRequiredRecord) == dict(enumerate(words, 1)) + _load(dialect_session, tmp_path, _DemoRequired, text) + assert _labels(dialect_session, DemoRequiredRecord) == dict(enumerate(words, 1)) @pytest.mark.parametrize("missing", ["", '""', " "]) -def test_required_text_empty_raises_and_loads_nothing(session, tmp_path, missing): +def test_required_text_empty_raises_and_loads_nothing(dialect_session, tmp_path, missing): text = f"demo_id,demo_label\n1,ok\n2,{missing}\n" with pytest.raises(IngestError, match="demo_label"): - _load(session, tmp_path, _DemoRequired, text) - session.rollback() - assert session.execute(sa.select(sa.func.count()).select_from(DemoRequiredRecord)).scalar_one() == 0 + _load(dialect_session, tmp_path, _DemoRequired, text) + assert dialect_session.execute(sa.select(sa.func.count()).select_from(DemoRequiredRecord)).scalar_one() == 0 -def test_required_empty_raises_atomically_with_batching(session, tmp_path): +def test_required_empty_raises_atomically_with_batching(dialect_session, tmp_path): text = "demo_id,demo_label\n1,a\n2,b\n3,\n" with pytest.raises(IngestError, match="demo_label"): - _load(session, tmp_path, _DemoRequired, text, merge_batch_size=1) - session.rollback() - assert session.execute(sa.select(sa.func.count()).select_from(DemoRequiredRecord)).scalar_one() == 0 + _load(dialect_session, tmp_path, _DemoRequired, text, merge_batch_size=1) + assert dialect_session.execute(sa.select(sa.func.count()).select_from(DemoRequiredRecord)).scalar_one() == 0 diff --git a/tests/loaders/test_parquet_loader.py b/tests/loaders/test_parquet_loader.py index 17d93ef..6d172a9 100644 --- a/tests/loaders/test_parquet_loader.py +++ b/tests/loaders/test_parquet_loader.py @@ -7,16 +7,15 @@ from sqlalchemy.orm import DeclarativeBase from typing import cast, Type -from orm_loader.tables.loadable_table import CSVLoadableTableInterface -from orm_loader.tables.typing import CSVTableProtocol -from orm_loader.loaders.loader_interface import ParquetLoader +from orm_loader.tables.loadable_table import FileLoadableTableInterface +from orm_loader.tables.typing import FileTableProtocol class Base(DeclarativeBase): pass -class ParquetTable(Base, CSVLoadableTableInterface): +class ParquetTable(Base, FileLoadableTableInterface): __tablename__ = "parquet_table" __table_args__ = {"schema": Role.PRIMARY.value} @@ -24,7 +23,7 @@ class ParquetTable(Base, CSVLoadableTableInterface): value: so.Mapped[int] = so.mapped_column(sa.Integer, nullable=False) -_ParquetTable = cast(Type[CSVTableProtocol], ParquetTable) +_ParquetTable = cast(Type[FileTableProtocol], ParquetTable) def test_parquet_loader(sqlite_session, sqlite_engine, tmp_path): @@ -40,10 +39,9 @@ def test_parquet_loader(sqlite_session, sqlite_engine, tmp_path): path = tmp_path / "parquet_table.parquet" pq.write_table(table, path) - inserted = _ParquetTable.load_csv( + inserted = _ParquetTable.load_file( sqlite_session, path, - loader=ParquetLoader(), dedupe=False, ) sqlite_session.commit() @@ -51,9 +49,8 @@ def test_parquet_loader(sqlite_session, sqlite_engine, tmp_path): assert inserted == 2 -class NullSentinelTable(Base, CSVLoadableTableInterface): - """Nullable text/number columns, for checking that the Arrow casting path - treats text null sentinels the same way the scalar path does.""" +class NullSentinelTable(Base, FileLoadableTableInterface): + """Nullable text columns, for checking that marker-like text survives a Parquet load.""" __tablename__ = "null_sentinel_table" __table_args__ = {"schema": Role.PRIMARY.value} @@ -63,82 +60,24 @@ class NullSentinelTable(Base, CSVLoadableTableInterface): role: so.Mapped[str | None] = so.mapped_column(sa.String, nullable=True) -_NullSentinelTable = cast(Type[CSVTableProtocol], NullSentinelTable) +_NullSentinelTable = cast(Type[FileTableProtocol], NullSentinelTable) -_SENTINELS = ["NULL", "null", " NULL ", "n/a", "N/A", "NaN", "none", "None", " ", ""] +_BLANKS = [" ", ""] +_MARKER_TEXT = ["NULL", "null", "n/a", "N/A", "NaN", "none", "None"] -def test_arrow_path_normalises_text_null_sentinels(sqlite_session, sqlite_engine, tmp_path): - # cast_scalar normalises these before dispatching; without the same - # treatment here the plain String branch wrote the literal text into - # staging, so the same file loaded via the two loaders disagreed. +def test_parquet_keeps_marker_text_and_nulls_blanks(sqlite_session, sqlite_engine, tmp_path): Base.metadata.create_all(sqlite_engine) - rows = [{"id": i, "name": s, "role": None} for i, s in enumerate(_SENTINELS, start=1)] + values = _MARKER_TEXT + _BLANKS + rows = [{"id": i, "name": v, "role": None} for i, v in enumerate(values, start=1)] rows.append({"id": 99, "name": " keep me ", "role": None}) path = tmp_path / "null_sentinel_table.parquet" pq.write_table(pa.Table.from_pandas(pd.DataFrame(rows)), path) - _NullSentinelTable.load_csv(sqlite_session, path, loader=ParquetLoader(), dedupe=False) + _NullSentinelTable.load_file(sqlite_session, path, dedupe=False) sqlite_session.commit() stored = dict(sqlite_session.execute(sa.select(NullSentinelTable.id, NullSentinelTable.name)).all()) - assert [stored[i] for i in range(1, len(_SENTINELS) + 1)] == [None] * len(_SENTINELS) + assert [stored[i] for i in range(1, len(values) + 1)] == _MARKER_TEXT + [None] * len(_BLANKS) assert stored[99] == "keep me" - - -def test_arrow_and_pandas_paths_agree_on_null_sentinels(sqlite_session, sqlite_engine, tmp_path): - from orm_loader.loaders.data.converters import cast_arrow_column, cast_scalar - - col = NullSentinelTable.__table__.c.name - for raw in _SENTINELS + [" keep me ", "plain"]: - arrow = cast_arrow_column(pa.array([raw], type=pa.string()), col)[0].as_py() - scalar = cast_scalar(raw, col.type) - assert arrow == scalar, f"{raw!r}: arrow={arrow!r} scalar={scalar!r}" - - -def test_arrow_column_rule_does_not_count_nulls_as_cast_failures(sqlite_session, sqlite_engine, tmp_path): - # The override branch resolved these to None either way, but recorded each - # one as a failure -- warning about rows that are simply null. - from enum import Enum - - from orm_loader.loaders.data.converters import ( - _COLUMN_CAST_RULES, - cast_arrow_column, - register_column_cast_rule, - ) - from orm_loader.loaders.data_classes import TableCastingStats - - class Role(str, Enum): - FIRST_AUTHOR = "first author" - - try: - register_column_cast_rule("null_sentinel_table", "role", enum_type=Role) - stats = TableCastingStats(table_name="null_sentinel_table") - arr = cast_arrow_column( - pa.array(_SENTINELS + ["first author", "BOGUS"], type=pa.string()), - NullSentinelTable.__table__.c.role, - stats=stats, - ) - assert arr.to_pylist() == [None] * len(_SENTINELS) + ["FIRST_AUTHOR", None] - # Only the genuinely unresolvable value is a failure. - assert stats.total_failures == 1 - assert stats.columns["role"].examples == ["BOGUS"] - finally: - _COLUMN_CAST_RULES.clear() - - -def test_arrow_path_nulls_float_nan(sqlite_session, sqlite_engine, tmp_path): - from orm_loader.loaders.data.converters import cast_arrow_column - - class FloatTable(Base, CSVLoadableTableInterface): - __tablename__ = "float_sentinel_table" - __table_args__ = {"schema": Role.PRIMARY.value} - id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) - score: so.Mapped[float | None] = so.mapped_column(sa.Float, nullable=True) - - arr = cast_arrow_column( - pa.array([1.5, float("nan"), None], type=pa.float64()), - FloatTable.__table__.c.score, - ) - assert arr.to_pylist() == [1.5, None, None] diff --git a/tests/loaders/test_pg_loader.py b/tests/loaders/test_pg_loader.py deleted file mode 100644 index 20f37cf..0000000 --- a/tests/loaders/test_pg_loader.py +++ /dev/null @@ -1,427 +0,0 @@ -import sqlalchemy as sa -import pandas as pd -import pytest -from orm_loader.backends import STAGING_SCHEMA, resolve_backend -from orm_loader.loaders.data.converters import _COLUMN_CAST_RULES, register_column_cast_rule -from orm_loader.loaders.loading_helpers import infer_encoding, infer_delim, check_line_ending, quick_load_pg - -from tests.models import EnumTable, Role, SimpleTable - - -def test_copy_into_staging_with_extra_identity_column(pg_session, tmp_path): - """COPY must succeed when the staging table has a _rownum identity column.""" - csv = tmp_path / "test_table.csv" - pd.DataFrame([{"id": 1, "name": "alpha"}, {"id": 2, "name": "beta"}]).to_csv(csv, index=False) - - SimpleTable.create_staging_table(pg_session, staging_schema_tag=STAGING_SCHEMA) - staging_name = resolve_backend(pg_session).staging_name_for_table(SimpleTable.__tablename__) - - cols = pg_session.execute(sa.text( - "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, schema=STAGING_SCHEMA) - assert total == 2 - - rownums = pg_session.execute( - 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" - - -def test_copy_and_orm_path_equivalence(pg_session, tmp_path): - csv = tmp_path / "test_table.csv" - - pd.DataFrame( - [ - {"id": 1, "name": "alpha"}, - {"id": 2, "name": "beta"}, - ] - ).to_csv(csv, index=False, sep="\t") - - SimpleTable.load_csv(pg_session, csv, staging_schema_tag=STAGING_SCHEMA) - pg_session.commit() - - rows = pg_session.execute(sa.select(SimpleTable).order_by(SimpleTable.id)).scalars().all() - assert [(r.id, r.name) for r in rows] == [ - (1, "alpha"), - (2, "beta"), - ] - - - -def test_postgres_copy_fast_path(pg_session, tmp_path): - csv = tmp_path / "test_table.csv" - pd.DataFrame([{"id": 1, "name": "alpha"}]).to_csv(csv, index=False) - - inserted = SimpleTable.load_csv(pg_session, csv, staging_schema_tag=STAGING_SCHEMA) - pg_session.commit() - - assert inserted == 1 - -def test_postgres_copy_fast_path_is_used(pg_session, tmp_path, monkeypatch): - csv = tmp_path / "test_table.csv" - pd.DataFrame([{"id": 1, "name": "alpha"}]).to_csv(csv, index=False) - - called = {"copy": False} - - def fake_quick_load_pg(*args, **kwargs): - called["copy"] = True - return 1 - - import orm_loader.loaders.loading_helpers as loading_helpers - monkeypatch.setattr(loading_helpers, "quick_load_pg", fake_quick_load_pg) - - inserted = SimpleTable.load_csv(pg_session, csv, staging_schema_tag=STAGING_SCHEMA) - pg_session.commit() - - assert called["copy"] is True - assert inserted == 1 - -def test_copy_failure_falls_back_to_orm(pg_session, tmp_path, monkeypatch): - csv = tmp_path / "test_table.csv" - pd.DataFrame([{"id": 1, "name": "alpha"}]).to_csv(csv, index=False) - - # See test_postgres_copy_fast_path_is_used above: `quick_load_pg` is - # imported lazily from `loaders.loading_helpers`, so that's what needs - # patching, not `backends.postgres`. - import orm_loader.loaders.loading_helpers as loading_helpers - - def broken_copy(*args, **kwargs): - raise RuntimeError("boom") - - monkeypatch.setattr(loading_helpers, "quick_load_pg", broken_copy) - - inserted = SimpleTable.load_csv(pg_session, csv, staging_schema_tag=STAGING_SCHEMA) - pg_session.commit() - - rows = pg_session.execute(sa.select(SimpleTable)).scalars().all() - assert inserted == 1 - assert [(r.id, r.name) for r in rows] == [(1, "alpha")] - - -def test_postgres_upsert_does_not_update(pg_session, tmp_path): - csv = tmp_path / "test_table.csv" - - pd.DataFrame([{"id": 1, "name": "alpha"}]).to_csv(csv, index=False) - SimpleTable.load_csv(pg_session, csv, staging_schema_tag=STAGING_SCHEMA) - pg_session.commit() - - pd.DataFrame([{"id": 1, "name": "alpha_updated"}]).to_csv(csv, index=False) - - SimpleTable.load_csv(pg_session, csv, merge_strategy="upsert", staging_schema_tag=STAGING_SCHEMA) - pg_session.commit() - - rows = pg_session.execute(sa.select(SimpleTable)).scalars().all() - assert [(r.id, r.name) for r in rows] == [(1, "alpha")] - - -def test_postgres_insert_if_empty(pg_session, tmp_path): - csv = tmp_path / "test_table.csv" - - pd.DataFrame( - [ - {"id": 1, "name": "alpha"}, - {"id": 2, "name": "beta"}, - ] - ).to_csv(csv, index=False) - - inserted = SimpleTable.load_csv( - pg_session, - csv, - merge_strategy="insert_if_empty", - staging_schema_tag=STAGING_SCHEMA, - ) - pg_session.commit() - - assert inserted == 2 - - rows = pg_session.execute(sa.select(SimpleTable).order_by(SimpleTable.id)).scalars().all() - assert [(r.id, r.name) for r in rows] == [ - (1, "alpha"), - (2, "beta"), - ] - - -def test_postgres_insert_if_empty_raises_on_non_empty_target(pg_session, tmp_path): - csv = tmp_path / "test_table.csv" - - pd.DataFrame([{"id": 1, "name": "alpha"}]).to_csv(csv, index=False) - SimpleTable.load_csv(pg_session, csv, staging_schema_tag=STAGING_SCHEMA) - pg_session.commit() - - pd.DataFrame([{"id": 2, "name": "beta"}]).to_csv(csv, index=False) - - with pytest.raises(ValueError, match="is not empty; cannot use merge strategy 'insert_if_empty'"): - SimpleTable.load_csv( - pg_session, - csv, - merge_strategy="insert_if_empty", - staging_schema_tag=STAGING_SCHEMA, - ) - - -def test_postgres_copy_large_batch(pg_session, tmp_path): - csv = tmp_path / "test_table.csv" - - df = pd.DataFrame( - [{"id": i, "name": f"name_{i}"} for i in range(1, 10_000)] - ) - df.to_csv(csv, index=False) - - inserted = SimpleTable.load_csv(pg_session, csv, staging_schema_tag=STAGING_SCHEMA) - pg_session.commit() - - count = pg_session.execute(sa.text('SELECT COUNT(*) FROM test_table')).scalar() - assert count == 9999 - assert inserted == 9999 - - -def test_staging_schema_matches_target(pg_session, tmp_path): - csv = tmp_path / "test_table.csv" - pd.DataFrame([{"id": 1, "name": "alpha"}]).to_csv(csv, index=False) - - SimpleTable.create_staging_table(pg_session, staging_schema_tag=STAGING_SCHEMA) - staging_name = resolve_backend(pg_session).staging_name_for_table(SimpleTable.__tablename__) - - cols = pg_session.execute(sa.text(""" - SELECT column_name, data_type - FROM information_schema.columns - WHERE table_schema = :s AND table_name = :table - ORDER BY ordinal_position - """), {"s": STAGING_SCHEMA, "table": staging_name}).all() - - data_cols = [(name, dtype) for name, dtype in cols if name != "_rownum"] - assert data_cols == [ - ("id", "integer"), - ("name", "character varying"), - ] - assert any(name == "_rownum" and dtype == "bigint" for name, dtype in cols) - - -def test_infer_encoding_ascii_promoted_to_utf8(tmp_path): - p = tmp_path / "ascii.csv" - p.write_bytes(b"id,name\n1,alpha\n") - - enc = infer_encoding(p) - assert enc["encoding"] == "utf-8" - -@pytest.mark.xfail(reason="chardet todo - finalise encoding inference") -def test_infer_encoding_utf8(tmp_path): - p = tmp_path / "utf8.csv" - p.write_text("id,name\n1,α\n", encoding="utf-8") - - enc = infer_encoding(p) - enc_str = enc["encoding"] - assert enc_str is not None - assert enc_str.lower().startswith("utf") - -def test_infer_delim_tab(tmp_path): - p = tmp_path / "tab.csv" - p.write_text("id\tname\n1\talpha\n") - - assert infer_delim(p) == "\t" - -def test_infer_delim_comma(tmp_path): - p = tmp_path / "comma.csv" - p.write_text("id,name\n1,alpha\n") - - assert infer_delim(p) == "," - -@pytest.mark.parametrize( - "raw,expected", - [ - ("a,b\r\n", "\r\n"), - ("a,b\n", "\n"), - ("a,b\r", "\r"), - ], -) -def test_check_line_ending(raw, expected): - assert check_line_ending(raw) == expected - - -def test_check_line_ending_unknown(caplog): - raw = "a,b" - out = check_line_ending(raw) - assert out == "\n" - assert "Unable to detect line ending" in caplog.text - - -def test_quick_load_pg_basic(pg_session, tmp_path): - csv = tmp_path / "test_table.csv" - csv.write_text("id,name\n1,alpha\n2,beta\n") - - total = quick_load_pg( - path=csv, - session=pg_session, - tablename="test_table", - ) - - rows = pg_session.execute(sa.text('SELECT id, name FROM test_table ORDER BY id')).all() - assert total == 2 - assert rows == [(1, "alpha"), (2, "beta")] - - -def test_quick_load_pg_lowercases_header(pg_session, tmp_path): - csv = tmp_path / "test_table.csv" - csv.write_text("ID,NAME\n1,alpha\n") - - total = quick_load_pg(path=csv, session=pg_session, tablename="test_table") - row = pg_session.execute(sa.text('SELECT id, name FROM test_table')).one() - - assert total == 1 - assert row == (1, "alpha") - - -def test_quick_load_pg_strips_literally_quoted_header(pg_session, tmp_path): - """A header row with literal quote characters around each column name - (a common CSV-export convention) used to round-trip into an invalid - quadruple-quoted COPY column list (`""id""`), failing with - 'zero-length delimited identifier'.""" - csv = tmp_path / "test_table.csv" - csv.write_text('"id","name"\n1,alpha\n2,beta\n') - - total = quick_load_pg(path=csv, session=pg_session, tablename="test_table") - rows = pg_session.execute(sa.text("SELECT id, name FROM test_table ORDER BY id")).all() - - assert total == 2 - assert rows == [(1, "alpha"), (2, "beta")] - - -def test_quick_load_pg_tab_delimiter(pg_session, tmp_path): - csv = tmp_path / "test_table.csv" - csv.write_text("id\tname\n1\talpha\n2\tbeta\n") - - total = quick_load_pg(path=csv, session=pg_session, tablename="test_table") - rows = pg_session.execute(sa.text("SELECT id, name FROM test_table ORDER BY id")).all() - - assert total == 2 - assert rows == [(1, "alpha"), (2, "beta")] - - -def test_quick_load_pg_rollback_on_error(pg_session, tmp_path): - csv = tmp_path / "test_table.csv" - csv.write_text("id,name\n1,alpha\n2,\n") # violates NOT NULL - - with pytest.raises(Exception): - quick_load_pg(path=csv, session=pg_session, tablename="test_table") - - rows = pg_session.execute(sa.text("SELECT COUNT(*) FROM test_table")).scalar_one() - assert rows == 0 - - -def test_quick_load_pg_equivalence_with_orm(pg_session, tmp_path): - csv = tmp_path / "test_table.csv" - csv.write_text("id,name\n1,alpha\n2,beta\n") - - quick_load_pg(path=csv, session=pg_session, tablename="test_table") - - rows_pg = pg_session.execute(sa.text("SELECT id, name FROM test_table ORDER BY id")).all() - - pg_session.execute(sa.text("TRUNCATE test_table")) - pg_session.commit() - - df = pd.read_csv(csv) - pg_session.execute(sa.text(""" - INSERT INTO test_table (id, name) - VALUES (:id, :name) - """), df.to_dict(orient="records")) - pg_session.commit() - - rows_orm = pg_session.execute(sa.text("SELECT id, name FROM test_table ORDER BY id")).all() - - assert rows_pg == rows_orm - - -def test_quick_load_pg_trailing_blank_lines(pg_session, tmp_path): - csv = tmp_path / "test_table.csv" - - csv.write_text( - "id,name\n" - "1,alpha\r\n" - "2,beta\r\n" - ) - - total = quick_load_pg(path=csv, session=pg_session, tablename="test_table") - - rows = pg_session.execute( - sa.text("SELECT id, name FROM test_table ORDER BY id") - ).all() - - assert total == 2 - assert rows == [(1, "alpha"), (2, "beta")] - -def test_copy_fails_with_raw_carriage_returns_but_succeeds_after_normalisation(pg_session, tmp_path): - csv = tmp_path / "test_table.csv" - - # Intentionally include CRLF to simulate Windows-originated files - raw = ( - b"id,name\r\n" - b"1,alpha\r\n" - b"2,beta\r\n" - ) - - csv.write_bytes(raw) - - # First: raw COPY (current behaviour) — this may fail on some systems - failed = False - try: - quick_load_pg(path=csv, session=pg_session, tablename="test_table") - pg_session.commit() - except Exception: - failed = True - - # Clean up table for second attempt - pg_session.execute(sa.text('TRUNCATE "test_table"')) - pg_session.commit() - - # Now normalise newlines exactly the way your fix would - normalised = raw.replace(b"\r\n", b"\n").replace(b"\r", b"\n") - csv.write_bytes(normalised) - - # Second: normalised COPY — should always succeed - total = quick_load_pg(path=csv, session=pg_session, tablename="test_table") - pg_session.commit() - - rows = pg_session.execute( - sa.text('SELECT id, name FROM test_table ORDER BY id') - ).all() - - assert total == 2 - assert rows == [(1, "alpha"), (2, "beta")] - - # This assertion is the diagnostic: on systems where COPY is strict, - # the first attempt should fail. - # On systems where it "works natively", you may see failed == False, - # which proves the environment difference. - print("Raw COPY failed:", failed) - - -@pytest.fixture(autouse=True) -def _clear_column_cast_rules(): - _COLUMN_CAST_RULES.clear() - yield - _COLUMN_CAST_RULES.clear() - - -def test_enum_column_cast_rule_round_trips_on_real_postgres(pg_session, tmp_path): - # The merge step that moves rows from staging to the target table is a - # plain SQL copy with no Python-level type translation, so whatever text - # a cast rule produces is exactly what lands in the target column. - # enum_type matches by .value but stores by .name -- sa.Enum's own - # default column-storage convention -- so EnumTable needs no - # customisation at all. Worth proving against a real native Postgres - # enum type specifically, not just SQLite (which has none). - register_column_cast_rule("enum_table", "role", enum_type=Role) - - csv = tmp_path / "enum_table.csv" - csv.write_text("id,role\n1,first author\n2,BOGUS\n3,last author\n") - - total = EnumTable.load_csv(pg_session, csv, quote_mode="csv") - pg_session.commit() - - assert total == 3 - rows = dict(pg_session.execute(sa.select(EnumTable.id, EnumTable.role)).all()) - assert rows == {1: Role.FIRST_AUTHOR, 2: None, 3: Role.LAST_AUTHOR} diff --git a/tests/loaders/test_pipeline.py b/tests/loaders/test_pipeline.py new file mode 100644 index 0000000..f059ce1 --- /dev/null +++ b/tests/loaders/test_pipeline.py @@ -0,0 +1,144 @@ +"""The whole pipeline, file to staging table, on every dialect and in every file format.""" + +import csv +import io +from datetime import date +from pathlib import Path +from typing import Type, cast + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +import sqlalchemy as sa + +from orm_loader.backends import STAGING_SCHEMA +from orm_loader.conversion import ConversionPolicy, ErrorPolicy +from orm_loader.helpers import ConversionError, RequiredValueError +from orm_loader.loaders import DelimitedTextReader, LoaderContext, ParquetReader, stage_file +from orm_loader.tables.typing import FileTableProtocol +from tests.models import DemoRecord, DemoRequiredRecord + +_Demo = cast(Type[FileTableProtocol], DemoRecord) +_DemoRequired = cast(Type[FileTableProtocol], DemoRequiredRecord) + +HEADER = ["demo_id", "demo_label", "demo_number", "demo_flag", "demo_date"] +REJECT = ConversionPolicy(errors=ErrorPolicy(on_error="reject_row")) + + +@pytest.fixture(params=["csv", "tsv", "parquet"]) +def fmt(request) -> str: + return request.param + + +def _write(tmp_path: Path, table, fmt: str, rows: list[tuple], header: list[str] = HEADER) -> Path: + """Write the same rows as comma or tab delimited text, or as Parquet text columns.""" + path = tmp_path / f"{table.__tablename__}.{'parquet' if fmt == 'parquet' else fmt}" + if fmt == "parquet": + pq.write_table(pa.table({name: pa.array([r[i] for r in rows], pa.string()) for i, name in enumerate(header)}), path) + else: + buffer = io.StringIO() + csv.writer(buffer, delimiter="," if fmt == "csv" else "\t", lineterminator="\n").writerows([header, *rows]) + path.write_text(buffer.getvalue(), encoding="utf-8", newline="") + return path + + +def _stage(session, table, path: Path, reader=None, **context): + tag = None if session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA + table.create_staging_table(session, staging_schema_tag=tag) + staging = table.get_staging_table(session, staging_schema_tag=tag) + ctx = LoaderContext( + tableclass=table, session=session, path=path, staging_table=staging, staging_schema_tag=tag, **context + ) + loaded = stage_file(ctx, reader or table._select_reader(ctx)) + return loaded, ctx, staging + + +def _staged(session, staging, *columns: str) -> list[tuple]: + order = staging.c._rownum if "_rownum" in staging.c else sa.literal_column("rowid") + return [tuple(r) for r in session.execute(sa.select(*(staging.c[c] for c in columns)).order_by(order)).all()] + + +def test_rows_are_converted_to_the_column_types(dialect_session, tmp_path, fmt): + rows = [("1", "alpha", "10", "true", "20200229"), ("2", " beta ", "", "no", "24-AUG-2017")] + loaded, ctx, staging = _stage(dialect_session, _Demo, _write(tmp_path, DemoRecord, fmt, rows)) + assert loaded == 2 and ctx.report.accounted + assert _staged(dialect_session, staging, *HEADER) == [ + (1, "alpha", 10, True, date(2020, 2, 29)), + (2, "beta", None, False, date(2017, 8, 24)), + ] + + +def test_marker_words_survive_and_only_blanks_are_null(dialect_session, tmp_path, fmt): + labels = ["NULL", "None", "N/A", "NaN", "", " ", " padded "] + rows = [(str(i), label, "", "", "") for i, label in enumerate(labels, start=1)] + _, _, staging = _stage(dialect_session, _Demo, _write(tmp_path, DemoRecord, fmt, rows)) + assert [r[0] for r in _staged(dialect_session, staging, "demo_label")] == [ + "NULL", "None", "N/A", "NaN", None, None, "padded", + ] + + +def test_missing_words_in_typed_columns_are_null(dialect_session, tmp_path, fmt): + rows = [("1", "x", "NA", "null", "N/A")] + _, _, staging = _stage(dialect_session, _Demo, _write(tmp_path, DemoRecord, fmt, rows)) + assert _staged(dialect_session, staging, "demo_number", "demo_flag", "demo_date") == [(None, None, None)] + + +def test_an_unreadable_value_raises_and_stages_nothing(dialect_session, tmp_path, fmt): + rows = [("1", "a", "1", "", ""), ("2", "b", "abc", "", "")] + path = _write(tmp_path, DemoRecord, fmt, rows) + with pytest.raises(ConversionError, match="demo_number"): + _stage(dialect_session, _Demo, path) + staging = _Demo.get_staging_table(dialect_session, staging_schema_tag=None if dialect_session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA) + assert _staged(dialect_session, staging, "demo_id") == [] + + +def test_a_missing_required_value_raises(dialect_session, tmp_path, fmt): + path = _write(tmp_path, DemoRequiredRecord, fmt, [("1", "ok"), ("2", "")], ["demo_id", "demo_label"]) + with pytest.raises(RequiredValueError): + _stage(dialect_session, _DemoRequired, path) + + +def test_reject_row_loads_the_rest_and_accounts_for_every_row(dialect_session, tmp_path, fmt): + rows = [("1", "a", "1", "", ""), ("2", "b", "abc", "", ""), ("3", "c", "3", "", "02/01/2020"), ("4", "d", "", "", "")] + loaded, ctx, staging = _stage( + dialect_session, _Demo, _write(tmp_path, DemoRecord, fmt, rows), policy=REJECT, rejects_dir=tmp_path / "rej" + ) + assert loaded == 2 + assert (ctx.report.rows_read, ctx.report.rows_loaded, ctx.report.rows_rejected) == (4, 2, 2) + assert ctx.report.accounted + assert [r[0] for r in _staged(dialect_session, staging, "demo_id")] == [1, 4] + assert ctx.report.rejects_path is not None and ctx.report.rejects_path.exists() + + +def test_dedupe_keeps_the_first_row_for_each_key(dialect_session, tmp_path, fmt): + rows = [("1", "first", "", "", ""), ("2", "x", "", "", ""), ("1", "second", "", "", "")] + loaded, ctx, staging = _stage(dialect_session, _Demo, _write(tmp_path, DemoRecord, fmt, rows), dedupe=True) + assert loaded == 2 and ctx.report.rows_deduplicated == 1 and ctx.report.accounted + assert sorted(_staged(dialect_session, staging, "demo_id", "demo_label")) == [(1, "first"), (2, "x")] + + +def test_a_file_in_many_batches_loads_every_row_in_order(dialect_session, tmp_path, fmt): + rows = [(str(i), f"label {i}", str(i), "", "") for i in range(1, 1501)] + path = _write(tmp_path, DemoRecord, fmt, rows) + columns = DemoRecord.model_columns() + reader = ParquetReader(path, columns, chunksize=100) if fmt == "parquet" else DelimitedTextReader(path, columns, block_size=1 << 10) + loaded, ctx, staging = _stage(dialect_session, _Demo, path, reader=reader) + assert loaded == 1500 and ctx.report.accounted + assert [r[0] for r in _staged(dialect_session, staging, "demo_id")] == list(range(1, 1501)) + + +def test_a_failing_reader_is_raised_and_stages_nothing(dialect_session, tmp_path, fmt): + path = _write(tmp_path, DemoRecord, fmt, [("1", "a", "", "", "")]) + reader = _Demo._select_reader( + LoaderContext(tableclass=_Demo, session=dialect_session, path=path, staging_table=cast(sa.Table, DemoRecord.__table__)) + ) + + def broken(): + raise RuntimeError("boom") + yield # pragma: no cover + + reader.batches = broken # type: ignore[method-assign] + with pytest.raises(RuntimeError, match="boom"): + _stage(dialect_session, _Demo, path, reader=reader) + tag = None if dialect_session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA + assert _staged(dialect_session, _Demo.get_staging_table(dialect_session, staging_schema_tag=tag), "demo_id") == [] diff --git a/tests/loaders/test_processor.py b/tests/loaders/test_processor.py new file mode 100644 index 0000000..97108fe --- /dev/null +++ b/tests/loaders/test_processor.py @@ -0,0 +1,140 @@ +"""BatchProcessor: raw batches in, typed rows, rejects and accounting out, no database.""" + +import csv +from datetime import date +from pathlib import Path + +import pyarrow as pa +import pytest + +from orm_loader.conversion import ConversionPolicy, DateHandler, ErrorPolicy +from orm_loader.helpers import ConversionError, IngestError, RejectLimitError, RequiredValueError +from orm_loader.loaders import LoadReport +from orm_loader.loaders.processor import BatchProcessor +from tests.models import DemoRecord, DemoRequiredRecord + +COLUMNS = ["demo_id", "demo_label", "demo_number", "demo_date"] +REJECT = ConversionPolicy(errors=ErrorPolicy(on_error="reject_row")) + + +def _batch(*rows: tuple, columns: list[str] = COLUMNS) -> pa.RecordBatch: + return pa.RecordBatch.from_pylist([dict(zip(columns, row)) for row in rows]) + + +def _processor(tmp_path: Path, policy: ConversionPolicy | None = None, table=DemoRecord, columns=COLUMNS, **kw): + report = LoadReport() + return BatchProcessor(table, columns, policy or ConversionPolicy(), report, tmp_path / "rejects.csv", **kw), report + + +def _rejects(tmp_path: Path) -> list[dict]: + with open(tmp_path / "rejects.csv", newline="") as handle: + return list(csv.DictReader(handle)) + + +def test_raw_text_is_converted_to_the_column_types(tmp_path): + processor, report = _processor(tmp_path) + table = processor.process(_batch(("1", "a", "10", "20200229"), ("2", " b ", "", ""))) + assert table.to_pylist() == [ + {"demo_id": 1, "demo_label": "a", "demo_number": 10, "demo_date": date(2020, 2, 29)}, + {"demo_id": 2, "demo_label": "b", "demo_number": None, "demo_date": None}, + ] + assert report.rows_read == 2 + + +def test_without_normalise_the_batch_passes_through_as_read(tmp_path): + processor, _ = _processor(tmp_path, normalise=False) + assert processor.process(_batch(("1", "a", "abc", "x"))).to_pylist() == [ + {"demo_id": "1", "demo_label": "a", "demo_number": "abc", "demo_date": "x"} + ] + + +def test_unreadable_values_raise_with_the_count_and_first_three_samples(tmp_path): + processor, _ = _processor(tmp_path) + batch = _batch(*[(str(i), "a", bad, "") for i, bad in enumerate(["1", "a", "b", "c", "d", "2"])]) + with pytest.raises(ConversionError) as exc: + processor.process(batch) + assert (exc.value.count, exc.value.samples) == (4, ["a", "b", "c"]) + assert "demo_record.demo_number" in str(exc.value) + + +def test_an_unreadable_date_names_the_option_that_declares_a_format(tmp_path): + processor, _ = _processor(tmp_path) + with pytest.raises(ConversionError, match="DateHandler"): + processor.process(_batch(("1", "a", "", "02/01/2020"))) + declared = ConversionPolicy(date=DateHandler(formats=("dd/mm/yyyy",))) + processor, _ = _processor(tmp_path, declared) + assert processor.process(_batch(("1", "a", "", "02/01/2020"))).column("demo_date").to_pylist() == [date(2020, 1, 2)] + + +def test_a_missing_required_value_raises_with_the_count_and_row_numbers(tmp_path): + processor, _ = _processor(tmp_path, table=DemoRequiredRecord, columns=["demo_id", "demo_label"]) + with pytest.raises(RequiredValueError) as exc: + processor.process(_batch(("1", "ok"), ("2", ""), ("3", None), columns=["demo_id", "demo_label"])) + assert (exc.value.count, exc.value.rows) == (2, [2, 3]) + + +def test_row_numbers_continue_across_batches(tmp_path): + processor, report = _processor(tmp_path, table=DemoRequiredRecord, columns=["demo_id", "demo_label"]) + processor.process(_batch(("1", "ok"), ("2", "ok"), columns=["demo_id", "demo_label"])) + with pytest.raises(RequiredValueError) as exc: + processor.process(_batch(("3", ""), columns=["demo_id", "demo_label"])) + assert exc.value.rows == [3] + assert report.rows_read == 3 + + +def test_reject_row_keeps_good_rows_and_writes_the_rest_with_raw_text_and_reason(tmp_path): + processor, report = _processor(tmp_path, REJECT) + kept = processor.process(_batch(("1", "a", "10", ""), ("2", "b", "x", ""), ("3", "c", "", "bad date"))) + processor.rejects.close() + assert kept.column("demo_id").to_pylist() == [1] + rejected = _rejects(tmp_path) + assert [(r["row"], r["column"], r["demo_number"], r["demo_date"]) for r in rejected] == [ + ("2", "demo_number", "x", ""), + ("3", "demo_date", "", "bad date"), + ] + assert all("cannot be read" in r["reason"] for r in rejected) + assert (report.rows_read, report.rows_rejected, report.rejects_path) == (3, 2, tmp_path / "rejects.csv") + + +def test_reject_row_also_rejects_missing_required_values(tmp_path): + processor, report = _processor(tmp_path, REJECT, table=DemoRequiredRecord, columns=["demo_id", "demo_label"]) + kept = processor.process(_batch(("1", "ok"), ("2", ""), columns=["demo_id", "demo_label"])) + processor.rejects.close() + assert kept.num_rows == 1 and report.rows_rejected == 1 + assert _rejects(tmp_path)[0]["reason"] == "required value is missing" + + +def test_a_clean_batch_writes_no_rejects_file(tmp_path): + processor, _ = _processor(tmp_path, REJECT) + processor.process(_batch(("1", "a", "10", ""))) + processor.rejects.close() + assert not (tmp_path / "rejects.csv").exists() + + +@pytest.mark.parametrize("limit, rows_bad, should_stop", [(1, 2, True), (2, 2, False), (0.5, 3, True), (0.5, 2, False)]) +def test_max_rejects_stops_a_load_over_a_count_or_a_fraction(tmp_path, limit, rows_bad, should_stop): + policy = ConversionPolicy(errors=ErrorPolicy(on_error="reject_row", max_rejects=limit)) + processor, _ = _processor(tmp_path, policy) + rows = [(str(i), "a", "x" if i <= rows_bad else "1", "") for i in range(1, 5)] + if should_stop: + with pytest.raises(RejectLimitError): + processor.process(_batch(*rows)) + else: + assert processor.process(_batch(*rows)).num_rows == 4 - rows_bad + + +def test_stream_skips_empty_results_and_closes_the_rejects_file(tmp_path): + processor, report = _processor(tmp_path, REJECT) + batches = [_batch(("1", "a", "x", "")), _batch(("2", "b", "5", ""))] + tables = list(processor.stream(batches)) + assert [t.column("demo_id").to_pylist() for t in tables] == [[2]] + assert processor.rejects._writer is None + assert [r["row"] for r in _rejects(tmp_path)] == ["1"] + assert report.rows_read == 2 + + +def test_an_unwritable_rejects_directory_fails_before_any_batch_is_read(tmp_path): + blocker = tmp_path / "file" + blocker.write_text("") + with pytest.raises(IngestError, match="not writable"): + BatchProcessor(DemoRecord, COLUMNS, REJECT, LoadReport(), blocker / "sub" / "rejects.csv") diff --git a/tests/loaders/test_readers.py b/tests/loaders/test_readers.py new file mode 100644 index 0000000..33faddb --- /dev/null +++ b/tests/loaders/test_readers.py @@ -0,0 +1,115 @@ +"""File readers: files in, Arrow record batches out, no database.""" + +from pathlib import Path + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest + +from orm_loader.conversion import NullPolicy +from orm_loader.loaders import DelimitedTextReader, ParquetReader, open_reader +from tests.models import DemoRecord + +COLUMNS = DemoRecord.model_columns() + + +def _write(tmp_path: Path, text: str, name: str = "demo_record.csv") -> Path: + path = tmp_path / name + path.write_text(text, encoding="utf-8", newline="") + return path + + +def _rows(reader) -> list[dict]: + return [row for batch in reader.batches() for row in batch.to_pylist()] + + +def test_open_reader_selects_the_reader_by_suffix(tmp_path): + pq.write_table(pa.table({"demo_id": [1]}), tmp_path / "demo_record.parquet") + assert isinstance(open_reader(tmp_path / "demo_record.parquet", COLUMNS), ParquetReader) + for name in ("demo_record.csv", "demo_record.tsv", "demo_record.txt"): + assert isinstance(open_reader(_write(tmp_path, "demo_id\n1\n", name), COLUMNS), DelimitedTextReader) + + +def test_delimited_text_is_read_as_raw_strings_of_the_model_columns_in_file_order(tmp_path): + path = _write(tmp_path, "demo_label,extra,demo_id\nalpha,x,1\n") + reader = DelimitedTextReader(path, COLUMNS) + assert reader.columns == ["demo_label", "demo_id"] + batch = next(iter(reader.batches())) + assert batch.schema.names == ["demo_label", "demo_id"] + assert batch.to_pylist() == [{"demo_label": "alpha", "demo_id": "1"}] + + +@pytest.mark.parametrize("delimiter", [",", "\t"]) +def test_the_delimiter_is_detected(tmp_path, delimiter): + path = _write(tmp_path, f"demo_id{delimiter}demo_label\n1{delimiter}alpha\n") + assert _rows(DelimitedTextReader(path, COLUMNS)) == [{"demo_id": "1", "demo_label": "alpha"}] + + +def test_headers_are_lowercased_and_stripped_of_quotes_and_hash_suffix(tmp_path): + path = _write(tmp_path, '"DEMO_ID","Demo_Label_hash"\n1,alpha\n') + assert DelimitedTextReader(path, COLUMNS).columns == ["demo_id", "demo_label"] + + +def test_an_unquoted_empty_field_is_null_and_a_quoted_one_is_an_empty_string(tmp_path): + path = _write(tmp_path, 'demo_id,demo_label\n1,\n2,""\n') + assert [r["demo_label"] for r in _rows(DelimitedTextReader(path, COLUMNS, quote_mode="csv"))] == [None, ""] + + +def test_embedded_newlines_and_tabs_survive(tmp_path): + path = _write(tmp_path, 'demo_id,demo_label\n1,"a\nb"\n2,"c\td"\n') + assert [r["demo_label"] for r in _rows(DelimitedTextReader(path, COLUMNS, quote_mode="csv"))] == ["a\nb", "c\td"] + + +def test_quote_modes_decide_what_double_quotes_mean(tmp_path): + path = _write(tmp_path, 'demo_id\tdemo_label\n1\t"x"\n', "demo_record.tsv") + label = lambda mode: _rows(DelimitedTextReader(path, COLUMNS, quote_mode=mode))[0]["demo_label"] # noqa: E731 + assert label("csv") == "x" + assert label("literal") == '"x"' + assert label("by_delimiter") == '"x"' + + +def test_a_file_is_read_in_several_batches_without_losing_rows(tmp_path): + lines = "".join(f"{i},label {i}\n" for i in range(1, 20_001)) + path = _write(tmp_path, "demo_id,demo_label\n" + lines) + reader = DelimitedTextReader(path, COLUMNS, block_size=64 << 10) + batches = list(reader.batches()) + assert len(batches) > 1 + assert [row["demo_id"] for batch in batches for row in batch.to_pylist()] == [str(i) for i in range(1, 20_001)] + + +def test_an_empty_file_has_no_content_and_a_header_only_file_has_no_rows(tmp_path): + empty = DelimitedTextReader(_write(tmp_path, "", "demo_record.csv"), COLUMNS) + assert empty.empty and list(empty.batches()) == [] + header_only = DelimitedTextReader(_write(tmp_path, "demo_id,demo_label\n", "other.csv"), COLUMNS) + assert not header_only.empty and list(header_only.batches()) == [] + + +def test_unquoted_markers_are_nulled_while_reading_but_quoted_ones_are_kept(tmp_path): + path = _write(tmp_path, 'demo_id,demo_label\n1,NULL\n2,"NULL"\n') + nulls = NullPolicy(text_markers=frozenset({"NULL"}), marker_quoting="unquoted_only") + labels = [r["demo_label"] for r in _rows(DelimitedTextReader(path, COLUMNS, nulls, quote_mode="csv"))] + assert labels == [None, "NULL"] + + +def test_parquet_keeps_stored_types_and_ignores_other_columns(tmp_path): + path = tmp_path / "demo_record.parquet" + pq.write_table(pa.table({"demo_id": [1, 2], "demo_label": ["a", None], "extra": [0, 0]}), path) + reader = ParquetReader(path, COLUMNS) + assert reader.columns == ["demo_id", "demo_label"] + batch = next(iter(reader.batches())) + assert batch.schema.types == [pa.int64(), pa.string()] + assert batch.to_pylist() == [{"demo_id": 1, "demo_label": "a"}, {"demo_id": 2, "demo_label": None}] + + +def test_parquet_is_read_in_batches_of_chunksize_rows(tmp_path): + path = tmp_path / "demo_record.parquet" + pq.write_table(pa.table({"demo_id": list(range(250))}), path) + sizes = [batch.num_rows for batch in ParquetReader(path, COLUMNS, chunksize=100).batches()] + assert sizes == [100, 100, 50] + + +def test_parquet_markers_are_nulled_while_reading(tmp_path): + path = tmp_path / "demo_record.parquet" + pq.write_table(pa.table({"demo_id": [1, 2], "demo_label": ["NULL", "kept"]}), path) + nulls = NullPolicy(text_markers=frozenset({"NULL"}), marker_quoting="unquoted_only") + assert [r["demo_label"] for r in _rows(ParquetReader(path, COLUMNS, nulls))] == [None, "kept"] diff --git a/tests/loaders/test_schema_translate_map.py b/tests/loaders/test_schema_translate_map.py index 9d6a15b..0ee4fa0 100644 --- a/tests/loaders/test_schema_translate_map.py +++ b/tests/loaders/test_schema_translate_map.py @@ -1,4 +1,4 @@ -"""End-to-end proof that load_csv() respects schema_translate_map with no +"""End-to-end proof that load_file() respects schema_translate_map with no caller-side workaround. This is the actual regression test for the bug this whole plan exists to fix, distinct from the backend unit tests in tests/backends/, which exercise the merge methods directly but never against @@ -23,13 +23,12 @@ from oa_configurator import Role as SchemaRole from orm_loader.backends import STAGING_SCHEMA -from orm_loader.loaders.loader_interface import PandasLoader from tests.conftest import schema_scoped_session from tests.models import SimpleTable, VocabSchemaTable -def test_load_csv_respects_non_default_schema_end_to_end(pg_db, tmp_path): +def test_load_file_respects_non_default_schema_end_to_end(pg_db, tmp_path): with schema_scoped_session(pg_db, SimpleTable.__table__, prefix="loader_primary") as (scoped, session): schema = scoped.schemas[SchemaRole.PRIMARY] csv_path = tmp_path / "test_table.csv" @@ -37,8 +36,8 @@ def test_load_csv_respects_non_default_schema_end_to_end(pg_db, tmp_path): [{"id": 1, "name": "alpha"}, {"id": 2, "name": "beta"}, {"id": 3, "name": "gamma"}] ).to_csv(csv_path, index=False, sep="\t") - inserted = SimpleTable.load_csv( - session, csv_path, dedupe=False, loader=PandasLoader(), staging_schema_tag=STAGING_SCHEMA + inserted = SimpleTable.load_file( + session, csv_path, dedupe=False, staging_schema_tag=STAGING_SCHEMA ) session.commit() @@ -58,7 +57,7 @@ def test_load_csv_respects_non_default_schema_end_to_end(pg_db, tmp_path): def test_replace_merge_respects_non_default_schema_end_to_end(pg_db, tmp_path): - """A second load_csv() call with merge_strategy="replace" against the + """A second load_file() call with merge_strategy="replace" against the same non-default schema. Proves the merge path itself, not just the initial insert-if-empty fast path, qualifies correctly.""" with schema_scoped_session(pg_db, SimpleTable.__table__, prefix="loader_replace") as (scoped, session): @@ -67,11 +66,10 @@ def test_replace_merge_respects_non_default_schema_end_to_end(pg_db, tmp_path): def _write_and_load(rows: list[dict], path_name: str) -> int: path = tmp_path / path_name pd.DataFrame(rows).to_csv(path, index=False, sep="\t") - return SimpleTable.load_csv( + return SimpleTable.load_file( session, path, dedupe=False, - loader=PandasLoader(), merge_strategy="replace", staging_schema_tag=STAGING_SCHEMA, ) @@ -88,7 +86,7 @@ def _write_and_load(rows: list[dict], path_name: str) -> int: assert rows == [(1, "alpha-updated"), (2, "beta")] -def test_load_csv_respects_non_primary_schema_tag_end_to_end(pg_db, tmp_path): +def test_load_file_respects_non_primary_schema_tag_end_to_end(pg_db, tmp_path): """Checks that the schema tag is correctly derived from the table's own schema (table.schema) for a non-primary tag.""" with schema_scoped_session( @@ -101,8 +99,8 @@ def test_load_csv_respects_non_primary_schema_tag_end_to_end(pg_db, tmp_path): csv_path, index=False, sep="\t" ) - inserted = VocabSchemaTable.load_csv( - session, csv_path, dedupe=False, loader=PandasLoader(), staging_schema_tag=STAGING_SCHEMA + inserted = VocabSchemaTable.load_file( + session, csv_path, dedupe=False, staging_schema_tag=STAGING_SCHEMA ) session.commit() diff --git a/tests/loaders/test_split_connection.py b/tests/loaders/test_split_connection.py index 3f183df..1529bd4 100644 --- a/tests/loaders/test_split_connection.py +++ b/tests/loaders/test_split_connection.py @@ -11,12 +11,11 @@ from sqlalchemy.exc import UnboundExecutionError from orm_loader.backends import staging_schema_claim -from orm_loader.loaders.loader_interface import PandasLoader from tests.conftest import schema_scoped_session from tests.models import SimpleTable, VocabSchemaTable -def test_vocab_csv_load_routes_every_session_bind_to_vocab_file(tmp_path): +def test_vocab_file_load_routes_every_session_bind_to_vocab_file(tmp_path): resolver = Resolver( StackConfig.for_session( connections={ @@ -59,12 +58,11 @@ def get_bind(self, mapper=None, clause=None, **kwargs): pd.DataFrame([{"id": 7, "name": "vocab-row"}]).to_csv(csv_path, index=False, sep="\t") sqlite_session = MapperRoutedSession() try: - VocabSchemaTable.load_csv( + VocabSchemaTable.load_file( sqlite_session, csv_path, merge_strategy="insert_if_empty", dedupe=False, - loader=PandasLoader(), staging_schema_tag=staging_schema_claim().schema_tag, ) sqlite_session.commit() @@ -81,7 +79,7 @@ def get_bind(self, mapper=None, clause=None, **kwargs): vocab_engine.dispose() -def test_load_csv_across_two_genuinely_separate_connections(pg_db, sqlite_session, tmp_path): +def test_load_file_across_two_genuinely_separate_connections(pg_db, sqlite_session, tmp_path): """``pg_db`` (real Postgres, primary connection) and ``sqlite_session`` (real SQLite, vocab connection, via the module-level ``sqlite_engine``/``sqlite_session`` fixtures) are two entirely different engines against two entirely @@ -97,10 +95,10 @@ def test_load_csv_across_two_genuinely_separate_connections(pg_db, sqlite_sessio # Interleaved on purpose: primary, then vocab, then primary again, so a # module-level cache keyed wrong (or reused across calls) would surface # as data landing in the wrong database. - SimpleTable.load_csv(primary_session, primary_csv, dedupe=False, loader=PandasLoader()) + SimpleTable.load_file(primary_session, primary_csv, dedupe=False) primary_session.commit() - VocabSchemaTable.load_csv(sqlite_session, vocab_csv, dedupe=False, loader=PandasLoader()) + VocabSchemaTable.load_file(sqlite_session, vocab_csv, dedupe=False) sqlite_session.commit() with scoped.engine.connect() as conn: diff --git a/tests/loaders/test_pandas_loader.py b/tests/loaders/test_stage_file.py similarity index 53% rename from tests/loaders/test_pandas_loader.py rename to tests/loaders/test_stage_file.py index b85e3c7..8e6d5e5 100644 --- a/tests/loaders/test_pandas_loader.py +++ b/tests/loaders/test_stage_file.py @@ -1,17 +1,18 @@ from pathlib import Path + import sqlalchemy as sa -from orm_loader.loaders.data_classes import LoaderContext -import pandas as pd -from orm_loader.loaders.loader_interface import PandasLoader -from tests.models import PandasLoaderTable + +from orm_loader.loaders import LoaderContext, stage_file +from tests.models import FileLoaderTable + def test_loader_context_fields(sqlite_session): ctx = LoaderContext( - tableclass=PandasLoaderTable, + tableclass=FileLoaderTable, session=sqlite_session, path=Path("file.csv"), - staging_table=PandasLoaderTable.__table__, + staging_table=FileLoaderTable.__table__, chunksize=100, normalise=False, dedupe=True, @@ -21,9 +22,9 @@ def test_loader_context_fields(sqlite_session): assert ctx.normalise is False assert ctx.dedupe is True -def test_pandas_loader_case_insensitive_headers(tmp_path, sqlite_session): +def test_stage_file_case_insensitive_headers(tmp_path, sqlite_session): """ - PandasLoader should accept mixed-case CSV headers and map them + The loader should accept mixed-case CSV headers and map them correctly to lowercase model columns. """ @@ -35,28 +36,28 @@ def test_pandas_loader_case_insensitive_headers(tmp_path, sqlite_session): ) ctx = LoaderContext( - tableclass=PandasLoaderTable, + tableclass=FileLoaderTable, session=sqlite_session, path=csv, - staging_table=PandasLoaderTable.__table__, + staging_table=FileLoaderTable.__table__, chunksize=100, normalise=True, dedupe=True, ) - n = PandasLoader.orm_file_load(ctx) + n = stage_file(ctx, ctx.tableclass._select_reader(ctx)) rows = sqlite_session.execute( - sa.text('SELECT id, value FROM test_pandas_loader ORDER BY id') + sa.text('SELECT id, value FROM test_file_loader ORDER BY id') ).all() assert n == 2 assert rows == [(1, "alpha"), (2, "beta")] -def test_pandas_loader_literal_mode_preserves_quotes(tmp_path, sqlite_session): +def test_stage_file_literal_mode_preserves_quotes(tmp_path, sqlite_session): """ - The pandas fallback must honour quote_mode. Under 'literal' it uses + The loader must honour quote_mode. Under 'literal' it uses QUOTE_NONE, so double-quotes are literal data and are preserved verbatim rather than being stripped as RFC-4180 field wrappers. """ @@ -69,31 +70,30 @@ def test_pandas_loader_literal_mode_preserves_quotes(tmp_path, sqlite_session): ) ctx = LoaderContext( - tableclass=PandasLoaderTable, + tableclass=FileLoaderTable, session=sqlite_session, path=csv, - staging_table=PandasLoaderTable.__table__, + staging_table=FileLoaderTable.__table__, chunksize=100, normalise=True, dedupe=True, quote_mode="literal", ) - n = PandasLoader.orm_file_load(ctx) + n = stage_file(ctx, ctx.tableclass._select_reader(ctx)) rows = sqlite_session.execute( - sa.text('SELECT id, value FROM test_pandas_loader ORDER BY id') + sa.text('SELECT id, value FROM test_file_loader ORDER BY id') ).all() assert n == 2 assert rows == [(1, '"hello"'), (2, '"world"')] -def test_pandas_loader_csv_mode_strips_quotes(tmp_path, sqlite_session): +def test_stage_file_csv_mode_strips_quotes(tmp_path, sqlite_session): """ - Contrast with the literal case above: the same file under 'csv' mode uses - QUOTE_MINIMAL, so the wrapping quotes are consumed. Proves quote_mode - actually flows through the fallback rather than being ignored. + Contrast with the literal case above: the same file under 'csv' mode reads + double quotes as field quoting, so the wrapping quotes are consumed. """ csv = tmp_path / "quoted.csv" @@ -104,27 +104,27 @@ def test_pandas_loader_csv_mode_strips_quotes(tmp_path, sqlite_session): ) ctx = LoaderContext( - tableclass=PandasLoaderTable, + tableclass=FileLoaderTable, session=sqlite_session, path=csv, - staging_table=PandasLoaderTable.__table__, + staging_table=FileLoaderTable.__table__, chunksize=100, normalise=True, dedupe=True, quote_mode="csv", ) - n = PandasLoader.orm_file_load(ctx) + n = stage_file(ctx, ctx.tableclass._select_reader(ctx)) rows = sqlite_session.execute( - sa.text('SELECT id, value FROM test_pandas_loader ORDER BY id') + sa.text('SELECT id, value FROM test_file_loader ORDER BY id') ).all() assert n == 2 assert rows == [(1, "hello"), (2, "world")] -def test_pandas_loader_by_delimiter_tab_preserves_quotes(tmp_path, sqlite_session): +def test_stage_file_by_delimiter_tab_preserves_quotes(tmp_path, sqlite_session): """ 'by_delimiter' derives the mode from the delimiter alone: a tab-delimited file resolves to 'literal', so embedded double-quotes are kept as data. @@ -138,41 +138,35 @@ def test_pandas_loader_by_delimiter_tab_preserves_quotes(tmp_path, sqlite_sessio ) ctx = LoaderContext( - tableclass=PandasLoaderTable, + tableclass=FileLoaderTable, session=sqlite_session, path=tsv, - staging_table=PandasLoaderTable.__table__, + staging_table=FileLoaderTable.__table__, chunksize=100, normalise=True, dedupe=True, quote_mode="by_delimiter", ) - n = PandasLoader.orm_file_load(ctx) + n = stage_file(ctx, ctx.tableclass._select_reader(ctx)) rows = sqlite_session.execute( - sa.text('SELECT id, value FROM test_pandas_loader ORDER BY id') + sa.text('SELECT id, value FROM test_file_loader ORDER BY id') ).all() assert n == 2 assert rows == [(1, '"hi"'), (2, "bye")] -def test_pandas_dedupe_internal(sqlite_session, tmp_path): - df = pd.DataFrame({ - "id": [1, 1, 2], - "value": [10, 10, 20], - }) - - ctx = LoaderContext( - tableclass=PandasLoaderTable, - session=sqlite_session, - path=Path("file.csv"), - staging_table=PandasLoaderTable.__table__, - chunksize=100, - normalise=False, - dedupe=True, - ) - - out = PandasLoader.dedupe(df, ctx) - assert len(out) == 2 +def test_file_without_data_rows_loads_nothing_whatever_its_header(tmp_path, sqlite_session): + for text in ("stub\n", "id,value\n", ""): + csv = tmp_path / "no_rows.csv" + csv.write_text(text) + ctx = LoaderContext( + tableclass=FileLoaderTable, + session=sqlite_session, + path=csv, + staging_table=FileLoaderTable.__table__, + ) + assert stage_file(ctx, ctx.tableclass._select_reader(ctx)) == 0 + assert ctx.report.rows_read == 0 diff --git a/tests/models.py b/tests/models.py index 3d6be90..dd00553 100644 --- a/tests/models.py +++ b/tests/models.py @@ -7,7 +7,7 @@ from oa_configurator import Role as SchemaRole from sqlalchemy.orm import declarative_base -from orm_loader.tables import CSVLoadableTableInterface +from orm_loader.tables import FileLoadableTableInterface Base = declarative_base() @@ -17,18 +17,15 @@ class Role(str, Enum): LAST_AUTHOR = "last author" -class Flag(str, Enum): - STANDARD = "S" - CLASSIFICATION = "C" -class PandasLoaderTable(CSVLoadableTableInterface, Base): - __tablename__ = "test_pandas_loader" +class FileLoaderTable(FileLoadableTableInterface, Base): + __tablename__ = "test_file_loader" __table_args__ = {"schema": SchemaRole.PRIMARY.value} id = sa.Column(sa.Integer, primary_key=True) value = sa.Column(sa.String, nullable=False) -class SimpleTable(Base, CSVLoadableTableInterface): +class SimpleTable(Base, FileLoadableTableInterface): __tablename__ = "test_table" __table_args__ = ( sa.Index("ix_test_table_name", "name"), @@ -39,7 +36,7 @@ class SimpleTable(Base, CSVLoadableTableInterface): name: so.Mapped[str] = so.mapped_column(sa.String, nullable=False) -class RequiredTable(Base, CSVLoadableTableInterface): +class RequiredTable(Base, FileLoadableTableInterface): __tablename__ = "required_table" __table_args__ = {"schema": SchemaRole.PRIMARY.value} @@ -47,7 +44,7 @@ class RequiredTable(Base, CSVLoadableTableInterface): name: so.Mapped[str] = so.mapped_column(sa.String, nullable=False) -class CompositeTable(Base, CSVLoadableTableInterface): +class CompositeTable(Base, FileLoadableTableInterface): __tablename__ = "composite_table" __table_args__ = {"schema": SchemaRole.PRIMARY.value} @@ -56,7 +53,7 @@ class CompositeTable(Base, CSVLoadableTableInterface): value: so.Mapped[str] = so.mapped_column(sa.String) -class CycleATable(Base, CSVLoadableTableInterface): +class CycleATable(Base, FileLoadableTableInterface): """Mutually-referencing pair with CycleBTable, for deferred-FK tests: DuckDB can't execute the ALTER TABLE ADD CONSTRAINT this cycle forces.""" @@ -69,7 +66,7 @@ class CycleATable(Base, CSVLoadableTableInterface): ) -class CycleBTable(Base, CSVLoadableTableInterface): +class CycleBTable(Base, FileLoadableTableInterface): __tablename__ = "cycle_b_table" __table_args__ = {"schema": SchemaRole.PRIMARY.value} @@ -79,13 +76,9 @@ class CycleBTable(Base, CSVLoadableTableInterface): ) -class EnumTable(Base, CSVLoadableTableInterface): - """A genuine sa.Enum column, unmodified -- register_column_cast_rule's - enum_type matches a raw value (e.g. "first author") by .value but stores - the matching member's .name, which is exactly sa.Enum's own default - column-storage convention -- no values_callable or other customisation - needed for this, the common case. - """ +class EnumTable(Base, FileLoadableTableInterface): + """A genuine sa.Enum column. Files hold the member names, which is what + sa.Enum stores by default.""" __tablename__ = "enum_table" __table_args__ = {"schema": SchemaRole.PRIMARY.value} @@ -94,11 +87,11 @@ class EnumTable(Base, CSVLoadableTableInterface): role: so.Mapped[Role | None] = so.mapped_column(sa.Enum(Role), nullable=True) -class ComputedColumnTable(Base, CSVLoadableTableInterface): +class ComputedColumnTable(Base, FileLoadableTableInterface): """A real, registered table with a computed column, for merge-method tests that need get_staging_table() to work. Unlike a bare __tablename__/__table__ pair, this actually implements - CSVLoadableTableInterface.""" + FileLoadableTableInterface.""" __tablename__ = "computed_column_table" __table_args__ = {"schema": SchemaRole.PRIMARY.value} @@ -109,20 +102,7 @@ class ComputedColumnTable(Base, CSVLoadableTableInterface): slug: so.Mapped[str | None] = so.mapped_column(sa.String, sa.Computed("lower(name)")) -class ImpliedEnumTable(Base, CSVLoadableTableInterface): - """A plain String column with no type-level enum signal at all -- the - OMOP CDM concept.standard_concept/invalid_reason shape register_column_cast_rule - exists to cover, since sa.Enum-keyed dispatch can never reach a column like this. - """ - - __tablename__ = "implied_enum_table" - __table_args__ = {"schema": SchemaRole.PRIMARY.value} - - id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) - flag: so.Mapped[str | None] = so.mapped_column(sa.String(1), nullable=True) - - -class ReferencedParentTable(Base, CSVLoadableTableInterface): +class ReferencedParentTable(Base, FileLoadableTableInterface): """Referenced by ReferencingChildTable *and* owns its own FK, for merge_replace's up-front DuckDB guard: a table in this shape can't be SET on its own FK column while still referenced, even with indexes @@ -138,7 +118,7 @@ class ReferencedParentTable(Base, CSVLoadableTableInterface): name: so.Mapped[str] = so.mapped_column(sa.String) -class UnreferencedGrandparentTable(Base, CSVLoadableTableInterface): +class UnreferencedGrandparentTable(Base, FileLoadableTableInterface): __tablename__ = "unreferenced_grandparent_table" __table_args__ = {"schema": SchemaRole.PRIMARY.value} @@ -146,7 +126,7 @@ class UnreferencedGrandparentTable(Base, CSVLoadableTableInterface): name: so.Mapped[str] = so.mapped_column(sa.String) -class ReferencingChildTable(Base, CSVLoadableTableInterface): +class ReferencingChildTable(Base, FileLoadableTableInterface): __tablename__ = "referencing_child_table" __table_args__ = {"schema": SchemaRole.PRIMARY.value} @@ -156,7 +136,7 @@ class ReferencingChildTable(Base, CSVLoadableTableInterface): ) -class VocabSchemaTable(Base, CSVLoadableTableInterface): +class VocabSchemaTable(Base, FileLoadableTableInterface): """A VOCAB-tagged table, so tests can prove the staging/index schema derivation actually resolves a non-primary schema tag correctly.""" @@ -170,7 +150,7 @@ class VocabSchemaTable(Base, CSVLoadableTableInterface): name: so.Mapped[str] = so.mapped_column(sa.String, nullable=False) -class VocabReferencingTable(Base, CSVLoadableTableInterface): +class VocabReferencingTable(Base, FileLoadableTableInterface): """PRIMARY-tagged table with a foreign key into VocabSchemaTable (VOCAB), for cross-schema foreign-key classification tests: DuckDB cannot enforce a key between two physical schemas, so oa-configurator omits @@ -186,7 +166,7 @@ class VocabReferencingTable(Base, CSVLoadableTableInterface): ) -class SelfReferentialTable(Base, CSVLoadableTableInterface): +class SelfReferentialTable(Base, FileLoadableTableInterface): """A table referencing its own rows, for dependency-ordered batching tests: a child row's parent may legitimately land in a later pagination batch of the same load. use_alter=True since DuckDB can't @@ -204,7 +184,7 @@ class SelfReferentialTable(Base, CSVLoadableTableInterface): ) -class DemoRecord(Base, CSVLoadableTableInterface): +class DemoRecord(Base, FileLoadableTableInterface): """Optional text label plus one column of each typed kind, for null-handling tests.""" __tablename__ = "demo_record" @@ -217,7 +197,7 @@ class DemoRecord(Base, CSVLoadableTableInterface): demo_date: so.Mapped[date | None] = so.mapped_column(sa.Date, nullable=True) -class DemoRequiredRecord(Base, CSVLoadableTableInterface): +class DemoRequiredRecord(Base, FileLoadableTableInterface): """Same shape as DemoRecord with a required text label.""" __tablename__ = "demo_required_record" diff --git a/tests/support.py b/tests/support.py new file mode 100644 index 0000000..3504912 --- /dev/null +++ b/tests/support.py @@ -0,0 +1,38 @@ +"""Helpers that run a column handler on a few values.""" + +from __future__ import annotations + +from typing import Any + +import pyarrow as pa +from sqlalchemy.types import TypeEngine + +from orm_loader.conversion import ConversionPolicy + + +def interpret( + values: pa.Array | pa.ChunkedArray | list[Any], + sa_type: TypeEngine[Any], + policy: ConversionPolicy | None = None, + *, + table_name: str | None = None, + column_name: str | None = None, +) -> tuple[list[Any], list[bool]]: + """Run the column's handler and return its values and invalid mask as lists.""" + policy = policy or ConversionPolicy() + arr = values if isinstance(values, (pa.Array, pa.ChunkedArray)) else pa.array(values) + handler = policy.for_column(sa_type, table_name, column_name) + out, invalid = handler.interpret(arr, sa_type, policy.nulls) + return out.to_pylist(), invalid.to_pylist() + + +def read(raw: Any, sa_type: TypeEngine[Any], policy: ConversionPolicy | None = None, **column: Any) -> Any: + """Return the value read from one raw value, failing if it is unreadable.""" + values, invalid = interpret([raw], sa_type, policy, **column) + assert not invalid[0], f"{raw!r} is unreadable as {sa_type!r}" + return values[0] + + +def unreadable(raw: Any, sa_type: TypeEngine[Any], policy: ConversionPolicy | None = None, **column: Any) -> bool: + """Return True when one raw value cannot be read as the column type.""" + return interpret([raw], sa_type, policy, **column)[1][0] diff --git a/tests/tables/test_typing_protocols.py b/tests/tables/test_typing_protocols.py index 0aa29c3..12ce537 100644 --- a/tests/tables/test_typing_protocols.py +++ b/tests/tables/test_typing_protocols.py @@ -2,10 +2,10 @@ import sqlalchemy.orm as so from orm_loader.tables.orm_table import ORMTableBase from orm_loader.tables.serialisable_table import SerialisableTableInterface -from orm_loader.tables.loadable_table import CSVLoadableTableInterface +from orm_loader.tables.loadable_table import FileLoadableTableInterface from orm_loader.tables.typing import ( ORMTableProtocol, - CSVTableProtocol, + FileTableProtocol, SerializedTableProtocol, ) @@ -38,13 +38,13 @@ class G(SerialisableTableInterface, ORMTableBase, Base): assert callable(obj.to_json) assert callable(obj.fingerprint) -def test_csv_table_protocol_conformance(): - class P(CSVLoadableTableInterface, ORMTableBase, Base): +def test_file_table_protocol_conformance(): + class P(FileLoadableTableInterface, ORMTableBase, Base): __tablename__ = "p" id = sa.Column(sa.Integer, primary_key=True) - assert isinstance(P, CSVTableProtocol) + assert isinstance(P, FileTableProtocol) # Required class-level API - assert callable(P.load_csv) + assert callable(P.load_file) assert callable(P.create_staging_table) assert callable(P.merge_from_staging) \ No newline at end of file diff --git a/uv.lock b/uv.lock index 439e964..5d15524 100644 --- a/uv.lock +++ b/uv.lock @@ -1,6 +1,10 @@ version = 1 -revision = 3 +revision = 5 requires-python = ">=3.12" +resolution-markers = [ + "python_full_version >= '3.14'", + "python_full_version < '3.14'", +] [[package]] name = "annotated-doc" @@ -152,6 +156,50 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, ] +[[package]] +name = "duckdb" +version = "1.5.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/59/0b/d65ea3be00ea79aa276a8388bec588a9cbf409ce637c6d306e5316210d15/duckdb-1.5.6.tar.gz", hash = "sha256:166a91dbfacfc0c9f08cc76c0243cb6d3d4296bfab5bad72a3cfb63140a5b7c8", size = 18032957, upload-time = "2026-09-28T13:38:37.978Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d9/d5/d0ab77a0a1702a43171c93874f44c1f6481e30038bd3987df0d77a16a5c6/duckdb-1.5.6-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:48d07d0651aaeac2c3974afd37599970154b7b79b54c18f27c319c14ccf98d9d", size = 32810486, upload-time = "2026-09-28T13:37:47.254Z" }, + { url = "https://files.pythonhosted.org/packages/9f/cd/b22201de5377faa3be6c38d5f3eaa504cb480392a448bed6a4d2239469b4/duckdb-1.5.6-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:79de3dfa8705b1ba0d59e7e3252e40ff399e0afd12f485502a6c7bf7c2fd809a", size = 17405278, upload-time = "2026-09-28T13:37:50.135Z" }, + { url = "https://files.pythonhosted.org/packages/9c/6d/f9cfb1493bbdc2f095693a402e42dce1192077f9e11573f00baed6a748de/duckdb-1.5.6-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:dcccce20965e6986cd083fdf192c461685ad0b93cd1ccd0b2a8207f1185f078b", size = 15532943, upload-time = "2026-09-28T13:37:52.927Z" }, + { url = "https://files.pythonhosted.org/packages/53/04/f65ccfaa5a833f2e570c4a140f03c8f95da416da9fe8ed08401f81f8242a/duckdb-1.5.6-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ce89a1025a5317ebe9c520876c48032b5247ac574865486648b1a004f6009875", size = 19454940, upload-time = "2026-09-28T13:37:55.732Z" }, + { url = "https://files.pythonhosted.org/packages/4c/99/be75c788a492f8d77b7a1cdc1b19939ae7be0007f2028691ad371a1a33ee/duckdb-1.5.6-cp312-cp312-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bc9619ed7d4ffa117b5155d84b44794366bb6635178d78ed5e13a6024845c757", size = 21568087, upload-time = "2026-09-28T13:37:58.191Z" }, + { url = "https://files.pythonhosted.org/packages/b5/95/889f8508960e47c0a7c75cc5bf57cde8512fc24f8db7b3129cca5388da42/duckdb-1.5.6-cp312-cp312-win_amd64.whl", hash = "sha256:09ff51b230219f0d8b47fc8a1e17fb595ba9fab0c3d96a6de4d00b8ff86b3cf1", size = 13190189, upload-time = "2026-09-28T13:38:00.407Z" }, + { url = "https://files.pythonhosted.org/packages/a4/c9/baab503364a68309f8368c88e77f5341e7d94927bdf3e6d703f0e5035f3e/duckdb-1.5.6-cp312-cp312-win_arm64.whl", hash = "sha256:b8d795c8b2d5634b3269f974aa97f1fdf878f62f032317a52252a151b693fb1e", size = 14021977, upload-time = "2026-09-28T13:38:02.682Z" }, + { url = "https://files.pythonhosted.org/packages/b1/5e/a476197fcba557738a588ec844747a19bc0a24b0e6f1809e308f29d68c0e/duckdb-1.5.6-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:ae352646374cacf48e9981cf031191c494865192fc436d13667a2531fc5d1da3", size = 32810376, upload-time = "2026-09-28T13:38:05.148Z" }, + { url = "https://files.pythonhosted.org/packages/0c/6d/5466a2b53ddd557644dfa47a763f68748efccdf282e6ae7c4f1bcfb3da69/duckdb-1.5.6-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:5a1261e90785e9d29953293e44f60fa073bd1137098924e8de21a037a861b051", size = 17405385, upload-time = "2026-09-28T13:38:07.363Z" }, + { url = "https://files.pythonhosted.org/packages/d4/a0/bf87071170835ee4a34fe764fc11c1c6e7040a0e021b36c1b6f834a4c22f/duckdb-1.5.6-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:97dd7a555b8f5298b76bc7d48a11cb2c64336e8de9bfde783cffb86ea9f54807", size = 15533132, upload-time = "2026-09-28T13:38:09.681Z" }, + { url = "https://files.pythonhosted.org/packages/31/e0/38095c8e140ecfbe847519ac07bcba94301b8fbb76b2870015e33e07f179/duckdb-1.5.6-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:364992ba1089a2b327391cfcb68fd0bd0ce9090cf293baef861a0ba6847abfee", size = 19454994, upload-time = "2026-09-28T13:38:11.836Z" }, + { url = "https://files.pythonhosted.org/packages/70/21/61dd2876bbaa69cf77d7b5c620e52e8b25faae7096f4d2e4a812b52095d7/duckdb-1.5.6-cp313-cp313-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:644f54ce99b3b61844bc9a3fe80e0aecb1ea4084b1fffc4396d1569db6111679", size = 21568700, upload-time = "2026-09-28T13:38:14.258Z" }, + { url = "https://files.pythonhosted.org/packages/4a/4a/100730e7785e85268be4d4d5bd62cfc8314e261d2f42efa208243eef35cb/duckdb-1.5.6-cp313-cp313-win_amd64.whl", hash = "sha256:ced693d33ddcee2e5345f077d342c87d2aaa80e41c514e64c9ff2d4e5963c251", size = 13190707, upload-time = "2026-09-28T13:38:16.875Z" }, + { url = "https://files.pythonhosted.org/packages/f3/2e/bc7f44eab4e89ee5c1cb427bb1168ad021d985042e6841ec0694c3d3d501/duckdb-1.5.6-cp313-cp313-win_arm64.whl", hash = "sha256:41ecc75bb9328d72d154a705c1a653d2c5c60f686a5c0c6578aa80020753c884", size = 14020962, upload-time = "2026-09-28T13:38:19.007Z" }, + { url = "https://files.pythonhosted.org/packages/fb/62/a8a30a4c6b94c0861d348ed5633b963f6745a5525527530f02f3c1a7c931/duckdb-1.5.6-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:aa21d2ad803b2524326e8622d7d96b2bb1ff1d5b60368e1978ee805df9c21fb3", size = 32828003, upload-time = "2026-09-28T13:38:21.414Z" }, + { url = "https://files.pythonhosted.org/packages/71/b7/1dcca0005eb8c67adf9fc06bf0cbb1d2bf4ea1974cc89e7a7c2ad66aac28/duckdb-1.5.6-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:8a1b2ad27d414068cbca06c55cfa802eece10f86ea4812ff082f8ab4cb25fc85", size = 17413912, upload-time = "2026-09-28T13:38:23.915Z" }, + { url = "https://files.pythonhosted.org/packages/93/b0/e3ac175443550f3464f2d95731a8b0aae9b4dc3875c3a186c352262b43c2/duckdb-1.5.6-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c79c6d222b1d015cde73b5139087186b00db65357fb4e2c94c2308fbbf465a72", size = 15543122, upload-time = "2026-09-28T13:38:26.317Z" }, + { url = "https://files.pythonhosted.org/packages/9d/08/cc510a7952aba69d5cdca17f3ef61c95713d86143f2ee9aa3e097d38f50b/duckdb-1.5.6-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1052b8050ef5696e2c0d8c836949c72f3dd11f0690466acbea739613e8e2750b", size = 19457946, upload-time = "2026-09-28T13:38:28.877Z" }, + { url = "https://files.pythonhosted.org/packages/ef/a5/6f8099d9a5a02ddff89e5c85875df3465054845b0920fb0703fbdf8dd2ec/duckdb-1.5.6-cp314-cp314-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:19c5e485e59613b8878d1670bcaa7a010f53c5a4da5ae8e08863e5e529ca6182", size = 21575132, upload-time = "2026-09-28T13:38:31.231Z" }, + { url = "https://files.pythonhosted.org/packages/9f/58/762f7159662d7859e201fa05ca29f306795daeabf84f3e087215a966b001/duckdb-1.5.6-cp314-cp314-win_amd64.whl", hash = "sha256:ebcbd09cd8578ab1093393e9b16289cda0e8f1791ac595bf00eb5bad75c3cf00", size = 13713963, upload-time = "2026-09-28T13:38:33.543Z" }, + { url = "https://files.pythonhosted.org/packages/46/69/64d165db322de13f5c3e75d377b6b9694df1821155ad1fa4b14b04601abc/duckdb-1.5.6-cp314-cp314-win_arm64.whl", hash = "sha256:820a8384faef11cd86068ea48c5da57ce2d8f1c7b3d2bdb9be3398317a7c3728", size = 14514368, upload-time = "2026-09-28T13:38:35.676Z" }, +] + +[[package]] +name = "duckdb-sqlalchemy" +version = "1.5.6.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "duckdb" }, + { name = "packaging" }, + { name = "pytz" }, + { name = "sqlalchemy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f8/fe/b2ced0727a5927aa893fae5fe95604ee030310604445a99e93dc237f0e41/duckdb_sqlalchemy-1.5.6.6.tar.gz", hash = "sha256:69fcb6a7a024dd2d49a9eeaebed587898ebdad26396c83a2dd78df7370d93e9b", size = 190836, upload-time = "2026-10-28T08:55:15.774Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/77/74/a0400aee4154a9bf2072eaaf6f35b8f4750d46784cb99063b356b4892951/duckdb_sqlalchemy-1.5.6.6-py3-none-any.whl", hash = "sha256:b4e4547ef3dcf46ea75964d56291d9f33195c3b03c0f3abe76fb7dfdbf5074ca", size = 76990, upload-time = "2026-10-28T08:55:14.278Z" }, +] + [[package]] name = "editorconfig" version = "0.17.1" @@ -559,7 +607,7 @@ wheels = [ [[package]] name = "oa-configurator" -version = "2.0.0" +version = "2.1.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, @@ -570,12 +618,15 @@ dependencies = [ { name = "typer" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/e9/62/64bf398b8fe381d138446199be69bc437ef93fdb70aab48cebf51467f6ea/oa_configurator-2.0.0.tar.gz", hash = "sha256:26406e90aeadd4c747646763b80f16025c4c6638afec687517788bc537d0d581", size = 278382, upload-time = "2026-10-10T22:48:42.275Z" } +sdist = { url = "https://files.pythonhosted.org/packages/5a/5a/e1061864a6f9abc9c49e26ab84a51dfc8f40616906c853e01ccbc022b477/oa_configurator-2.1.0.tar.gz", hash = "sha256:89e33ab58bd428081b1c1c8ce1b4ad7d4a7d4b4627bcbc20d8cd6502811f627c", size = 298778, upload-time = "2026-10-26T21:01:07.87Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/06/9b/99f0073d97de3545d48cf28c8f39b724d543baf83e8af8d0ba5c8bffbb01/oa_configurator-2.0.0-py3-none-any.whl", hash = "sha256:2186beb5567850ba5bfe4950c2820b909d3746c6fa233a6074a85396a343af35", size = 118027, upload-time = "2026-10-10T22:48:40.571Z" }, + { url = "https://files.pythonhosted.org/packages/59/44/38cdd910fe7544b118d7a5fc741417f1be0a4a95fdb754d8d5f3a016b9f9/oa_configurator-2.1.0-py3-none-any.whl", hash = "sha256:44c2fc43cf01398e598fc4d0fd4b17f9a3090cc35b6a378cdaf8e40fbfab0a0a", size = 126679, upload-time = "2026-10-26T21:01:05.893Z" }, ] [package.optional-dependencies] +duckdb = [ + { name = "duckdb-sqlalchemy" }, +] postgres = [ { name = "psycopg", extra = ["binary"] }, ] @@ -586,7 +637,6 @@ source = { editable = "." } dependencies = [ { name = "chardet" }, { name = "oa-configurator" }, - { name = "pandas" }, { name = "pyarrow" }, { name = "sqlalchemy" }, ] @@ -597,7 +647,8 @@ dev = [ { name = "mkdocs-material" }, { name = "mkdocs-mermaid2-plugin" }, { name = "mkdocstrings-python" }, - { name = "oa-configurator", extra = ["postgres"] }, + { name = "oa-configurator", extra = ["duckdb", "postgres"] }, + { name = "pandas" }, { name = "pygments" }, { name = "pytest" }, { name = "python-dotenv" }, @@ -605,6 +656,9 @@ dev = [ { name = "ruff" }, { name = "ty" }, ] +duckdb = [ + { name = "duckdb-sqlalchemy" }, +] postgres = [ { name = "psycopg", extra = ["binary"] }, ] @@ -612,13 +666,14 @@ postgres = [ [package.metadata] requires-dist = [ { name = "chardet", specifier = ">=5.2.0" }, + { name = "duckdb-sqlalchemy", marker = "extra == 'duckdb'", specifier = ">=1.5" }, { name = "mkdocs", marker = "extra == 'dev'", specifier = ">=1.6.1" }, { name = "mkdocs-material", marker = "extra == 'dev'", specifier = ">=9.7.7" }, { name = "mkdocs-mermaid2-plugin", marker = "extra == 'dev'" }, { name = "mkdocstrings-python", marker = "extra == 'dev'", specifier = ">=2.0.1" }, - { name = "oa-configurator", specifier = ">=2.0.0,<3" }, - { name = "oa-configurator", extras = ["postgres"], marker = "extra == 'dev'", specifier = ">=2.0.0,<3" }, - { name = "pandas", specifier = ">=2.3.3" }, + { name = "oa-configurator", specifier = ">=2.1.0,<3" }, + { name = "oa-configurator", extras = ["duckdb", "postgres"], marker = "extra == 'dev'", specifier = ">=2.0.0,<3" }, + { name = "pandas", marker = "extra == 'dev'", specifier = ">=2.3.3" }, { name = "psycopg", extras = ["binary"], marker = "extra == 'postgres'", specifier = ">=3.2" }, { name = "pyarrow", specifier = ">=23.0.1" }, { name = "pygments", marker = "extra == 'dev'", specifier = ">=2.20.0" }, @@ -629,7 +684,7 @@ requires-dist = [ { name = "sqlalchemy", specifier = ">=2.0.45" }, { name = "ty", marker = "extra == 'dev'", specifier = ">=0.0.59" }, ] -provides-extras = ["dev", "postgres"] +provides-extras = ["dev", "duckdb", "postgres"] [[package]] name = "packaging" From 1f8916087b05450eecf2898d334f14a69d1efda9 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Thu, 29 Oct 2026 01:33:06 +0000 Subject: [PATCH 04/21] Remove deprecated bulk module --- src/orm_loader/helpers/__init__.py | 3 -- src/orm_loader/helpers/bulk.py | 46 ------------------------------ 2 files changed, 49 deletions(-) delete mode 100644 src/orm_loader/helpers/bulk.py diff --git a/src/orm_loader/helpers/__init__.py b/src/orm_loader/helpers/__init__.py index 55533f9..be24464 100644 --- a/src/orm_loader/helpers/__init__.py +++ b/src/orm_loader/helpers/__init__.py @@ -1,7 +1,6 @@ from .errors import ConversionError, IngestError, RejectLimitError, RequiredValueError, ValidationError from .bootstrap import bootstrap, create_db, create_tables from .engines import dispose_engines -from .bulk import bulk_load_context, engine_with_replica_role from .metadata import Base from .discovery import get_model_by_tablename @@ -15,8 +14,6 @@ "create_db", "create_tables", "dispose_engines", - "bulk_load_context", - "engine_with_replica_role", "Base", "get_model_by_tablename", ] diff --git a/src/orm_loader/helpers/bulk.py b/src/orm_loader/helpers/bulk.py deleted file mode 100644 index d425b09..0000000 --- a/src/orm_loader/helpers/bulk.py +++ /dev/null @@ -1,46 +0,0 @@ -from contextlib import contextmanager -from sqlalchemy import Engine -from sqlalchemy.orm import Session -from typing import Generator - - -@contextmanager -def bulk_load_context( - session: Session, - *, - disable_fk: bool = True, - no_autoflush: bool = True, -) -> Generator[None, None, None]: - """ - Wrap a trusted bulk operation in backend-aware session settings. - - Foreign-key checks are disabled and autoflush is suspended for the - duration, then restored. See ``DatabaseBackend.bulk_load_context``. - """ - from ..backends.resolve import resolve_backend - backend = resolve_backend(session) - with backend.bulk_load_context( - session, - disable_fk=disable_fk, - no_autoflush=no_autoflush, - ): - yield - - -@contextmanager -def engine_with_replica_role(engine: Engine) -> Generator[Engine, None, None]: - """ - Force ``session_replication_role=replica`` on PostgreSQL engine connections. - - 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: - raise NotImplementedError( - f"Backend '{backend.name}' does not support replica-role engine contexts" - ) - with method(engine) as wrapped: - yield wrapped From e0b4b6f5de71b356ac00d14bb5310f883f6166a7 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Thu, 29 Oct 2026 02:27:54 +0000 Subject: [PATCH 05/21] Foreign key validation routine that is used after loading of large files --- docs/loaders/staging.md | 9 ++ src/orm_loader/backends/__init__.py | 2 + src/orm_loader/backends/base.py | 159 ++++++++++++++++++++++-- src/orm_loader/backends/duckdb.py | 71 ++--------- src/orm_loader/tables/loadable_table.py | 50 +++++++- src/orm_loader/tables/typing.py | 10 ++ tests/backends/test_foreign_keys.py | 81 ++++++++++++ 7 files changed, 305 insertions(+), 77 deletions(-) create mode 100644 tests/backends/test_foreign_keys.py diff --git a/docs/loaders/staging.md b/docs/loaders/staging.md index 645d27a..9a81bae 100644 --- a/docs/loaders/staging.md +++ b/docs/loaders/staging.md @@ -34,3 +34,12 @@ A failure in the sink is raised. There is no fallback to another loading path. With `dedupe` set, the staging table keeps the first row for each primary key, in source order. The backend removes the others after the load and reports a `DedupeResult` on `LoadReport` with the rows before and the rows removed. The merge step decides what happens when incoming rows overlap with existing target rows. No schema changes are performed at the loader layer. + +--- + +## Referential integrity + +On PostgreSQL and SQLite a load relaxes foreign keys while it merges, so a row loaded before the row it references is accepted. After every table is loaded, `Table.validate_foreign_keys(session)` checks every foreign key of the table with an anti-join and raises `ForeignKeyViolationError` with up to five unmatched keys. Null references are not violations, and the load order does not matter. + +DuckDB enforces the keys it creates, so load parent tables first. `load_file` already validates the keys DuckDB never enforces (cross-schema, cyclic and self-referential) after each table's merge, and `validate_foreign_keys` repeats that check, so it works the same way on every backend. A DuckDB table with foreign keys needs `resolved=` for both. + diff --git a/src/orm_loader/backends/__init__.py b/src/orm_loader/backends/__init__.py index 4aae24a..c076d13 100644 --- a/src/orm_loader/backends/__init__.py +++ b/src/orm_loader/backends/__init__.py @@ -7,6 +7,7 @@ DatabaseBackend, DedupeResult, Dialect, + ForeignKeyViolationError, STAGING_SCHEMA, staging_schema_claim, ) @@ -24,6 +25,7 @@ "DedupeResult", "Dialect", "DuckDBBackend", + "ForeignKeyViolationError", "MaterializationError", "MaterializationFailure", "MaterializationOperation", diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index f151d5a..3670489 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -5,12 +5,22 @@ from contextlib import AbstractContextManager, contextmanager from dataclasses import dataclass from functools import wraps -from typing import TYPE_CHECKING, Any, Callable, ParamSpec, Type, TypeVar, cast +from typing import TYPE_CHECKING, Any, Callable, ClassVar, ParamSpec, Type, TypeVar, cast import pyarrow as pa import sqlalchemy as sa import sqlalchemy.orm as so -from oa_configurator import Dialect, ResolvedDatabase, SchemaClaim, open_connection, qualified +from oa_configurator import ( + Dialect, + ResolvedCDMDatabase, + ResolvedDatabase, + Role, + SchemaClaim, + is_ephemeral_url, + open_connection, + qualified, + unenforced_foreign_keys, +) from sqlalchemy.engine import Connection, Engine from sqlalchemy.sql.compiler import IdentifierPreparer @@ -59,22 +69,36 @@ class BackendCapabilities: supports_replace_with_kept_indexes: bool = False -class DeferredForeignKeyViolationError(RuntimeError): - """A foreign key this backend can't enforce at the DB level is violated. +class ForeignKeyViolationError(RuntimeError): + """Rows reference a row that does not exist in the referred table. + + Parameters + ---------- + constraint + The violated foreign key. + orphan_keys + Up to five referencing key values that have no match. """ + qualifier: ClassVar[str] = "" + def __init__(self, constraint: sa.ForeignKeyConstraint, orphan_keys: list[tuple[Any, ...]]) -> None: self.constraint = constraint self.orphan_keys = orphan_keys local_cols = ", ".join(c.name for c in constraint.columns) referred_table = next(iter(constraint.elements)).column.table.name super().__init__( - f"{constraint.table.name}.({local_cols}) references {referred_table!r} via a " - f"constraint this backend can't enforce at the DB level; {len(orphan_keys)} " - f"row(s) violate it (showing up to 5): {orphan_keys}" + f"{constraint.table.name}.({local_cols}) references {referred_table!r}{self.qualifier}; " + f"{len(orphan_keys)} row(s) violate it (showing up to 5): {orphan_keys}" ) +class DeferredForeignKeyViolationError(ForeignKeyViolationError): + """A foreign key this backend can't enforce at the DB level is violated.""" + + qualifier = " via a constraint this backend can't enforce at the DB level" + + STAGING_SCHEMA: str = "staging" @@ -355,24 +379,133 @@ def validate_deferred_fks( referenced table. """ for fk in self.deferred_fk_constraints(table_cls): - self._validate_one_deferred_fk(fk, session) + self._validate_one_fk(fk, session, DeferredForeignKeyViolationError) + + def validate_foreign_keys( + self, + table_cls: Type["FileTableProtocol"], + session: so.Session, + ) -> None: + """Raise if any row in *table_cls* references a row that does not exist. + + Checks every foreign key of the table with an anti-join, whether or + not the database enforces it. Loads relax foreign keys while merging + on backends that can toggle them, so this is how referential + integrity is verified once all tables are loaded, in any load order. + Null references are not violations. + + Raises + ------ + ForeignKeyViolationError + If any row's FK columns are non-null with no matching row in the + referred table. + """ + own_fks = tuple(table_cls.__table__.foreign_key_constraints) + cross_database: set[sa.ForeignKeyConstraint] = set() + if own_fks and self.resolved is not None: + metadata = table_cls.__table__.metadata + cross_database = { + entry.constraint + for entry in unenforced_foreign_keys( + tuple(metadata.tables.values()), resolved=self.resolved, physical_schemas=self.physical_schemas + ) + if entry.reason == "cross_database" + } + for fk in own_fks: + if fk in cross_database: + self._validate_cross_database_fk(fk, session) + else: + self._validate_one_fk(fk, session) + + def _validate_cross_database_fk( + self, + fk: sa.ForeignKeyConstraint, + session: so.Session, + error: type[ForeignKeyViolationError] = ForeignKeyViolationError, + ) -> None: + """Anti-join fk across two connections by comparing both sides in Python. + + Pulls fk's distinct non-null local values from *session*'s own + connection, then opens the referred table's own connection (via + self.resolved) to check which of them actually exist there. One + bulk query per side instead of one query spanning both. + + Raises + ------ + TypeError + self.resolved is not a ResolvedCDMDatabase: a cross-database FK + can only arise when PRIMARY/VOCAB are split across connections, + which only a CDM-aware resolved can describe. + ValueError + The remote role's connection is ephemeral (DuckDB ``:memory:``): + a freshly built engine would get its own, empty database rather + than the one already holding the remote data. + ForeignKeyViolationError + A local value has no match on the referred connection. + """ + if not isinstance(self.resolved, ResolvedCDMDatabase): + raise TypeError( + f"{fk.name} was classified cross_database, which requires split PRIMARY/VOCAB " + f"connections, but self.resolved is {type(self.resolved).__name__}, not " + "ResolvedCDMDatabase." + ) + local_cols = list(fk.columns) + remote_cols = [element.column for element in fk.elements] + not_null = sa.and_(*(col.isnot(None) for col in local_cols)) + local_values = session.execute(sa.select(*local_cols).where(not_null).distinct()).all() + if not local_values: + return + + remote_tag = remote_cols[0].table.schema + remote_role = Role.VOCAB if remote_tag == Role.VOCAB.value else Role.PRIMARY + remote_connection = self.resolved.connection_for_role(remote_role) + if is_ephemeral_url(remote_connection.safe_url): + raise ValueError( + f"{fk.name} crosses to {remote_connection.name!r}, an ephemeral connection " + "(DuckDB ':memory:'): a freshly built engine would see an empty, disconnected " + "database instead of the data already loaded there. Cross-database FK " + "validation needs a real, file-backed or server-backed remote connection." + ) + from ..helpers.engines import dispose_engines + + primary_engine, vocab_engine = self.resolved.create_engines(register_claims=False) + try: + remote_engine = vocab_engine if remote_tag == Role.VOCAB.value else primary_engine + with remote_engine.connect() as remote_conn: + if len(remote_cols) == 1: + condition = remote_cols[0].in_(value for (value,) in local_values) + else: + condition = sa.tuple_(*remote_cols).in_(tuple(value) for value in local_values) + existing = {tuple(row) for row in remote_conn.execute(sa.select(*remote_cols).where(condition))} + finally: + dispose_engines(primary_engine, vocab_engine) + + missing = [tuple(value) for value in local_values if tuple(value) not in existing] + if missing: + raise error(fk, missing[:5]) @staticmethod - def _validate_one_deferred_fk(fk: sa.ForeignKeyConstraint, session: so.Session) -> None: + def _validate_one_fk( + fk: sa.ForeignKeyConstraint, + session: so.Session, + error: type[ForeignKeyViolationError] = ForeignKeyViolationError, + ) -> None: """Raise if any row of fk's local table has a non-null value with no match in the referred table. Parameters ---------- fk : sa.ForeignKeyConstraint - The deferred FK constraint to check. + The FK constraint to check. session : so.Session Session to run the anti-join query against. + error + Exception raised on a violation. Raises ------ - DeferredForeignKeyViolationError + ForeignKeyViolationError If any row's FK column is non-null with no matching row in the - referenced table. + referenced table (``error`` is raised, a subclass). Notes ----- @@ -394,7 +527,7 @@ def _validate_one_deferred_fk(fk: sa.ForeignKeyConstraint, session: so.Session) sa.select(*local_cols).where(not_null, ~orphan_exists).limit(5) ).all() if orphans: - raise DeferredForeignKeyViolationError(fk, [tuple(row) for row in orphans]) + raise error(fk, [tuple(row) for row in orphans]) @abstractmethod def create_staging_table( diff --git a/src/orm_loader/backends/duckdb.py b/src/orm_loader/backends/duckdb.py index 02f494e..36c9e42 100644 --- a/src/orm_loader/backends/duckdb.py +++ b/src/orm_loader/backends/duckdb.py @@ -7,10 +7,7 @@ import sqlalchemy as sa import sqlalchemy.orm as so from oa_configurator import ( - ResolvedCDMDatabase, - Role, UnenforcedForeignKey, - is_ephemeral_url, physical_schema_of, qualified, unenforced_foreign_keys, @@ -18,7 +15,6 @@ from sqlalchemy.schema import CreateSequence, DropSequence, DropTable from sqlalchemy.sql.compiler import IdentifierPreparer -from ..helpers.engines import dispose_engines from .base import BackendCapabilities, DatabaseBackend, DeferredForeignKeyViolationError, Dialect if TYPE_CHECKING: @@ -116,69 +112,18 @@ def validate_deferred_fks(self, table_cls: Type["FileTableProtocol"], session: s if entry.constraint not in own_fks: continue if entry.reason == "cross_database": - self._validate_cross_database_fk(entry.constraint, session) + self._validate_cross_database_fk(entry.constraint, session, DeferredForeignKeyViolationError) else: - self._validate_one_deferred_fk(entry.constraint, session) + self._validate_one_fk(entry.constraint, session, DeferredForeignKeyViolationError) - def _validate_cross_database_fk(self, fk: sa.ForeignKeyConstraint, session: so.Session) -> None: - """Anti-join fk across two connections by comparing both sides in Python. + def validate_foreign_keys(self, table_cls: Type["FileTableProtocol"], session: so.Session) -> None: + """Raise if any row in *table_cls* references a row that does not exist. - Pulls fk's distinct non-null local values from *session*'s own - connection, then opens the referred table's own connection (via - self.resolved) to check which of them actually exist there. One - bulk query per side instead of one query spanning both. - - Raises - ------ - TypeError - self.resolved is not a ResolvedCDMDatabase: a cross-database FK - can only arise when PRIMARY/VOCAB are split across connections, - which only a CDM-aware resolved can describe. - ValueError - The remote role's connection is ephemeral (DuckDB ``:memory:``): - a freshly built engine would get its own, empty database rather - than the one already holding the remote data. - DeferredForeignKeyViolationError - A local value has no match on the referred connection. + DuckDB enforces the foreign keys it creates, so only the ones it + never enforces need an anti-join, which is + :meth:`validate_deferred_fks`. """ - if not isinstance(self.resolved, ResolvedCDMDatabase): - raise TypeError( - f"{fk.name} was classified cross_database, which requires split PRIMARY/VOCAB " - f"connections, but self.resolved is {type(self.resolved).__name__}, not " - "ResolvedCDMDatabase." - ) - local_cols = list(fk.columns) - remote_cols = [element.column for element in fk.elements] - not_null = sa.and_(*(col.isnot(None) for col in local_cols)) - local_values = session.execute(sa.select(*local_cols).where(not_null).distinct()).all() - if not local_values: - return - - remote_tag = remote_cols[0].table.schema - remote_role = Role.VOCAB if remote_tag == Role.VOCAB.value else Role.PRIMARY - remote_connection = self.resolved.connection_for_role(remote_role) - if is_ephemeral_url(remote_connection.safe_url): - raise ValueError( - f"{fk.name} crosses to {remote_connection.name!r}, an ephemeral connection " - "(DuckDB ':memory:'): a freshly built engine would see an empty, disconnected " - "database instead of the data already loaded there. Cross-database FK " - "validation needs a real, file-backed or server-backed remote connection." - ) - primary_engine, vocab_engine = self.resolved.create_engines(register_claims=False) - try: - remote_engine = vocab_engine if remote_tag == Role.VOCAB.value else primary_engine - with remote_engine.connect() as remote_conn: - if len(remote_cols) == 1: - condition = remote_cols[0].in_(value for (value,) in local_values) - else: - condition = sa.tuple_(*remote_cols).in_(tuple(value) for value in local_values) - existing = {tuple(row) for row in remote_conn.execute(sa.select(*remote_cols).where(condition))} - finally: - dispose_engines(primary_engine, vocab_engine) - - missing = [tuple(value) for value in local_values if tuple(value) not in existing] - if missing: - raise DeferredForeignKeyViolationError(fk, missing[:5]) + self.validate_deferred_fks(table_cls, session) def _tables_referencing(self, target: sa.Table) -> set[sa.Table]: """Other tables with an enforced FK pointing at target. diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index c6261d9..c277c58 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -385,7 +385,8 @@ def load_file( Also run before every intermediate commit a batched merge makes, so a violation is caught before that batch commits rather than only after the whole load, when an earlier batch's commit can - no longer be rolled back. + no longer be rolled back. Every other key is verified by + ``validate_foreign_keys``. resolved Topology this table's engine(s) were resolved from. Required on DuckDB whenever table_cls has any foreign key at all, to @@ -511,6 +512,53 @@ def _target_has_rows( return row is not None + @classmethod + def validate_foreign_keys( + cls: Type[FileTableProtocol], + session: so.Session, + *, + staging_schema_tag: str | None = None, + resolved: ResolvedDatabase | None = None, + physical_schemas: Mapping[str, str | None] | None = None, + ) -> None: + """ + Raise if any row references a row that does not exist. + + Loads relax foreign keys during the merge on PostgreSQL and SQLite, + so a row loaded before the row it references is accepted. Call this + once every table is loaded to verify referential integrity, + whatever the load order. DuckDB enforces most of its keys itself, so + only the keys it cannot enforce are checked. + + Parameters + ---------- + session + An active SQLAlchemy session. + staging_schema_tag + schema_translate_map tag used to resolve the backend, as in + ``load_file``. + resolved + Topology this table's engine(s) were resolved from. Required on + DuckDB for a table with foreign keys. + physical_schemas + A schema tag's physical name beyond resolved's own configured + ones. + + Raises + ------ + ForeignKeyViolationError + A non-null reference has no matching row in the referred table. + """ + _require_bind(session, cls) + backend = resolve_backend( + session, + staging_schema_tag=staging_schema_tag, + mapper=cls, + resolved=resolved, + physical_schemas=physical_schemas, + ) + backend.validate_foreign_keys(cls, session) + @classmethod def merge_from_staging( cls: Type[FileTableProtocol], diff --git a/src/orm_loader/tables/typing.py b/src/orm_loader/tables/typing.py index c03e21d..12cd45f 100644 --- a/src/orm_loader/tables/typing.py +++ b/src/orm_loader/tables/typing.py @@ -118,6 +118,16 @@ def merge_from_staging( @classmethod def drop_staging_table(cls, session: so.Session, *, staging_schema_tag: str | None = None) -> None: ... + @classmethod + def validate_foreign_keys( + cls, + session: so.Session, + *, + staging_schema_tag: str | None = None, + resolved: Optional["ResolvedDatabase"] = None, + physical_schemas: Optional[Mapping[str, str | None]] = None, + ) -> None: ... + @classmethod def _target_has_rows(cls, session: so.Session) -> bool: ... diff --git a/tests/backends/test_foreign_keys.py b/tests/backends/test_foreign_keys.py new file mode 100644 index 0000000..21c60b3 --- /dev/null +++ b/tests/backends/test_foreign_keys.py @@ -0,0 +1,81 @@ +"""Referential integrity after a load: the same check on every dialect.""" + +from pathlib import Path +from typing import Type, cast + +import pytest + +from oa_configurator import Dialect +from orm_loader.backends import STAGING_SCHEMA, ForeignKeyViolationError +from orm_loader.tables.typing import FileTableProtocol +from tests.models import ( + ReferencedParentTable, + ReferencingChildTable, + SelfReferentialTable, + SimpleTable, + UnreferencedGrandparentTable, +) + +_Child = cast(Type[FileTableProtocol], ReferencingChildTable) +_Parent = cast(Type[FileTableProtocol], ReferencedParentTable) +_Grandparent = cast(Type[FileTableProtocol], UnreferencedGrandparentTable) +_Self = cast(Type[FileTableProtocol], SelfReferentialTable) +_Simple = cast(Type[FileTableProtocol], SimpleTable) + + +@pytest.fixture +def options(dialect_session, request): + """Keyword arguments every load and validation on this dialect needs.""" + if dialect_session.get_bind().dialect.name == "sqlite": + return {"staging_schema_tag": None} + options = {"staging_schema_tag": STAGING_SCHEMA} + if dialect_session.get_bind().dialect.name == Dialect.DUCKDB: + options["resolved"] = request.getfixturevalue("duckdb_db").resolved + return options + + +def _load(session, table, tmp_path: Path, text: str, options: dict, **extra) -> None: + path = tmp_path / f"{table.__tablename__}.csv" + path.write_text(text) + table.load_file(session, path, dedupe=False, **options, **extra) + session.commit() + + +def _validate(session, table, options: dict) -> None: + table.validate_foreign_keys(session, **options) + + +def test_every_reference_resolving_passes(dialect_session, tmp_path, options): + _load(dialect_session, _Grandparent, tmp_path, "id,name\n1,g\n", options) + _load(dialect_session, _Parent, tmp_path, "id,grandparent_id,name\n1,1,p\n", options) + _load(dialect_session, _Child, tmp_path, "id,parent_id\n1,1\n2,1\n", options) + for table in (_Grandparent, _Parent, _Child): + _validate(dialect_session, table, options) + + +def test_a_table_without_foreign_keys_passes(dialect_session, tmp_path, options): + _load(dialect_session, _Simple, tmp_path, "id,name\n1,a\n", options) + _validate(dialect_session, _Simple, options) + + +def test_null_references_are_not_violations(dialect_session, tmp_path, options): + _load(dialect_session, _Self, tmp_path, "id,parent_id\n1,\n2,1\n", options) + _validate(dialect_session, _Self, options) + + +def test_an_orphan_is_reported_with_its_key(dialect_session, tmp_path, options): + _load(dialect_session, _Self, tmp_path, "id,parent_id\n1,\n2,99\n", options, validate_deferred_fks=False) + with pytest.raises(ForeignKeyViolationError, match=r"self_referential_table.*\[\(99,\)\]"): + _validate(dialect_session, _Self, options) + + +def test_tables_loaded_child_first_validate_once_the_parents_are_loaded(dialect_session, tmp_path, options): + if dialect_session.get_bind().dialect.name == Dialect.DUCKDB: + pytest.skip("DuckDB enforces this key itself, so a child cannot be loaded first") + _load(dialect_session, _Child, tmp_path, "id,parent_id\n1,1\n", options) + with pytest.raises(ForeignKeyViolationError, match="referencing_child_table"): + _validate(dialect_session, _Child, options) + _load(dialect_session, _Parent, tmp_path, "id,grandparent_id,name\n1,1,p\n", options) + _load(dialect_session, _Grandparent, tmp_path, "id,name\n1,g\n", options) + for table in (_Grandparent, _Parent, _Child): + _validate(dialect_session, table, options) From 0c25ff85b3b6e978605b45d4b83bc044ff074ece Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Thu, 29 Oct 2026 04:31:48 +0000 Subject: [PATCH 06/21] Address first PR review: integer width, column defaults, typed input checks, quote mode, malformed lines --- docs/loaders/context.md | 2 +- docs/loaders/conversion.md | 30 +++-- docs/loaders/readers.md | 7 +- docs/loaders/staging.md | 3 +- pyproject.toml | 6 +- src/orm_loader/backends/base.py | 34 +++-- src/orm_loader/backends/duckdb.py | 39 +++--- src/orm_loader/backends/postgres.py | 29 +++-- src/orm_loader/backends/sqlite.py | 15 +-- src/orm_loader/conversion/__init__.py | 2 + .../conversion/handlers/__init__.py | 3 +- src/orm_loader/conversion/handlers/base.py | 15 ++- src/orm_loader/conversion/handlers/dates.py | 58 +++++---- src/orm_loader/conversion/handlers/numbers.py | 68 ++++++++-- src/orm_loader/conversion/handlers/text.py | 16 ++- src/orm_loader/conversion/policy.py | 43 +++--- src/orm_loader/conversion/values.py | 52 +++----- src/orm_loader/helpers/__init__.py | 3 +- src/orm_loader/helpers/errors.py | 25 ++++ src/orm_loader/loaders/__init__.py | 2 +- src/orm_loader/loaders/loader.py | 26 ++-- src/orm_loader/loaders/processor.py | 43 ++++-- src/orm_loader/loaders/readers/__init__.py | 8 +- src/orm_loader/loaders/readers/base.py | 23 +++- src/orm_loader/loaders/readers/delimited.py | 36 ++--- src/orm_loader/loaders/readers/detection.py | 112 ++++------------ src/orm_loader/loaders/readers/parquet.py | 10 +- src/orm_loader/tables/loadable_table.py | 17 +-- src/orm_loader/tables/orm_table.py | 37 +++++- src/orm_loader/tables/typing.py | 5 +- tests/backends/test_sinks.py | 4 +- tests/conversion/test_date_handlers.py | 34 ++++- tests/conversion/test_number_handlers.py | 91 ++++++++++++- tests/conversion/test_policy.py | 35 +++-- tests/conversion/test_text_handler.py | 23 +++- tests/conversion/test_typed_input.py | 74 +++++++++++ tests/loaders/test_defaults.py | 69 ++++++++++ tests/loaders/test_helpers.py | 63 ++------- tests/loaders/test_load_policy.py | 16 +-- tests/loaders/test_loader_e2e.py | 2 + tests/loaders/test_malformed_files.py | 123 ++++++++++++++++++ tests/loaders/test_pipeline.py | 96 +++++++++++++- tests/loaders/test_readers.py | 7 + tests/loaders/test_stage_file.py | 25 ++++ tests/models.py | 33 ++++- tests/support.py | 3 +- 46 files changed, 1063 insertions(+), 404 deletions(-) create mode 100644 tests/conversion/test_typed_input.py create mode 100644 tests/loaders/test_defaults.py create mode 100644 tests/loaders/test_malformed_files.py diff --git a/docs/loaders/context.md b/docs/loaders/context.md index 00f728d..b5ef6a0 100644 --- a/docs/loaders/context.md +++ b/docs/loaders/context.md @@ -25,7 +25,7 @@ on globals or implicit configuration. | `chunksize` | Rows per batch for the default reader of formats that batch by rows | | `normalise` | Whether to interpret values under the policy; when false, values are written as read | | `dedupe` | Whether the staging table keeps only the first row for each primary key | -| `quote_mode` | Quoting mode for delimited text: `csv`, `literal`, `by_delimiter` or `auto` | +| `quote_mode` | Quoting mode for delimited text: `csv`, `literal` or `by_delimiter` (default) | | `policy` | `ConversionPolicy`: how values are interpreted and what happens to unreadable ones | | `rejects_dir` | Directory that rejected rows are written to | | `report` | `LoadReport` row accounting | diff --git a/docs/loaders/conversion.md b/docs/loaders/conversion.md index b279256..7506843 100644 --- a/docs/loaders/conversion.md +++ b/docs/loaders/conversion.md @@ -31,7 +31,7 @@ policy = ConversionPolicy( nulls=NullPolicy(text_markers=frozenset({"NULL"})), errors=ErrorPolicy(on_error="reject_row", max_rejects=0.01), date=DateHandler(formats=("dd/mm/yyyy",)), - columns={("concept", "concept_name"): TextHandler(markers=frozenset({"N/A"}))}, + columns={("concept", "concept_code"): TextHandler(strip=False)}, ) ``` @@ -42,7 +42,7 @@ policy = ConversionPolicy( ## NullPolicy - Text columns keep every nonempty value verbatim. Only empty and whitespace-only text is null. -- `text_markers` adds exact, case-sensitive words that are null in text columns. `TextHandler(markers=...)` replaces them for the columns it serves. +- `text_markers` adds exact, case-sensitive words that are null in every text column of the load. To treat a word as missing in some columns only, subclass `TextHandler`, override `missing()` and register it in `ConversionPolicy.columns`. - `marker_quoting="unquoted_only"` matches a marker only when the source did not quote it. - Typed columns accept a valid value, a missing word (`null`, `none`, `na`, `n/a`, `nan` and similar) or a real null, and fail otherwise. @@ -60,22 +60,34 @@ policy = ConversionPolicy( ## Handlers -The handler for a column is the first match of: the entry in `ConversionPolicy.columns` and the handler for the most specific class of the column type. A column type with no handler is passed through unchanged. +The handler for a column is the first match of: the entry in `ConversionPolicy.columns` and the handler for the most specific class of the column type. A `TypeDecorator` is handled as the type it stores. A column type with no handler is an error that names the table, column and type, raised before any row is processed. Loading such a column unchanged is an explicit choice: `columns={("table", "column"): PassthroughHandler()}`. | Handler | Column types | Accepted text | Settings | |------|------|------|------| -| `TextHandler` | `String`, `Text`, `Enum` | Any nonempty text; longer than the column is an error | `strip`, `markers` | -| `IntegerHandler` | `Integer` | Optional sign and digits | none | -| `FloatHandler` | `Float` | Decimal or exponent form, no `Infinity` | none | -| `NumericHandler` | `Numeric` | Decimal form, converted to the declared precision and scale without rounding | none | +| `TextHandler` | `String`, `Text` | Any nonempty text; longer than the column is an error | `strip` | +| `EnumHandler` | `Enum` | One of the labels the column stores | `strip` | +| `IntegerHandler` | `Integer`, `SmallInteger`, `BigInteger` | Optional sign and digits, with an optional all-zero fraction (`1.0`); a value outside the 16, 32 or 64 bits of the column type is an error | none | +| `FloatHandler` | `Float` | Decimal or exponent form; `Infinity` and values that overflow to infinity are errors | none | +| `NumericHandler` | `Numeric` | Decimal form, converted to the declared precision and scale without rounding; a precision without a scale has scale 0 and a column without a precision reads as `decimal128(38, 18)` | none | | `BooleanHandler` | `Boolean` | `true`, `t`, `yes`, `y`, `1`, `on` and their opposites | `true_words`, `false_words` | -| `DateHandler` | `Date` | One of `formats`, validated against the calendar | `formats`, `parser` | -| `DateTimeHandler` | `DateTime` | ISO date and time, or a date in one of `formats` | `formats`, `parser` | +| `DateHandler` | `Date` | One of `formats`, validated against the calendar; year 0 is an error | `formats`, `parser` | +| `DateTimeHandler` | `DateTime` | ISO date and time, or a date in one of `formats`; see time zones below | `formats`, `parser` | Dates default to `yyyymmdd`, `dd-mon-yyyy` and `yyyy-mm-dd`. Ambiguous forms such as `02/01/2020` are never guessed; declare `DateHandler(formats=("dd/mm/yyyy",))` or `("mm/dd/yyyy",)`. A timestamp column reads date-only values with its own `formats`. `DateHandler(parser=...)` replaces the date parser with a function that has the signature of `parse_dates`. +### Time zones + +The column type decides, so no value is converted by a convention the schema did not ask for. + +- A `DateTime(timezone=True)` column requires an offset (`Z` or `+hh:mm`) and stores UTC. A value without an offset and a date-only value are errors. +- A `DateTime` column without a time zone rejects a value with an offset. + +### Typed input + +Input that already has a type, such as a Parquet column, is cast to the column type and checked like parsed text. A value that does not come back unchanged when cast back to its input type is an error, so an integer `2` in a `Boolean` column, a timestamp with a time in a `Date` column and a decimal that loses digits are all rejected. Float `NaN` is missing. + ### Adding a column type Subclass `ConvertingHandler`, implement `arrow_type` and `parse`, and pass the handler in `ConversionPolicy.handlers` for the type class, or in `ConversionPolicy.columns` for a single column. `clean`, `missing` and `check` can be overridden. diff --git a/docs/loaders/readers.md b/docs/loaders/readers.md index ab41898..4343a47 100644 --- a/docs/loaders/readers.md +++ b/docs/loaders/readers.md @@ -21,7 +21,7 @@ A new format is a new `FileReader` subclass that implements `batches()` and sets ## DelimitedTextReader -Reads raw strings in 4 MB blocks, so memory does not grow with the row count. Header names are lowercased and stripped of quotes and a `_hash` suffix, and file columns that are not model columns are ignored. An unquoted empty field is null and a quoted empty field is an empty string. +Reads raw strings in 4 MB blocks, so memory does not grow with the row count. Header names are lowercased and stripped of quotes, padding and a trailing `_hash`, and file columns that are not model columns are ignored with a warning. Parquet column names are normalised the same way. An unquoted empty field is null and a quoted empty field is an empty string. `quote_mode` selects how double quotes are read: @@ -29,8 +29,7 @@ Reads raw strings in 4 MB blocks, so memory does not grow with the row count. He |------|------| | `csv` | Double quotes quote fields | | `literal` | Double quotes are data | -| `by_delimiter` | `literal` for tab-delimited files, `csv` otherwise | -| `auto` | Samples the file and picks the mode that gives consistent column counts | +| `by_delimiter` (default) | `literal` for tab-delimited files, `csv` otherwise | ::: orm_loader.loaders.readers.delimited.DelimitedTextReader @@ -50,7 +49,7 @@ Reads raw strings in 4 MB blocks, so memory does not grow with the row count. He ## Detection -The delimited reader detects the encoding with `chardet` (ASCII is normalised to UTF-8, and only a small prefix of the file is inspected), the delimiter from the first line (tabs win over commas when more frequent), and the quote mode when `quote_mode="auto"`. +The delimited reader detects the encoding with `chardet` (ASCII is normalised to UTF-8, and only a small prefix of the file is inspected), the delimiter from the first line (tabs win over commas when more frequent). ::: orm_loader.loaders.readers.detection.infer_encoding ::: orm_loader.loaders.readers.detection.infer_delim diff --git a/docs/loaders/staging.md b/docs/loaders/staging.md index 9a81bae..a50986c 100644 --- a/docs/loaders/staging.md +++ b/docs/loaders/staging.md @@ -10,7 +10,8 @@ - Values are interpreted by column type, as described in [Conversion](conversion.md). - With the default policy an unreadable value or a missing required value raises and nothing is loaded. -- With `ErrorPolicy(on_error="reject_row")` those rows are written to `/.rejects.csv` and the rest load. `rejects_dir` defaults to a `rejects` directory inside the input file's directory. `LoadReport` accounts for every row read. +- With `ErrorPolicy(on_error="reject_row")` those rows, and delimited lines with the wrong number of fields, are written to `/.rejects.csv` and the rest load. A malformed line has no row number, and its text is in the reason. `rejects_dir` defaults to a `rejects` directory inside the input file's directory, and the file of an earlier load is removed. `LoadReport` accounts for every row read, checked against the rows counted in the staging table. +- A column of a table that has a default which SQL can apply (a server default or a scalar Python default) and is not nullable takes it when the value is empty or the column is not in the file. A column with a callable default stays required. - With `normalise=False` raw text is passed to the staging table unchanged. --- diff --git a/pyproject.toml b/pyproject.toml index 3ed6f36..001065b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -91,9 +91,13 @@ addopts = "-ra -m 'not db_dialect'" reportMissingTypeStubs = false # Temporary ty override until #43 picks it up in a follow-up -# with pyarrow-stubs +# with pyarrow-stubs, which type the compute functions too narrowly [[tool.ty.overrides]] include = ["src/orm_loader/conversion/**", "src/orm_loader/loaders/processor.py", "src/orm_loader/loaders/readers/**"] [tool.ty.overrides.rules] unresolved-attribute = "ignore" +invalid-argument-type = "ignore" +invalid-return-type = "ignore" +invalid-assignment = "ignore" +no-matching-overload = "ignore" diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 3670489..abc718f 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -350,6 +350,23 @@ def _insertable_column_names( """ return [c.name for c in table_cls.__table__.columns if c.computed is None] + def _staged_columns( + self, + table_cls: Type["FileTableProtocol"], + staging: sa.Table, + ) -> dict[str, sa.ColumnElement[Any]]: + """ + Return the staging value of each insertable column by column name. + + A null in a non-nullable column takes the column's SQL default, so an + empty value and an absent column mean the default. + """ + defaults = table_cls.sql_defaults() + return { + name: sa.func.coalesce(staging.c[name], defaults[name]) if name in defaults else staging.c[name] + for name in self._insertable_column_names(table_cls) + } + def deferred_fk_constraints( self, table_cls: Type["FileTableProtocol"], @@ -551,7 +568,7 @@ def insert_batches( ctx: "LoaderContext", batches: Iterable[pa.Table], columns: Sequence[str], - ) -> int: + ) -> None: """ Write typed batches into the load's staging table. @@ -566,11 +583,6 @@ def insert_batches( columns Names of the columns the batches carry. - Returns - ------- - int - Number of rows written. - Notes ----- A backend may commit between batches to bound memory, which also ends @@ -602,6 +614,11 @@ def _staging_reference(ctx: "LoaderContext", connection: Connection) -> str: schema = translate.get(staging.schema, staging.schema) return qualified(connection, staging.name, physical_schema=schema) + def count_staging(self, ctx: "LoaderContext") -> int: + """Return the number of rows in the load's staging table.""" + count = sa.select(sa.func.count()).select_from(ctx.staging_table) + return ctx.session.execute(count, bind_arguments={"mapper": ctx.tableclass}).scalar_one() + def deduplicate_staging( self, ctx: "LoaderContext", @@ -635,10 +652,9 @@ def deduplicate_staging( .label("position"), ).subquery() duplicates = sa.select(ranked.c.row_key).where(ranked.c.position > 1) - count = sa.select(sa.func.count()).select_from(staging) - before = ctx.session.execute(count, bind_arguments=arguments).scalar_one() + before = self.count_staging(ctx) ctx.session.execute(sa.delete(staging).where(row_key.in_(duplicates)), bind_arguments=arguments) - after = ctx.session.execute(count, bind_arguments=arguments).scalar_one() + after = self.count_staging(ctx) return DedupeResult(tuple(key_columns), before, before - after) @staticmethod diff --git a/src/orm_loader/backends/duckdb.py b/src/orm_loader/backends/duckdb.py index 36c9e42..4b8a7fd 100644 --- a/src/orm_loader/backends/duckdb.py +++ b/src/orm_loader/backends/duckdb.py @@ -338,25 +338,21 @@ def insert_batches( ctx: "LoaderContext", batches: Iterable[pa.Table], columns: Sequence[str], - ) -> int: + ) -> None: """Insert the batches into the staging table straight from Arrow. Notes ----- - - The fastest bulk insertion method is DuckDB's ``INSERT ... SELECT`` from a - temporary view over the Arrow memory. - - The staging table is committed every ``STAGING_COMMIT_EVERY`` batches, - because DuckDB keeps uncommitted rows in memory and one transaction would - hold the whole file. - - Notes - ----- - Committing the staging table also ends the caller's open transaction. + Each batch is registered as a temporary view over the Arrow memory and + copied with ``INSERT ... SELECT``, which is DuckDB's fastest bulk path. + The staging table is committed every ``STAGING_COMMIT_EVERY`` batches, + because DuckDB keeps uncommitted rows in memory and one transaction + would hold the whole file. The commit also ends the caller's open + transaction. """ batch = sa.table("_orm_loader_batch", *(sa.column(c) for c in columns)) insert = ctx.staging_table.insert().from_select(list(columns), sa.select(*batch.c)) bind_arguments = {"mapper": ctx.tableclass} - total = 0 for index, table in enumerate(batches, start=1): native = ctx.session.connection(bind_arguments=bind_arguments).connection native.register("_orm_loader_batch", table) @@ -364,10 +360,8 @@ def insert_batches( ctx.session.execute(insert, bind_arguments=bind_arguments) finally: native.unregister("_orm_loader_batch") - total += table.num_rows if index % self.STAGING_COMMIT_EVERY == 0: ctx.session.commit() - return total def check_replace_feasible( self, @@ -439,7 +433,8 @@ def merge_replace( self.check_replace_feasible(table_cls, pk_cols, session) target = table_cls.__table__ staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) - update_cols = [c for c in self._insertable_column_names(table_cls) if c not in pk_cols] + staged = self._staged_columns(table_cls, staging) + update_cols = [c for c in staged if c not in pk_cols] pk_join = sa.and_(*(target.c[c] == staging.c[c] for c in pk_cols)) def _batch_keys(extra: sa.ColumnElement[bool] | None = None) -> sa.Select[Any]: @@ -454,7 +449,7 @@ def _pk_in(keys: sa.Select[Any]) -> sa.ColumnElement[bool]: return sa.tuple_(*(target.c[c] for c in pk_cols)).in_(keys) def _update(extra: sa.ColumnElement[bool] | None = None) -> None: - stmt = sa.update(target).values({c: staging.c[c] for c in update_cols}).where(pk_join) + stmt = sa.update(target).values({c: staged[c] for c in update_cols}).where(pk_join) if extra is not None: stmt = stmt.where(_pk_in(_batch_keys(extra))) session.execute(stmt) @@ -477,8 +472,8 @@ def _drop_matched_from_staging(extra: sa.ColumnElement[bool] | None = None) -> N self_referential_fks = self._self_referential_deferred_fks(table_cls) if self_referential_fks: - insertable_cols = self._insertable_column_names(table_cls) - non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) + insertable_cols = list(staged) + non_paginated_select = sa.select(*staged.values()) def _insert(select_: Any) -> sa.Insert: return sa.insert(target).from_select(insertable_cols, select_) @@ -522,7 +517,8 @@ def merge_upsert( self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) - insertable_cols = self._insertable_column_names(table_cls) + staged = self._staged_columns(table_cls, staging) + insertable_cols = list(staged) def _upsert(select_: Any) -> sa.Insert: return ( @@ -531,7 +527,7 @@ def _upsert(select_: Any) -> sa.Insert: .on_conflict_do_nothing(index_elements=pk_cols) ) - non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) + non_paginated_select = sa.select(*staged.values()) self_referential = bool(self._self_referential_deferred_fks(table_cls)) for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): @@ -556,8 +552,9 @@ def merge_insert( self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) - insertable_cols = self._insertable_column_names(table_cls) - non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) + staged = self._staged_columns(table_cls, staging) + insertable_cols = list(staged) + non_paginated_select = sa.select(*staged.values()) def _insert(select_: Any) -> sa.Insert: return sa.insert(target).from_select(insertable_cols, select_) diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index 25efbd6..e0b42a6 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -109,6 +109,12 @@ def create_staging_table( bind_arguments=bind_arguments, ) + for name in table_cls.sql_defaults(): + session.execute( + sa.text(f'ALTER TABLE {staging_ref} ALTER COLUMN {preparer.quote_identifier(name)} DROP NOT NULL;'), + bind_arguments=bind_arguments, + ) + computed_cols = [c.name for c in table.columns if c.computed is not None] for col in computed_cols: session.execute( @@ -143,7 +149,7 @@ def insert_batches( ctx: "LoaderContext", batches: Iterable[pa.Table], columns: Sequence[str], - ) -> int: + ) -> None: """Stream the batches into the staging table with one COPY, as CSV written by Arrow. An unquoted empty field is NULL and a quoted one an empty string, so @@ -151,23 +157,20 @@ def insert_batches( any batch aborts the COPY. """ connection = ctx.session.connection(bind_arguments={"mapper": ctx.tableclass}) - raw_conn = connection.connection - if not hasattr(raw_conn, "cursor"): - raise RuntimeError("Expected DB-API connection for COPY") + cursor = connection.connection.cursor() + if not hasattr(cursor, "copy"): + cursor.close() + raise RuntimeError("The PostgreSQL sink needs a psycopg 3 connection, whose cursors provide copy()") table_ref = self._staging_reference(ctx, connection) column_list = ", ".join(self.identifier_preparer.quote(c) for c in columns) options = pacsv.WriteOptions(include_header=False, quoting_style="needed") - total = 0 - cursor = raw_conn.cursor() try: with cursor.copy(f"COPY {table_ref} ({column_list}) FROM STDIN WITH (FORMAT csv)") as copy: for table in batches: buffer = io.BytesIO() pacsv.write_csv(table, buffer, options) copy.write(buffer.getvalue()) - total += table.num_rows ctx.session.flush() - return total except Exception: ctx.session.rollback() raise @@ -271,7 +274,8 @@ def merge_upsert( self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) - insertable_cols = self._insertable_column_names(table_cls) + staged = self._staged_columns(table_cls, staging) + insertable_cols = list(staged) def _upsert(select_: Any) -> sa.Insert: # sa.insert() has no .on_conflict_do_nothing() @@ -281,7 +285,7 @@ def _upsert(select_: Any) -> sa.Insert: .on_conflict_do_nothing(index_elements=pk_cols) ) - non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) + non_paginated_select = sa.select(*staged.values()) for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): select_ = non_paginated_select if batch_range is None else non_paginated_select.where(batch_range) @@ -306,8 +310,9 @@ def merge_insert( self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) - insertable_cols = self._insertable_column_names(table_cls) - non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) + staged = self._staged_columns(table_cls, staging) + insertable_cols = list(staged) + non_paginated_select = sa.select(*staged.values()) def _insert(select_: Any) -> sa.Insert: return sa.insert(target).from_select(insertable_cols, select_) diff --git a/src/orm_loader/backends/sqlite.py b/src/orm_loader/backends/sqlite.py index a0fae64..592184e 100644 --- a/src/orm_loader/backends/sqlite.py +++ b/src/orm_loader/backends/sqlite.py @@ -121,7 +121,7 @@ def insert_batches( ctx: "LoaderContext", batches: Iterable[pa.Table], columns: Sequence[str], - ) -> int: + ) -> None: """Insert the batches with one native ``executemany`` per batch. The statement goes straight to the DB-API without per-value bind @@ -136,12 +136,9 @@ def insert_batches( f"INSERT INTO {table_ref} ({', '.join(preparer.quote(c) for c in columns)}) " f"VALUES ({', '.join('?' for _ in columns)})" ) - total = 0 for table in batches: values = [self._storable(table[name]) for name in columns] connection.exec_driver_sql(sql, list(zip(*values))) - total += table.num_rows - return total @staticmethod def _storable(column: pa.ChunkedArray) -> list[Any]: @@ -282,7 +279,8 @@ def merge_upsert( self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) - insertable_cols = self._insertable_column_names(table_cls) + staged = self._staged_columns(table_cls, staging) + insertable_cols = list(staged) def _upsert(select_: Any) -> sa.Insert: return ( @@ -291,7 +289,7 @@ def _upsert(select_: Any) -> sa.Insert: .on_conflict_do_nothing(index_elements=pk_cols) ) - non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) + non_paginated_select = sa.select(*staged.values()) for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): @@ -309,8 +307,9 @@ def merge_insert( self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) - insertable_cols = self._insertable_column_names(table_cls) - non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) + staged = self._staged_columns(table_cls, staging) + insertable_cols = list(staged) + non_paginated_select = sa.select(*staged.values()) def _insert(select_: Any) -> sa.Insert: return sa.insert(target).from_select(insertable_cols, select_) diff --git a/src/orm_loader/conversion/__init__.py b/src/orm_loader/conversion/__init__.py index 69da4fe..5c02151 100644 --- a/src/orm_loader/conversion/__init__.py +++ b/src/orm_loader/conversion/__init__.py @@ -5,6 +5,7 @@ DateFormats, DateHandler, DateTimeHandler, + EnumHandler, FloatHandler, IntegerHandler, NumberHandler, @@ -24,6 +25,7 @@ "DateFormats", "DateHandler", "DateTimeHandler", + "EnumHandler", "ErrorPolicy", "FloatHandler", "IntegerHandler", diff --git a/src/orm_loader/conversion/handlers/__init__.py b/src/orm_loader/conversion/handlers/__init__.py index 2230205..09b3b22 100644 --- a/src/orm_loader/conversion/handlers/__init__.py +++ b/src/orm_loader/conversion/handlers/__init__.py @@ -2,7 +2,7 @@ from .booleans import BooleanHandler from .dates import DateFormats, DateHandler, DateTimeHandler, parse_dates from .numbers import FloatHandler, IntegerHandler, NumberHandler, NumericHandler -from .text import TextHandler +from .text import EnumHandler, TextHandler __all__ = [ "BooleanHandler", @@ -11,6 +11,7 @@ "DateFormats", "DateHandler", "DateTimeHandler", + "EnumHandler", "FloatHandler", "IntegerHandler", "NumberHandler", diff --git a/src/orm_loader/conversion/handlers/base.py b/src/orm_loader/conversion/handlers/base.py index d8dc910..c05b15d 100644 --- a/src/orm_loader/conversion/handlers/base.py +++ b/src/orm_loader/conversion/handlers/base.py @@ -9,7 +9,7 @@ import pyarrow.compute as pc from sqlalchemy.types import TypeEngine -from ..values import Interpretation, as_array, is_text_like, no_failures, NULL_STRING, safe_cast +from ..values import NULL_STRING, Interpretation, as_array, is_text_like, no_failures, safe_cast if TYPE_CHECKING: from ..policy import NullPolicy @@ -62,7 +62,9 @@ class ConvertingHandler(ColumnHandler): """Converts a column in four stages: clean, detect missing, parse, check. Text input follows the stages. Input that already has a type, such as a - Parquet column, is cast directly with floating NaN read as missing. + Parquet column, is cast directly with floating NaN read as missing. A + value that does not survive the cast back to its input type is lossy and + unreadable, and ``check`` runs on the cast values. Subclasses supply ``arrow_type`` and ``parse`` and may override ``clean``, ``missing`` and ``check``. @@ -132,6 +134,13 @@ def _interpret_text(self, raw: pa.Array, sa_type: TypeEngine[Any], nulls: NullPo return Interpretation(values, invalid if rejected is None else pc.or_(invalid, rejected)) def _interpret_typed(self, arr: pa.Array, sa_type: TypeEngine[Any]) -> Interpretation: + target = self.arrow_type(sa_type) + if pa.types.is_null(arr.type): + return Interpretation(pa.nulls(len(arr), target), no_failures(len(arr))) if pa.types.is_floating(arr.type): arr = pc.if_else(pc.fill_null(pc.is_nan(arr), False), pa.scalar(None, arr.type), arr) - return Interpretation(*safe_cast(arr, self.arrow_type(sa_type))) + values, invalid = safe_cast(arr, target) + round_trip, _ = safe_cast(values, arr.type) + invalid = pc.or_(invalid, pc.fill_null(pc.not_equal(round_trip, arr), False)) + rejected = self.check(values, sa_type) + return Interpretation(values, invalid if rejected is None else pc.or_(invalid, rejected)) diff --git a/src/orm_loader/conversion/handlers/dates.py b/src/orm_loader/conversion/handlers/dates.py index cf19fc1..f51e2f3 100644 --- a/src/orm_loader/conversion/handlers/dates.py +++ b/src/orm_loader/conversion/handlers/dates.py @@ -10,7 +10,7 @@ import pyarrow.compute as pc from sqlalchemy.types import TypeEngine -from ..values import NULL_STRING, Interpretation, cast_each, safe_cast +from ..values import NULL_STRING, Interpretation, no_failures, safe_cast from .base import ConvertingHandler DateShape = Callable[[pa.Array], tuple[pa.Array, pa.Array]] @@ -31,7 +31,7 @@ class DateFormats: MONTHS Lowercase month abbreviations used by ``dd-mon-yyyy``. SHAPES - Shape function by format name. + Shape function by format name; the one place a format is registered. """ DEFAULT: ClassVar[tuple[str, ...]] = ("yyyymmdd", "dd-mon-yyyy", "yyyy-mm-dd") @@ -69,21 +69,23 @@ def mm_dd_yyyy(values: pa.Array) -> tuple[pa.Array, pa.Array]: cut = pc.utf8_slice_codeunits return ok, pc.binary_join_element_wise(cut(values, 6, 10), cut(values, 0, 2), cut(values, 3, 5), "-") + SHAPES: ClassVar[dict[str, DateShape]] = { + "yyyymmdd": yyyymmdd, + "yyyy-mm-dd": yyyy_mm_dd, + "dd-mon-yyyy": dd_mon_yyyy, + "dd/mm/yyyy": dd_mm_yyyy, + "mm/dd/yyyy": mm_dd_yyyy, + } + @classmethod def shape(cls, name: str) -> DateShape: """Return the shape function for a format name.""" - return { - "yyyymmdd": cls.yyyymmdd, - "yyyy-mm-dd": cls.yyyy_mm_dd, - "dd-mon-yyyy": cls.dd_mon_yyyy, - "dd/mm/yyyy": cls.dd_mm_yyyy, - "mm/dd/yyyy": cls.mm_dd_yyyy, - }[name] + return cls.SHAPES[name] @classmethod def names(cls) -> tuple[str, ...]: """Return every supported format name.""" - return ("yyyymmdd", "yyyy-mm-dd", "dd-mon-yyyy", "dd/mm/yyyy", "mm/dd/yyyy") + return tuple(cls.SHAPES) def parse_dates(values: pa.Array, formats: tuple[str, ...]) -> tuple[pa.Array, pa.Array]: @@ -107,7 +109,7 @@ def parse_dates(values: pa.Array, formats: tuple[str, ...]) -> tuple[pa.Array, p boolean array that is true where a value is invalid. """ iso = pa.nulls(len(values), pa.string()) - matched = pa.array([False] * len(values), pa.bool_()) + matched = no_failures(len(values)) for name in formats: ok, candidate = DateFormats.shape(name)(values) ok = pc.and_(pc.fill_null(ok, False), pc.is_valid(candidate)) @@ -115,11 +117,8 @@ def parse_dates(values: pa.Array, formats: tuple[str, ...]) -> tuple[pa.Array, p iso = pc.if_else(ok, candidate, iso) matched = pc.or_(matched, ok) invalid = pc.and_(pc.is_valid(values), pc.invert(matched)) - try: - return pc.cast(iso, pa.date32()), invalid - except pa.ArrowInvalid: - dates, calendar_failed = cast_each(iso, pa.date32()) - return dates, pc.or_(invalid, calendar_failed) + dates, calendar_failed = safe_cast(iso, pa.date32()) + return dates, pc.or_(invalid, calendar_failed) @dataclass(frozen=True) @@ -154,6 +153,9 @@ def arrow_type(self, sa_type: TypeEngine[Any]) -> pa.DataType: def parse(self, cleaned: pa.Array, sa_type: TypeEngine[Any]) -> Interpretation: return Interpretation(*self.parser(cleaned, self.formats)) + def check(self, values: pa.Array, sa_type: TypeEngine[Any]) -> pa.Array | None: + return pc.fill_null(pc.less(pc.year(values), 1), False) + @dataclass(frozen=True) class DateTimeHandler(DateHandler): @@ -161,22 +163,34 @@ class DateTimeHandler(DateHandler): A timestamp has date and time separated by ``T`` or a space, with optional seconds and up to six fractional digits. A date-only value - becomes midnight. + becomes midnight. The column type decides about time zones: a + ``DateTime(timezone=True)`` column requires an offset (``Z`` or + ``+hh:mm``) and stores UTC, and a column without a time zone rejects + an offset. A date-only value is not read for a time zone column. """ distinct: ClassVar[bool] = False - ISO: ClassVar[str] = r"^\d{4}-\d{2}-\d{2}[T ]\d{2}:\d{2}(:\d{2}(\.\d{1,6})?)?$" + _STAMP = r"^\d{4}-\d{2}-\d{2}[T ]\d{2}:\d{2}(:\d{2}(\.\d{1,6})?)?" + NAIVE: ClassVar[str] = _STAMP + "$" + AWARE: ClassVar[str] = _STAMP + r"(Z|[+-]\d{2}:\d{2})$" def arrow_type(self, sa_type: TypeEngine[Any]) -> pa.DataType: - return pa.timestamp("us") + return pa.timestamp("us", "UTC" if getattr(sa_type, "timezone", False) else None) + + def expected(self, sa_type: TypeEngine[Any]) -> str: + return "DateTime with an offset (Z or +hh:mm)" if getattr(sa_type, "timezone", False) else "DateTime without an offset" def parse(self, cleaned: pa.Array, sa_type: TypeEngine[Any]) -> Interpretation: - is_iso = pc.fill_null(pc.match_substring_regex(cleaned, self.ISO), False) - stamps, stamp_failed = safe_cast(pc.if_else(is_iso, cleaned, NULL_STRING), pa.timestamp("us")) + target = self.arrow_type(sa_type) + aware = target.tz is not None + is_iso = pc.fill_null(pc.match_substring_regex(cleaned, self.AWARE if aware else self.NAIVE), False) + stamps, stamp_failed = safe_cast(pc.if_else(is_iso, cleaned, NULL_STRING), target) + if aware: + return Interpretation(stamps, pc.or_(stamp_failed, pc.and_(pc.is_valid(cleaned), pc.invert(is_iso)))) rest = pc.dictionary_encode(pc.if_else(is_iso, NULL_STRING, cleaned)) dates, invalid = self.parser(rest.dictionary, self.formats) dates = pc.take(dates, rest.indices) invalid = pc.fill_null(pc.take(invalid, rest.indices), False) - values = pc.coalesce(stamps, pc.cast(dates, pa.timestamp("us"))) + values = pc.coalesce(stamps, pc.cast(dates, target)) return Interpretation(values, pc.or_(invalid, stamp_failed)) diff --git a/src/orm_loader/conversion/handlers/numbers.py b/src/orm_loader/conversion/handlers/numbers.py index 315335f..479624c 100644 --- a/src/orm_loader/conversion/handlers/numbers.py +++ b/src/orm_loader/conversion/handlers/numbers.py @@ -7,7 +7,7 @@ import pyarrow as pa import pyarrow.compute as pc -from sqlalchemy.types import TypeEngine +from sqlalchemy.types import BigInteger, SmallInteger, TypeEngine from ..values import NULL_STRING, Interpretation, safe_cast from .base import ConvertingHandler @@ -20,14 +20,23 @@ class NumberHandler(ConvertingHandler): Attributes ---------- PATTERN - Regular expression a trimmed value must match. + Regular expression a trimmed value must match, unless ``pattern`` + is overridden. + STRIPPED + Regular expression for the parts of a matching value removed before + the cast. """ PATTERN: ClassVar[str] + STRIPPED: ClassVar[str] = r"^\+" + + def pattern(self, sa_type: TypeEngine[Any]) -> str: + """Return the regular expression a trimmed value must match for the column.""" + return self.PATTERN def parse(self, cleaned: pa.Array, sa_type: TypeEngine[Any]) -> Interpretation: - fits = pc.fill_null(pc.match_substring_regex(cleaned, self.PATTERN), False) - candidate = pc.replace_substring_regex(pc.if_else(fits, cleaned, NULL_STRING), r"^\+", "") + fits = pc.fill_null(pc.match_substring_regex(cleaned, self.pattern(sa_type)), False) + candidate = pc.replace_substring_regex(pc.if_else(fits, cleaned, NULL_STRING), self.STRIPPED, "") values, cast_failed = safe_cast(candidate, self.arrow_type(sa_type)) invalid = pc.and_(pc.is_valid(cleaned), pc.invert(fits)) return Interpretation(values, pc.or_(invalid, cast_failed)) @@ -35,36 +44,67 @@ def parse(self, cleaned: pa.Array, sa_type: TypeEngine[Any]) -> Interpretation: @dataclass(frozen=True) class IntegerHandler(NumberHandler): - """Accepts an optionally signed run of digits.""" + """Accepts an optionally signed run of digits, with an optional all-zero fraction. + + A value is unreadable when it does not fit the column: 16 bits for + ``SmallInteger``, 64 for ``BigInteger`` and 32 for any other integer. + """ - PATTERN: ClassVar[str] = r"^[+-]?\d+$" + PATTERN: ClassVar[str] = r"^[+-]?\d+(\.0*)?$" + STRIPPED: ClassVar[str] = r"^\+|\.0*$" def arrow_type(self, sa_type: TypeEngine[Any]) -> pa.DataType: return pa.int64() + def check(self, values: pa.Array, sa_type: TypeEngine[Any]) -> pa.Array | None: + low, high = self.bounds(sa_type) + return pc.fill_null(pc.or_(pc.less(values, low), pc.greater(values, high)), False) + + @staticmethod + def bounds(sa_type: TypeEngine[Any]) -> tuple[int, int]: + """Return the smallest and largest value the column type holds.""" + bits = 64 if isinstance(sa_type, BigInteger) else 16 if isinstance(sa_type, SmallInteger) else 32 + return -(1 << (bits - 1)), (1 << (bits - 1)) - 1 + @dataclass(frozen=True) class FloatHandler(NumberHandler): - """Accepts decimal and exponent forms; infinity is not accepted.""" + """Accepts decimal and exponent forms; infinity, written or overflowing, is not accepted.""" PATTERN: ClassVar[str] = r"^[+-]?(\d+\.?\d*|\.\d+)([eE][+-]?\d+)?$" def arrow_type(self, sa_type: TypeEngine[Any]) -> pa.DataType: return pa.float64() + def check(self, values: pa.Array, sa_type: TypeEngine[Any]) -> pa.Array | None: + return pc.fill_null(pc.is_inf(values), False) + @dataclass(frozen=True) class NumericHandler(NumberHandler): """Accepts decimal forms and keeps the column's declared precision and scale. - A value with more digits than the column holds is unreadable. A column - without a declared precision and scale is read as a float. + A precision without a scale has scale 0. A column without a precision is + read as ``decimal128(38, 18)``. A value with more digits than the column + holds is unreadable; trailing zeros beyond the scale are accepted. """ - PATTERN: ClassVar[str] = r"^[+-]?(\d+\.?\d*|\.\d+)$" + @staticmethod + def precision_scale(sa_type: TypeEngine[Any]) -> tuple[int, int]: + """Return the precision and scale a column is read with.""" + precision = getattr(sa_type, "precision", None) + if precision is None: + return 38, 18 + return precision, getattr(sa_type, "scale", None) or 0 def arrow_type(self, sa_type: TypeEngine[Any]) -> pa.DataType: - precision, scale = getattr(sa_type, "precision", None), getattr(sa_type, "scale", None) - if precision is not None and scale is not None: - return pa.decimal128(precision, scale) - return pa.float64() + return pa.decimal128(*self.precision_scale(sa_type)) + + def pattern(self, sa_type: TypeEngine[Any]) -> str: + """Return the decimal forms whose integer part fits the column. + + Arrow wraps around instead of failing when a string has more digits + than a decimal holds, so the integer digits are bounded here. + """ + precision, scale = self.precision_scale(sa_type) + return rf"^[+-]?(0*\d{{0,{precision - scale}}}(\.\d*)?|\.\d+)$" diff --git a/src/orm_loader/conversion/handlers/text.py b/src/orm_loader/conversion/handlers/text.py index e2758db..58300ff 100644 --- a/src/orm_loader/conversion/handlers/text.py +++ b/src/orm_loader/conversion/handlers/text.py @@ -24,14 +24,9 @@ class TextHandler(ConvertingHandler): ---------- strip Surrounding whitespace is removed from stored text. - markers - Exact, case-sensitive values treated as missing, ignoring surrounding - whitespace. Replaces ``NullPolicy.text_markers`` for the columns this - handler serves; None uses the policy markers. """ strip: bool = True - markers: frozenset[str] | None = None def arrow_type(self, sa_type: TypeEngine[Any]) -> pa.DataType: return pa.string() @@ -40,7 +35,7 @@ def clean(self, raw: pa.Array, trimmed: pa.Array) -> pa.Array: return trimmed if self.strip else raw def missing(self, trimmed: pa.Array, nulls: NullPolicy) -> pa.Array: - return nulls.text_missing(trimmed, self.markers) + return nulls.text_missing(trimmed) def parse(self, cleaned: pa.Array, sa_type: TypeEngine[Any]) -> Interpretation: return Interpretation(cleaned, no_failures(len(cleaned))) @@ -50,3 +45,12 @@ def check(self, values: pa.Array, sa_type: TypeEngine[Any]) -> pa.Array | None: if not length: return None return pc.fill_null(pc.greater(pc.utf8_length(values), length), False) + + +@dataclass(frozen=True) +class EnumHandler(TextHandler): + """Keeps text verbatim and rejects values that are not a label the column stores.""" + + def check(self, values: pa.Array, sa_type: TypeEngine[Any]) -> pa.Array | None: + known = pc.is_in(values, value_set=pa.array(sa_type.enums, pa.string())) + return pc.and_(pc.is_valid(values), pc.invert(known)) diff --git a/src/orm_loader/conversion/policy.py b/src/orm_loader/conversion/policy.py index bd8f343..c512bc4 100644 --- a/src/orm_loader/conversion/policy.py +++ b/src/orm_loader/conversion/policy.py @@ -4,6 +4,7 @@ from collections.abc import Mapping from dataclasses import dataclass, field +from types import MappingProxyType from typing import Any, ClassVar, Literal import pyarrow as pa @@ -12,6 +13,7 @@ Boolean, Date, DateTime, + Enum, Float, Integer, Numeric, @@ -19,18 +21,19 @@ TypeEngine, ) +from ..helpers import PolicyError from .handlers import ( BooleanHandler, ColumnHandler, DateHandler, DateTimeHandler, + EnumHandler, FloatHandler, IntegerHandler, NumericHandler, - PassthroughHandler, TextHandler, ) -from .values import no_failures +from .values import no_failures, underlying_type @dataclass(frozen=True) @@ -68,24 +71,20 @@ def __post_init__(self) -> None: if self.marker_quoting not in ("ignore", "unquoted_only"): raise ValueError(f"marker_quoting must be 'ignore' or 'unquoted_only', not {self.marker_quoting!r}") - def text_missing(self, trimmed: pa.Array, markers: frozenset[str] | None = None) -> pa.Array: + def text_missing(self, trimmed: pa.Array) -> pa.Array: """Return a mask of trimmed text values that are missing in a text column. Parameters ---------- trimmed String array without surrounding whitespace. - markers - Markers that replace ``text_markers``, or None to use them. Returns ------- pa.Array Boolean array that is true where the value is missing. """ - words = self.text_markers if markers is None else markers - if self.marker_quoting != "ignore": - words = frozenset() + words = self.text_markers if self.marker_quoting == "ignore" else frozenset() return self._mask(trimmed, words, empty=self.empty_is_null, casefold=False) def typed_missing(self, trimmed: pa.Array) -> pa.Array: @@ -131,8 +130,9 @@ class ConversionPolicy: """The settings of one load, shared by every column. The handler for a column is, in order, the entry in ``columns`` and the - handler for the most specific class of the column type among - ``handlers`` and the built-in fields. + handler for the most specific class of the column type (the stored type + of a ``TypeDecorator``) among ``handlers`` and the built-in fields. A + column type with no handler is an error. Attributes ---------- @@ -142,6 +142,8 @@ class ConversionPolicy: What happens to unreadable and missing required values. text Handler for string and text columns. + enum + Handler for enumerated columns. integer Handler for integer columns. floating @@ -164,6 +166,7 @@ class ConversionPolicy: nulls: NullPolicy = field(default_factory=NullPolicy) errors: ErrorPolicy = field(default_factory=ErrorPolicy) text: TextHandler = field(default_factory=TextHandler) + enum: EnumHandler = field(default_factory=EnumHandler) integer: IntegerHandler = field(default_factory=IntegerHandler) floating: FloatHandler = field(default_factory=FloatHandler) numeric: NumericHandler = field(default_factory=NumericHandler) @@ -174,10 +177,8 @@ class ConversionPolicy: columns: Mapping[tuple[str, str], ColumnHandler] = field(default_factory=dict) def __post_init__(self) -> None: - if self.nulls.marker_quoting == "unquoted_only": - texts = [self.text, *self.columns.values(), *self.handlers.values()] - if any(isinstance(h, TextHandler) and h.markers is not None for h in texts): - raise ValueError("marker_quoting='unquoted_only' supports NullPolicy.text_markers only") + object.__setattr__(self, "handlers", MappingProxyType(dict(self.handlers))) + object.__setattr__(self, "columns", MappingProxyType(dict(self.columns))) def for_column( self, sa_type: TypeEngine[Any], table_name: str | None = None, column_name: str | None = None @@ -196,8 +197,13 @@ def for_column( Returns ------- ColumnHandler - The handler for the column; a passthrough handler for a column - type with none. + The handler for the column. + + Raises + ------ + PolicyError + No handler matches the column type. Passthrough is chosen + explicitly with ``columns``. """ if table_name and column_name: handler = self.columns.get((table_name, column_name)) @@ -206,6 +212,7 @@ def for_column( # Dispatch default first and then overwrite with any registered handlers by_type: dict[type, ColumnHandler] = { String: self.text, + Enum: self.enum, Integer: self.integer, Float: self.floating, Numeric: self.numeric, @@ -214,7 +221,7 @@ def for_column( DateTime: self.datetime, **self.handlers, } - for cls in type(sa_type).__mro__: + for cls in type(underlying_type(sa_type)).__mro__: if cls in by_type: return by_type[cls] - return PassthroughHandler() + raise PolicyError(table_name=table_name, column_name=column_name, column_type=repr(sa_type)) diff --git a/src/orm_loader/conversion/values.py b/src/orm_loader/conversion/values.py index 3bd2056..553b2fd 100644 --- a/src/orm_loader/conversion/values.py +++ b/src/orm_loader/conversion/values.py @@ -2,10 +2,11 @@ from __future__ import annotations -from typing import Any, NamedTuple +from typing import NamedTuple import pyarrow as pa import pyarrow.compute as pc +from sqlalchemy.types import TypeDecorator, TypeEngine NULL_STRING = pa.scalar(None, pa.string()) @@ -25,6 +26,13 @@ class Interpretation(NamedTuple): invalid: pa.Array +def underlying_type(sa_type: TypeEngine) -> TypeEngine: + """Return the type a ``TypeDecorator`` stores its values as, or the type itself.""" + while isinstance(sa_type, TypeDecorator): + sa_type = sa_type.impl + return sa_type + + def as_array(arr: pa.Array | pa.ChunkedArray) -> pa.Array: """Return a chunked array as one contiguous array.""" if isinstance(arr, pa.ChunkedArray): @@ -41,39 +49,14 @@ def is_text_like(arr_type: pa.DataType) -> bool: def no_failures(length: int) -> pa.Array: """Return a boolean array of ``length`` false values.""" - return pa.array([False] * length, pa.bool_()) - - -def cast_each(values: pa.Array, target: pa.DataType) -> tuple[pa.Array, pa.Array]: - """Cast one element at a time. - - Parameters - ---------- - values - Array to cast. - target - Arrow type to cast to. - - Returns - ------- - tuple[pa.Array, pa.Array] - The cast array with nulls where a cast failed, and a boolean array - that is true where it failed. - """ - out: list[Any] = [] - failed: list[bool] = [] - for i in range(len(values)): - try: - out.append(pc.cast(values[i : i + 1], target)[0].as_py()) - failed.append(False) - except (pa.ArrowInvalid, pa.ArrowNotImplementedError): - out.append(None) - failed.append(True) - return pa.array(out, target), pa.array(failed, pa.bool_()) + return pa.repeat(False, length) def safe_cast(values: pa.Array, target: pa.DataType) -> tuple[pa.Array, pa.Array]: - """Cast the whole array, falling back to element-wise casts on failure. + """Cast an array, isolating the values that cannot be cast. + + The whole array is cast first. When that fails it is split in halves and + each half is cast the same way, so few unreadable values cost few casts. Parameters ---------- @@ -91,4 +74,9 @@ def safe_cast(values: pa.Array, target: pa.DataType) -> tuple[pa.Array, pa.Array try: return pc.cast(values, target), no_failures(len(values)) except (pa.ArrowInvalid, pa.ArrowNotImplementedError): - return cast_each(values, target) + if len(values) == 1: + return pa.nulls(1, target), pa.array([True]) + middle = len(values) // 2 + head, head_failed = safe_cast(values.slice(0, middle), target) + tail, tail_failed = safe_cast(values.slice(middle), target) + return pa.concat_arrays([head, tail]), pa.concat_arrays([head_failed, tail_failed]) diff --git a/src/orm_loader/helpers/__init__.py b/src/orm_loader/helpers/__init__.py index be24464..255d98e 100644 --- a/src/orm_loader/helpers/__init__.py +++ b/src/orm_loader/helpers/__init__.py @@ -1,4 +1,4 @@ -from .errors import ConversionError, IngestError, RejectLimitError, RequiredValueError, ValidationError +from .errors import ConversionError, IngestError, PolicyError, RejectLimitError, RequiredValueError, ValidationError from .bootstrap import bootstrap, create_db, create_tables from .engines import dispose_engines from .metadata import Base @@ -6,6 +6,7 @@ __all__ = [ "ConversionError", + "PolicyError", "RejectLimitError", "RequiredValueError", "IngestError", diff --git a/src/orm_loader/helpers/errors.py b/src/orm_loader/helpers/errors.py index 665f226..71edd51 100644 --- a/src/orm_loader/helpers/errors.py +++ b/src/orm_loader/helpers/errors.py @@ -6,6 +6,31 @@ class ValidationError(Exception): """Raised when schema or specification validation fails.""" +class PolicyError(IngestError): + """Raised when the conversion policy cannot handle a column. + + Parameters + ---------- + table_name + Table the column belongs to. + column_name + The column. + column_type + Name of the column type that has no handler. + """ + + def __init__(self, *, table_name: str | None, column_name: str | None, column_type: str) -> None: + self.table_name = table_name + self.column_name = column_name + self.column_type = column_type + target = ".".join(part for part in (table_name, column_name) if part) or "column" + super().__init__( + f"{target}: no handler for column type {column_type}. Register one in " + f"ConversionPolicy(handlers=...) or ConversionPolicy(columns=...), " + f"or PassthroughHandler() to load the values unchanged." + ) + + class ConversionError(IngestError): """Raised when values cannot be read as their column type. diff --git a/src/orm_loader/loaders/__init__.py b/src/orm_loader/loaders/__init__.py index d9705ca..2d62ef8 100644 --- a/src/orm_loader/loaders/__init__.py +++ b/src/orm_loader/loaders/__init__.py @@ -1,4 +1,4 @@ -from .loader import LoadReport, LoaderContext, stage_file +from .loader import LoaderContext, LoadReport, stage_file from .readers import DelimitedTextReader, FileReader, ParquetReader, open_reader __all__ = [ diff --git a/src/orm_loader/loaders/loader.py b/src/orm_loader/loaders/loader.py index 6ca55bd..8fadb47 100644 --- a/src/orm_loader/loaders/loader.py +++ b/src/orm_loader/loaders/loader.py @@ -10,10 +10,10 @@ import logging from collections.abc import Iterable -from itertools import chain from dataclasses import dataclass, field +from itertools import chain from pathlib import Path -from typing import TYPE_CHECKING, Type +from typing import TYPE_CHECKING import sqlalchemy as sa import sqlalchemy.orm as so @@ -42,8 +42,9 @@ class LoadReport: rows_loaded Rows written to the staging table. rows_rejected - Rows dropped because a value was unreadable or a required value was - missing, with ``ErrorPolicy.on_error`` set to ``"reject_row"``. + Rows dropped because a value was unreadable, a required value was + missing or a line was malformed, with ``ErrorPolicy.on_error`` set to + ``"reject_row"``. dedupe Result of deduplicating the staging table, if it was requested. rejects_path @@ -92,7 +93,7 @@ class LoaderContext: Whether the staging table keeps only the first row for each primary key. quote_mode Quoting mode of the default reader for delimited text: ``"csv"``, - ``"literal"``, ``"by_delimiter"`` or ``"auto"``. + ``"literal"`` or ``"by_delimiter"``. staging_schema_tag schema_translate_map tag the staging table lives in, passed to resolve_backend() so every backend resolution within this load @@ -106,7 +107,7 @@ class LoaderContext: Row accounting, filled in as the load runs. """ - tableclass: Type["FileTableProtocol"] + tableclass: type[FileTableProtocol] session: so.Session path: Path staging_table: sa.Table @@ -114,14 +115,14 @@ class LoaderContext: chunksize: int | None = None normalise: bool = True dedupe: bool = True - quote_mode: str = "auto" + quote_mode: str = "by_delimiter" staging_schema_tag: str | None = None policy: ConversionPolicy = field(default_factory=ConversionPolicy) rejects_dir: Path | None = None report: LoadReport = field(default_factory=LoadReport) -def _require_columns_present(tableclass: Type["FileTableProtocol"], present_columns: Iterable[str]) -> None: +def _require_columns_present(tableclass: type[FileTableProtocol], present_columns: Iterable[str]) -> None: """Raise before staging if a required column is entirely absent from the file. Partial-column loads are not supported, so every merge strategy @@ -183,9 +184,12 @@ def stage_file(ctx: LoaderContext, reader: FileReader) -> int: table_name = ctx.tableclass.__tablename__ batches = reader.batches() first = None if reader.empty else next(batches, None) - if first is None: + if first is None and not reader.skipped: logger.info("File %s has no data rows; skipping load for %s", ctx.path.name, table_name) return 0 + ignored = [name for name in reader.file_columns if name not in reader.columns] + if ignored: + logger.warning("File %s has columns that %s does not have, which are ignored: %s", ctx.path.name, table_name, ignored) if ctx.normalise: _require_columns_present(ctx.tableclass, reader.columns) rejects_dir = ctx.rejects_dir or ctx.path.parent / "rejects" @@ -196,9 +200,11 @@ def stage_file(ctx: LoaderContext, reader: FileReader) -> int: ctx.report, rejects_dir / f"{ctx.path.stem}.rejects.csv", ctx.normalise, + reader.skipped, ) backend = resolve_backend(ctx.session, staging_schema_tag=ctx.staging_schema_tag, mapper=ctx.tableclass) - total = backend.insert_batches(ctx, processor.stream(chain([first], batches)), reader.columns) + backend.insert_batches(ctx, processor.stream(chain([] if first is None else [first], batches)), reader.columns) + total = backend.count_staging(ctx) keys = [c for c in ctx.tableclass.pk_names() if c in reader.columns] if ctx.dedupe and keys: ctx.report.dedupe = backend.deduplicate_staging(ctx, keys) diff --git a/src/orm_loader/loaders/processor.py b/src/orm_loader/loaders/processor.py index d6272fc..4166c72 100644 --- a/src/orm_loader/loaders/processor.py +++ b/src/orm_loader/loaders/processor.py @@ -13,7 +13,7 @@ import pyarrow.csv as pv from ..conversion import ConversionPolicy -from ..conversion.values import no_failures +from ..conversion.values import no_failures, underlying_type from ..helpers import ConversionError, IngestError, RejectLimitError, RequiredValueError if TYPE_CHECKING: @@ -74,7 +74,8 @@ class BatchProcessor: A batch with an unreadable value or a missing required value either raises or has those rows written to the rejects file, depending on - ``ErrorPolicy.on_error``. + ``ErrorPolicy.on_error``. A rejects file left by an earlier load is + removed. Parameters ---------- @@ -90,6 +91,9 @@ class BatchProcessor: File that rejected rows are written to. normalise Whether to interpret values. When false, batches pass through. + skipped + Lines the reader left out as malformed, each with the reason. They + are written to the rejects file and count as rows read and rejected. """ def __init__( @@ -100,28 +104,32 @@ def __init__( report: LoadReport, rejects_path: Path, normalise: bool = True, + skipped: list[str] | None = None, ) -> None: self.columns = columns + self.skipped = skipped if skipped is not None else [] self.policy = policy self.report = report self.normalise = normalise self.table_name = tableclass.__tablename__ - self.model_columns = tableclass.model_columns() - self.handlers = { - name: policy.for_column(self.model_columns[name].type, self.table_name, name) for name in columns - } + model_columns = tableclass.model_columns() + self.types = {name: underlying_type(model_columns[name].type) for name in columns} + self.handlers = {name: policy.for_column(self.types[name], self.table_name, name) for name in columns} self.required = [c for c in tableclass.required_columns() if c in columns] self.rejects = RejectsWriter(rejects_path) if policy.errors.on_error == "reject_row": self.rejects.check_writable() + rejects_path.unlink(missing_ok=True) def stream(self, batches: Iterable[pa.RecordBatch]) -> Iterator[pa.Table]: """Yield typed tables for the raw batches, closing the rejects file afterwards.""" try: for raw in batches: table = self.process(raw) + self._reject_skipped() if table.num_rows: yield table + self._reject_skipped() finally: self.rejects.close() @@ -153,7 +161,7 @@ def process(self, raw: pa.RecordBatch) -> pa.Table: if not self.normalise: return pa.Table.from_batches([raw]) results = { - name: self.handlers[name].interpret(raw.column(name), self.model_columns[name].type, self.policy.nulls) + name: self.handlers[name].interpret(raw.column(name), self.types[name], self.policy.nulls) for name in self.columns } failures = self._failures(results) @@ -168,7 +176,7 @@ def _failures(self, results: dict[str, Any]) -> dict[str, tuple[pa.Array, str]]: for name in self.columns: invalid = results[name].invalid if pc.any(invalid).as_py(): - found[name] = (invalid, f"cannot be read as {self.handlers[name].expected(self.model_columns[name].type)}") + found[name] = (invalid, f"cannot be read as {self.handlers[name].expected(self.types[name])}") for name in self.required: missing = pc.and_(pc.is_null(results[name].values), pc.invert(results[name].invalid)) if pc.any(missing).as_py(): @@ -202,7 +210,7 @@ def _conversion_error(self, name: str, raw: pa.Array | pa.ChunkedArray, invalid: return ConversionError( table_name=self.table_name, column_name=name, - expected=self.handlers[name].expected(self.model_columns[name].type), + expected=self.handlers[name].expected(self.types[name]), count=count, samples=raw.filter(invalid).slice(0, 3).to_pylist(), ) @@ -249,6 +257,23 @@ def _reject( keep = pc.invert(rejected) return pa.table({n: results[n].values for n in self.columns}).filter(keep) + def _reject_skipped(self) -> None: + """Write the lines the reader left out to the rejects file, without a row number.""" + count = len(self.skipped) + if not count: + return + text, empty = pa.string(), pa.nulls(count, pa.string()) + rows = pa.table( + {"row": pa.nulls(count, pa.int64()), "column": empty, "reason": pa.array(self.skipped, text)} + | {name: empty for name in self.columns} + ) + self.skipped.clear() + self.rejects.write(rows) + self.report.rows_read += count + self.report.rows_rejected += count + self.report.rejects_path = self.rejects.path + self._check_limit() + def _check_limit(self) -> None: limit = self.policy.errors.max_rejects if limit is None: diff --git a/src/orm_loader/loaders/readers/__init__.py b/src/orm_loader/loaders/readers/__init__.py index 8eeeb22..5df851f 100644 --- a/src/orm_loader/loaders/readers/__init__.py +++ b/src/orm_loader/loaders/readers/__init__.py @@ -15,8 +15,9 @@ def open_reader( model_columns: Mapping[str, sa.ColumnElement[Any]], *, nulls: NullPolicy | None = None, - quote_mode: str = "auto", + quote_mode: str = "by_delimiter", chunksize: int | None = None, + skip_malformed: bool = False, ) -> FileReader: """Open a file with the reader for its format. @@ -35,6 +36,9 @@ def open_reader( Quote handling for delimited text. chunksize Rows per batch for formats that batch by rows. + skip_malformed + Leave out lines with the wrong number of fields and record them in + the reader's ``skipped`` list. Delimited text only. Returns ------- @@ -43,7 +47,7 @@ def open_reader( """ if path.suffix.lower() == ".parquet": return ParquetReader(path, model_columns, nulls, chunksize) - return DelimitedTextReader(path, model_columns, nulls, quote_mode) + return DelimitedTextReader(path, model_columns, nulls, quote_mode, skip_malformed=skip_malformed) diff --git a/src/orm_loader/loaders/readers/base.py b/src/orm_loader/loaders/readers/base.py index 050d4a3..5008522 100644 --- a/src/orm_loader/loaders/readers/base.py +++ b/src/orm_loader/loaders/readers/base.py @@ -3,7 +3,7 @@ from __future__ import annotations from abc import ABC, abstractmethod -from collections.abc import Iterator, Mapping +from collections.abc import Iterable, Iterator, Mapping from pathlib import Path from typing import Any @@ -12,6 +12,7 @@ import sqlalchemy as sa from ...conversion import NullPolicy +from ...helpers import IngestError class FileReader(ABC): @@ -36,8 +37,13 @@ class FileReader(ABC): ---------- columns Model columns present in the file, in file order. + file_columns + Every column of the file under its normalised name, in file order. empty True when the file has no content at all. + skipped + Lines the reader left out as malformed, each with the reason, filled + in while the batches are read. """ def __init__( @@ -50,7 +56,22 @@ def __init__( self.model_columns = model_columns self.nulls = nulls or NullPolicy() self.columns: list[str] = [] + self.file_columns: list[str] = [] self.empty = False + self.skipped: list[str] = [] + + def normalise_headers(self, names: Iterable[str]) -> list[str]: + """Return header names lowercased, stripped of quotes and padding and of a ``_hash`` suffix. + + Raises + ------ + IngestError + Two columns have the same normalised name. + """ + normalised = [name.strip().strip('"').strip().lower().removesuffix("_hash") for name in names] + if len(set(normalised)) != len(normalised): + raise IngestError(f"{self.path.name}: duplicate column names in header {normalised}") + return normalised @abstractmethod def batches(self) -> Iterator[pa.RecordBatch]: diff --git a/src/orm_loader/loaders/readers/delimited.py b/src/orm_loader/loaders/readers/delimited.py index 836634e..1f199d8 100644 --- a/src/orm_loader/loaders/readers/delimited.py +++ b/src/orm_loader/loaders/readers/delimited.py @@ -6,7 +6,7 @@ import logging from collections.abc import Iterator, Mapping from pathlib import Path -from typing import Any +from typing import Any, Literal import pyarrow as pa import pyarrow.csv as pv @@ -14,8 +14,8 @@ from ...conversion import NullPolicy from ...helpers import IngestError -from .detection import infer_delim, infer_encoding, resolve_quote_mode from .base import FileReader +from .detection import infer_delim, infer_encoding, resolve_quote_mode logger = logging.getLogger(__name__) @@ -26,9 +26,9 @@ class DelimitedTextReader(FileReader): """A delimited text file read as raw strings in bounded blocks. Date columns are dictionary-encoded. An unquoted empty field is null and - a quoted empty field is an empty string. Header names are lowercased and - stripped of quotes and a ``_hash`` suffix. File columns that are not - model columns are ignored. + a quoted empty field is an empty string. Header names are normalised as + described by ``normalise_headers``. File columns that are not model + columns are ignored. Parameters ---------- @@ -41,9 +41,12 @@ class DelimitedTextReader(FileReader): quote_mode ``"csv"`` reads double quotes as field quoting, ``"literal"`` keeps them as data, ``"by_delimiter"`` chooses ``"literal"`` for tab - delimited files and ``"csv"`` otherwise, ``"auto"`` samples the file. + delimited files and ``"csv"`` otherwise. block_size Bytes per block; memory use follows this, not the row count. + skip_malformed + Leave out a line with the wrong number of fields and record it in + ``skipped``, instead of failing the read. """ def __init__( @@ -51,11 +54,13 @@ def __init__( path: Path, model_columns: Mapping[str, sa.ColumnElement[Any]], nulls: NullPolicy | None = None, - quote_mode: str = "auto", + quote_mode: str = "by_delimiter", block_size: int = 4 << 20, + skip_malformed: bool = False, ) -> None: super().__init__(path, model_columns, nulls) self.block_size = block_size + self.skip_malformed = skip_malformed if path.stat().st_size == 0: self.empty = True return @@ -65,7 +70,7 @@ def __init__( codecs.lookup(self.encoding) except LookupError as exc: raise IngestError(f"{path.name}: unsupported encoding {self.encoding!r}") from exc - self.quote = resolve_quote_mode(quote_mode, path, self.delimiter, self.encoding) + self.quote = resolve_quote_mode(quote_mode, self.delimiter) self.file_columns = self._header_names() self.columns = [name for name in self.file_columns if name in model_columns] if self.nulls.marker_quoting == "unquoted_only" and self.nulls.text_markers: @@ -74,13 +79,11 @@ def __init__( def _header_names(self) -> list[str]: with open(self.path, "rb") as handle: header = handle.readline().decode(self.encoding, errors="replace") - header = header.lstrip("").rstrip("\r\n") - names = [ - name.strip().strip('"').strip().lower().replace("_hash", "") for name in header.split(self.delimiter) - ] - if len(set(names)) != len(names): - raise IngestError(f"{self.path.name}: duplicate column names in header {names}") - return names + return self.normalise_headers(header.lstrip("\ufeff").rstrip("\r\n").split(self.delimiter)) + + def _skip_row(self, row: pv.InvalidRow) -> Literal["skip"]: + self.skipped.append(f"expected {row.expected_columns} columns, got {row.actual_columns}: {row.text.rstrip()}") + return "skip" def _warn_if_everything_is_quoted(self) -> None: with open(self.path, "rb") as handle: @@ -117,6 +120,7 @@ def batches(self) -> Iterator[pa.RecordBatch]: quote_char='"' if self.quote == "csv" else False, escape_char=False, newlines_in_values=self.quote == "csv", + invalid_row_handler=self._skip_row if self.skip_malformed else None, ), convert_options=pv.ConvertOptions( column_types=column_types, @@ -130,7 +134,7 @@ def batches(self) -> Iterator[pa.RecordBatch]: for batch in reader: if batch.num_rows: yield batch - except pa.ArrowInvalid as exc: + except (pa.ArrowInvalid, UnicodeDecodeError) as exc: if "Empty CSV file" in str(exc): return hint = ( diff --git a/src/orm_loader/loaders/readers/detection.py b/src/orm_loader/loaders/readers/detection.py index 2814cdf..95eb0e4 100644 --- a/src/orm_loader/loaders/readers/detection.py +++ b/src/orm_loader/loaders/readers/detection.py @@ -2,10 +2,7 @@ from __future__ import annotations -import csv as _csv -import io import logging -from pathlib import Path import chardet @@ -19,101 +16,36 @@ def infer_encoding(file): return encoding def infer_delim(file): - with open(file, 'r') as infile: + with open(file, 'rb') as infile: line = infile.readline() - tabs = line.count('\t') - commas = line.count(',') - if tabs > commas: - return '\t' - return ',' - - -def infer_quote_mode( - path: Path, - delimiter: str, - encoding: str = "utf-8", - sample_rows: int = 2000, -) -> str: - """Return 'csv' or 'literal' by comparing column-count consistency under both - quoting interpretations across a sample of rows. - - - 'csv' → standard RFC-4180 quoting; surrounding double-quotes are stripped - and embedded delimiters/newlines inside quotes are preserved. - - 'literal' → double-quote has no special meaning; every byte is stored as-is. - - Defaults to 'csv' when both modes produce identical output (no quoting in play) - or when the evidence is tied. Callers can always override by passing an - explicit value instead of relying on auto-detection. - """ - with open(path, encoding=encoding, errors="replace", newline="") as f: - lines = [f.readline() for _ in range(sample_rows + 1)] - - raw = "".join(ln for ln in lines if ln) - if not raw: - return "csv" - - try: - rows_csv = list(_csv.reader(io.StringIO(raw), delimiter=delimiter)) - except _csv.Error: - return "literal" - - try: - rows_lit = list( - _csv.reader(io.StringIO(raw), delimiter=delimiter, quoting=_csv.QUOTE_NONE) - ) - except _csv.Error: - return "csv" + return '\t' if line.count(b'\t') > line.count(b',') else ',' - if not rows_csv: - return "csv" - ncols = len(rows_csv[0]) - if ncols <= 1: - return "csv" - - # No difference between modes → no quoting is active, csv is the safe default - if rows_csv == rows_lit: - return "csv" - - data_csv = rows_csv[1:] - data_lit = rows_lit[1:] if len(rows_lit) > 1 else [] - - if not data_csv: - return "csv" - - csv_ok = sum(1 for r in data_csv if len(r) == ncols) - lit_ok = sum(1 for r in data_lit if len(r) == ncols) - - # Prefer csv on a tie; only choose literal when it is strictly more consistent - return "literal" if lit_ok > csv_ok else "csv" - - -def resolve_quote_mode( - quote_mode: str, - path: Path, - delimiter: str, - encoding: str = "utf-8", -) -> str: +def resolve_quote_mode(quote_mode: str, delimiter: str) -> str: """Resolve a requested quote_mode to a concrete ``"csv"`` or ``"literal"``. - Recognised inputs: - - - ``"csv"`` / ``"literal"`` — returned unchanged (explicit override). - - ``"auto"`` — legacy content sniff via :func:`infer_quote_mode`. Compares - column-count consistency under both interpretations across a sample of - rows. Supports RFC-4180 quoting (embedded delimiters/newlines) but can - guess wrong when the discriminating rows fall outside the sample. - - ``"by_delimiter"`` — derive the mode from the delimiter alone, no content - scan: tab-delimited input never embeds the delimiter in a field so any - double-quote is literal data (``"literal"``); for non-tab delimiters this - resolves to ``"csv"`` where RFC-4180 quoting can be meaningful. Deterministic - and immune to the sampling gamble; use it when the source guarantees tab - fields never contain embedded tabs/newlines. + Parameters + ---------- + quote_mode + ``"csv"`` reads double quotes as RFC 4180 field quoting. ``"literal"`` + keeps them as data. ``"by_delimiter"`` chooses ``"literal"`` for tab + delimited files, where a field never embeds the delimiter, and + ``"csv"`` otherwise. + delimiter + Delimiter of the file. + + Returns + ------- + str + ``"csv"`` or ``"literal"``. + + Raises + ------ + ValueError + ``quote_mode`` is not one of the above. """ if quote_mode in ("csv", "literal"): return quote_mode if quote_mode == "by_delimiter": return "literal" if delimiter == "\t" else "csv" - if quote_mode == "auto": - return infer_quote_mode(path, delimiter=delimiter, encoding=encoding) raise ValueError(f"Unknown quote_mode: {quote_mode!r}") diff --git a/src/orm_loader/loaders/readers/parquet.py b/src/orm_loader/loaders/readers/parquet.py index 8fe2563..7e0ba50 100644 --- a/src/orm_loader/loaders/readers/parquet.py +++ b/src/orm_loader/loaders/readers/parquet.py @@ -17,6 +17,8 @@ class ParquetReader(FileReader): """A Parquet file read with its stored types. + Column names are normalised as described by ``normalise_headers``. + Parameters ---------- path @@ -41,12 +43,16 @@ def __init__( super().__init__(path, model_columns, nulls) self.chunksize = chunksize or 64_000 self._dataset = ds.dataset(path, format="parquet") - self.columns = [name for name in self._dataset.schema.names if name in model_columns] + self.file_columns = self.normalise_headers(self._dataset.schema.names) + stored = dict(zip(self.file_columns, self._dataset.schema.names)) + self.columns = [name for name in self.file_columns if name in model_columns] + self._stored = [stored[name] for name in self.columns] def batches(self) -> Iterator[pa.RecordBatch]: """Yield the stored rows in batches of ``chunksize`` rows.""" markers = self.reader_markers() - for batch in self._dataset.to_batches(columns=self.columns, batch_size=self.chunksize): + for batch in self._dataset.to_batches(columns=self._stored, batch_size=self.chunksize): + batch = batch.rename_columns(self.columns) if markers: batch = self.null_markers(batch, markers) if batch.num_rows: diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index c277c58..a16ed3c 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -312,6 +312,7 @@ def _select_reader(cls: Type[FileTableProtocol], ctx: LoaderContext) -> FileRead nulls=ctx.policy.nulls, quote_mode=ctx.quote_mode, chunksize=ctx.chunksize, + skip_malformed=ctx.policy.errors.on_error == "reject_row", ) @classmethod @@ -325,7 +326,7 @@ def load_file( dedupe: bool = False, chunksize: int | None = None, merge_strategy: str = "replace", - quote_mode: str = "auto", + quote_mode: str = "by_delimiter", index_strategy: str = "auto", merge_batch_size: int | None = None, staging_schema_tag: str | None = None, @@ -370,7 +371,7 @@ def load_file( quote_mode How double quotes in delimited text are read: ``"csv"`` (field quoting), ``"literal"`` (kept as data), ``"by_delimiter"`` - (literal for tab-delimited files) or ``"auto"`` (sample the file). + (literal for tab-delimited files, field quoting otherwise). index_strategy Index handling strategy during merge. Use ``"auto"`` to let the backend choose a sensible default. @@ -463,13 +464,13 @@ def load_file( report=report or LoadReport(), ) - # Load to staging (Indices are already excluded via updated create_staging_table) - logger.info(f"Table `{cls.__tablename__}`: Loading data into staging table") - total = cls.load_staging(loader_context, reader or cls._select_reader(loader_context)) - - # Merge staging to target (Wrapped in our index dropper!) - logger.info(f"Table `{cls.__tablename__}`: Merging staging data into target table") try: + # Load to staging (Indices are already excluded via updated create_staging_table) + logger.info(f"Table `{cls.__tablename__}`: Loading data into staging table") + total = cls.load_staging(loader_context, reader or cls._select_reader(loader_context)) + + # Merge staging to target (Wrapped in our index dropper!) + logger.info(f"Table `{cls.__tablename__}`: Merging staging data into target table") with cls.manage_indices(session, index_strategy=index_strategy, staging_schema_tag=staging_schema_tag): cls.merge_from_staging( session, diff --git a/src/orm_loader/tables/orm_table.py b/src/orm_loader/tables/orm_table.py index e015eb4..a5b9d31 100644 --- a/src/orm_loader/tables/orm_table.py +++ b/src/orm_loader/tables/orm_table.py @@ -141,8 +141,7 @@ def required_columns(cls) -> set[str]: A column is considered required if it is: - non-nullable - - has no Python-side default - - has no server-side default + - has no default that SQL can apply (see ``sql_defaults``) - not an autoincrementing column (the database assigns its value; it's never expected in inbound data at all) @@ -155,15 +154,41 @@ def required_columns(cls) -> set[str]: A set of required column names. """ mapper = cls.mapper_for() + defaults = cls.sql_defaults() return { c.key for c in mapper.columns - if not c.nullable - and not c.default - and not c.server_default - and not _resolves_to_autoincrement(c) + if not c.nullable and c.key not in defaults and not _resolves_to_autoincrement(c) } + @classmethod + def sql_defaults(cls) -> dict[str, sa.ColumnElement[Any]]: + """ + Return the default of each non-nullable column as a SQL expression. + + A server default given as text is cast to the column type, one given + as a SQL expression is used as written, and a scalar Python default + is a literal. A callable or context-sensitive default has no SQL + form, so its column is not listed. + + Returns + ------- + dict[str, sqlalchemy.ColumnElement] + A mapping of column name to the expression that supplies its default. + """ + defaults: dict[str, sa.ColumnElement[Any]] = {} + for c in cls.mapper_for().columns: + if c.nullable: + continue + server_default = getattr(c.server_default, "arg", None) + if isinstance(server_default, str): + defaults[c.name] = sa.cast(sa.literal(server_default), c.type) + elif server_default is not None: + defaults[c.name] = server_default + elif c.default is not None and c.default.is_scalar: + defaults[c.name] = sa.literal(c.default.arg, c.type) # ty: ignore[unresolved-attribute] + return defaults + @classmethod def max_id(cls, session: so.Session) -> int: """ diff --git a/src/orm_loader/tables/typing.py b/src/orm_loader/tables/typing.py index 12cd45f..c5ee018 100644 --- a/src/orm_loader/tables/typing.py +++ b/src/orm_loader/tables/typing.py @@ -51,6 +51,9 @@ def model_columns(cls) -> dict[str, sa.ColumnElement[Any]]: ... @classmethod def required_columns(cls) -> set[str]: ... + @classmethod + def sql_defaults(cls) -> dict[str, sa.ColumnElement[Any]]: ... + @runtime_checkable class FileTableProtocol(ORMTableProtocol, Protocol): """ @@ -87,7 +90,7 @@ def load_file( dedupe: bool = False, chunksize: int | None = None, merge_strategy: str = "replace", - quote_mode: str = "auto", + quote_mode: str = "by_delimiter", index_strategy: str = "auto", merge_batch_size: int | None = None, staging_schema_tag: str | None = None, diff --git a/tests/backends/test_sinks.py b/tests/backends/test_sinks.py index 32e074c..4f9e93d 100644 --- a/tests/backends/test_sinks.py +++ b/tests/backends/test_sinks.py @@ -45,7 +45,7 @@ def _stored(session, staging): def test_typed_values_and_nulls_are_stored_across_more_batches_than_a_commit_interval(dialect_session, tmp_path): backend, ctx, staging = _sink(dialect_session, tmp_path) batches = [_batch(i, f"row {i}" if i % 2 else None) for i in range(1, 13)] - assert backend.insert_batches(ctx, batches, COLUMNS) == 12 + backend.insert_batches(ctx, batches, COLUMNS) stored = _stored(dialect_session, staging) assert [r.demo_id for r in stored] == list(range(1, 13)) assert stored[0].demo_label == "row 1" and stored[1].demo_label is None @@ -62,7 +62,7 @@ def test_rows_keep_the_order_they_were_written_in(dialect_session, tmp_path): def test_no_batches_write_nothing(dialect_session, tmp_path): backend, ctx, staging = _sink(dialect_session, tmp_path) - assert backend.insert_batches(ctx, [], COLUMNS) == 0 + backend.insert_batches(ctx, [], COLUMNS) assert _stored(dialect_session, staging) == [] diff --git a/tests/conversion/test_date_handlers.py b/tests/conversion/test_date_handlers.py index bd457d0..c32f427 100644 --- a/tests/conversion/test_date_handlers.py +++ b/tests/conversion/test_date_handlers.py @@ -1,4 +1,4 @@ -from datetime import date, datetime +from datetime import date, datetime, timezone import numpy as np import pyarrow as pa @@ -137,6 +137,36 @@ def test_default_formats_are_mutually_distinguishable(): def test_typed_input_is_cast_without_parsing(): assert read(date(2020, 5, 17), sa.Date()) == date(2020, 5, 17) - assert read(datetime(2020, 5, 17, 14, 30), sa.Date()) == date(2020, 5, 17) + assert read(datetime(2020, 5, 17), sa.Date()) == date(2020, 5, 17) assert read(datetime(2020, 5, 17, 14, 30, 5), sa.DateTime()) == datetime(2020, 5, 17, 14, 30, 5) assert read(date(2020, 5, 17), sa.DateTime()) == datetime(2020, 5, 17) + + +def test_typed_timestamp_with_a_time_is_not_a_date(): + assert unreadable(datetime(2020, 5, 17, 14, 30), sa.Date()) + + +@pytest.mark.parametrize("sa_type", [sa.Date(), sa.DateTime()]) +def test_year_zero_is_unreadable(sa_type): + assert unreadable("0000-01-01", sa_type) + assert read("0001-01-01", sa_type) is not None + + +def test_a_timezone_column_requires_an_offset_and_stores_utc(): + aware = sa.DateTime(timezone=True) + assert read("2020-01-01T10:00:00+10:00", aware) == datetime(2020, 1, 1, tzinfo=timezone.utc) + assert read("2020-01-01 10:00:00Z", aware) == datetime(2020, 1, 1, 10, tzinfo=timezone.utc) + assert unreadable("2020-01-01 10:00:00", aware) + assert unreadable("20200101", aware) + + +def test_a_column_without_a_time_zone_rejects_an_offset(): + assert read("2020-01-01 10:00:00", sa.DateTime()) == datetime(2020, 1, 1, 10) + assert unreadable("2020-01-01T10:00:00+10:00", sa.DateTime()) + assert unreadable("2020-01-01T10:00:00Z", sa.DateTime()) + + +def test_a_batch_with_one_bad_date_isolates_it(): + values, invalid = interpret(["20200101"] * 1000 + ["20200231"] + ["20200102"] * 1000, sa.Date()) + assert [i for i, bad in enumerate(invalid) if bad] == [1000] + assert values[0] == date(2020, 1, 1) and values[-1] == date(2020, 1, 2) diff --git a/tests/conversion/test_number_handlers.py b/tests/conversion/test_number_handlers.py index afcdad3..23c3dc3 100644 --- a/tests/conversion/test_number_handlers.py +++ b/tests/conversion/test_number_handlers.py @@ -1,12 +1,14 @@ +import time from decimal import Decimal +import pyarrow as pa import pytest import sqlalchemy as sa -from tests.support import read, unreadable +from tests.support import interpret, read, unreadable -@pytest.mark.parametrize("raw", ["12abc", "0x1F", "1.0", "1e3", "--1", "Infinity"]) +@pytest.mark.parametrize("raw", ["12abc", "0x1F", "1.5", "1e3", "--1", "Infinity", "1.0.0"]) def test_integer_rejects_non_integer_text(raw): assert unreadable(raw, sa.Integer()) @@ -38,3 +40,88 @@ def test_integer_trims_padding_and_reads_blank_and_missing_words_as_null(): def test_typed_input_is_cast_without_parsing(): assert read(42.0, sa.Integer()) == 42 assert unreadable(42.5, sa.Integer()) + + +@pytest.mark.parametrize("raw", ["1.0", "1.00", "1.", "+2.0", "-3.000"]) +def test_integer_accepts_integer_valued_decimals(raw): + assert read(raw, sa.Integer()) == int(float(raw)) + + +@pytest.mark.parametrize( + "sa_type, low, high", + [ + (sa.SmallInteger(), -(2**15), 2**15 - 1), + (sa.Integer(), -(2**31), 2**31 - 1), + (sa.BigInteger(), -(2**63), 2**63 - 1), + ], +) +def test_integer_width_follows_the_column_type(sa_type, low, high): + assert read(str(low), sa_type) == low + assert read(str(high), sa_type) == high + assert unreadable(str(low - 1), sa_type) + assert unreadable(str(high + 1), sa_type) + + +def test_integer_far_beyond_64_bits_is_unreadable(): + assert unreadable("1" + "0" * 40, sa.BigInteger()) + + +def test_float_rejects_infinity_from_text_overflow(): + assert unreadable("1e999", sa.Float()) + assert read("1e300", sa.Float()) == 1e300 + + +def test_numeric_precision_without_scale_has_scale_zero(): + assert read("12345678901234567890", sa.Numeric(30)) == Decimal("12345678901234567890") + assert unreadable("1.5", sa.Numeric(30)) + + +def test_numeric_without_precision_is_exact_and_rejects_what_does_not_fit(): + assert read("12345678901234567890.123456789", sa.Numeric()) == Decimal("12345678901234567890.123456789000000000") + assert unreadable("1" + "0" * 30, sa.Numeric()) + assert unreadable("0.1234567890123456789", sa.Numeric()) + + +def test_numeric_digits_beyond_the_precision_do_not_wrap_around(): + assert unreadable("9" * 39, sa.Numeric(38, 0)) + assert unreadable("9" * 45, sa.Numeric(38, 0)) + assert read("9" * 38, sa.Numeric(38, 0)) == Decimal("9" * 38) + + +def _reference_integer(raw: str, bits: int) -> int | None: + text = raw.strip() + sign = -1 if text.startswith("-") else 1 + digits = text[1:] if text[:1] in ("+", "-") else text + whole, _, fraction = digits.partition(".") + if not (whole.isascii() and whole.isdigit()) or (fraction and set(fraction) != {"0"}) or text.count(".") > 1: + return None + value = sign * int(whole) + return value if -(2 ** (bits - 1)) <= value < 2 ** (bits - 1) else None + + +EDGE_INTEGERS = [ + "0", "-0", "+0", "007", "-007", "1.0", "1.", ".5", "5.", "1.50", "+7", "--1", "1e3", "", + "32767", "32768", "-32768", "-32769", "2147483647", "2147483648", "-2147483649", + "9223372036854775807", "9223372036854775808", "-9223372036854775808", "-9223372036854775809", + "0" * 30 + "12", "9" * 25, "12abc", " 7 ", "1_000", "١٢", +] + + +@pytest.mark.parametrize("sa_type, bits", [(sa.SmallInteger(), 16), (sa.Integer(), 32), (sa.BigInteger(), 64)]) +def test_integer_parsing_matches_a_per_element_reference(sa_type, bits): + values, invalid = interpret(EDGE_INTEGERS, sa_type) + for raw, value, bad in zip(EDGE_INTEGERS, values, invalid): + if raw.strip() == "": + assert value is None and not bad, raw + continue + expected = _reference_integer(raw, bits) + assert (value, bad) == ((expected, False) if expected is not None else (None, True)), raw + + +def test_one_bad_value_in_a_large_batch_does_not_slow_the_whole_batch(): + raw = pa.array([str(i) for i in range(200_000)]) + raw = pa.concat_arrays([raw.slice(0, 100_000), pa.array(["9" * 25]), raw.slice(100_000)]) + started = time.perf_counter() + values, invalid = interpret(raw, sa.BigInteger()) + assert time.perf_counter() - started < 0.5 + assert sum(invalid) == 1 and invalid[100_000] diff --git a/tests/conversion/test_policy.py b/tests/conversion/test_policy.py index 34ab58b..4b97414 100644 --- a/tests/conversion/test_policy.py +++ b/tests/conversion/test_policy.py @@ -8,6 +8,7 @@ ConvertingHandler, DateHandler, DateTimeHandler, + EnumHandler, ErrorPolicy, FloatHandler, IntegerHandler, @@ -17,6 +18,7 @@ PassthroughHandler, TextHandler, ) +from orm_loader.helpers import PolicyError from tests.support import interpret, read, unreadable @@ -29,7 +31,7 @@ class _Money(sa.Numeric): [ (sa.String(), TextHandler), (sa.Text(), TextHandler), - (sa.Enum("a", "b"), TextHandler), + (sa.Enum("a", "b"), EnumHandler), (sa.Integer(), IntegerHandler), (sa.BigInteger(), IntegerHandler), (sa.Float(), FloatHandler), @@ -39,16 +41,30 @@ class _Money(sa.Numeric): (sa.Date(), DateHandler), (sa.DateTime(), DateTimeHandler), (sa.TIMESTAMP(), DateTimeHandler), - (sa.LargeBinary(), PassthroughHandler), ], ) def test_the_most_specific_type_class_selects_the_handler(sa_type, handler): assert type(ConversionPolicy().for_column(sa_type)) is handler -def test_a_type_without_a_handler_is_returned_unchanged(): +class _Wrapped(sa.TypeDecorator): + impl = sa.Integer + cache_ok = True + + +def test_a_type_without_a_handler_is_an_error_naming_the_column(): + with pytest.raises(PolicyError, match="t.c: no handler for column type LargeBinary"): + ConversionPolicy().for_column(sa.LargeBinary(), "t", "c") + + +def test_passthrough_is_an_explicit_choice_for_a_column(): + policy = ConversionPolicy(columns={("t", "c"): PassthroughHandler()}) raw = pa.array([b"a", None]) - assert interpret(raw, sa.LargeBinary())[0] == [b"a", None] + assert interpret(raw, sa.LargeBinary(), policy, table_name="t", column_name="c")[0] == [b"a", None] + + +def test_a_type_decorator_is_handled_as_the_type_it_stores(): + assert type(ConversionPolicy().for_column(_Wrapped())) is IntegerHandler class _Hex(ConvertingHandler): @@ -90,13 +106,10 @@ def test_column_handlers_win_over_type_handlers(): assert read("7", sa.String(), policy, table_name="t", column_name="other") == "7" -def test_text_markers_in_a_text_handler_replace_the_policy_markers(): - policy = ConversionPolicy( - nulls=NullPolicy(text_markers=frozenset({"NULL"})), - text=TextHandler(markers=frozenset({"-"})), - ) - assert read("-", sa.String(), policy) is None - assert read("NULL", sa.String(), policy) == "NULL" +def test_the_handler_and_column_mappings_cannot_be_changed_afterwards(): + policy = ConversionPolicy(columns={("t", "c"): IntegerHandler()}) + with pytest.raises(TypeError): + policy.columns[("t", "d")] = IntegerHandler() # type: ignore[index] def test_unknown_error_mode_is_rejected(): diff --git a/tests/conversion/test_text_handler.py b/tests/conversion/test_text_handler.py index dc923da..3982dbc 100644 --- a/tests/conversion/test_text_handler.py +++ b/tests/conversion/test_text_handler.py @@ -1,8 +1,9 @@ +import pyarrow.compute as pc import pytest import sqlalchemy as sa from orm_loader.conversion import ConversionPolicy, NullPolicy, TextHandler -from tests.support import read, unreadable +from tests.support import interpret, read, unreadable MARKER = ConversionPolicy(nulls=NullPolicy(text_markers=frozenset({"NULL"}))) @@ -30,15 +31,20 @@ def test_default_policy_has_no_text_markers(): assert read("NULL", sa.String()) == "NULL" -def test_column_markers_replace_the_import_markers_for_their_column(): +class _DashIsMissing(TextHandler): + def missing(self, trimmed, nulls): + return pc.equal(trimmed, "-") + + +def test_a_text_handler_subclass_can_choose_its_own_missing_values(): policy = ConversionPolicy( nulls=NullPolicy(text_markers=frozenset({"NULL"})), - columns={("t", "a"): TextHandler(markers=frozenset({"N/A"}))}, + columns={("t", "a"): _DashIsMissing()}, ) + assert read("-", sa.String(), policy, table_name="t", column_name="a") is None assert read("NULL", sa.String(), policy, table_name="t", column_name="a") == "NULL" - assert read("N/A", sa.String(), policy, table_name="t", column_name="a") is None assert read("NULL", sa.String(), policy, table_name="t", column_name="b") is None - assert read("N/A", sa.String(), policy, table_name="t", column_name="b") == "N/A" + assert read("-", sa.String(), policy, table_name="t", column_name="b") == "-" def test_policies_do_not_leak_between_conversions(): @@ -58,3 +64,10 @@ def test_blank_text_is_kept_as_empty_when_empty_is_not_null(): def test_text_longer_than_the_column_raises(): assert unreadable("ab", sa.String(1)) + + +def test_an_enum_column_accepts_only_its_labels(): + enum = sa.Enum("red", "green", name="colour") + values, invalid = interpret(["red", " green ", "blue", ""], enum) + assert values == ["red", "green", None, None] + assert invalid == [False, False, True, False] diff --git a/tests/conversion/test_typed_input.py b/tests/conversion/test_typed_input.py new file mode 100644 index 0000000..37a17f6 --- /dev/null +++ b/tests/conversion/test_typed_input.py @@ -0,0 +1,74 @@ +"""Input that already has a type, such as a Parquet column, is cast and checked like parsed text.""" + +from datetime import date, datetime, timezone +from decimal import Decimal + +import pyarrow as pa +import sqlalchemy as sa + +from tests.support import interpret + + +def _result(values: pa.Array, sa_type): + return interpret(values, sa_type) + + +def test_a_boolean_column_accepts_only_zero_and_one_from_integers(): + values, invalid = _result(pa.array([0, 1, 2, None]), sa.Boolean()) + assert values == [False, True, None, None] + assert invalid == [False, False, True, False] + + +def test_a_date_column_accepts_only_midnight_timestamps(): + midnight, noon = datetime(2020, 5, 17), datetime(2020, 5, 17, 12) + values, invalid = _result(pa.array([midnight, noon], pa.timestamp("us")), sa.Date()) + assert values == [date(2020, 5, 17), None] + assert invalid == [False, True] + + +def test_the_length_of_text_cast_from_another_type_is_checked(): + values, invalid = _result(pa.array([123, 123456]), sa.String(5)) + assert values == ["123", None] + assert invalid == [False, True] + + +def test_integer_width_is_checked_on_typed_input(): + _, invalid = _result(pa.array([70_000, 1], pa.int64()), sa.SmallInteger()) + assert invalid == [True, False] + + +def test_floats_cannot_be_truncated_to_integers(): + values, invalid = _result(pa.array([1.0, 1.5]), sa.Integer()) + assert values == [1, None] + assert invalid == [False, True] + + +def test_float_nan_is_missing_and_infinity_is_unreadable(): + values, invalid = _result(pa.array([float("nan"), float("inf"), 2.5]), sa.Float()) + assert values == [None, None, 2.5] + assert invalid == [False, True, False] + + +def test_decimals_that_lose_digits_are_unreadable(): + values, invalid = _result(pa.array([Decimal("1.25"), Decimal("1.234")], pa.decimal128(10, 3)), sa.Numeric(10, 2)) + assert values == [Decimal("1.25"), None] + assert invalid == [False, True] + + +def test_an_all_null_column_of_unknown_type_is_missing_values(): + values, invalid = _result(pa.array([None, None]), sa.Integer()) + assert values == [None, None] + assert invalid == [False, False] + + +def test_timestamps_with_and_without_a_time_zone_match_their_column(): + aware = pa.array([datetime(2020, 1, 1, 10, tzinfo=timezone.utc)], pa.timestamp("us", "UTC")) + naive = pa.array([datetime(2020, 1, 1, 10)], pa.timestamp("us")) + assert _result(aware, sa.DateTime(timezone=True)) == ([datetime(2020, 1, 1, 10, tzinfo=timezone.utc)], [False]) + assert _result(naive, sa.DateTime()) == ([datetime(2020, 1, 1, 10)], [False]) + + +def test_year_zero_dates_are_unreadable_on_typed_input(): + year_zero, year_two_thousand = -719_528, 10_957 + _, invalid = _result(pa.array([year_zero, year_two_thousand], pa.date32()), sa.Date()) + assert invalid == [True, False] diff --git a/tests/loaders/test_defaults.py b/tests/loaders/test_defaults.py new file mode 100644 index 0000000..cd54f41 --- /dev/null +++ b/tests/loaders/test_defaults.py @@ -0,0 +1,69 @@ +"""A non-nullable column with a default takes it when the value is empty or the column is absent.""" + +from pathlib import Path +from typing import Type, cast + +import pytest +import sqlalchemy as sa + +from orm_loader.backends import STAGING_SCHEMA +from orm_loader.conversion import ConversionPolicy, ErrorPolicy +from orm_loader.helpers import IngestError, RequiredValueError +from orm_loader.tables.typing import FileTableProtocol +from tests.models import DemoDefaultedRecord + +_Table = cast(Type[FileTableProtocol], DemoDefaultedRecord) +FULL = "demo_id,demo_status,demo_priority,demo_stamp,demo_note\n" + + +def _load(session, tmp_path: Path, text: str, **kwargs) -> int: + path = tmp_path / "demo_defaulted_record.csv" + path.write_text(text, encoding="utf-8") + tag = None if session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA + loaded = _Table.load_file(session, path, staging_schema_tag=tag, **kwargs) + session.commit() + session.expire_all() + return loaded + + +def _rows(session) -> list[tuple]: + model = DemoDefaultedRecord + columns = (model.demo_id, model.demo_status, model.demo_priority, model.demo_stamp, model.demo_note) + return [tuple(r) for r in session.execute(sa.select(*columns).order_by(model.demo_id))] + + +def test_only_a_default_that_sql_can_apply_makes_a_column_optional(): + assert _Table.required_columns() == {"demo_stamp"} + assert set(_Table.sql_defaults()) == {"demo_status", "demo_priority"} + + +def test_an_empty_value_takes_the_default(dialect_session, tmp_path): + _load(dialect_session, tmp_path, FULL + "1,,,s,\n2,closed,9,s,kept\n") + assert _rows(dialect_session) == [(1, "open", 5, "s", None), (2, "closed", 9, "s", "kept")] + + +def test_an_absent_column_takes_the_default(dialect_session, tmp_path): + _load(dialect_session, tmp_path, "demo_id,demo_stamp\n1,s\n2,t\n") + assert _rows(dialect_session) == [(1, "open", 5, "s", None), (2, "open", 5, "t", None)] + + +def test_a_column_with_a_callable_default_stays_required(dialect_session, tmp_path): + with pytest.raises(RequiredValueError, match="demo_stamp"): + _load(dialect_session, tmp_path, FULL + "1,,,,\n") + with pytest.raises(IngestError, match="demo_stamp"): + _load(dialect_session, tmp_path, "demo_id,demo_status\n1,open\n") + + +def test_a_missing_callable_default_can_be_rejected_by_row(dialect_session, tmp_path): + policy = ConversionPolicy(errors=ErrorPolicy(on_error="reject_row")) + _load(dialect_session, tmp_path, FULL + "1,,,,\n2,,,s,\n", policy=policy) + assert [r[0] for r in _rows(dialect_session)] == [2] + + +@pytest.mark.parametrize("strategy", ["replace", "upsert"]) +def test_a_reload_applies_the_default_to_an_existing_row(dialect_session, tmp_path, strategy): + _load(dialect_session, tmp_path, FULL + "1,closed,9,s,kept\n") + _load(dialect_session, tmp_path, FULL + "1,,,t,\n2,,,u,\n", merge_strategy=strategy) + expected = {"replace": [(1, "open", 5, "t", None), (2, "open", 5, "u", None)], + "upsert": [(1, "closed", 9, "s", "kept"), (2, "open", 5, "u", None)]} + assert _rows(dialect_session) == expected[strategy] diff --git a/tests/loaders/test_helpers.py b/tests/loaders/test_helpers.py index cd2dade..615d23f 100644 --- a/tests/loaders/test_helpers.py +++ b/tests/loaders/test_helpers.py @@ -1,11 +1,8 @@ -from pathlib import Path - import pytest from orm_loader.loaders.readers.detection import ( infer_delim, infer_encoding, - infer_quote_mode, resolve_quote_mode, ) @@ -33,78 +30,36 @@ def test_infer_encoding_utf8(tmp_path): assert enc.lower() in {"utf-8", "ascii"} -def test_infer_quote_mode_unquoted_tsv_returns_csv(tmp_path): - # No quotes anywhere: both modes identical, csv is the safe default - p = tmp_path / "x.csv" - p.write_text("id\tname\tvalue\n1\tAlice\t10\n2\tBob\t20\n") - assert infer_quote_mode(p, delimiter="\t") == "csv" - - -def test_infer_quote_mode_rfc4180_quoted_field_returns_csv(tmp_path): - # Athena-style: quoted concept_name at the varchar(255) boundary, - # no embedded delimiter — the column-count tie-break must favour csv - p = tmp_path / "x.csv" - long_name = "A" * 255 - p.write_text(f'id\tname\n1\t"{long_name}"\n2\tnormal\n') - assert infer_quote_mode(p, delimiter="\t") == "csv" - - -def test_infer_quote_mode_embedded_delimiter_in_quoted_field_returns_csv(tmp_path): - # Quoted field contains the delimiter: csv mode keeps column count consistent, - # literal mode splits on the embedded tab and produces ragged rows - p = tmp_path / "x.csv" - p.write_text('id\tname\tval\n1\t"foo\tbar"\t99\n2\tbaz\t0\n') - assert infer_quote_mode(p, delimiter="\t") == "csv" - - -def test_infer_quote_mode_unbalanced_quote_returns_literal(tmp_path): - # Unbalanced leading quote breaks CSV parsing: literal mode produces - # consistent 2-column rows while csv mode does not - p = tmp_path / "x.csv" - p.write_text('id\tname\n1\t"open\n2\t"open\n3\t"open\n') - assert infer_quote_mode(p, delimiter="\t") == "literal" - - # --------------------------------------------------------------------------- # resolve_quote_mode # --------------------------------------------------------------------------- def test_resolve_quote_mode_explicit_csv_passthrough(): # Explicit modes are returned unchanged and never read the file - assert resolve_quote_mode("csv", Path("does-not-exist.csv"), ",") == "csv" + assert resolve_quote_mode("csv", ",") == "csv" def test_resolve_quote_mode_explicit_literal_passthrough(): - assert resolve_quote_mode("literal", Path("does-not-exist.csv"), "\t") == "literal" + assert resolve_quote_mode("literal", "\t") == "literal" def test_resolve_quote_mode_by_delimiter_tab_is_literal(): # Tab-delimited input never embeds the delimiter in a field, so any # double-quote is literal data — derived from the delimiter, no file scan - assert resolve_quote_mode("by_delimiter", Path("does-not-exist.tsv"), "\t") == "literal" + assert resolve_quote_mode("by_delimiter", "\t") == "literal" def test_resolve_quote_mode_by_delimiter_comma_is_csv(): # Comma-delimited input must wrap fields containing the delimiter in quotes, # so quoting is meaningful - assert resolve_quote_mode("by_delimiter", Path("does-not-exist.csv"), ",") == "csv" + assert resolve_quote_mode("by_delimiter", ",") == "csv" -def test_resolve_quote_mode_auto_delegates_to_infer(tmp_path): - # "auto" defers to the content sniff: an unbalanced leading quote is only - # consistent under literal parsing - p = tmp_path / "x.csv" - p.write_text('id\tname\n1\t"open\n2\t"open\n3\t"open\n') - assert resolve_quote_mode("auto", p, delimiter="\t") == "literal" - - -def test_resolve_quote_mode_auto_clean_file_returns_csv(tmp_path): - # No quotes anywhere: the sniff falls back to the csv default - p = tmp_path / "x.csv" - p.write_text("id\tname\n1\tAlice\n2\tBob\n") - assert resolve_quote_mode("auto", p, delimiter="\t") == "csv" +def test_resolve_quote_mode_auto_is_not_a_mode(): + with pytest.raises(ValueError, match="Unknown quote_mode"): + resolve_quote_mode("auto", "\t") -def test_resolve_quote_mode_unknown_raises(tmp_path): +def test_resolve_quote_mode_unknown_raises(): with pytest.raises(ValueError, match="Unknown quote_mode"): - resolve_quote_mode("bogus", tmp_path / "x.csv", ",") + resolve_quote_mode("bogus", ",") diff --git a/tests/loaders/test_load_policy.py b/tests/loaders/test_load_policy.py index 5f3f2da..486ebd2 100644 --- a/tests/loaders/test_load_policy.py +++ b/tests/loaders/test_load_policy.py @@ -12,7 +12,7 @@ import sqlalchemy as sa from orm_loader.backends import STAGING_SCHEMA -from orm_loader.conversion import ConversionPolicy, DateHandler, ErrorPolicy, NullPolicy, TextHandler +from orm_loader.conversion import ConversionPolicy, DateHandler, ErrorPolicy, NullPolicy from orm_loader.helpers import ConversionError, IngestError, RejectLimitError, RequiredValueError from orm_loader.backends import DedupeResult from orm_loader.loaders import LoadReport @@ -62,12 +62,6 @@ def test_marker_in_a_required_column_raises(dialect_session, tmp_path): _load(dialect_session, tmp_path, _DemoRequired, text, policy=MARKER) -def test_column_markers_apply_only_to_their_column(dialect_session, tmp_path): - policy = ConversionPolicy(columns={("demo_record", "demo_label"): TextHandler(markers=frozenset({"N/A"}))}) - _load(dialect_session, tmp_path, _Demo, "demo_id,demo_label\n1,N/A\n2,NULL\n", policy=policy) - assert _labels(dialect_session, DemoRecord) == {1: None, 2: "NULL"} - - def test_quote_sensitive_marker_matches_only_unquoted(dialect_session, tmp_path): policy = ConversionPolicy(nulls=NullPolicy(text_markers=frozenset({"NULL"}), marker_quoting="unquoted_only")) _load(dialect_session, tmp_path, _Demo, SEVEN_ROWS, policy=policy) @@ -85,14 +79,6 @@ def test_quote_sensitive_marker_warns_when_every_field_is_quoted(dialect_session assert _labels(dialect_session, DemoRecord) == {1: "NULL", 2: "x"} -def test_quote_sensitive_markers_reject_column_markers(): - with pytest.raises(ValueError, match="unquoted_only"): - ConversionPolicy( - nulls=NullPolicy(marker_quoting="unquoted_only"), - columns={("t", "c"): TextHandler(markers=frozenset({"x"}))}, - ) - - def test_both_athena_date_formats_and_iso_load_in_one_file(dialect_session, tmp_path): text = "demo_id,demo_date\n1,20200229\n2,01-JAN-2020\n3,2020-01-02\n4,\n5,NA\n" _load(dialect_session, tmp_path, _Demo, text) diff --git a/tests/loaders/test_loader_e2e.py b/tests/loaders/test_loader_e2e.py index 725a72d..3ea4b6a 100644 --- a/tests/loaders/test_loader_e2e.py +++ b/tests/loaders/test_loader_e2e.py @@ -485,6 +485,7 @@ class TextTable(Base, FileLoadableTableInterface): sqlite_session, csv, dedupe=False, + quote_mode="csv", ) sqlite_session.commit() @@ -510,6 +511,7 @@ class TextTable2(Base, FileLoadableTableInterface): sqlite_session, csv, dedupe=False, + quote_mode="csv", ) sqlite_session.commit() diff --git a/tests/loaders/test_malformed_files.py b/tests/loaders/test_malformed_files.py new file mode 100644 index 0000000..19df9d3 --- /dev/null +++ b/tests/loaders/test_malformed_files.py @@ -0,0 +1,123 @@ +"""Files that are broken as a whole: ragged lines, undecodable bytes, stale rejects and ignored columns.""" + +import logging +from pathlib import Path +from typing import Type, cast + +import pyarrow as pa +import pyarrow.csv as pv +import pyarrow.parquet as pq +import pytest +import sqlalchemy as sa + +from orm_loader.backends import STAGING_SCHEMA, resolve_backend +from orm_loader.conversion import ConversionPolicy, ErrorPolicy +from orm_loader.helpers import IngestError +from orm_loader.loaders import DelimitedTextReader, LoadReport, ParquetReader +from orm_loader.tables.typing import FileTableProtocol +from tests.models import DemoRecord + +_Demo = cast(Type[FileTableProtocol], DemoRecord) +REJECT = ConversionPolicy(errors=ErrorPolicy(on_error="reject_row")) +COLUMNS = DemoRecord.model_columns() + + +def _load(session, tmp_path: Path, text: str, **kwargs) -> int: + path = tmp_path / "demo_record.csv" + path.write_text(text, encoding="utf-8", newline="") + tag = None if session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA + loaded = _Demo.load_file(session, path, staging_schema_tag=tag, **kwargs) + session.commit() + return loaded + + +def _ids(session) -> list[int]: + return [r[0] for r in session.execute(sa.select(DemoRecord.demo_id).order_by(DemoRecord.demo_id))] + + +def _rejects(path: Path) -> list[dict]: + options = pv.ConvertOptions(column_types={"row": pa.int64(), "column": pa.string(), "reason": pa.string()}) + return pv.read_csv(path, convert_options=options).to_pylist() + + +RAGGED = "demo_id,demo_label\n1,a\n2\n3,c,extra\n4,d\n" + + +def test_a_ragged_line_fails_the_read_by_default(dialect_session, tmp_path): + with pytest.raises(IngestError, match="demo_record.csv.*Expected 2 columns"): + _load(dialect_session, tmp_path, RAGGED) + + +def test_ragged_lines_are_rejected_with_their_text_under_reject_row(dialect_session, tmp_path): + report = LoadReport() + assert _load(dialect_session, tmp_path, RAGGED, policy=REJECT, report=report) == 2 + assert _ids(dialect_session) == [1, 4] + assert report.rows_read == 4 and report.rows_rejected == 2 and report.accounted + rejects = _rejects(report.rejects_path) + assert [r["reason"] for r in rejects] == [ + "expected 2 columns, got 1: 2", + "expected 2 columns, got 3: 3,c,extra", + ] + assert all(r["row"] is None and not r["column"] for r in rejects) + + +def test_a_file_of_only_ragged_lines_is_rejected_not_silently_empty(dialect_session, tmp_path): + report = LoadReport() + assert _load(dialect_session, tmp_path, "demo_id,demo_label\n1\n2\n", policy=REJECT, report=report) == 0 + assert report.rows_read == 2 and report.rows_rejected == 2 and report.accounted + + +def test_undecodable_bytes_fail_with_the_file_name(tmp_path): + path = tmp_path / "demo_record.csv" + path.write_bytes(b"demo_id,demo_label\n1,a\x81\n") + reader = DelimitedTextReader(path, COLUMNS) + reader.encoding = "cp1252" + with pytest.raises(IngestError, match="demo_record.csv"): + list(reader.batches()) + + +def test_a_clean_reload_removes_the_rejects_file_of_the_previous_load(dialect_session, tmp_path): + first = LoadReport() + _load(dialect_session, tmp_path, RAGGED, policy=REJECT, report=first) + assert first.rejects_path.exists() + second = LoadReport() + _load(dialect_session, tmp_path, "demo_id,demo_label\n1,a\n", policy=REJECT, report=second) + assert second.rejects_path is None + assert not first.rejects_path.exists() + + +def test_a_failed_read_leaves_no_staging_table_behind(dialect_session, tmp_path): + with pytest.raises(IngestError): + _load(dialect_session, tmp_path, RAGGED) + backend = resolve_backend(dialect_session, mapper=DemoRecord) + name = backend.staging_name_for_table(DemoRecord.__tablename__) + schema = None if dialect_session.get_bind().dialect.name == "sqlite" else backend.staging_schema + assert not sa.inspect(dialect_session.get_bind()).has_table(name, schema=schema) + + +def test_file_columns_the_table_lacks_are_reported(dialect_session, tmp_path, caplog): + with caplog.at_level(logging.WARNING): + _load(dialect_session, tmp_path, "demo_id,demo_labl,demo_label\n1,typo,a\n") + assert "which are ignored: ['demo_labl']" in caplog.text + + +def test_parquet_column_names_are_normalised_like_delimited_headers(tmp_path): + path = tmp_path / "demo_record.parquet" + pq.write_table(pa.table({"Demo_ID": [1], '"DEMO_LABEL_hash"': ["a"], "Other": [1]}), path) + reader = ParquetReader(path, COLUMNS) + assert reader.columns == ["demo_id", "demo_label"] + assert reader.file_columns == ["demo_id", "demo_label", "other"] + assert [b.schema.names for b in reader.batches()] == [["demo_id", "demo_label"]] + + +def test_a_hash_suffix_is_stripped_only_at_the_end_of_a_name(tmp_path): + path = tmp_path / "demo_record.csv" + path.write_text("demo_id,my_hash_value,demo_label_hash\n1,x,a\n") + assert DelimitedTextReader(path, COLUMNS).file_columns == ["demo_id", "my_hash_value", "demo_label"] + + +def test_duplicate_names_after_normalisation_are_an_error(tmp_path): + path = tmp_path / "demo_record.csv" + path.write_text("demo_id,Demo_ID\n1,2\n") + with pytest.raises(IngestError, match="duplicate column names"): + DelimitedTextReader(path, COLUMNS) diff --git a/tests/loaders/test_pipeline.py b/tests/loaders/test_pipeline.py index f059ce1..c2daaf9 100644 --- a/tests/loaders/test_pipeline.py +++ b/tests/loaders/test_pipeline.py @@ -2,7 +2,8 @@ import csv import io -from datetime import date +from datetime import date, datetime, timezone +from decimal import Decimal from pathlib import Path from typing import Type, cast @@ -16,10 +17,11 @@ from orm_loader.helpers import ConversionError, RequiredValueError from orm_loader.loaders import DelimitedTextReader, LoaderContext, ParquetReader, stage_file from orm_loader.tables.typing import FileTableProtocol -from tests.models import DemoRecord, DemoRequiredRecord +from tests.models import DemoRecord, DemoRequiredRecord, DemoTypedRecord _Demo = cast(Type[FileTableProtocol], DemoRecord) _DemoRequired = cast(Type[FileTableProtocol], DemoRequiredRecord) +_Typed = cast(Type[FileTableProtocol], DemoTypedRecord) HEADER = ["demo_id", "demo_label", "demo_number", "demo_flag", "demo_date"] REJECT = ConversionPolicy(errors=ErrorPolicy(on_error="reject_row")) @@ -142,3 +144,93 @@ def broken(): _stage(dialect_session, _Demo, path, reader=reader) tag = None if dialect_session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA assert _staged(dialect_session, _Demo.get_staging_table(dialect_session, staging_schema_tag=tag), "demo_id") == [] + + +TYPED_HEADER = [ + "typed_id", "typed_small", "typed_float", "typed_numeric", "typed_flag", + "typed_date", "typed_stamp", "typed_stamp_utc", "typed_label", +] +UTC = timezone.utc + + +def _typed_file(tmp_path: Path, **overrides: pa.Array) -> Path: + columns = { + "typed_id": pa.array([1, 2], pa.int64()), + "typed_small": pa.array([7, None], pa.int32()), + "typed_float": pa.array([1.5, float("nan")], pa.float32()), + "typed_numeric": pa.array([Decimal("12.34"), None], pa.decimal128(10, 2)), + "typed_flag": pa.array([True, None], pa.bool_()), + "typed_date": pa.array([date(2020, 2, 29), None], pa.date32()), + "typed_stamp": pa.array([datetime(2020, 1, 1, 10, 30), None], pa.timestamp("us")), + "typed_stamp_utc": pa.array([datetime(2020, 1, 1, 10, 30, tzinfo=UTC), None], pa.timestamp("us", "UTC")), + "typed_label": pa.array(["alpha", None], pa.string()), + **overrides, + } + path = tmp_path / "demo_typed_record.parquet" + pq.write_table(pa.table(columns), path) + return path + + +def test_typed_parquet_columns_keep_their_values_on_every_dialect(dialect_session, tmp_path): + loaded, ctx, staging = _stage(dialect_session, _Typed, _typed_file(tmp_path)) + assert loaded == 2 and ctx.report.accounted + rows = _staged(dialect_session, staging, *TYPED_HEADER) + assert rows[0] == ( + 1, 7, 1.5, Decimal("12.34"), True, date(2020, 2, 29), + datetime(2020, 1, 1, 10, 30), datetime(2020, 1, 1, 10, 30, tzinfo=UTC), "alpha", + ) + assert rows[1] == (2, None, None, None, None, None, None, None, None) + + +@pytest.mark.parametrize( + "column, bad", + [ + ("typed_flag", pa.array([1, 2], pa.int64())), + ("typed_date", pa.array([datetime(2020, 1, 1), datetime(2020, 1, 1, 12)], pa.timestamp("us"))), + ("typed_small", pa.array([1, 70_000], pa.int64())), + ("typed_label", pa.array([1, 12345678901], pa.int64())), + ("typed_numeric", pa.array([Decimal("1.25"), Decimal("1.234")], pa.decimal128(10, 3))), + ], +) +def test_typed_parquet_values_that_would_change_are_unreadable(dialect_session, tmp_path, column, bad): + with pytest.raises(ConversionError, match=column): + _stage(dialect_session, _Typed, _typed_file(tmp_path, **{column: bad})) + loaded, ctx, _ = _stage(dialect_session, _Typed, _typed_file(tmp_path, **{column: bad}), policy=REJECT) + assert loaded == 1 and ctx.report.rows_rejected == 1 and ctx.report.accounted + + +def _typed_rows(**cells: str) -> list[tuple]: + return [tuple(cells.get(name, "1" if name == "typed_id" else "") for name in TYPED_HEADER)] + + +def test_offsets_are_stored_as_utc_in_a_time_zone_column_and_rejected_in_a_naive_one(dialect_session, tmp_path): + rows = _typed_rows(typed_stamp_utc="2020-01-01T10:00:00+10:00", typed_stamp="2020-01-01 10:00:00") + _, _, staging = _stage(dialect_session, _Typed, _write(tmp_path, DemoTypedRecord, "csv", rows, TYPED_HEADER)) + assert _staged(dialect_session, staging, "typed_stamp", "typed_stamp_utc") == [ + (datetime(2020, 1, 1, 10), datetime(2020, 1, 1, 0, tzinfo=UTC)) + ] + with pytest.raises(ConversionError, match="typed_stamp.*without an offset"): + bad = _typed_rows(typed_stamp="2020-01-01T10:00:00+10:00") + _stage(dialect_session, _Typed, _write(tmp_path, DemoTypedRecord, "csv", bad, TYPED_HEADER)) + with pytest.raises(ConversionError, match="typed_stamp_utc.*with an offset"): + bad = _typed_rows(typed_stamp_utc="2020-01-01 10:00:00") + _stage(dialect_session, _Typed, _write(tmp_path, DemoTypedRecord, "csv", bad, TYPED_HEADER)) + + +@pytest.mark.parametrize("policy", [None, REJECT], ids=["raise", "reject_row"]) +def test_an_integer_too_wide_for_the_column_never_reaches_the_database(dialect_session, tmp_path, policy): + rows = [*_typed_rows(typed_small="32767"), *_typed_rows(typed_id="2", typed_small="32768")] + path = _write(tmp_path, DemoTypedRecord, "csv", rows, TYPED_HEADER) + if policy is None: + with pytest.raises(ConversionError, match="typed_small"): + _stage(dialect_session, _Typed, path) + return + loaded, ctx, staging = _stage(dialect_session, _Typed, path, policy=policy) + assert loaded == 1 and ctx.report.rows_rejected == 1 + assert _staged(dialect_session, staging, "typed_small") == [(32767,)] + + +def test_year_zero_is_rejected_on_every_dialect(dialect_session, tmp_path): + rows = _typed_rows(typed_date="0000-01-01", typed_stamp="0000-01-01 00:00:00") + with pytest.raises(ConversionError, match="typed_date"): + _stage(dialect_session, _Typed, _write(tmp_path, DemoTypedRecord, "csv", rows, TYPED_HEADER)) diff --git a/tests/loaders/test_readers.py b/tests/loaders/test_readers.py index 33faddb..25eb2b3 100644 --- a/tests/loaders/test_readers.py +++ b/tests/loaders/test_readers.py @@ -68,6 +68,13 @@ def test_quote_modes_decide_what_double_quotes_mean(tmp_path): assert label("by_delimiter") == '"x"' +def test_a_quote_far_into_a_tab_delimited_file_is_read_literally_by_default(tmp_path): + lines = "".join(f"{i}\tlabel {i}\n" for i in range(1, 5001)) + path = _write(tmp_path, "demo_id\tdemo_label\n" + lines + '5001\t"Sick" headaches\n5002\tnext\n', "demo_record.tsv") + rows = _rows(DelimitedTextReader(path, COLUMNS)) + assert [r["demo_label"] for r in rows[-2:]] == ['"Sick" headaches', "next"] + + def test_a_file_is_read_in_several_batches_without_losing_rows(tmp_path): lines = "".join(f"{i},label {i}\n" for i in range(1, 20_001)) path = _write(tmp_path, "demo_id,demo_label\n" + lines) diff --git a/tests/loaders/test_stage_file.py b/tests/loaders/test_stage_file.py index 8e6d5e5..2e93c8c 100644 --- a/tests/loaders/test_stage_file.py +++ b/tests/loaders/test_stage_file.py @@ -1,7 +1,10 @@ from pathlib import Path +import pytest import sqlalchemy as sa +from orm_loader.backends import resolve_backend +from orm_loader.helpers import IngestError from orm_loader.loaders import LoaderContext, stage_file from tests.models import FileLoaderTable @@ -170,3 +173,25 @@ def test_file_without_data_rows_loads_nothing_whatever_its_header(tmp_path, sqli ) assert stage_file(ctx, ctx.tableclass._select_reader(ctx)) == 0 assert ctx.report.rows_read == 0 + + +def test_accounting_is_checked_against_the_staging_table(tmp_path, sqlite_session, monkeypatch): + """A sink that silently loses rows is caught, because the staged rows are counted in the database.""" + csv = tmp_path / "test_file_loader.csv" + csv.write_text("id,value\n1,a\n2,b\n3,c\n") + ctx = LoaderContext( + tableclass=FileLoaderTable, + session=sqlite_session, + path=csv, + staging_table=FileLoaderTable.__table__, + dedupe=False, + ) + backend = resolve_backend(sqlite_session, mapper=FileLoaderTable) + insert = type(backend).insert_batches + + def losing(self, ctx, batches, columns): + insert(self, ctx, list(batches)[:-1] or [], columns) + + monkeypatch.setattr(type(backend), "insert_batches", losing) + with pytest.raises(IngestError, match="read 3 rows but loaded 0"): + stage_file(ctx, ctx.tableclass._select_reader(ctx)) diff --git a/tests/models.py b/tests/models.py index dd00553..3061d40 100644 --- a/tests/models.py +++ b/tests/models.py @@ -1,5 +1,6 @@ -from datetime import date +from datetime import date, datetime +from decimal import Decimal from enum import Enum import sqlalchemy as sa @@ -205,3 +206,33 @@ class DemoRequiredRecord(Base, FileLoadableTableInterface): demo_id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) demo_label: so.Mapped[str] = so.mapped_column(sa.String, nullable=False) + + +class DemoDefaultedRecord(Base, FileLoadableTableInterface): + """Non-nullable columns with a server default, a scalar default and a callable default.""" + + __tablename__ = "demo_defaulted_record" + __table_args__ = {"schema": SchemaRole.PRIMARY.value} + + demo_id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) + demo_status: so.Mapped[str] = so.mapped_column(sa.String, nullable=False, server_default="open") + demo_priority: so.Mapped[int] = so.mapped_column(sa.Integer, nullable=False, default=5) + demo_stamp: so.Mapped[str] = so.mapped_column(sa.String, nullable=False, default=lambda: "made") + demo_note: so.Mapped[str | None] = so.mapped_column(sa.String, nullable=True, default="n") + + +class DemoTypedRecord(Base, FileLoadableTableInterface): + """One column of each kind a typed (Parquet) input can fill.""" + + __tablename__ = "demo_typed_record" + __table_args__ = {"schema": SchemaRole.PRIMARY.value} + + typed_id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) + typed_small: so.Mapped[int | None] = so.mapped_column(sa.SmallInteger, nullable=True) + typed_float: so.Mapped[float | None] = so.mapped_column(sa.Float, nullable=True) + typed_numeric: so.Mapped[Decimal | None] = so.mapped_column(sa.Numeric(10, 2), nullable=True) + typed_flag: so.Mapped[bool | None] = so.mapped_column(sa.Boolean, nullable=True) + typed_date: so.Mapped[date | None] = so.mapped_column(sa.Date, nullable=True) + typed_stamp: so.Mapped[datetime | None] = so.mapped_column(sa.DateTime, nullable=True) + typed_stamp_utc: so.Mapped[datetime | None] = so.mapped_column(sa.DateTime(timezone=True), nullable=True) + typed_label: so.Mapped[str | None] = so.mapped_column(sa.String(10), nullable=True) diff --git a/tests/support.py b/tests/support.py index 3504912..0f38faf 100644 --- a/tests/support.py +++ b/tests/support.py @@ -5,6 +5,7 @@ from typing import Any import pyarrow as pa +import pyarrow.compute as pc from sqlalchemy.types import TypeEngine from orm_loader.conversion import ConversionPolicy @@ -23,7 +24,7 @@ def interpret( arr = values if isinstance(values, (pa.Array, pa.ChunkedArray)) else pa.array(values) handler = policy.for_column(sa_type, table_name, column_name) out, invalid = handler.interpret(arr, sa_type, policy.nulls) - return out.to_pylist(), invalid.to_pylist() + return pc.if_else(invalid, pa.scalar(None, out.type), out).to_pylist(), invalid.to_pylist() def read(raw: Any, sa_type: TypeEngine[Any], policy: ConversionPolicy | None = None, **column: Any) -> Any: From d2ac0f5d79cdd195f72a184c4cfcf6c21e3e7495 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Thu, 29 Oct 2026 05:24:34 +0000 Subject: [PATCH 07/21] Address second PR review: row numbers, numeric fractions, defaults on all dialects, typed time zones --- src/orm_loader/backends/postgres.py | 16 ++++-- src/orm_loader/conversion/handlers/base.py | 4 +- src/orm_loader/conversion/handlers/dates.py | 23 ++++++-- src/orm_loader/conversion/handlers/numbers.py | 16 ++++-- src/orm_loader/conversion/values.py | 3 +- src/orm_loader/loaders/processor.py | 55 ++++++++++++++----- src/orm_loader/loaders/readers/base.py | 7 ++- src/orm_loader/loaders/readers/delimited.py | 9 ++- src/orm_loader/tables/orm_table.py | 21 +++---- tests/backends/test_sinks.py | 11 ++++ tests/conversion/test_number_handlers.py | 15 +++++ tests/conversion/test_typed_input.py | 9 +++ tests/loaders/test_defaults.py | 31 ++++++++--- tests/loaders/test_malformed_files.py | 46 +++++++++++++++- tests/models.py | 28 ++++++++++ 15 files changed, 234 insertions(+), 60 deletions(-) diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index e0b42a6..a179e67 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -109,12 +109,6 @@ def create_staging_table( bind_arguments=bind_arguments, ) - for name in table_cls.sql_defaults(): - session.execute( - sa.text(f'ALTER TABLE {staging_ref} ALTER COLUMN {preparer.quote_identifier(name)} DROP NOT NULL;'), - bind_arguments=bind_arguments, - ) - computed_cols = [c.name for c in table.columns if c.computed is not None] for col in computed_cols: session.execute( @@ -122,6 +116,16 @@ def create_staging_table( bind_arguments=bind_arguments, ) + # staging only lands the file, so a null reaches the merge, where the column's default + # or the required check deals with it; LIKE would keep the NOT NULL of the target + nullable = ", ".join( + f"ALTER COLUMN {preparer.quote_identifier(c.name)} DROP NOT NULL" + for c in table.columns + if c.computed is None and not c.nullable + ) + if nullable: + session.execute(sa.text(f"ALTER TABLE {staging_ref} {nullable};"), bind_arguments=bind_arguments) + # 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( diff --git a/src/orm_loader/conversion/handlers/base.py b/src/orm_loader/conversion/handlers/base.py index c05b15d..2dc36a7 100644 --- a/src/orm_loader/conversion/handlers/base.py +++ b/src/orm_loader/conversion/handlers/base.py @@ -39,8 +39,8 @@ def interpret( Returns ------- Interpretation - Converted values with nulls where a value is missing or - unreadable, and a mask that is true where a value was unreadable. + Converted values with nulls where a value is missing, and a mask + that is true where a value was unreadable. """ def expected(self, sa_type: TypeEngine[Any]) -> str: diff --git a/src/orm_loader/conversion/handlers/dates.py b/src/orm_loader/conversion/handlers/dates.py index f51e2f3..7e2aeee 100644 --- a/src/orm_loader/conversion/handlers/dates.py +++ b/src/orm_loader/conversion/handlers/dates.py @@ -161,12 +161,16 @@ def check(self, values: pa.Array, sa_type: TypeEngine[Any]) -> pa.Array | None: class DateTimeHandler(DateHandler): """Reads ISO timestamps, and date-only values in the declared date formats. - A timestamp has date and time separated by ``T`` or a space, with - optional seconds and up to six fractional digits. A date-only value - becomes midnight. The column type decides about time zones: a - ``DateTime(timezone=True)`` column requires an offset (``Z`` or - ``+hh:mm``) and stores UTC, and a column without a time zone rejects - an offset. A date-only value is not read for a time zone column. + - A timestamp has date and time separated by ``T`` or a space, with + optional seconds and up to six fractional digits. + - A date-only value becomes midnight. + - The column type decides about time zones: + - ``DateTime(timezone=True)`` column requires an offset (``Z`` or + ``+hh:mm``) and stores UTC, and + - a column without a time zone rejects an offset. + - A date-only value is not read for a time zone column. + - Typed input follows the same rule: a time zone aware timestamp does not fill a + naive column and a naive one does not fill an aware column. """ distinct: ClassVar[bool] = False @@ -181,6 +185,13 @@ def arrow_type(self, sa_type: TypeEngine[Any]) -> pa.DataType: def expected(self, sa_type: TypeEngine[Any]) -> str: return "DateTime with an offset (Z or +hh:mm)" if getattr(sa_type, "timezone", False) else "DateTime without an offset" + def _interpret_typed(self, arr: pa.Array, sa_type: TypeEngine[Any]) -> Interpretation: + target = self.arrow_type(sa_type) + stored_aware = pa.types.is_timestamp(arr.type) and arr.type.tz is not None + if not pa.types.is_null(arr.type) and stored_aware != (target.tz is not None): + return Interpretation(pa.nulls(len(arr), target), pc.is_valid(arr)) + return super()._interpret_typed(arr, sa_type) + def parse(self, cleaned: pa.Array, sa_type: TypeEngine[Any]) -> Interpretation: target = self.arrow_type(sa_type) aware = target.tz is not None diff --git a/src/orm_loader/conversion/handlers/numbers.py b/src/orm_loader/conversion/handlers/numbers.py index 479624c..5b78166 100644 --- a/src/orm_loader/conversion/handlers/numbers.py +++ b/src/orm_loader/conversion/handlers/numbers.py @@ -85,8 +85,8 @@ class NumericHandler(NumberHandler): """Accepts decimal forms and keeps the column's declared precision and scale. A precision without a scale has scale 0. A column without a precision is - read as ``decimal128(38, 18)``. A value with more digits than the column - holds is unreadable; trailing zeros beyond the scale are accepted. + read as ``decimal128(38, 18)``. Trailing zeros of the fraction are + dropped, and a value with more digits than the column holds is unreadable. """ @staticmethod @@ -100,11 +100,17 @@ def precision_scale(sa_type: TypeEngine[Any]) -> tuple[int, int]: def arrow_type(self, sa_type: TypeEngine[Any]) -> pa.DataType: return pa.decimal128(*self.precision_scale(sa_type)) + def clean(self, raw: pa.Array, trimmed: pa.Array) -> pa.Array: + return pc.replace_substring_regex(trimmed, r"(\.\d*?)0+$", r"\1") + def pattern(self, sa_type: TypeEngine[Any]) -> str: - """Return the decimal forms whose integer part fits the column. + """Return the decimal forms that fit the column's integer and fraction digits. Arrow wraps around instead of failing when a string has more digits - than a decimal holds, so the integer digits are bounded here. + than a decimal holds, so both parts are bounded here. """ precision, scale = self.precision_scale(sa_type) - return rf"^[+-]?(0*\d{{0,{precision - scale}}}(\.\d*)?|\.\d+)$" + forms = [rf"0*\d{{0,{precision - scale}}}(\.\d{{0,{scale}}})?"] + if scale: + forms.append(rf"\.\d{{1,{scale}}}") + return rf"^[+-]?({'|'.join(forms)})$" diff --git a/src/orm_loader/conversion/values.py b/src/orm_loader/conversion/values.py index 553b2fd..b5e198b 100644 --- a/src/orm_loader/conversion/values.py +++ b/src/orm_loader/conversion/values.py @@ -17,7 +17,8 @@ class Interpretation(NamedTuple): Attributes ---------- values - Converted values, null where a value is missing or unreadable. + Converted values, null where a value is missing. Where ``invalid`` is + true the value is not meaningful and is not used. invalid Boolean array that is true where a value could not be read. """ diff --git a/src/orm_loader/loaders/processor.py b/src/orm_loader/loaders/processor.py index 4166c72..787cbe7 100644 --- a/src/orm_loader/loaders/processor.py +++ b/src/orm_loader/loaders/processor.py @@ -4,6 +4,7 @@ import logging import os +from bisect import bisect_left from collections.abc import Iterable, Iterator from pathlib import Path from typing import TYPE_CHECKING, Any @@ -92,8 +93,9 @@ class BatchProcessor: normalise Whether to interpret values. When false, batches pass through. skipped - Lines the reader left out as malformed, each with the reason. They - are written to the rejects file and count as rows read and rejected. + Lines the reader left out as malformed, each as its line number in the + file and the reason. They are written to the rejects file and count as + rows read and rejected. """ def __init__( @@ -104,17 +106,21 @@ def __init__( report: LoadReport, rejects_path: Path, normalise: bool = True, - skipped: list[str] | None = None, + skipped: list[tuple[int | None, str]] | None = None, ) -> None: self.columns = columns self.skipped = skipped if skipped is not None else [] + self._parsed = 0 + self._shifts: list[int] = [] self.policy = policy self.report = report self.normalise = normalise self.table_name = tableclass.__tablename__ model_columns = tableclass.model_columns() self.types = {name: underlying_type(model_columns[name].type) for name in columns} - self.handlers = {name: policy.for_column(self.types[name], self.table_name, name) for name in columns} + self.handlers = ( + {name: policy.for_column(self.types[name], self.table_name, name) for name in columns} if normalise else {} + ) self.required = [c for c in tableclass.required_columns() if c in columns] self.rejects = RejectsWriter(rejects_path) if policy.errors.on_error == "reject_row": @@ -125,8 +131,8 @@ def stream(self, batches: Iterable[pa.RecordBatch]) -> Iterator[pa.Table]: """Yield typed tables for the raw batches, closing the rejects file afterwards.""" try: for raw in batches: - table = self.process(raw) self._reject_skipped() + table = self.process(raw) if table.num_rows: yield table self._reject_skipped() @@ -156,7 +162,8 @@ def process(self, raw: pa.RecordBatch) -> pa.Table: RejectLimitError Rejected rows exceed ``max_rejects``. """ - first_row = self.report.rows_read + 1 + first = self._parsed + self._parsed += raw.num_rows self.report.rows_read += raw.num_rows if not self.normalise: return pa.Table.from_batches([raw]) @@ -167,8 +174,8 @@ def process(self, raw: pa.RecordBatch) -> pa.Table: failures = self._failures(results) if failures: if self.policy.errors.on_error == "raise": - self._raise(raw, results, first_row) - return self._reject(raw, results, failures, first_row) + self._raise(raw, results, first) + return self._reject(raw, results, failures, first) return pa.table({name: results[name].values for name in self.columns}) def _failures(self, results: dict[str, Any]) -> dict[str, tuple[pa.Array, str]]: @@ -215,7 +222,7 @@ def _conversion_error(self, name: str, raw: pa.Array | pa.ChunkedArray, invalid: samples=raw.filter(invalid).slice(0, 3).to_pylist(), ) - def _raise(self, raw: pa.RecordBatch, results: dict[str, Any], first_row: int) -> None: + def _raise(self, raw: pa.RecordBatch, results: dict[str, Any], first: int) -> None: for name in self.columns: error = self._conversion_error(name, raw.column(name), results[name].invalid) if error is not None: @@ -224,7 +231,7 @@ def _raise(self, raw: pa.RecordBatch, results: dict[str, Any], first_row: int) - missing = pc.is_null(results[name].values) count = pc.sum(pc.cast(missing, pa.int64())).as_py() or 0 if count: - rows = [first_row + index for index in pc.indices_nonzero(missing).slice(0, 3).to_pylist()] + rows = self._row_numbers(first, raw.num_rows).take(pc.indices_nonzero(missing).slice(0, 3)).to_pylist() raise RequiredValueError( table_name=self.table_name, column_name=name, count=count, rows=rows ) @@ -234,7 +241,7 @@ def _reject( raw: pa.RecordBatch, results: dict[str, Any], failures: dict[str, tuple[pa.Array, str]], - first_row: int, + first: int, ) -> pa.Table: rejected = no_failures(raw.num_rows) column = pa.nulls(raw.num_rows, pa.string()) @@ -244,7 +251,7 @@ def _reject( column = pc.if_else(first_hit, pa.scalar(name), column) reason = pc.if_else(first_hit, pa.scalar(why), reason) rejected = pc.or_(rejected, mask) - row_numbers = pc.add(pa.array(range(raw.num_rows), pa.int64()), first_row) + row_numbers = self._row_numbers(first, raw.num_rows) plain = [c.dictionary_decode() if pa.types.is_dictionary(c.type) else c for c in raw.columns] listing = pa.table( [row_numbers, column, reason, *plain], names=["row", "column", "reason", *raw.schema.names] @@ -257,16 +264,36 @@ def _reject( keep = pc.invert(rejected) return pa.table({n: results[n].values for n in self.columns}).filter(keep) + def _row_numbers(self, first: int, count: int) -> pa.Array: + """Return the data row numbers in the file of ``count`` parsed rows, counting malformed lines. + + Parameters + ---------- + first + Position among the parsed rows of the first row (starting at 0). + count + Number of rows. + """ + positions = range(first + 1, first + count + 1) + return pa.array([p + bisect_left(self._shifts, p) for p in positions], pa.int64()) + def _reject_skipped(self) -> None: - """Write the lines the reader left out to the rejects file, without a row number.""" + """Write the lines the reader left out to the rejects file and remember where they were.""" count = len(self.skipped) if not count: return text, empty = pa.string(), pa.nulls(count, pa.string()) rows = pa.table( - {"row": pa.nulls(count, pa.int64()), "column": empty, "reason": pa.array(self.skipped, text)} + { + "row": pa.array([None if n is None else n - 1 for n, _ in self.skipped], pa.int64()), + "column": empty, + "reason": pa.array([why for _, why in self.skipped], text), + } | {name: empty for name in self.columns} ) + for number, _ in self.skipped: + if number is not None: + self._shifts.append(number - 1 - (len(self._shifts) + 1)) self.skipped.clear() self.rejects.write(rows) self.report.rows_read += count diff --git a/src/orm_loader/loaders/readers/base.py b/src/orm_loader/loaders/readers/base.py index 5008522..8bc360a 100644 --- a/src/orm_loader/loaders/readers/base.py +++ b/src/orm_loader/loaders/readers/base.py @@ -42,8 +42,9 @@ class FileReader(ABC): empty True when the file has no content at all. skipped - Lines the reader left out as malformed, each with the reason, filled - in while the batches are read. + Lines the reader left out as malformed, each as the number of the + line in the file (None when unknown) and the reason, filled in while + the batches are read. """ def __init__( @@ -58,7 +59,7 @@ def __init__( self.columns: list[str] = [] self.file_columns: list[str] = [] self.empty = False - self.skipped: list[str] = [] + self.skipped: list[tuple[int | None, str]] = [] def normalise_headers(self, names: Iterable[str]) -> list[str]: """Return header names lowercased, stripped of quotes and padding and of a ``_hash`` suffix. diff --git a/src/orm_loader/loaders/readers/delimited.py b/src/orm_loader/loaders/readers/delimited.py index 1f199d8..cbde7f1 100644 --- a/src/orm_loader/loaders/readers/delimited.py +++ b/src/orm_loader/loaders/readers/delimited.py @@ -82,7 +82,7 @@ def _header_names(self) -> list[str]: return self.normalise_headers(header.lstrip("\ufeff").rstrip("\r\n").split(self.delimiter)) def _skip_row(self, row: pv.InvalidRow) -> Literal["skip"]: - self.skipped.append(f"expected {row.expected_columns} columns, got {row.actual_columns}: {row.text.rstrip()}") + self.skipped.append((row.number, f"expected {row.expected_columns} columns, got {row.actual_columns}: {row.text.rstrip()}")) return "skip" def _warn_if_everything_is_quoted(self) -> None: @@ -114,6 +114,7 @@ def batches(self) -> Iterator[pa.RecordBatch]: skip_rows=1, block_size=self.block_size, encoding=self.encoding, + use_threads=not self.skip_malformed, ), parse_options=pv.ParseOptions( delimiter=self.delimiter, @@ -134,12 +135,14 @@ def batches(self) -> Iterator[pa.RecordBatch]: for batch in reader: if batch.num_rows: yield batch - except (pa.ArrowInvalid, UnicodeDecodeError) as exc: + except UnicodeDecodeError as exc: + raise IngestError(f"{self.path.name}: cannot decode the file as {self.encoding}: {exc.reason}.") from exc + except pa.ArrowInvalid as exc: if "Empty CSV file" in str(exc): return hint = ( " The file may contain literal double quotes; try quote_mode='literal' or 'by_delimiter'." - if self.quote == "csv" + if self.quote == "csv" and str(exc).startswith("CSV parse error") else "" ) raise IngestError(f"{self.path.name}: {exc}.{hint}") from exc diff --git a/src/orm_loader/tables/orm_table.py b/src/orm_loader/tables/orm_table.py index a5b9d31..25a3950 100644 --- a/src/orm_loader/tables/orm_table.py +++ b/src/orm_loader/tables/orm_table.py @@ -1,8 +1,10 @@ +import logging +from typing import Any + import sqlalchemy as sa import sqlalchemy.orm as so from sqlalchemy.exc import NoInspectionAvailable -from typing import Any -import logging + from .allocators import IdAllocator logger = logging.getLogger(__name__) @@ -164,12 +166,13 @@ def required_columns(cls) -> set[str]: @classmethod def sql_defaults(cls) -> dict[str, sa.ColumnElement[Any]]: """ - Return the default of each non-nullable column as a SQL expression. + Return the default of each column that has one SQL can apply. - A server default given as text is cast to the column type, one given - as a SQL expression is used as written, and a scalar Python default - is a literal. A callable or context-sensitive default has no SQL - form, so its column is not listed. + A server default given as text is written inline, as in the DDL, so + the database reads it as the column's type. One given as a SQL + expression is used as written, and a scalar Python default is a typed + literal. A callable or context-sensitive default has no SQL form, so + its column is not listed. Returns ------- @@ -178,11 +181,9 @@ def sql_defaults(cls) -> dict[str, sa.ColumnElement[Any]]: """ defaults: dict[str, sa.ColumnElement[Any]] = {} for c in cls.mapper_for().columns: - if c.nullable: - continue server_default = getattr(c.server_default, "arg", None) if isinstance(server_default, str): - defaults[c.name] = sa.cast(sa.literal(server_default), c.type) + defaults[c.name] = sa.literal(server_default, sa.String()).render_literal_execute() elif server_default is not None: defaults[c.name] = server_default elif c.default is not None and c.default.is_scalar: diff --git a/tests/backends/test_sinks.py b/tests/backends/test_sinks.py index 4f9e93d..f7c1e68 100644 --- a/tests/backends/test_sinks.py +++ b/tests/backends/test_sinks.py @@ -85,3 +85,14 @@ def batches(): with pytest.raises(RuntimeError, match="boom"): backend.insert_batches(ctx, batches(), COLUMNS) + + +def test_a_staging_table_accepts_a_null_in_every_column(dialect_session): + from tests.models import DemoRequiredRecord + + table = cast(Type[FileTableProtocol], DemoRequiredRecord) + tag = None if dialect_session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA + table.create_staging_table(dialect_session, staging_schema_tag=tag) + staging = table.get_staging_table(dialect_session, staging_schema_tag=tag) + assert not DemoRequiredRecord.__table__.c.demo_label.nullable + assert [c.name for c in staging.c if not c.nullable and c.name != "_rownum"] == [] diff --git a/tests/conversion/test_number_handlers.py b/tests/conversion/test_number_handlers.py index 23c3dc3..aed12bf 100644 --- a/tests/conversion/test_number_handlers.py +++ b/tests/conversion/test_number_handlers.py @@ -125,3 +125,18 @@ def test_one_bad_value_in_a_large_batch_does_not_slow_the_whole_batch(): values, invalid = interpret(raw, sa.BigInteger()) assert time.perf_counter() - started < 0.5 assert sum(invalid) == 1 and invalid[100_000] + + +def test_numeric_fractions_longer_than_the_decimal_do_not_wrap_around(): + values = ["1.5" + "0" * 45, "1.50" + "0" * 40, "0." + "1" * 50, "1.234", "1.2300"] + assert interpret(values, sa.Numeric(10, 2)) == ( + [Decimal("1.50"), Decimal("1.50"), None, None, Decimal("1.23")], + [False, False, True, True, False], + ) + + +def test_numeric_with_scale_zero_accepts_a_zero_fraction_only(): + assert interpret(["2", "2.0", "2.5", "."], sa.Numeric(30)) == ( + [Decimal("2"), Decimal("2"), None, None], + [False, False, True, True], + ) diff --git a/tests/conversion/test_typed_input.py b/tests/conversion/test_typed_input.py index 37a17f6..fce9678 100644 --- a/tests/conversion/test_typed_input.py +++ b/tests/conversion/test_typed_input.py @@ -72,3 +72,12 @@ def test_year_zero_dates_are_unreadable_on_typed_input(): year_zero, year_two_thousand = -719_528, 10_957 _, invalid = _result(pa.array([year_zero, year_two_thousand], pa.date32()), sa.Date()) assert invalid == [True, False] + + +def test_the_time_zone_rule_of_the_column_applies_to_typed_timestamps(): + aware = pa.array([datetime(2020, 1, 1, 10, tzinfo=timezone.utc), None], pa.timestamp("us", "UTC")) + naive = pa.array([datetime(2020, 1, 1, 10), None], pa.timestamp("us")) + assert _result(aware, sa.DateTime())[1] == [True, False] + assert _result(naive, sa.DateTime(timezone=True))[1] == [True, False] + assert _result(pa.array([date(2020, 1, 1)], pa.date32()), sa.DateTime(timezone=True))[1] == [True] + assert _result(pa.array([date(2020, 1, 1)], pa.date32()), sa.DateTime())[1] == [False] diff --git a/tests/loaders/test_defaults.py b/tests/loaders/test_defaults.py index cd54f41..970dba6 100644 --- a/tests/loaders/test_defaults.py +++ b/tests/loaders/test_defaults.py @@ -1,5 +1,6 @@ -"""A non-nullable column with a default takes it when the value is empty or the column is absent.""" +"""A column with a default takes it when the value is empty or the column is absent.""" +from datetime import date, datetime from pathlib import Path from typing import Type, cast @@ -10,7 +11,7 @@ from orm_loader.conversion import ConversionPolicy, ErrorPolicy from orm_loader.helpers import IngestError, RequiredValueError from orm_loader.tables.typing import FileTableProtocol -from tests.models import DemoDefaultedRecord +from tests.models import DemoDefaultedRecord, DemoDefaultKinds _Table = cast(Type[FileTableProtocol], DemoDefaultedRecord) FULL = "demo_id,demo_status,demo_priority,demo_stamp,demo_note\n" @@ -34,17 +35,17 @@ def _rows(session) -> list[tuple]: def test_only_a_default_that_sql_can_apply_makes_a_column_optional(): assert _Table.required_columns() == {"demo_stamp"} - assert set(_Table.sql_defaults()) == {"demo_status", "demo_priority"} + assert set(_Table.sql_defaults()) == {"demo_status", "demo_priority", "demo_note"} def test_an_empty_value_takes_the_default(dialect_session, tmp_path): _load(dialect_session, tmp_path, FULL + "1,,,s,\n2,closed,9,s,kept\n") - assert _rows(dialect_session) == [(1, "open", 5, "s", None), (2, "closed", 9, "s", "kept")] + assert _rows(dialect_session) == [(1, "open", 5, "s", "n"), (2, "closed", 9, "s", "kept")] def test_an_absent_column_takes_the_default(dialect_session, tmp_path): _load(dialect_session, tmp_path, "demo_id,demo_stamp\n1,s\n2,t\n") - assert _rows(dialect_session) == [(1, "open", 5, "s", None), (2, "open", 5, "t", None)] + assert _rows(dialect_session) == [(1, "open", 5, "s", "n"), (2, "open", 5, "t", "n")] def test_a_column_with_a_callable_default_stays_required(dialect_session, tmp_path): @@ -64,6 +65,22 @@ def test_a_missing_callable_default_can_be_rejected_by_row(dialect_session, tmp_ def test_a_reload_applies_the_default_to_an_existing_row(dialect_session, tmp_path, strategy): _load(dialect_session, tmp_path, FULL + "1,closed,9,s,kept\n") _load(dialect_session, tmp_path, FULL + "1,,,t,\n2,,,u,\n", merge_strategy=strategy) - expected = {"replace": [(1, "open", 5, "t", None), (2, "open", 5, "u", None)], - "upsert": [(1, "closed", 9, "s", "kept"), (2, "open", 5, "u", None)]} + expected = {"replace": [(1, "open", 5, "t", "n"), (2, "open", 5, "u", "n")], + "upsert": [(1, "closed", 9, "s", "kept"), (2, "open", 5, "u", "n")]} assert _rows(dialect_session) == expected[strategy] + + +def test_every_kind_of_default_is_applied_and_read_back_on_every_dialect(dialect_session, tmp_path): + path = tmp_path / "demo_default_kinds.csv" + path.write_text("kind_id\n1\n2\n") + tag = None if dialect_session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA + cast(Type[FileTableProtocol], DemoDefaultKinds).load_file(dialect_session, path, staging_schema_tag=tag) + dialect_session.commit() + dialect_session.expire_all() + rows = dialect_session.execute(sa.select(DemoDefaultKinds)).scalars().all() + assert len(rows) == 2 + for row in rows: + assert (row.kind_text, row.kind_int, row.kind_bool) == ("it's text", 7, True) + assert (row.kind_date, row.kind_stamp) == (date(2000, 1, 2), datetime(2000, 1, 2, 10, 30)) + assert (row.kind_expression, row.kind_py_int, row.kind_py_date) == ("expr", 5, date(2001, 2, 3)) + assert isinstance(row.kind_now, datetime) diff --git a/tests/loaders/test_malformed_files.py b/tests/loaders/test_malformed_files.py index 19df9d3..9c1cbab 100644 --- a/tests/loaders/test_malformed_files.py +++ b/tests/loaders/test_malformed_files.py @@ -12,12 +12,13 @@ from orm_loader.backends import STAGING_SCHEMA, resolve_backend from orm_loader.conversion import ConversionPolicy, ErrorPolicy -from orm_loader.helpers import IngestError +from orm_loader.helpers import IngestError, PolicyError from orm_loader.loaders import DelimitedTextReader, LoadReport, ParquetReader from orm_loader.tables.typing import FileTableProtocol -from tests.models import DemoRecord +from tests.models import DemoBinaryRecord, DemoRecord, DemoRequiredRecord _Demo = cast(Type[FileTableProtocol], DemoRecord) +_DemoRequired = cast(Type[FileTableProtocol], DemoRequiredRecord) REJECT = ConversionPolicy(errors=ErrorPolicy(on_error="reject_row")) COLUMNS = DemoRecord.model_columns() @@ -58,7 +59,26 @@ def test_ragged_lines_are_rejected_with_their_text_under_reject_row(dialect_sess "expected 2 columns, got 1: 2", "expected 2 columns, got 3: 3,c,extra", ] - assert all(r["row"] is None and not r["column"] for r in rejects) + assert [r["row"] for r in rejects] == [2, 3] and not any(r["column"] for r in rejects) + + +def test_row_numbers_count_the_malformed_lines_before_them(dialect_session, tmp_path): + report = LoadReport() + text = "demo_id,demo_label,demo_number\n1,a,1\n2,b\n3,c,3\n4,d,x\n5,e,5\n" + assert _load(dialect_session, tmp_path, text, policy=REJECT, report=report) == 3 + assert [(r["row"], r["column"]) for r in _rejects(report.rejects_path)] == [(2, ""), (4, "demo_number")] + + +def test_a_required_value_error_names_the_row_in_the_file(dialect_session, tmp_path): + from orm_loader.helpers import RequiredValueError + from tests.models import DemoRequiredRecord + + path = tmp_path / "demo_required_record.csv" + path.write_text("demo_id,demo_label\n1,a\n2\n3,\n") + reader = DelimitedTextReader(path, DemoRequiredRecord.model_columns(), skip_malformed=True) + with pytest.raises(RequiredValueError) as error: + _DemoRequired.load_file(dialect_session, path, reader=reader, policy=ConversionPolicy()) + assert error.value.rows == [3] def test_a_file_of_only_ragged_lines_is_rejected_not_silently_empty(dialect_session, tmp_path): @@ -121,3 +141,23 @@ def test_duplicate_names_after_normalisation_are_an_error(tmp_path): path.write_text("demo_id,Demo_ID\n1,2\n") with pytest.raises(IngestError, match="duplicate column names"): DelimitedTextReader(path, COLUMNS) + + +def test_invalid_utf8_is_reported_without_the_quote_mode_hint(tmp_path): + path = tmp_path / "demo_record.csv" + path.write_bytes(b"demo_id,demo_label\n1,a\xff\n") + reader = DelimitedTextReader(path, COLUMNS) + reader.encoding = "utf-8" + with pytest.raises(IngestError, match="demo_record.csv") as error: + list(reader.batches()) + assert "quote_mode" not in str(error.value) + + +def test_a_column_type_without_a_handler_is_fine_when_values_are_not_interpreted(dialect_session, tmp_path): + path = tmp_path / "demo_binary_record.csv" + path.write_text("demo_id,demo_blob\n1,abc\n") + tag = None if dialect_session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA + table = cast(Type[FileTableProtocol], DemoBinaryRecord) + with pytest.raises(PolicyError, match="demo_blob"): + table.load_file(dialect_session, path, staging_schema_tag=tag) + assert table.load_file(dialect_session, path, staging_schema_tag=tag, normalise=False) == 1 diff --git a/tests/models.py b/tests/models.py index 3061d40..7bbd68f 100644 --- a/tests/models.py +++ b/tests/models.py @@ -236,3 +236,31 @@ class DemoTypedRecord(Base, FileLoadableTableInterface): typed_stamp: so.Mapped[datetime | None] = so.mapped_column(sa.DateTime, nullable=True) typed_stamp_utc: so.Mapped[datetime | None] = so.mapped_column(sa.DateTime(timezone=True), nullable=True) typed_label: so.Mapped[str | None] = so.mapped_column(sa.String(10), nullable=True) + + +class DemoDefaultKinds(Base, FileLoadableTableInterface): + """One column for each kind of default a column can declare.""" + + __tablename__ = "demo_default_kinds" + __table_args__ = {"schema": SchemaRole.PRIMARY.value} + + kind_id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) + kind_text: so.Mapped[str] = so.mapped_column(sa.String, nullable=False, server_default="it's text") + kind_int: so.Mapped[int] = so.mapped_column(sa.Integer, nullable=False, server_default="7") + kind_bool: so.Mapped[bool] = so.mapped_column(sa.Boolean, nullable=False, server_default=sa.true()) + kind_date: so.Mapped[date] = so.mapped_column(sa.Date, nullable=False, server_default="2000-01-02") + kind_stamp: so.Mapped[datetime] = so.mapped_column(sa.DateTime, nullable=False, server_default="2000-01-02 10:30:00") + kind_expression: so.Mapped[str] = so.mapped_column(sa.String, nullable=False, server_default=sa.text("'expr'")) + kind_now: so.Mapped[datetime | None] = so.mapped_column(sa.DateTime, nullable=True, server_default=sa.func.current_timestamp()) + kind_py_int: so.Mapped[int | None] = so.mapped_column(sa.Integer, nullable=True, default=5) + kind_py_date: so.Mapped[date | None] = so.mapped_column(sa.Date, nullable=True, default=date(2001, 2, 3)) + + +class DemoBinaryRecord(Base, FileLoadableTableInterface): + """A column type that the conversion policy has no handler for.""" + + __tablename__ = "demo_binary_record" + __table_args__ = {"schema": SchemaRole.PRIMARY.value} + + demo_id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) + demo_blob: so.Mapped[str | None] = so.mapped_column(sa.LargeBinary, nullable=True) From 12c360c0886dee08018acb0d83ba3bb5b3650fdc Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Thu, 29 Oct 2026 21:36:44 +0000 Subject: [PATCH 08/21] Rework staging and defaults: one staging definition per backend, staging passed as an argument, defaults filled by the processor --- docs/loaders/context.md | 1 - docs/loaders/staging.md | 12 ++- src/orm_loader/backends/base.py | 136 +++++++++++++++++------- src/orm_loader/backends/duckdb.py | 92 +++------------- src/orm_loader/backends/postgres.py | 96 +++++------------ src/orm_loader/backends/sqlite.py | 51 +++------ src/orm_loader/loaders/loader.py | 16 +-- src/orm_loader/loaders/processor.py | 83 +++++++++++++-- src/orm_loader/tables/loadable_table.py | 63 +++++------ src/orm_loader/tables/orm_table.py | 45 ++++---- src/orm_loader/tables/typing.py | 11 +- tests/backends/test_base_backend.py | 4 +- tests/backends/test_duckdb_backend.py | 18 +++- tests/backends/test_postgres_backend.py | 9 -- tests/backends/test_sinks.py | 56 +++++++--- tests/backends/test_sqlite_backend.py | 15 +-- tests/loaders/test_defaults.py | 61 ++++++++++- tests/loaders/test_pipeline.py | 10 +- tests/loaders/test_stage_file.py | 23 ++-- tests/models.py | 20 ++++ 20 files changed, 460 insertions(+), 362 deletions(-) diff --git a/docs/loaders/context.md b/docs/loaders/context.md index b5ef6a0..d2c61a9 100644 --- a/docs/loaders/context.md +++ b/docs/loaders/context.md @@ -21,7 +21,6 @@ on globals or implicit configuration. | `tableclass` | Target ORM table class | | `session` | Active SQLAlchemy session | | `path` | Path to the input file | -| `staging_table` | SQLAlchemy Table used for staging | | `chunksize` | Rows per batch for the default reader of formats that batch by rows | | `normalise` | Whether to interpret values under the policy; when false, values are written as read | | `dedupe` | Whether the staging table keeps only the first row for each primary key | diff --git a/docs/loaders/staging.md b/docs/loaders/staging.md index a50986c..0af150d 100644 --- a/docs/loaders/staging.md +++ b/docs/loaders/staging.md @@ -1,6 +1,6 @@ # Staging and Sinks -`stage_file(ctx, reader)` connects a reader, the batch processor and the sink of the session's backend. It loads into the staging table, respects the `LoaderContext` flags, returns the number of rows left in staging, and leaves the merge to the table layer. +`stage_file(ctx, reader, staging)` connects a reader, the batch processor and the sink of the session's backend. It loads into the given staging table, respects the `LoaderContext` flags, returns the number of rows left in staging, and leaves the merge to the table layer. ::: orm_loader.loaders.loader.stage_file @@ -10,12 +10,18 @@ - Values are interpreted by column type, as described in [Conversion](conversion.md). - With the default policy an unreadable value or a missing required value raises and nothing is loaded. -- With `ErrorPolicy(on_error="reject_row")` those rows, and delimited lines with the wrong number of fields, are written to `/.rejects.csv` and the rest load. A malformed line has no row number, and its text is in the reason. `rejects_dir` defaults to a `rejects` directory inside the input file's directory, and the file of an earlier load is removed. `LoadReport` accounts for every row read, checked against the rows counted in the staging table. -- A column of a table that has a default which SQL can apply (a server default or a scalar Python default) and is not nullable takes it when the value is empty or the column is not in the file. A column with a callable default stays required. +- With `ErrorPolicy(on_error="reject_row")` those rows, and delimited lines with the wrong number of fields, are written to `/.rejects.csv` and the rest load. A malformed line is listed with its row number in the file, and its text is in the reason. `rejects_dir` defaults to a `rejects` directory inside the input file's directory, and the file of an earlier load is removed. `LoadReport` accounts for every row read, checked against the rows counted in the staging table. +- A blank value, and a column the file lacks, take the column's default when it has one: a server default written as text, a scalar Python default, or a SQL expression such as `current_date`. The column's handler reads the default like a value from the file, and an expression is evaluated by the database for each batch that needs it. A column with a callable default stays required. - With `normalise=False` raw text is passed to the staging table unchanged. --- +## Staging table + +Staging mirrors the target: its columns, their types and their nullability, with no key, index, default or constraint, and a `_rownum` column that orders the rows (SQLite uses its rowid). The processor completes every value before it is written, so a NOT NULL column never receives NULL. A key the file provides is a staging column like any other. When the file lacks an autoincrement key, staging leaves it out and the merge leaves it out of its insert, so the database assigns it. + +--- + ## Sinks Each backend uses the fastest bulk interface its engine offers. diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index abc718f..22c3e9e 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -1,7 +1,7 @@ from __future__ import annotations from abc import ABC, abstractmethod -from collections.abc import Generator, Iterable, Iterator, Mapping, Sequence +from collections.abc import Collection, Generator, Iterable, Iterator, Mapping, Sequence from contextlib import AbstractContextManager, contextmanager from dataclasses import dataclass from functools import wraps @@ -341,31 +341,21 @@ def _as_connection( def _insertable_column_names( self, table_cls: Type["FileTableProtocol"], + staging: sa.Table | None = None, ) -> list[str]: """ Return column names safe to include in generic insert statements. Computed columns are excluded because backend loaders and merge helpers - should not attempt to write to them directly. + should not attempt to write to them directly. With a staging table, + only the columns it has are returned, so a column the load left out + (an autoincrement key) is assigned by the database. """ - return [c.name for c in table_cls.__table__.columns if c.computed is None] - - def _staged_columns( - self, - table_cls: Type["FileTableProtocol"], - staging: sa.Table, - ) -> dict[str, sa.ColumnElement[Any]]: - """ - Return the staging value of each insertable column by column name. - - A null in a non-nullable column takes the column's SQL default, so an - empty value and an absent column mean the default. - """ - defaults = table_cls.sql_defaults() - return { - name: sa.func.coalesce(staging.c[name], defaults[name]) if name in defaults else staging.c[name] - for name in self._insertable_column_names(table_cls) - } + return [ + c.name + for c in table_cls.__table__.columns + if c.computed is None and (staging is None or c.name in staging.c) + ] def deferred_fk_constraints( self, @@ -546,26 +536,95 @@ def _validate_one_fk( if orphans: raise error(fk, [tuple(row) for row in orphans]) - @abstractmethod def create_staging_table( self, table_cls: Type["FileTableProtocol"], session: so.Session, + *, + without: Collection[str] = (), ) -> None: - """Create a staging table for the supplied ORM table class.""" + """ + Create the staging table for a table class, replacing an existing one. + + The staging table has the target's columns with their types and + nullability, and no key, index, default or constraint, plus the row + key column of the backend. The processor completes every value before + it is written, so the target's NOT NULL holds in staging. + + Parameters + ---------- + table_cls + Table the staging table is for. + session + Session whose connection creates the table. + without + Names of target columns the staging table leaves out. The merge + then leaves them out of its insert. + """ + self.drop_staging_table(table_cls, session) + columns = [ + sa.Column(c.name, c.type, nullable=c.nullable) + for c in table_cls.__table__.columns + if c.computed is None and c.name not in without + ] + staging = self._staging_table(table_cls, columns) + staging.create(bind=session.connection(bind_arguments={"mapper": table_cls}), checkfirst=True) + session.commit() - @abstractmethod def drop_staging_table( self, table_cls: Type["FileTableProtocol"], session: so.Session, ) -> None: - """Drop a staging table if it exists.""" + """Drop a staging table and what belongs to its row key, if they exist.""" + staging = self._staging_table(table_cls, []) + staging.drop(bind=session.connection(bind_arguments={"mapper": table_cls}), checkfirst=True) + + def _staging_table(self, table_cls: Type["FileTableProtocol"], columns: Sequence[sa.Column[Any]]) -> sa.Table: + """ + Return the definition of the staging table, which create and drop both use. + + Parameters + ---------- + table_cls + Table the staging table is for. + columns + The target's columns the staging table holds, none when dropping. + + Returns + ------- + sqlalchemy.Table + The columns plus the backend's row key column. The default has no + row key, for a backend whose tables already carry one. + """ + return sa.Table( + self.staging_name_for_table(table_cls.__tablename__), + sa.MetaData(), + *columns, + schema=self.staging_schema, + ) + + def evaluate_default( + self, + session: so.Session, + table_cls: Type["FileTableProtocol"], + expression: sa.ColumnElement[Any], + column_type: Any, + ) -> Any: + """ + Return the value of a SQL expression default, as a column of the given type stores it. + + The database evaluates the expression and casts the result to the + column type, as it would for the default of an inserted row. + """ + select = sa.select(sa.cast(expression, column_type)) + return session.execute(select, bind_arguments={"mapper": table_cls}).scalar_one() @abstractmethod def insert_batches( self, ctx: "LoaderContext", + staging: sa.Table, batches: Iterable[pa.Table], columns: Sequence[str], ) -> None: @@ -577,7 +636,9 @@ def insert_batches( Parameters ---------- ctx - Context of the load; its staging table receives the rows. + Context of the load. + staging + Staging table that receives the rows. batches Typed tables, one per batch. columns @@ -590,7 +651,7 @@ def insert_batches( """ @staticmethod - def _staging_reference(ctx: "LoaderContext", connection: Connection) -> str: + def _staging_reference(staging: sa.Table, connection: Connection) -> str: """ Return the quoted, schema-qualified name of the load's staging table. @@ -599,8 +660,8 @@ def _staging_reference(ctx: "LoaderContext", connection: Connection) -> str: Parameters ---------- - ctx - Context of the load. + staging + The staging table. connection Connection the SQL runs on. @@ -609,23 +670,23 @@ def _staging_reference(ctx: "LoaderContext", connection: Connection) -> str: str The staging table name as the connection resolves it. """ - staging = ctx.staging_table translate = connection.get_execution_options().get("schema_translate_map") or {} schema = translate.get(staging.schema, staging.schema) return qualified(connection, staging.name, physical_schema=schema) - def count_staging(self, ctx: "LoaderContext") -> int: - """Return the number of rows in the load's staging table.""" - count = sa.select(sa.func.count()).select_from(ctx.staging_table) + def count_staging(self, ctx: "LoaderContext", staging: sa.Table) -> int: + """Return the number of rows in the staging table.""" + count = sa.select(sa.func.count()).select_from(staging) return ctx.session.execute(count, bind_arguments={"mapper": ctx.tableclass}).scalar_one() def deduplicate_staging( self, ctx: "LoaderContext", + staging: sa.Table, key_columns: Sequence[str], ) -> "DedupeResult": """ - Keep the first row for each primary key in the load's staging table. + Keep the first row for each primary key in the staging table. The first row is the one with the lowest row key, which follows the order of the source file. @@ -633,7 +694,9 @@ def deduplicate_staging( Parameters ---------- ctx - Context of the load; its staging table is deduplicated. + Context of the load. + staging + Staging table to deduplicate. key_columns Primary key columns that identify a duplicate. @@ -642,7 +705,6 @@ def deduplicate_staging( DedupeResult How many rows the staging table held and how many were removed. """ - staging = ctx.staging_table arguments = {"mapper": ctx.tableclass} row_key = self._pagination_column(ctx.tableclass, staging) ranked = sa.select( @@ -652,9 +714,9 @@ def deduplicate_staging( .label("position"), ).subquery() duplicates = sa.select(ranked.c.row_key).where(ranked.c.position > 1) - before = self.count_staging(ctx) + before = self.count_staging(ctx, staging) ctx.session.execute(sa.delete(staging).where(row_key.in_(duplicates)), bind_arguments=arguments) - after = self.count_staging(ctx) + after = self.count_staging(ctx, staging) return DedupeResult(tuple(key_columns), before, before - after) @staticmethod diff --git a/src/orm_loader/backends/duckdb.py b/src/orm_loader/backends/duckdb.py index 4b8a7fd..d106266 100644 --- a/src/orm_loader/backends/duckdb.py +++ b/src/orm_loader/backends/duckdb.py @@ -9,10 +9,8 @@ from oa_configurator import ( UnenforcedForeignKey, physical_schema_of, - qualified, unenforced_foreign_keys, ) -from sqlalchemy.schema import CreateSequence, DropSequence, DropTable from sqlalchemy.sql.compiler import IdentifierPreparer from .base import BackendCapabilities, DatabaseBackend, DeferredForeignKeyViolationError, Dialect @@ -156,68 +154,12 @@ def _is_referenced_by_existing_rows(self, target: sa.Table, session: so.Session) for other in self._tables_referencing(target) ) - def _staging_rownum_sequence_ref(self, tablename: str) -> str: - """Fully schema-qualified, quoted name of tablename's staging _rownum sequence. - - DuckDB can't ALTER TABLE ADD a GENERATED ALWAYS AS IDENTITY column - after creation, so pagination uses a sequence-backed DEFAULT - instead for a stable, monotonic key. - """ - seq_name = f"{self.staging_name_for_table(tablename)}_rownum_seq" - return qualified(self.identifier_preparer, seq_name, physical_schema=self.staging_schema) - - def create_staging_table( - self, - table_cls: type["FileTableProtocol"], - session: so.Session, - ) -> None: - bind_arguments = {"mapper": table_cls} - staging_name = self.staging_name_for_table(table_cls.__tablename__) - staging_ref = self.qualified_staging_name(table_cls.__tablename__) - seq_ref = self._staging_rownum_sequence_ref(table_cls.__tablename__) - sequence = sa.Sequence(f"{staging_name}_rownum_seq", schema=self.staging_schema) - - session.execute( - DropTable(sa.Table(staging_name, sa.MetaData(), schema=self.staging_schema), if_exists=True), - bind_arguments=bind_arguments, - ) - session.execute(DropSequence(sequence, if_exists=True), bind_arguments=bind_arguments) - - metadata = sa.MetaData() - staging_columns = [ - sa.Column(col.name, col.type, nullable=True) - for col in table_cls.__table__.columns - ] - staging_table = sa.Table(staging_name, metadata, *staging_columns, schema=self.staging_schema) - metadata.create_all( - bind=session.connection(bind_arguments=bind_arguments), tables=[staging_table] - ) - - session.execute(CreateSequence(sequence), bind_arguments=bind_arguments) - # ALTER TABLE ADD COLUMN has no Core DDL construct, so this stays raw. - session.execute( - sa.text(f"ALTER TABLE {staging_ref} ADD COLUMN _rownum BIGINT DEFAULT nextval('{seq_ref}');"), - bind_arguments=bind_arguments, - ) - session.commit() - - def drop_staging_table( - self, - table_cls: type["FileTableProtocol"], - session: so.Session, - ) -> None: - bind_arguments = {"mapper": table_cls} - staging_name = self.staging_name_for_table(table_cls.__tablename__) - session.execute( - DropTable(sa.Table(staging_name, sa.MetaData(), schema=self.staging_schema), if_exists=True), - bind_arguments=bind_arguments, - ) - session.execute( - DropSequence( - sa.Sequence(f"{staging_name}_rownum_seq", schema=self.staging_schema), if_exists=True - ), - bind_arguments=bind_arguments, - ) + def _staging_table(self, table_cls: type["FileTableProtocol"], columns: Sequence[sa.Column[Any]]) -> sa.Table: + """Number the rows with a sequence, which DuckDB creates and drops with the table.""" + name = self.staging_name_for_table(table_cls.__tablename__) + sequence = sa.Sequence(f"{name}_rownum_seq", schema=self.staging_schema) + row_key = sa.Column("_rownum", sa.BigInteger, sequence, server_default=sequence.next_value()) + return sa.Table(name, sa.MetaData(), *columns, row_key, schema=self.staging_schema) @staticmethod def _normalize_fk_check_state(previous_state: str | int) -> str | int: @@ -336,6 +278,7 @@ def _commits_between_pagination_batches(self) -> bool: def insert_batches( self, ctx: "LoaderContext", + staging: sa.Table, batches: Iterable[pa.Table], columns: Sequence[str], ) -> None: @@ -351,7 +294,7 @@ def insert_batches( transaction. """ batch = sa.table("_orm_loader_batch", *(sa.column(c) for c in columns)) - insert = ctx.staging_table.insert().from_select(list(columns), sa.select(*batch.c)) + insert = staging.insert().from_select(list(columns), sa.select(*batch.c)) bind_arguments = {"mapper": ctx.tableclass} for index, table in enumerate(batches, start=1): native = ctx.session.connection(bind_arguments=bind_arguments).connection @@ -433,8 +376,7 @@ def merge_replace( self.check_replace_feasible(table_cls, pk_cols, session) target = table_cls.__table__ staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) - staged = self._staged_columns(table_cls, staging) - update_cols = [c for c in staged if c not in pk_cols] + update_cols = [c for c in self._insertable_column_names(table_cls, staging) if c not in pk_cols] pk_join = sa.and_(*(target.c[c] == staging.c[c] for c in pk_cols)) def _batch_keys(extra: sa.ColumnElement[bool] | None = None) -> sa.Select[Any]: @@ -449,7 +391,7 @@ def _pk_in(keys: sa.Select[Any]) -> sa.ColumnElement[bool]: return sa.tuple_(*(target.c[c] for c in pk_cols)).in_(keys) def _update(extra: sa.ColumnElement[bool] | None = None) -> None: - stmt = sa.update(target).values({c: staged[c] for c in update_cols}).where(pk_join) + stmt = sa.update(target).values({c: staging.c[c] for c in update_cols}).where(pk_join) if extra is not None: stmt = stmt.where(_pk_in(_batch_keys(extra))) session.execute(stmt) @@ -472,8 +414,8 @@ def _drop_matched_from_staging(extra: sa.ColumnElement[bool] | None = None) -> N self_referential_fks = self._self_referential_deferred_fks(table_cls) if self_referential_fks: - insertable_cols = list(staged) - non_paginated_select = sa.select(*staged.values()) + insertable_cols = self._insertable_column_names(table_cls, staging) + non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) def _insert(select_: Any) -> sa.Insert: return sa.insert(target).from_select(insertable_cols, select_) @@ -517,8 +459,7 @@ def merge_upsert( self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) - staged = self._staged_columns(table_cls, staging) - insertable_cols = list(staged) + insertable_cols = self._insertable_column_names(table_cls, staging) def _upsert(select_: Any) -> sa.Insert: return ( @@ -527,7 +468,7 @@ def _upsert(select_: Any) -> sa.Insert: .on_conflict_do_nothing(index_elements=pk_cols) ) - non_paginated_select = sa.select(*staged.values()) + non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) self_referential = bool(self._self_referential_deferred_fks(table_cls)) for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): @@ -552,9 +493,8 @@ def merge_insert( self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) - staged = self._staged_columns(table_cls, staging) - insertable_cols = list(staged) - non_paginated_select = sa.select(*staged.values()) + insertable_cols = self._insertable_column_names(table_cls, staging) + non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) def _insert(select_: Any) -> sa.Insert: return sa.insert(target).from_select(insertable_cols, select_) diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index a179e67..c108a7f 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -13,12 +13,10 @@ from oa_configurator import ( Dialect, autocommit_connection, - physical_schema_of, qualified, ) from sqlalchemy.dialects import postgresql from sqlalchemy.exc import OperationalError -from sqlalchemy.schema import DropTable from sqlalchemy.sql.compiler import IdentifierPreparer from ..mappers.materialised_view_errors import ( @@ -81,76 +79,34 @@ def capabilities(self) -> BackendCapabilities: supports_replace_with_kept_indexes=True, ) - def create_staging_table( + def evaluate_default( self, - table_cls: type["FileTableProtocol"], session: so.Session, - ) -> None: - table = table_cls.__table__ - preparer = self.identifier_preparer - staging_name = self.staging_name_for_table(table_cls.__tablename__) - staging_ref = self.qualified_staging_name(table_cls.__tablename__) - bind = session.get_bind(mapper=table_cls) - source_ref = qualified( - bind, table.name, physical_schema=physical_schema_of(bind, schema_tag=table.schema) - ) - bind_arguments = {"mapper": table_cls} - session.execute( - DropTable(sa.Table(staging_name, sa.MetaData(), schema=self.staging_schema), if_exists=True), - bind_arguments=bind_arguments, - ) - session.execute( - sa.text( - f''' - CREATE UNLOGGED TABLE {staging_ref} - (LIKE {source_ref} INCLUDING DEFAULTS INCLUDING CONSTRAINTS); - ''' - ), - bind_arguments=bind_arguments, - ) - - 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_ref} DROP COLUMN {preparer.quote_identifier(col)};'), - bind_arguments=bind_arguments, - ) - - # staging only lands the file, so a null reaches the merge, where the column's default - # or the required check deals with it; LIKE would keep the NOT NULL of the target - nullable = ", ".join( - f"ALTER COLUMN {preparer.quote_identifier(c.name)} DROP NOT NULL" - for c in table.columns - if c.computed is None and not c.nullable - ) - if nullable: - session.execute(sa.text(f"ALTER TABLE {staging_ref} {nullable};"), bind_arguments=bind_arguments) - - # 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_ref} ADD COLUMN _rownum BIGINT' - f" GENERATED ALWAYS AS IDENTITY (CACHE 1000);" - ), - bind_arguments=bind_arguments, - ) - - session.commit() - - def drop_staging_table( - self, table_cls: type["FileTableProtocol"], - session: so.Session, - ) -> None: - staging_name = self.staging_name_for_table(table_cls.__tablename__) - session.execute( - DropTable(sa.Table(staging_name, sa.MetaData(), schema=self.staging_schema), if_exists=True), - bind_arguments={"mapper": table_cls}, + expression: sa.ColumnElement[Any], + column_type: Any, + ) -> Any: + """Evaluate on a separate connection, because the session's is busy with the COPY.""" + bind = session.get_bind(mapper=table_cls) + engine = bind.engine if isinstance(bind, sa.Connection) else bind + with engine.connect() as connection: + return connection.execute(sa.select(sa.cast(expression, column_type))).scalar_one() + + def _staging_table(self, table_cls: type["FileTableProtocol"], columns: Sequence[sa.Column[Any]]) -> sa.Table: + """Unlogged, with an identity column, which paginates the merge in O(N log N) without an index.""" + return sa.Table( + self.staging_name_for_table(table_cls.__tablename__), + sa.MetaData(), + *columns, + sa.Column("_rownum", sa.BigInteger, sa.Identity(always=True, cache=1000)), + schema=self.staging_schema, + prefixes=["UNLOGGED"], ) def insert_batches( self, ctx: "LoaderContext", + staging: sa.Table, batches: Iterable[pa.Table], columns: Sequence[str], ) -> None: @@ -165,7 +121,7 @@ def insert_batches( if not hasattr(cursor, "copy"): cursor.close() raise RuntimeError("The PostgreSQL sink needs a psycopg 3 connection, whose cursors provide copy()") - table_ref = self._staging_reference(ctx, connection) + table_ref = self._staging_reference(staging, connection) column_list = ", ".join(self.identifier_preparer.quote(c) for c in columns) options = pacsv.WriteOptions(include_header=False, quoting_style="needed") try: @@ -278,8 +234,7 @@ def merge_upsert( self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) - staged = self._staged_columns(table_cls, staging) - insertable_cols = list(staged) + insertable_cols = self._insertable_column_names(table_cls, staging) def _upsert(select_: Any) -> sa.Insert: # sa.insert() has no .on_conflict_do_nothing() @@ -289,7 +244,7 @@ def _upsert(select_: Any) -> sa.Insert: .on_conflict_do_nothing(index_elements=pk_cols) ) - non_paginated_select = sa.select(*staged.values()) + non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): select_ = non_paginated_select if batch_range is None else non_paginated_select.where(batch_range) @@ -314,9 +269,8 @@ def merge_insert( self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) - staged = self._staged_columns(table_cls, staging) - insertable_cols = list(staged) - non_paginated_select = sa.select(*staged.values()) + insertable_cols = self._insertable_column_names(table_cls, staging) + non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) def _insert(select_: Any) -> sa.Insert: return sa.insert(target).from_select(insertable_cols, select_) diff --git a/src/orm_loader/backends/sqlite.py b/src/orm_loader/backends/sqlite.py index 592184e..017a831 100644 --- a/src/orm_loader/backends/sqlite.py +++ b/src/orm_loader/backends/sqlite.py @@ -13,7 +13,6 @@ from sqlalchemy import event, text from sqlalchemy.dialects import sqlite as sqlite_dialect from sqlalchemy.exc import IntegrityError -from sqlalchemy.schema import DropTable from sqlalchemy.sql.compiler import IdentifierPreparer from .base import BackendCapabilities, DatabaseBackend, Dialect @@ -119,6 +118,7 @@ def default_index_strategy(self) -> str: def insert_batches( self, ctx: "LoaderContext", + staging: sa.Table, batches: Iterable[pa.Table], columns: Sequence[str], ) -> None: @@ -131,7 +131,7 @@ def insert_batches( """ connection = ctx.session.connection(bind_arguments={"mapper": ctx.tableclass}) preparer = connection.dialect.identifier_preparer - table_ref = self._staging_reference(ctx, connection) + table_ref = self._staging_reference(staging, connection) sql = ( f"INSERT INTO {table_ref} ({', '.join(preparer.quote(c) for c in columns)}) " f"VALUES ({', '.join('?' for _ in columns)})" @@ -148,40 +148,15 @@ def _storable(column: pa.ChunkedArray) -> list[Any]: column = pc.cast(column, pa.string()) return column.to_pylist() - def create_staging_table( + def evaluate_default( self, - table_cls: type["FileTableProtocol"], session: so.Session, - ) -> None: - staging_name = self.staging_name_for_table(table_cls.__tablename__) - session.execute( - DropTable(sa.Table(staging_name, sa.MetaData()), if_exists=True), - bind_arguments={"mapper": table_cls}, - ) - - metadata = sa.MetaData() - staging_columns = [ - sa.Column(col.name, col.type, nullable=True) - for col in table_cls.__table__.columns - ] - staging_table = sa.Table( - staging_name, metadata, *staging_columns, schema=table_cls.__table__.schema - ) - metadata.create_all( - bind=session.connection(bind_arguments={"mapper": table_cls}), tables=[staging_table] - ) - session.commit() - - def drop_staging_table( - self, table_cls: type["FileTableProtocol"], - session: so.Session, - ) -> None: - staging_name = self.staging_name_for_table(table_cls.__tablename__) - session.execute( - DropTable(sa.Table(staging_name, sa.MetaData()), if_exists=True), - bind_arguments={"mapper": table_cls}, - ) + expression: sa.ColumnElement[Any], + column_type: Any, + ) -> Any: + """Evaluate without a cast: SQLite casts a date to the integer of its year.""" + return session.execute(sa.select(expression), bind_arguments={"mapper": table_cls}).scalar_one() def disable_fk_check(self, session: so.Session, *, mapper: Any | None = None) -> str | int: execute_options = {"bind_arguments": {"mapper": mapper}} if mapper is not None else {} @@ -279,8 +254,7 @@ def merge_upsert( self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) - staged = self._staged_columns(table_cls, staging) - insertable_cols = list(staged) + insertable_cols = self._insertable_column_names(table_cls, staging) def _upsert(select_: Any) -> sa.Insert: return ( @@ -289,7 +263,7 @@ def _upsert(select_: Any) -> sa.Insert: .on_conflict_do_nothing(index_elements=pk_cols) ) - non_paginated_select = sa.select(*staged.values()) + non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): @@ -307,9 +281,8 @@ def merge_insert( self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) - staged = self._staged_columns(table_cls, staging) - insertable_cols = list(staged) - non_paginated_select = sa.select(*staged.values()) + insertable_cols = self._insertable_column_names(table_cls, staging) + non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) def _insert(select_: Any) -> sa.Insert: return sa.insert(target).from_select(insertable_cols, select_) diff --git a/src/orm_loader/loaders/loader.py b/src/orm_loader/loaders/loader.py index 8fadb47..69a165c 100644 --- a/src/orm_loader/loaders/loader.py +++ b/src/orm_loader/loaders/loader.py @@ -83,8 +83,6 @@ class LoaderContext: Active SQLAlchemy session. path Path to the source data file. - staging_table - SQLAlchemy Table object representing the staging table. chunksize Rows per batch for the default reader of formats that batch by rows. normalise @@ -110,7 +108,6 @@ class LoaderContext: tableclass: type[FileTableProtocol] session: so.Session path: Path - staging_table: sa.Table chunksize: int | None = None normalise: bool = True @@ -149,7 +146,7 @@ def _require_columns_present(tableclass: type[FileTableProtocol], present_column ) -def stage_file(ctx: LoaderContext, reader: FileReader) -> int: +def stage_file(ctx: LoaderContext, reader: FileReader, staging: sa.Table) -> int: """Load a file into the staging table. The reader streams the file in bounded batches, so memory does not grow @@ -163,6 +160,8 @@ def stage_file(ctx: LoaderContext, reader: FileReader) -> int: Loader context. reader Reader that yields the file as Arrow batches. + staging + Staging table the rows are written to. Returns ------- @@ -192,6 +191,7 @@ def stage_file(ctx: LoaderContext, reader: FileReader) -> int: logger.warning("File %s has columns that %s does not have, which are ignored: %s", ctx.path.name, table_name, ignored) if ctx.normalise: _require_columns_present(ctx.tableclass, reader.columns) + backend = resolve_backend(ctx.session, staging_schema_tag=ctx.staging_schema_tag, mapper=ctx.tableclass) rejects_dir = ctx.rejects_dir or ctx.path.parent / "rejects" processor = BatchProcessor( ctx.tableclass, @@ -201,13 +201,13 @@ def stage_file(ctx: LoaderContext, reader: FileReader) -> int: rejects_dir / f"{ctx.path.stem}.rejects.csv", ctx.normalise, reader.skipped, + lambda expression, column_type: backend.evaluate_default(ctx.session, ctx.tableclass, expression, column_type), ) - backend = resolve_backend(ctx.session, staging_schema_tag=ctx.staging_schema_tag, mapper=ctx.tableclass) - backend.insert_batches(ctx, processor.stream(chain([] if first is None else [first], batches)), reader.columns) - total = backend.count_staging(ctx) + backend.insert_batches(ctx, staging, processor.stream(chain([] if first is None else [first], batches)), processor.out_columns) + total = backend.count_staging(ctx, staging) keys = [c for c in ctx.tableclass.pk_names() if c in reader.columns] if ctx.dedupe and keys: - ctx.report.dedupe = backend.deduplicate_staging(ctx, keys) + ctx.report.dedupe = backend.deduplicate_staging(ctx, staging, keys) total = ctx.report.dedupe.rows_after ctx.report.rows_loaded += total if not ctx.report.accounted: diff --git a/src/orm_loader/loaders/processor.py b/src/orm_loader/loaders/processor.py index 787cbe7..70331e2 100644 --- a/src/orm_loader/loaders/processor.py +++ b/src/orm_loader/loaders/processor.py @@ -5,17 +5,24 @@ import logging import os from bisect import bisect_left -from collections.abc import Iterable, Iterator +from collections.abc import Callable, Iterable, Iterator from pathlib import Path from typing import TYPE_CHECKING, Any import pyarrow as pa import pyarrow.compute as pc import pyarrow.csv as pv +import sqlalchemy as sa -from ..conversion import ConversionPolicy +from ..conversion import ConversionPolicy, Interpretation from ..conversion.values import no_failures, underlying_type -from ..helpers import ConversionError, IngestError, RejectLimitError, RequiredValueError +from ..helpers import ( + ConversionError, + IngestError, + PolicyError, + RejectLimitError, + RequiredValueError, +) if TYPE_CHECKING: from ..tables.typing import FileTableProtocol @@ -96,6 +103,23 @@ class BatchProcessor: Lines the reader left out as malformed, each as its line number in the file and the reason. They are written to the rejects file and count as rows read and rejected. + evaluate + Function that returns the value of a SQL expression default as a + column of the given type stores it. + + Attributes + ---------- + out_columns + Columns of the typed batches: the file's columns, then the columns + the file lacks that have a default. + + Notes + ----- + A missing value takes the column's default, and a column the file lacks + is filled with it. The default is read by the column's handler like a + value from the file. A SQL expression default is evaluated for each batch + that needs it. A missing value left in a non-nullable column is a + required value that is missing. """ def __init__( @@ -107,6 +131,7 @@ def __init__( rejects_path: Path, normalise: bool = True, skipped: list[tuple[int | None, str]] | None = None, + evaluate: Callable[[Any, Any], Any] | None = None, ) -> None: self.columns = columns self.skipped = skipped if skipped is not None else [] @@ -117,11 +142,17 @@ def __init__( self.normalise = normalise self.table_name = tableclass.__tablename__ model_columns = tableclass.model_columns() - self.types = {name: underlying_type(model_columns[name].type) for name in columns} + self.defaults = tableclass.column_defaults() if normalise else {} + self.out_columns = [*columns, *(name for name in self.defaults if name not in columns)] + self.types = {name: underlying_type(model_columns[name].type) for name in self.out_columns} self.handlers = ( - {name: policy.for_column(self.types[name], self.table_name, name) for name in columns} if normalise else {} + {name: policy.for_column(self.types[name], self.table_name, name) for name in self.out_columns} + if normalise + else {} ) - self.required = [c for c in tableclass.required_columns() if c in columns] + self.required = [name for name in self.out_columns if not model_columns[name].nullable] + self.evaluate = evaluate + self._literals: dict[str, pa.Scalar] = {} self.rejects = RejectsWriter(rejects_path) if policy.errors.on_error == "reject_row": self.rejects.check_writable() @@ -171,12 +202,48 @@ def process(self, raw: pa.RecordBatch) -> pa.Table: name: self.handlers[name].interpret(raw.column(name), self.types[name], self.policy.nulls) for name in self.columns } + self._fill_defaults(results, raw.num_rows) failures = self._failures(results) if failures: if self.policy.errors.on_error == "raise": self._raise(raw, results, first) return self._reject(raw, results, failures, first) - return pa.table({name: results[name].values for name in self.columns}) + return pa.table({name: results[name].values for name in self.out_columns}) + + def _fill_defaults(self, results: dict[str, Interpretation], count: int) -> None: + """Give a missing value, and a column the file lacks, the column's default.""" + for name, default in self.defaults.items(): + if name not in results: + results[name] = Interpretation(pa.repeat(self._fill_value(name, default), count), no_failures(count)) + continue + values, invalid = results[name] + missing = pc.and_(pc.is_null(values), pc.invert(invalid)) + if pc.any(missing).as_py(): + filled = pc.if_else(missing, self._fill_value(name, default), values) + results[name] = Interpretation(filled, invalid) + + def _fill_value(self, name: str, default: Any) -> pa.Scalar: + """Return the scalar that fills a column: a literal is read once, an expression is evaluated for each batch.""" + if isinstance(default, sa.ClauseElement): + if self.evaluate is None: + raise PolicyError(table_name=self.table_name, column_name=name, column_type="SQL expression default") + return self._read_default(name, self.evaluate(default, self.types[name])) + if name not in self._literals: + self._literals[name] = self._read_default(name, default) + return self._literals[name] + + def _read_default(self, name: str, value: Any) -> pa.Scalar: + """Read a default value the way the column's handler reads a value from the file.""" + values, invalid = self.handlers[name].interpret(pa.array([value]), self.types[name], self.policy.nulls) + if invalid[0].as_py(): + raise ConversionError( + table_name=self.table_name, + column_name=name, + expected=self.handlers[name].expected(self.types[name]), + count=1, + samples=[value], + ) + return values[0] def _failures(self, results: dict[str, Any]) -> dict[str, tuple[pa.Array, str]]: found: dict[str, tuple[pa.Array, str]] = {} @@ -262,7 +329,7 @@ def _reject( self.report.rejects_path = self.rejects.path self._check_limit() keep = pc.invert(rejected) - return pa.table({n: results[n].values for n in self.columns}).filter(keep) + return pa.table({n: results[n].values for n in self.out_columns}).filter(keep) def _row_numbers(self, first: int, count: int) -> pa.Array: """Return the data row numbers in the file of ``count`` parsed rows, counting malformed lines. diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index a16ed3c..8836fab 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -7,7 +7,7 @@ and index handling. """ import logging -from collections.abc import Mapping +from collections.abc import Collection, Mapping from contextlib import contextmanager from pathlib import Path from time import perf_counter @@ -15,7 +15,7 @@ import sqlalchemy as sa import sqlalchemy.orm as so -from oa_configurator import Dialect, ResolvedDatabase, physical_schema_of +from oa_configurator import ResolvedDatabase, physical_schema_of from sqlalchemy.exc import InvalidRequestError, UnboundExecutionError from ..backends.resolve import resolve_backend @@ -64,12 +64,14 @@ def create_staging_table( session: so.Session, *, staging_schema_tag: str | None = None, + without: Collection[str] = (), ): """ Create a fresh staging table for ingestion. Any existing staging table with the same name is dropped first. - The staging table schema mirrors the target table schema. + The staging table has the target's columns with their types and + nullability. Parameters ---------- @@ -79,6 +81,9 @@ def create_staging_table( schema_translate_map tag the staging table should be created under. ``None`` means no schema qualification (backend-default behavior). + without + Names of target columns the staging table leaves out. The merge + then leaves them out of its insert. Raises ------ @@ -89,7 +94,7 @@ def create_staging_table( """ _require_bind(session, cls) backend = resolve_backend(session, staging_schema_tag=staging_schema_tag, mapper=cls) - backend.create_staging_table(cls, session) + backend.create_staging_table(cls, session, without=without) @classmethod @contextmanager @@ -214,7 +219,7 @@ def get_staging_table( staging_schema_tag: str | None = None, ) -> sa.Table: """ - Return the reflected staging table, creating it if necessary. + Return the reflected staging table. Parameters ---------- @@ -229,38 +234,26 @@ def get_staging_table( sqlalchemy.Table The reflected staging table. + Raises + ------ + sqlalchemy.exc.NoSuchTableError + The staging table has not been created. + Notes ----- - Inspects and reflects via ``session.connection()``, not the bare - engine. Confirmed empirically: on SQLite's SingletonThreadPool, a - second connection opened straight from the engine is the same - underlying DBAPI connection, and closing that second wrapper resets - its perceived transaction state, silently discarding the session's - own uncommitted work. Using the session's own already-open - connection avoids ever opening a second one. Fetched fresh both - before and after the possible ``create_staging_table()`` call below, - since that call commits, which can invalidate an earlier reference. + Reflects via ``session.connection()``, not the bare engine. On + SQLite's SingletonThreadPool a second connection opened straight + from the engine is the same underlying DBAPI connection, and closing + that second wrapper resets its perceived transaction state, silently + discarding the session's own uncommitted work. """ _require_bind(session, cls) backend = resolve_backend(session, staging_schema_tag=staging_schema_tag, mapper=cls) - connection = session.connection(bind_arguments={"mapper": cls}) - inspector = sa.inspect(connection) - staging_name = backend.staging_name_for_table(cls.__tablename__) - - staging_schema = backend.staging_schema - if staging_schema is None and backend.dialect == Dialect.SQLITE: - staging_schema = cls.__table__.schema - 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, staging_schema_tag=staging_schema_tag) - - connection = session.connection(bind_arguments={"mapper": cls}) - return sa.Table( - staging_name, + backend.staging_name_for_table(cls.__tablename__), sa.MetaData(), # throwaway — keeps staging table out of Base.metadata - autoload_with=connection, - schema=staging_schema, + autoload_with=session.connection(bind_arguments={"mapper": cls}), + schema=backend.staging_schema, ) @classmethod @@ -284,9 +277,12 @@ def load_staging( int Number of rows loaded into the staging table. """ - _require_bind(loader_context.session, cls) - cls.create_staging_table(loader_context.session, staging_schema_tag=loader_context.staging_schema_tag) - return stage_file(loader_context, reader) + session, tag = loader_context.session, loader_context.staging_schema_tag + _require_bind(session, cls) + left_out = set() if reader.empty else cls.autoincrement_columns() - set(reader.columns) + cls.create_staging_table(session, staging_schema_tag=tag, without=left_out) + staging = cls.get_staging_table(session, staging_schema_tag=tag) + return stage_file(loader_context, reader, staging) @classmethod def _select_reader(cls: Type[FileTableProtocol], ctx: LoaderContext) -> FileReader: @@ -453,7 +449,6 @@ def load_file( tableclass=cls, session=session, path=path, - staging_table=cls.get_staging_table(session, staging_schema_tag=staging_schema_tag), chunksize=chunksize, normalise=normalise, dedupe=dedupe, diff --git a/src/orm_loader/tables/orm_table.py b/src/orm_loader/tables/orm_table.py index 25a3950..a299cdc 100644 --- a/src/orm_loader/tables/orm_table.py +++ b/src/orm_loader/tables/orm_table.py @@ -143,7 +143,7 @@ def required_columns(cls) -> set[str]: A column is considered required if it is: - non-nullable - - has no default that SQL can apply (see ``sql_defaults``) + - has no default that can be supplied (see ``column_defaults``) - not an autoincrementing column (the database assigns its value; it's never expected in inbound data at all) @@ -155,39 +155,48 @@ def required_columns(cls) -> set[str]: set[str] A set of required column names. """ - mapper = cls.mapper_for() - defaults = cls.sql_defaults() + defaults = cls.column_defaults() return { c.key - for c in mapper.columns + for c in cls.mapper_for().columns if not c.nullable and c.key not in defaults and not _resolves_to_autoincrement(c) } @classmethod - def sql_defaults(cls) -> dict[str, sa.ColumnElement[Any]]: + def autoincrement_columns(cls) -> set[str]: + """ + Return the names of the columns whose value the database assigns. + + Returns + ------- + set[str] + A set of column names. + """ + return {c.key for c in cls.mapper_for().columns if _resolves_to_autoincrement(c)} + + @classmethod + def column_defaults(cls) -> dict[str, Any]: """ - Return the default of each column that has one SQL can apply. + Return the default of each column that has one that can be supplied. - A server default given as text is written inline, as in the DDL, so - the database reads it as the column's type. One given as a SQL - expression is used as written, and a scalar Python default is a typed - literal. A callable or context-sensitive default has no SQL form, so - its column is not listed. + A server default given as text and a scalar Python default are + returned as the value itself. A server default given as a SQL + expression is returned as that expression, for the database to + evaluate. A callable or context-sensitive default cannot be supplied + for a whole file, so its column is not listed. Returns ------- - dict[str, sqlalchemy.ColumnElement] - A mapping of column name to the expression that supplies its default. + dict[str, Any] + A mapping of column name to the value, or SQL expression, of its default. """ - defaults: dict[str, sa.ColumnElement[Any]] = {} + defaults: dict[str, Any] = {} for c in cls.mapper_for().columns: server_default = getattr(c.server_default, "arg", None) - if isinstance(server_default, str): - defaults[c.name] = sa.literal(server_default, sa.String()).render_literal_execute() - elif server_default is not None: + if server_default is not None: defaults[c.name] = server_default elif c.default is not None and c.default.is_scalar: - defaults[c.name] = sa.literal(c.default.arg, c.type) # ty: ignore[unresolved-attribute] + defaults[c.name] = c.default.arg # ty: ignore[unresolved-attribute] return defaults @classmethod diff --git a/src/orm_loader/tables/typing.py b/src/orm_loader/tables/typing.py index c5ee018..df548d7 100644 --- a/src/orm_loader/tables/typing.py +++ b/src/orm_loader/tables/typing.py @@ -1,7 +1,7 @@ from typing import Protocol, ClassVar, runtime_checkable, TYPE_CHECKING, Optional, Type, Dict, Any, Unpack, TypedDict import sqlalchemy.orm as so import sqlalchemy as sa -from collections.abc import Mapping +from collections.abc import Collection, Mapping from pathlib import Path from contextlib import AbstractContextManager if TYPE_CHECKING: @@ -52,7 +52,10 @@ def model_columns(cls) -> dict[str, sa.ColumnElement[Any]]: ... def required_columns(cls) -> set[str]: ... @classmethod - def sql_defaults(cls) -> dict[str, sa.ColumnElement[Any]]: ... + def autoincrement_columns(cls) -> set[str]: ... + + @classmethod + def column_defaults(cls) -> dict[str, Any]: ... @runtime_checkable class FileTableProtocol(ORMTableProtocol, Protocol): @@ -74,7 +77,9 @@ class FileTableProtocol(ORMTableProtocol, Protocol): def _select_reader(cls, ctx: "LoaderContext") -> "FileReader": ... @classmethod - def create_staging_table(cls, session: so.Session, *, staging_schema_tag: str | None = None) -> None: ... + def create_staging_table( + cls, session: so.Session, *, staging_schema_tag: str | None = None, without: Collection[str] = () + ) -> None: ... @classmethod def load_staging(cls: Type["FileTableProtocol"], loader_context: "LoaderContext", reader: "FileReader") -> int: ... diff --git a/tests/backends/test_base_backend.py b/tests/backends/test_base_backend.py index 3a6537e..69ac326 100644 --- a/tests/backends/test_base_backend.py +++ b/tests/backends/test_base_backend.py @@ -39,10 +39,10 @@ class FakeBackend(DatabaseBackend): def __init__(self) -> None: self.calls: list[tuple[str, object]] = [] - def insert_batches(self, ctx, batches, columns) -> int: + def insert_batches(self, ctx, staging, batches, columns) -> int: total = 0 for table in batches: - ctx.session.execute(ctx.staging_table.insert(), table.to_pylist()) + ctx.session.execute(staging.insert(), table.to_pylist()) total += table.num_rows return total diff --git a/tests/backends/test_duckdb_backend.py b/tests/backends/test_duckdb_backend.py index 519232c..8c8b7ee 100644 --- a/tests/backends/test_duckdb_backend.py +++ b/tests/backends/test_duckdb_backend.py @@ -191,8 +191,9 @@ def test_duckdb_backend_create_staging_table(duckdb_session): inspector = sa.inspect(duckdb_session.get_bind()) assert inspector.has_table(_STAGING_TABLE) is True cols = inspector.get_columns(_STAGING_TABLE) - assert [c["name"] for c in cols] == ["id", "name", "slug", "_rownum"] - assert all(c["nullable"] is True for c in cols if c["name"] != "_rownum") + assert [c["name"] for c in cols] == ["id", "name", "_rownum"] + target = {c.name: c.nullable for c in ComputedColumnTable.__table__.columns} + assert {c["name"]: c["nullable"] for c in cols if c["name"] != "_rownum"} == {name: target[name] for name in ("id", "name")} @pytest.mark.duckdb @@ -229,6 +230,19 @@ def test_duckdb_backend_drop_staging_table_also_drops_rownum_sequence(duckdb_ses assert not any(f"{_STAGING_TABLE}_rownum_seq" in name for name in sequences_after_drop) +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_creates_staging_again_when_only_its_sequence_was_left_behind(duckdb_session): + backend = DuckDBBackend() + backend.create_staging_table(_ComputedTableCls, duckdb_session) + duckdb_session.execute(sa.text(f"DROP TABLE {_STAGING_TABLE}")) + duckdb_session.commit() + + backend.create_staging_table(_ComputedTableCls, duckdb_session) + + assert sa.inspect(duckdb_session.get_bind()).has_table(_STAGING_TABLE) + + @pytest.mark.duckdb @pytest.mark.db_dialect def test_duckdb_backend_rownum_populates_monotonically_on_insert(duckdb_session): diff --git a/tests/backends/test_postgres_backend.py b/tests/backends/test_postgres_backend.py index 8bb8843..5823233 100644 --- a/tests/backends/test_postgres_backend.py +++ b/tests/backends/test_postgres_backend.py @@ -124,15 +124,6 @@ def test_postgres_backend_create_staging_table_drops_computed_columns(pg_session assert cols == {"id", "name", "_rownum"} # slug is computed, excluded -def test_postgres_backend_drop_staging_table(): - backend = PostgresBackend(staging_schema=STAGING_SCHEMA) - session = _FakeSession() - - backend.drop_staging_table(_ComputedTableCls, _sess(session)) - - assert session.statements == [f'DROP TABLE IF EXISTS {_STAGING_TABLE_WITH_SCHEMA}'] - - def test_postgres_backend_fk_methods_emit_expected_sql(): backend = PostgresBackend() session = _FakeSession() diff --git a/tests/backends/test_sinks.py b/tests/backends/test_sinks.py index f7c1e68..ec4a7ec 100644 --- a/tests/backends/test_sinks.py +++ b/tests/backends/test_sinks.py @@ -20,9 +20,7 @@ def _sink(session, tmp_path): tag = None if session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA _Demo.create_staging_table(session, staging_schema_tag=tag) staging = _Demo.get_staging_table(session, staging_schema_tag=tag) - ctx = LoaderContext( - tableclass=_Demo, session=session, path=tmp_path / "demo_record.csv", staging_table=staging, staging_schema_tag=tag - ) + ctx = LoaderContext(tableclass=_Demo, session=session, path=tmp_path / "demo_record.csv", staging_schema_tag=tag) return resolve_backend(session, staging_schema_tag=tag, mapper=_Demo), ctx, staging @@ -45,7 +43,7 @@ def _stored(session, staging): def test_typed_values_and_nulls_are_stored_across_more_batches_than_a_commit_interval(dialect_session, tmp_path): backend, ctx, staging = _sink(dialect_session, tmp_path) batches = [_batch(i, f"row {i}" if i % 2 else None) for i in range(1, 13)] - backend.insert_batches(ctx, batches, COLUMNS) + backend.insert_batches(ctx, staging, batches, COLUMNS) stored = _stored(dialect_session, staging) assert [r.demo_id for r in stored] == list(range(1, 13)) assert stored[0].demo_label == "row 1" and stored[1].demo_label is None @@ -56,13 +54,13 @@ def test_typed_values_and_nulls_are_stored_across_more_batches_than_a_commit_int def test_rows_keep_the_order_they_were_written_in(dialect_session, tmp_path): backend, ctx, staging = _sink(dialect_session, tmp_path) ids = [5, 3, 9, 1, 7] - backend.insert_batches(ctx, [_batch(i, None) for i in ids], COLUMNS) + backend.insert_batches(ctx, staging, [_batch(i, None) for i in ids], COLUMNS) assert [r.demo_id for r in _stored(dialect_session, staging)] == ids def test_no_batches_write_nothing(dialect_session, tmp_path): backend, ctx, staging = _sink(dialect_session, tmp_path) - backend.insert_batches(ctx, [], COLUMNS) + backend.insert_batches(ctx, staging, [], COLUMNS) assert _stored(dialect_session, staging) == [] @@ -71,28 +69,56 @@ def test_no_batches_write_nothing(dialect_session, tmp_path): ) def test_text_is_stored_verbatim_and_an_empty_string_stays_distinct_from_null(dialect_session, tmp_path, label): backend, ctx, staging = _sink(dialect_session, tmp_path) - backend.insert_batches(ctx, [_batch(1, label), _batch(2, None)], COLUMNS) + backend.insert_batches(ctx, staging, [_batch(1, label), _batch(2, None)], COLUMNS) stored = _stored(dialect_session, staging) assert (stored[0].demo_label, stored[1].demo_label) == (label, None) def test_a_failing_batch_source_is_raised_not_swallowed(dialect_session, tmp_path): - backend, ctx, _ = _sink(dialect_session, tmp_path) + backend, ctx, staging = _sink(dialect_session, tmp_path) def batches(): yield _batch(1, "a") raise RuntimeError("boom") with pytest.raises(RuntimeError, match="boom"): - backend.insert_batches(ctx, batches(), COLUMNS) + backend.insert_batches(ctx, staging, batches(), COLUMNS) + + +def _staging_for(session, table, **kwargs): + tag = None if session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA + table.create_staging_table(session, staging_schema_tag=tag, **kwargs) + return table.get_staging_table(session, staging_schema_tag=tag) -def test_a_staging_table_accepts_a_null_in_every_column(dialect_session): +def test_a_staging_table_mirrors_the_columns_and_nullability_of_the_target(dialect_session): from tests.models import DemoRequiredRecord - table = cast(Type[FileTableProtocol], DemoRequiredRecord) + staging = _staging_for(dialect_session, cast(Type[FileTableProtocol], DemoRequiredRecord)) + target = DemoRequiredRecord.__table__ + assert [c.name for c in staging.c if c.name != "_rownum"] == [c.name for c in target.c] + assert {c.name: c.nullable for c in staging.c if c.name != "_rownum"} == {c.name: c.nullable for c in target.c} + assert not target.c.demo_label.nullable + + +def test_a_staging_table_has_no_key_and_no_default(dialect_session): + from tests.models import DemoDefaultedRecord + + staging = _staging_for(dialect_session, cast(Type[FileTableProtocol], DemoDefaultedRecord)) + assert staging.primary_key.columns.keys() == [] + assert all(c.server_default is None for c in staging.c if c.name != "_rownum") + + +def test_columns_can_be_left_out_of_a_staging_table(dialect_session): + staging = _staging_for(dialect_session, _Demo, without={"demo_label"}) + assert "demo_label" not in staging.c and "demo_id" in staging.c + + +def test_a_staging_table_can_be_dropped_twice_and_created_again(dialect_session): + staging = _staging_for(dialect_session, _Demo) tag = None if dialect_session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA - table.create_staging_table(dialect_session, staging_schema_tag=tag) - staging = table.get_staging_table(dialect_session, staging_schema_tag=tag) - assert not DemoRequiredRecord.__table__.c.demo_label.nullable - assert [c.name for c in staging.c if not c.nullable and c.name != "_rownum"] == [] + for _ in range(2): + _Demo.drop_staging_table(dialect_session, staging_schema_tag=tag) + with pytest.raises(sa.exc.NoSuchTableError): + _Demo.get_staging_table(dialect_session, staging_schema_tag=tag) + assert _staging_for(dialect_session, _Demo).name == staging.name diff --git a/tests/backends/test_sqlite_backend.py b/tests/backends/test_sqlite_backend.py index 45af16d..f2c1e19 100644 --- a/tests/backends/test_sqlite_backend.py +++ b/tests/backends/test_sqlite_backend.py @@ -82,18 +82,9 @@ def test_sqlite_backend_create_staging_table(sqlite_session, sqlite_engine): inspector = sa.inspect(sqlite_engine) 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) - - -def test_sqlite_backend_drop_staging_table(): - backend = SQLiteBackend() - sqlite_session = _FakeSession() - - backend.drop_staging_table(_ComputedTableCls, _sess(sqlite_session)) - - assert sqlite_session.statements == [f"DROP TABLE IF EXISTS {_STAGING_TABLE}"] - + assert [c["name"] for c in cols] == ["id", "name"] + target = {c.name: c.nullable for c in ComputedColumnTable.__table__.columns} + assert {c["name"]: c["nullable"] for c in cols} == {name: target[name] for name in ("id", "name")} def test_sqlite_backend_disable_fk_reads_then_sets(): diff --git a/tests/loaders/test_defaults.py b/tests/loaders/test_defaults.py index 970dba6..b615f57 100644 --- a/tests/loaders/test_defaults.py +++ b/tests/loaders/test_defaults.py @@ -33,9 +33,9 @@ def _rows(session) -> list[tuple]: return [tuple(r) for r in session.execute(sa.select(*columns).order_by(model.demo_id))] -def test_only_a_default_that_sql_can_apply_makes_a_column_optional(): +def test_only_a_default_that_can_be_supplied_makes_a_column_optional(): assert _Table.required_columns() == {"demo_stamp"} - assert set(_Table.sql_defaults()) == {"demo_status", "demo_priority", "demo_note"} + assert set(_Table.column_defaults()) == {"demo_status", "demo_priority", "demo_note"} def test_an_empty_value_takes_the_default(dialect_session, tmp_path): @@ -84,3 +84,60 @@ def test_every_kind_of_default_is_applied_and_read_back_on_every_dialect(dialect assert (row.kind_date, row.kind_stamp) == (date(2000, 1, 2), datetime(2000, 1, 2, 10, 30)) assert (row.kind_expression, row.kind_py_int, row.kind_py_date) == ("expr", 5, date(2001, 2, 3)) assert isinstance(row.kind_now, datetime) + + +def _keyed(session, tmp_path: Path, text: str, **kwargs) -> list[tuple]: + from tests.models import DemoKeyedRecord + + path = tmp_path / "demo_keyed_record.csv" + path.write_text(text) + tag = None if session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA + cast(Type[FileTableProtocol], DemoKeyedRecord).load_file(session, path, staging_schema_tag=tag, **kwargs) + session.commit() + columns = (DemoKeyedRecord.keyed_id, DemoKeyedRecord.keyed_name) + return [tuple(r) for r in session.execute(sa.select(*columns).order_by(DemoKeyedRecord.keyed_id))] + + +def test_a_key_the_file_lacks_is_assigned_by_the_database(dialect_session, tmp_path): + assert _keyed(dialect_session, tmp_path, "keyed_name\na\nb\n") == [(1, "a"), (2, "b")] + + +def test_a_key_in_the_file_is_loaded_and_may_not_be_blank(dialect_session, tmp_path): + assert _keyed(dialect_session, tmp_path, "keyed_id,keyed_name\n10,a\n20,b\n") == [(10, "a"), (20, "b")] + with pytest.raises(RequiredValueError, match="keyed_id"): + _keyed(dialect_session, tmp_path, "keyed_id,keyed_name\n,a\n") + + +def test_an_expression_default_is_evaluated_for_each_batch_that_needs_it(dialect_session, tmp_path, monkeypatch): + from orm_loader.backends import resolve_backend + from orm_loader.loaders import DelimitedTextReader + + calls: list[object] = [] + backend_class = type(resolve_backend(dialect_session, mapper=DemoDefaultKinds)) + original = backend_class.evaluate_default + + def counting(self, session, table_cls, expression, column_type): + calls.append(expression) + return original(self, session, table_cls, expression, column_type) + + monkeypatch.setattr(backend_class, "evaluate_default", counting) + path = tmp_path / "demo_default_kinds.csv" + path.write_text("kind_id,kind_now\n" + "".join(f"{i},\n" for i in range(1, 30_001))) + tag = None if dialect_session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA + table = cast(Type[FileTableProtocol], DemoDefaultKinds) + reader = DelimitedTextReader(path, DemoDefaultKinds.model_columns(), block_size=64 << 10) + table.load_file(dialect_session, path, staging_schema_tag=tag, reader=reader) + batches = len(list(DelimitedTextReader(path, DemoDefaultKinds.model_columns(), block_size=64 << 10).batches())) + now_calls = [c for c in calls if isinstance(c, sa.sql.functions.FunctionElement)] + assert batches > 1 and len(now_calls) == batches + + +def test_a_default_the_handler_cannot_read_is_an_error_naming_the_column(dialect_session, tmp_path): + from orm_loader.helpers import ConversionError + from tests.models import DemoBadDefaultRecord + + path = tmp_path / "demo_bad_default_record.csv" + path.write_text("bad_id\n1\n") + tag = None if dialect_session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA + with pytest.raises(ConversionError, match="bad_count"): + cast(Type[FileTableProtocol], DemoBadDefaultRecord).load_file(dialect_session, path, staging_schema_tag=tag) diff --git a/tests/loaders/test_pipeline.py b/tests/loaders/test_pipeline.py index c2daaf9..82688fb 100644 --- a/tests/loaders/test_pipeline.py +++ b/tests/loaders/test_pipeline.py @@ -48,10 +48,8 @@ def _stage(session, table, path: Path, reader=None, **context): tag = None if session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA table.create_staging_table(session, staging_schema_tag=tag) staging = table.get_staging_table(session, staging_schema_tag=tag) - ctx = LoaderContext( - tableclass=table, session=session, path=path, staging_table=staging, staging_schema_tag=tag, **context - ) - loaded = stage_file(ctx, reader or table._select_reader(ctx)) + ctx = LoaderContext(tableclass=table, session=session, path=path, staging_schema_tag=tag, **context) + loaded = stage_file(ctx, reader or table._select_reader(ctx), staging) return loaded, ctx, staging @@ -131,9 +129,7 @@ def test_a_file_in_many_batches_loads_every_row_in_order(dialect_session, tmp_pa def test_a_failing_reader_is_raised_and_stages_nothing(dialect_session, tmp_path, fmt): path = _write(tmp_path, DemoRecord, fmt, [("1", "a", "", "", "")]) - reader = _Demo._select_reader( - LoaderContext(tableclass=_Demo, session=dialect_session, path=path, staging_table=cast(sa.Table, DemoRecord.__table__)) - ) + reader = _Demo._select_reader(LoaderContext(tableclass=_Demo, session=dialect_session, path=path)) def broken(): raise RuntimeError("boom") diff --git a/tests/loaders/test_stage_file.py b/tests/loaders/test_stage_file.py index 2e93c8c..cb630a3 100644 --- a/tests/loaders/test_stage_file.py +++ b/tests/loaders/test_stage_file.py @@ -15,7 +15,6 @@ def test_loader_context_fields(sqlite_session): tableclass=FileLoaderTable, session=sqlite_session, path=Path("file.csv"), - staging_table=FileLoaderTable.__table__, chunksize=100, normalise=False, dedupe=True, @@ -42,13 +41,12 @@ def test_stage_file_case_insensitive_headers(tmp_path, sqlite_session): tableclass=FileLoaderTable, session=sqlite_session, path=csv, - staging_table=FileLoaderTable.__table__, chunksize=100, normalise=True, dedupe=True, ) - n = stage_file(ctx, ctx.tableclass._select_reader(ctx)) + n = stage_file(ctx, ctx.tableclass._select_reader(ctx), FileLoaderTable.__table__) rows = sqlite_session.execute( sa.text('SELECT id, value FROM test_file_loader ORDER BY id') @@ -76,14 +74,13 @@ def test_stage_file_literal_mode_preserves_quotes(tmp_path, sqlite_session): tableclass=FileLoaderTable, session=sqlite_session, path=csv, - staging_table=FileLoaderTable.__table__, chunksize=100, normalise=True, dedupe=True, quote_mode="literal", ) - n = stage_file(ctx, ctx.tableclass._select_reader(ctx)) + n = stage_file(ctx, ctx.tableclass._select_reader(ctx), FileLoaderTable.__table__) rows = sqlite_session.execute( sa.text('SELECT id, value FROM test_file_loader ORDER BY id') @@ -110,14 +107,13 @@ def test_stage_file_csv_mode_strips_quotes(tmp_path, sqlite_session): tableclass=FileLoaderTable, session=sqlite_session, path=csv, - staging_table=FileLoaderTable.__table__, chunksize=100, normalise=True, dedupe=True, quote_mode="csv", ) - n = stage_file(ctx, ctx.tableclass._select_reader(ctx)) + n = stage_file(ctx, ctx.tableclass._select_reader(ctx), FileLoaderTable.__table__) rows = sqlite_session.execute( sa.text('SELECT id, value FROM test_file_loader ORDER BY id') @@ -144,14 +140,13 @@ def test_stage_file_by_delimiter_tab_preserves_quotes(tmp_path, sqlite_session): tableclass=FileLoaderTable, session=sqlite_session, path=tsv, - staging_table=FileLoaderTable.__table__, chunksize=100, normalise=True, dedupe=True, quote_mode="by_delimiter", ) - n = stage_file(ctx, ctx.tableclass._select_reader(ctx)) + n = stage_file(ctx, ctx.tableclass._select_reader(ctx), FileLoaderTable.__table__) rows = sqlite_session.execute( sa.text('SELECT id, value FROM test_file_loader ORDER BY id') @@ -169,9 +164,8 @@ def test_file_without_data_rows_loads_nothing_whatever_its_header(tmp_path, sqli tableclass=FileLoaderTable, session=sqlite_session, path=csv, - staging_table=FileLoaderTable.__table__, ) - assert stage_file(ctx, ctx.tableclass._select_reader(ctx)) == 0 + assert stage_file(ctx, ctx.tableclass._select_reader(ctx), FileLoaderTable.__table__) == 0 assert ctx.report.rows_read == 0 @@ -183,15 +177,14 @@ def test_accounting_is_checked_against_the_staging_table(tmp_path, sqlite_sessio tableclass=FileLoaderTable, session=sqlite_session, path=csv, - staging_table=FileLoaderTable.__table__, dedupe=False, ) backend = resolve_backend(sqlite_session, mapper=FileLoaderTable) insert = type(backend).insert_batches - def losing(self, ctx, batches, columns): - insert(self, ctx, list(batches)[:-1] or [], columns) + def losing(self, ctx, staging, batches, columns): + insert(self, ctx, staging, list(batches)[:-1] or [], columns) monkeypatch.setattr(type(backend), "insert_batches", losing) with pytest.raises(IngestError, match="read 3 rows but loaded 0"): - stage_file(ctx, ctx.tableclass._select_reader(ctx)) + stage_file(ctx, ctx.tableclass._select_reader(ctx), FileLoaderTable.__table__) diff --git a/tests/models.py b/tests/models.py index 7bbd68f..d7b15c2 100644 --- a/tests/models.py +++ b/tests/models.py @@ -264,3 +264,23 @@ class DemoBinaryRecord(Base, FileLoadableTableInterface): demo_id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) demo_blob: so.Mapped[str | None] = so.mapped_column(sa.LargeBinary, nullable=True) + + +class DemoKeyedRecord(Base, FileLoadableTableInterface): + """A key the database assigns.""" + + __tablename__ = "demo_keyed_record" + __table_args__ = {"schema": SchemaRole.PRIMARY.value} + + keyed_id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True, autoincrement=True) + keyed_name: so.Mapped[str] = so.mapped_column(sa.String, nullable=False) + + +class DemoBadDefaultRecord(Base, FileLoadableTableInterface): + """A default that cannot be read as the type of its column.""" + + __tablename__ = "demo_bad_default_record" + __table_args__ = {"schema": SchemaRole.PRIMARY.value} + + bad_id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) + bad_count: so.Mapped[int | None] = so.mapped_column(sa.Integer, nullable=True, default="abc") From 82bb14ad6f03b7183cc190ed15b9005d215edeaf Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Thu, 29 Oct 2026 22:26:16 +0000 Subject: [PATCH 09/21] Pass the mapper straight through to the foreign key toggles --- src/orm_loader/backends/base.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 22c3e9e..382f7a5 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -949,11 +949,7 @@ def bulk_load_context( try: if disable_fk: self._require_capability("supports_fk_toggle", "foreign key toggling") - raw_state = ( - self.disable_fk_check(session) - if mapper is None - else self.disable_fk_check(session, mapper=mapper) - ) + raw_state = self.disable_fk_check(session, mapper=mapper) previous_fk_state = self._normalize_fk_check_state(raw_state) if no_autoflush: @@ -968,10 +964,7 @@ def bulk_load_context( finally: if previous_fk_state is not None: - if mapper is None: - self.restore_fk_check(session, previous_fk_state) - else: - self.restore_fk_check(session, previous_fk_state, mapper=mapper) + self.restore_fk_check(session, previous_fk_state, mapper=mapper) @requires_capability("supports_materialized_views", "materialized views") def create_materialized_view( From f37776c542e77f082f2b0bc1ccd709b8f1c93bb4 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Thu, 29 Oct 2026 22:39:01 +0000 Subject: [PATCH 10/21] Replace bulk_load_context and merge_context with relaxed_foreign_keys that logs the toggling --- docs/loaders/staging.md | 4 +- src/orm_loader/backends/base.py | 49 ++++++++++++++----------- src/orm_loader/tables/loadable_table.py | 31 ++++------------ tests/backends/test_base_backend.py | 37 +++++++------------ tests/backends/test_duckdb_backend.py | 17 +++++++-- tests/backends/test_sqlite_backend.py | 9 +++++ tests/loaders/test_loader_e2e.py | 4 +- 7 files changed, 74 insertions(+), 77 deletions(-) diff --git a/docs/loaders/staging.md b/docs/loaders/staging.md index 0af150d..d3137ed 100644 --- a/docs/loaders/staging.md +++ b/docs/loaders/staging.md @@ -46,7 +46,9 @@ No schema changes are performed at the loader layer. ## Referential integrity -On PostgreSQL and SQLite a load relaxes foreign keys while it merges, so a row loaded before the row it references is accepted. After every table is loaded, `Table.validate_foreign_keys(session)` checks every foreign key of the table with an anti-join and raises `ForeignKeyViolationError` with up to five unmatched keys. Null references are not violations, and the load order does not matter. +On PostgreSQL and SQLite a load relaxes foreign keys while it merges, so the caller needs no foreign key handling and a row loaded before the row it references is accepted. To verify the result, call `Table.validate_foreign_keys(session)` after every table is loaded. It checks every foreign key of the table with an anti-join and raises `ForeignKeyViolationError` with up to five unmatched keys. Null references are not violations, and the load order does not matter. + +To insert ORM objects built in Python that reference each other, wrap the inserts in `backend.relaxed_foreign_keys(session)`, which relaxes the keys on PostgreSQL and SQLite and leaves them enforced on DuckDB. DuckDB enforces the keys it creates, so load parent tables first. `load_file` already validates the keys DuckDB never enforces (cross-schema, cyclic and self-referential) after each table's merge, and `validate_foreign_keys` repeats that check, so it works the same way on every backend. A DuckDB table with foreign keys needs `resolved=` for both. diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 382f7a5..948c1d5 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -1,10 +1,12 @@ from __future__ import annotations +import logging from abc import ABC, abstractmethod from collections.abc import Collection, Generator, Iterable, Iterator, Mapping, Sequence -from contextlib import AbstractContextManager, contextmanager +from contextlib import contextmanager from dataclasses import dataclass from functools import wraps +from time import perf_counter from typing import TYPE_CHECKING, Any, Callable, ClassVar, ParamSpec, Type, TypeVar, cast import pyarrow as pa @@ -29,6 +31,8 @@ from ..mappers.materialised_view_contracts import MaterializedViewIndex from ..tables.typing import FileTableProtocol +logger = logging.getLogger(__name__) + @dataclass(frozen=True) class DedupeResult: @@ -917,40 +921,39 @@ def _self_referential_deferred_fks( target = table_cls.__table__ return tuple(fk for fk in self.deferred_fk_constraints(table_cls) if fk.referred_table is target) - def merge_context( - self, - table_cls: Type["FileTableProtocol"], - session: so.Session, - ) -> AbstractContextManager[None]: - """Disable backend FK checks for a table's merge and restore them afterward. - - A backend without supports_fk_toggle has no state to disable in the - first place, so this skips the attempt entirely rather than raising; - such a backend's own merge methods must avoid operations (like - deleting a referenced row) that would need the toggle to be safe. - """ - return self.bulk_load_context( - session, mapper=table_cls, disable_fk=self.capabilities.supports_fk_toggle, no_autoflush=False - ) - @contextmanager - def bulk_load_context( + def relaxed_foreign_keys( self, session: so.Session, *, mapper: Any | None = None, - disable_fk: bool = True, no_autoflush: bool = True, ): """ - Generic bulk-load context that defers FK semantics to the backend. + Relax foreign key checks for the block where the backend can, and restore them on exit. + + A backend without ``supports_fk_toggle`` has nothing to relax, so its + constraints stay enforced and rows must be inserted parents first. + The session is rolled back when the block raises. + + Parameters + ---------- + session + Session whose connection the foreign key state belongs to. + mapper + Mapped class selecting the session's bind, for a session with + more than one. + no_autoflush + Whether to suspend the session's autoflush inside the block. """ + scope = f"Table `{mapper.__tablename__}`: " if hasattr(mapper, "__tablename__") else "" previous_fk_state: str | int | None = None try: - if disable_fk: - self._require_capability("supports_fk_toggle", "foreign key toggling") + if self.capabilities.supports_fk_toggle: + started = perf_counter() raw_state = self.disable_fk_check(session, mapper=mapper) previous_fk_state = self._normalize_fk_check_state(raw_state) + logger.info(f"{scope}Foreign key checks disabled in {perf_counter() - started:.2f}s.") if no_autoflush: with session.no_autoflush: @@ -964,7 +967,9 @@ def bulk_load_context( finally: if previous_fk_state is not None: + started = perf_counter() self.restore_fk_check(session, previous_fk_state, mapper=mapper) + logger.info(f"{scope}Foreign key checks restored in {perf_counter() - started:.2f}s.") @requires_capability("supports_materialized_views", "materialized views") def create_materialized_view( diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index 8836fab..6f104ca 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -143,38 +143,21 @@ def manage_indices( f"{_format_elapsed(perf_counter() - commit_started)}." ) - fk_restore_started: float | None = None try: - logger.info(f"Table `{table_name}`: Disabling foreign key checks before merge.") - fk_disable_started = perf_counter() - with backend.merge_context(cls, session): + with backend.relaxed_foreign_keys(session, mapper=cls, no_autoflush=False): + yield + logger.info(f"Table `{table_name}`: Committing merged rows.") + commit_started = perf_counter() + session.commit() logger.info( - f"Table `{table_name}`: Foreign key checks disabled in " - f"{_format_elapsed(perf_counter() - fk_disable_started)}." + f"Table `{table_name}`: Merge commit completed in " + f"{_format_elapsed(perf_counter() - commit_started)}." ) - try: - yield - logger.info(f"Table `{table_name}`: Committing merged rows.") - commit_started = perf_counter() - session.commit() - logger.info( - f"Table `{table_name}`: Merge commit completed in " - f"{_format_elapsed(perf_counter() - commit_started)}." - ) - finally: - logger.info(f"Table `{table_name}`: Restoring foreign key checks.") - fk_restore_started = perf_counter() - except Exception as e: session.rollback() logger.error(f"Table `{table_name}`: Merge operation failed - {e}") raise finally: - if fk_restore_started is not None: - logger.info( - f"Table `{table_name}`: Foreign key checks restored in " - f"{_format_elapsed(perf_counter() - fk_restore_started)}." - ) if indices: logger.info(f"Table `{table_name}`: Verifying/Rebuilding indices.") rebuild_started = perf_counter() diff --git a/tests/backends/test_base_backend.py b/tests/backends/test_base_backend.py index 69ac326..c33069e 100644 --- a/tests/backends/test_base_backend.py +++ b/tests/backends/test_base_backend.py @@ -125,15 +125,17 @@ def _commits_between_pagination_batches(self) -> bool: def _normalize_fk_check_state(previous_state: str | int) -> str | int: return previous_state - def disable_fk_check(self, session: so.Session) -> str | int: + def disable_fk_check(self, session: so.Session, *, mapper: Any | None = None) -> str | int: self.calls.append(("disable_fk_check", session)) return "enabled" - def enable_fk_check(self, session: so.Session) -> str | int: + def enable_fk_check(self, session: so.Session, *, mapper: Any | None = None) -> str | int: self.calls.append(("enable_fk_check", session)) return "disabled" - def restore_fk_check(self, session: so.Session, previous_state: str | int) -> None: + def restore_fk_check( + self, session: so.Session, previous_state: str | int, *, mapper: Any | None = None + ) -> None: self.calls.append(("restore_fk_check", previous_state)) def create_materialized_view( @@ -197,9 +199,6 @@ def test_fake_backend_can_implement_contract(): assert backend.resolve_index_strategy("auto") == "drop_rebuild" assert backend.resolve_index_strategy("keep") == "keep" - with backend.merge_context(cast("Type[FileTableProtocol]", None), cast(so.Session, None)): - pass - def test_require_capability_passes_for_supported_feature(): backend = FakeBackend() @@ -234,10 +233,10 @@ def test_insertable_column_names_exclude_computed_columns(): assert backend._insertable_column_names(_ComputedTableCls) == ["id", "name"] -def test_bulk_load_context_toggles_fk_and_restores(sqlite_session): +def test_relaxed_foreign_keys_toggles_fk_and_restores(sqlite_session): backend = FakeBackend() - with backend.bulk_load_context(sqlite_session): + with backend.relaxed_foreign_keys(sqlite_session): pass assert backend.calls == [ @@ -246,16 +245,7 @@ def test_bulk_load_context_toggles_fk_and_restores(sqlite_session): ] -def test_bulk_load_context_without_fk_toggle(sqlite_session): - backend = FakeBackend() - - with backend.bulk_load_context(sqlite_session, disable_fk=False): - pass - - assert backend.calls == [] - - -def test_bulk_load_context_raises_when_capability_missing(sqlite_session): +def test_relaxed_foreign_keys_leaves_keys_enforced_when_backend_cannot_toggle(sqlite_session): class NoFKBackend(FakeBackend): @property def capabilities(self) -> BackendCapabilities: @@ -263,16 +253,17 @@ def capabilities(self) -> BackendCapabilities: backend = NoFKBackend() - with pytest.raises(NotImplementedError, match="does not support foreign key toggling"): - with backend.bulk_load_context(sqlite_session): - pass + with backend.relaxed_foreign_keys(sqlite_session): + pass + + assert backend.calls == [] -def test_bulk_load_context_rolls_back_and_restores(sqlite_session): +def test_relaxed_foreign_keys_rolls_back_and_restores(sqlite_session): backend = FakeBackend() with pytest.raises(RuntimeError, match="boom"): - with backend.bulk_load_context(sqlite_session): + with backend.relaxed_foreign_keys(sqlite_session): raise RuntimeError("boom") assert backend.calls == [ diff --git a/tests/backends/test_duckdb_backend.py b/tests/backends/test_duckdb_backend.py index 8c8b7ee..00ba145 100644 --- a/tests/backends/test_duckdb_backend.py +++ b/tests/backends/test_duckdb_backend.py @@ -271,15 +271,24 @@ def test_duckdb_backend_fk_toggle_methods_raise(): @pytest.mark.duckdb @pytest.mark.db_dialect -def test_duckdb_backend_merge_context_runs_without_toggling_fk(duckdb_session): - """merge_context() must not require supports_fk_toggle: DuckDB has none.""" +def test_duckdb_backend_relaxed_foreign_keys_runs_without_toggling_fk(duckdb_session): + """DuckDB has no foreign key toggle, so the context leaves its keys enforced.""" backend = DuckDBBackend() - backend.create_staging_table(_ComputedTableCls, duckdb_session) - with backend.merge_context(_ComputedTableCls, duckdb_session): + with backend.relaxed_foreign_keys(duckdb_session, mapper=_ComputedTableCls): pass +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_a_merge_on_duckdb_does_not_log_foreign_key_toggling(duckdb_session, caplog): + with caplog.at_level("INFO", logger="orm_loader.backends.base"): + with ComputedColumnTable.manage_indices(duckdb_session): # ty: ignore[invalid-argument-type] + pass + + assert "foreign key" not in caplog.text.lower() + + @pytest.mark.duckdb @pytest.mark.db_dialect def test_duckdb_backend_materialized_view_methods_raise(duckdb_session): diff --git a/tests/backends/test_sqlite_backend.py b/tests/backends/test_sqlite_backend.py index f2c1e19..5854b69 100644 --- a/tests/backends/test_sqlite_backend.py +++ b/tests/backends/test_sqlite_backend.py @@ -297,3 +297,12 @@ def test_sqlite_backend_fk_toggle_round_trip(sqlite_session): backend.restore_fk_check(sqlite_session, previous) assert sqlite_session.execute(sa.text("PRAGMA foreign_keys")).scalar() == 1 + + +def test_a_merge_on_sqlite_logs_foreign_key_toggling(sqlite_session, caplog): + with caplog.at_level("INFO", logger="orm_loader.backends.base"): + with ComputedColumnTable.manage_indices(sqlite_session): # ty: ignore[invalid-argument-type] + pass + + assert "Foreign key checks disabled" in caplog.text + assert "Foreign key checks restored" in caplog.text diff --git a/tests/loaders/test_loader_e2e.py b/tests/loaders/test_loader_e2e.py index 3ea4b6a..9ec7046 100644 --- a/tests/loaders/test_loader_e2e.py +++ b/tests/loaders/test_loader_e2e.py @@ -607,7 +607,7 @@ def test_drop_rebuild_logging_shows_merge_phases(sqlite_session, tmp_path, caplo ] ).to_csv(csv_path, index=False, sep="\t") - caplog.set_level(logging.INFO, logger="orm_loader.tables.loadable_table") + caplog.set_level(logging.INFO, logger="orm_loader") _SimpleTable.load_file( sqlite_session, @@ -624,7 +624,6 @@ def test_drop_rebuild_logging_shows_merge_phases(sqlite_session, tmp_path, caplo assert any("Finished dropping 1 active indices in " in message for message in messages) assert any("Committing after index drop." in message for message in messages) assert any("Commit after index drop completed in " in message for message in messages) - assert any("Disabling foreign key checks before merge." in message for message in messages) assert any("Foreign key checks disabled in " in message for message in messages) assert any("Merging staging data into target table" in message for message in messages) assert any("Merge replace delete phase starting." in message for message in messages) @@ -634,7 +633,6 @@ def test_drop_rebuild_logging_shows_merge_phases(sqlite_session, tmp_path, caplo assert any("Merging staging data into target table" in message for message in messages) assert any("Committing merged rows." in message for message in messages) assert any("Merge commit completed in " in message for message in messages) - assert any("Restoring foreign key checks." in message for message in messages) assert any("Foreign key checks restored in " in message for message in messages) assert any("Verifying/Rebuilding indices." in message for message in messages) assert any("Restoring missing index: ix_test_table_name" in message for message in messages) From 16bdbbc4e635aa6d706deae3f8743b0e00219ad1 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Fri, 30 Oct 2026 01:04:23 +0000 Subject: [PATCH 11/21] Run merge strategies as named: remove the empty-target shortcut and database-assigned keys, rename insert_if_empty to insert, move reader options off LoaderContext --- docs/loaders/context.md | 2 - docs/loaders/readers.md | 2 +- docs/loaders/staging.md | 18 +- src/orm_loader/backends/base.py | 12 +- src/orm_loader/backends/duckdb.py | 7 +- src/orm_loader/loaders/loader.py | 7 - src/orm_loader/loaders/readers/__init__.py | 4 +- src/orm_loader/loaders/readers/delimited.py | 20 +- src/orm_loader/tables/loadable_table.py | 240 +++++++++----------- src/orm_loader/tables/orm_table.py | 35 +-- src/orm_loader/tables/typing.py | 13 +- tests/backends/test_postgres_backend.py | 8 +- tests/backends/test_sinks.py | 9 +- tests/loaders/test_defaults.py | 7 +- tests/loaders/test_loader_e2e.py | 63 +++-- tests/loaders/test_malformed_files.py | 9 + tests/loaders/test_merge_semantics.py | 79 +++++++ tests/loaders/test_pipeline.py | 6 +- tests/loaders/test_schema_translate_map.py | 2 +- tests/loaders/test_split_connection.py | 2 +- tests/loaders/test_stage_file.py | 31 +-- 21 files changed, 309 insertions(+), 267 deletions(-) create mode 100644 tests/loaders/test_merge_semantics.py diff --git a/docs/loaders/context.md b/docs/loaders/context.md index d2c61a9..56e7752 100644 --- a/docs/loaders/context.md +++ b/docs/loaders/context.md @@ -21,10 +21,8 @@ on globals or implicit configuration. | `tableclass` | Target ORM table class | | `session` | Active SQLAlchemy session | | `path` | Path to the input file | -| `chunksize` | Rows per batch for the default reader of formats that batch by rows | | `normalise` | Whether to interpret values under the policy; when false, values are written as read | | `dedupe` | Whether the staging table keeps only the first row for each primary key | -| `quote_mode` | Quoting mode for delimited text: `csv`, `literal` or `by_delimiter` (default) | | `policy` | `ConversionPolicy`: how values are interpreted and what happens to unreadable ones | | `rejects_dir` | Directory that rejected rows are written to | | `report` | `LoadReport` row accounting | diff --git a/docs/loaders/readers.md b/docs/loaders/readers.md index 4343a47..0f275df 100644 --- a/docs/loaders/readers.md +++ b/docs/loaders/readers.md @@ -2,7 +2,7 @@ A reader yields a file as Arrow record batches holding the model columns present in the file. Text formats yield raw strings, so the conversion policy sees exactly what the file contains. Typed formats yield their stored types. -`open_reader` picks the reader from the file suffix: `.parquet` is read as Parquet, any other suffix as delimited text. A table chooses its reader in `_select_reader`, and `load_file(reader=...)` passes one explicitly. +`open_reader` picks the reader from the file suffix: `.parquet` is read as Parquet, any other suffix as delimited text. `load_file` opens the reader with it, taking `quote_mode` and `chunksize` as arguments, and `load_file(reader=...)` passes a configured reader explicitly. | Reader | Input | |------|------| diff --git a/docs/loaders/staging.md b/docs/loaders/staging.md index d3137ed..77f9b54 100644 --- a/docs/loaders/staging.md +++ b/docs/loaders/staging.md @@ -11,14 +11,14 @@ - Values are interpreted by column type, as described in [Conversion](conversion.md). - With the default policy an unreadable value or a missing required value raises and nothing is loaded. - With `ErrorPolicy(on_error="reject_row")` those rows, and delimited lines with the wrong number of fields, are written to `/.rejects.csv` and the rest load. A malformed line is listed with its row number in the file, and its text is in the reason. `rejects_dir` defaults to a `rejects` directory inside the input file's directory, and the file of an earlier load is removed. `LoadReport` accounts for every row read, checked against the rows counted in the staging table. +- With `normalise=False` raw text is passed to the staging table unchanged, without defaults, required checks or type conversion. What the database does with a value that does not fit its column then depends on the dialect. - A blank value, and a column the file lacks, take the column's default when it has one: a server default written as text, a scalar Python default, or a SQL expression such as `current_date`. The column's handler reads the default like a value from the file, and an expression is evaluated by the database for each batch that needs it. A column with a callable default stays required. -- With `normalise=False` raw text is passed to the staging table unchanged. --- ## Staging table -Staging mirrors the target: its columns, their types and their nullability, with no key, index, default or constraint, and a `_rownum` column that orders the rows (SQLite uses its rowid). The processor completes every value before it is written, so a NOT NULL column never receives NULL. A key the file provides is a staging column like any other. When the file lacks an autoincrement key, staging leaves it out and the merge leaves it out of its insert, so the database assigns it. +Staging mirrors the target: its columns, their types and their nullability, with no key, index, default or constraint, and a `_rownum` column that orders the rows (SQLite uses its rowid). The processor completes every value before it is written, so a NOT NULL column never receives NULL. Every column that is required and has no default must be in the file, primary key columns included, so a file without its key fails before anything is loaded. --- @@ -40,6 +40,20 @@ A failure in the sink is raised. There is no fallback to another loading path. With `dedupe` set, the staging table keeps the first row for each primary key, in source order. The backend removes the others after the load and reports a `DedupeResult` on `LoadReport` with the rows before and the rows removed. The merge step decides what happens when incoming rows overlap with existing target rows. +--- + +## Merge strategies + +Each strategy does what its name says, whatever the target already holds. + +| Strategy | Effect | +|------|------| +| `replace` | Deletes target rows whose primary key is in the file, then inserts the file | +| `upsert` | Inserts the file and skips rows whose primary key already exists, keeping the first of a repeated key | +| `insert` | Inserts the file without any key check, so a key that already exists fails with the database's duplicate key error | + +`insert` suits a first load into an empty table, or any file whose keys are new. With `dedupe=True` it is several times faster than `upsert`, because `upsert` checks every row against the key index. A load with `upsert` into an empty table logs this hint and runs unchanged. With `merge_batch_size` set, batches commit separately on PostgreSQL, so a duplicate key partway through leaves the earlier batches in the table. + No schema changes are performed at the loader layer. --- diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 948c1d5..f3ad4e0 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -2,7 +2,7 @@ import logging from abc import ABC, abstractmethod -from collections.abc import Collection, Generator, Iterable, Iterator, Mapping, Sequence +from collections.abc import Generator, Iterable, Iterator, Mapping, Sequence from contextlib import contextmanager from dataclasses import dataclass from functools import wraps @@ -352,8 +352,7 @@ def _insertable_column_names( Computed columns are excluded because backend loaders and merge helpers should not attempt to write to them directly. With a staging table, - only the columns it has are returned, so a column the load left out - (an autoincrement key) is assigned by the database. + only the columns it has are returned. """ return [ c.name @@ -544,8 +543,6 @@ def create_staging_table( self, table_cls: Type["FileTableProtocol"], session: so.Session, - *, - without: Collection[str] = (), ) -> None: """ Create the staging table for a table class, replacing an existing one. @@ -561,15 +558,12 @@ def create_staging_table( Table the staging table is for. session Session whose connection creates the table. - without - Names of target columns the staging table leaves out. The merge - then leaves them out of its insert. """ self.drop_staging_table(table_cls, session) columns = [ sa.Column(c.name, c.type, nullable=c.nullable) for c in table_cls.__table__.columns - if c.computed is None and c.name not in without + if c.computed is None ] staging = self._staging_table(table_cls, columns) staging.create(bind=session.connection(bind_arguments={"mapper": table_cls}), checkfirst=True) diff --git a/src/orm_loader/backends/duckdb.py b/src/orm_loader/backends/duckdb.py index d106266..85c462a 100644 --- a/src/orm_loader/backends/duckdb.py +++ b/src/orm_loader/backends/duckdb.py @@ -397,12 +397,7 @@ def _update(extra: sa.ColumnElement[bool] | None = None) -> None: session.execute(stmt) def _drop_matched_from_staging(extra: sa.ColumnElement[bool] | None = None) -> None: - matched = sa.select(*(target.c[c] for c in pk_cols)).where(pk_join) - if len(pk_cols) == 1: - condition: sa.ColumnElement[bool] = staging.c[pk_cols[0]].in_(matched) - else: - condition = sa.tuple_(*(staging.c[c] for c in pk_cols)).in_(matched) - stmt = sa.delete(staging).where(condition) + stmt = sa.delete(staging).where(sa.exists().where(pk_join)) if extra is not None: stmt = stmt.where(extra) session.execute(stmt) diff --git a/src/orm_loader/loaders/loader.py b/src/orm_loader/loaders/loader.py index 69a165c..3a729d6 100644 --- a/src/orm_loader/loaders/loader.py +++ b/src/orm_loader/loaders/loader.py @@ -83,15 +83,10 @@ class LoaderContext: Active SQLAlchemy session. path Path to the source data file. - chunksize - Rows per batch for the default reader of formats that batch by rows. normalise Whether to interpret values under the policy. dedupe Whether the staging table keeps only the first row for each primary key. - quote_mode - Quoting mode of the default reader for delimited text: ``"csv"``, - ``"literal"`` or ``"by_delimiter"``. staging_schema_tag schema_translate_map tag the staging table lives in, passed to resolve_backend() so every backend resolution within this load @@ -109,10 +104,8 @@ class LoaderContext: session: so.Session path: Path - chunksize: int | None = None normalise: bool = True dedupe: bool = True - quote_mode: str = "by_delimiter" staging_schema_tag: str | None = None policy: ConversionPolicy = field(default_factory=ConversionPolicy) rejects_dir: Path | None = None diff --git a/src/orm_loader/loaders/readers/__init__.py b/src/orm_loader/loaders/readers/__init__.py index 5df851f..a2ad040 100644 --- a/src/orm_loader/loaders/readers/__init__.py +++ b/src/orm_loader/loaders/readers/__init__.py @@ -15,7 +15,7 @@ def open_reader( model_columns: Mapping[str, sa.ColumnElement[Any]], *, nulls: NullPolicy | None = None, - quote_mode: str = "by_delimiter", + quote_mode: str | None = None, chunksize: int | None = None, skip_malformed: bool = False, ) -> FileReader: @@ -33,7 +33,7 @@ def open_reader( nulls Null policy of the load. quote_mode - Quote handling for delimited text. + Quote handling for delimited text. ``None`` uses the reader's default. chunksize Rows per batch for formats that batch by rows. skip_malformed diff --git a/src/orm_loader/loaders/readers/delimited.py b/src/orm_loader/loaders/readers/delimited.py index cbde7f1..f987de0 100644 --- a/src/orm_loader/loaders/readers/delimited.py +++ b/src/orm_loader/loaders/readers/delimited.py @@ -41,20 +41,28 @@ class DelimitedTextReader(FileReader): quote_mode ``"csv"`` reads double quotes as field quoting, ``"literal"`` keeps them as data, ``"by_delimiter"`` chooses ``"literal"`` for tab - delimited files and ``"csv"`` otherwise. + delimited files and ``"csv"`` otherwise. ``None`` uses + ``DEFAULT_QUOTE_MODE``. block_size Bytes per block; memory use follows this, not the row count. skip_malformed Leave out a line with the wrong number of fields and record it in ``skipped``, instead of failing the read. + + Attributes + ---------- + DEFAULT_QUOTE_MODE + Quote mode used when none is given. """ + DEFAULT_QUOTE_MODE = "by_delimiter" + def __init__( self, path: Path, model_columns: Mapping[str, sa.ColumnElement[Any]], nulls: NullPolicy | None = None, - quote_mode: str = "by_delimiter", + quote_mode: str | None = None, block_size: int = 4 << 20, skip_malformed: bool = False, ) -> None: @@ -70,7 +78,7 @@ def __init__( codecs.lookup(self.encoding) except LookupError as exc: raise IngestError(f"{path.name}: unsupported encoding {self.encoding!r}") from exc - self.quote = resolve_quote_mode(quote_mode, self.delimiter) + self.quote_mode = resolve_quote_mode(quote_mode or self.DEFAULT_QUOTE_MODE, self.delimiter) self.file_columns = self._header_names() self.columns = [name for name in self.file_columns if name in model_columns] if self.nulls.marker_quoting == "unquoted_only" and self.nulls.text_markers: @@ -118,9 +126,9 @@ def batches(self) -> Iterator[pa.RecordBatch]: ), parse_options=pv.ParseOptions( delimiter=self.delimiter, - quote_char='"' if self.quote == "csv" else False, + quote_char='"' if self.quote_mode == "csv" else False, escape_char=False, - newlines_in_values=self.quote == "csv", + newlines_in_values=self.quote_mode == "csv", invalid_row_handler=self._skip_row if self.skip_malformed else None, ), convert_options=pv.ConvertOptions( @@ -142,7 +150,7 @@ def batches(self) -> Iterator[pa.RecordBatch]: return hint = ( " The file may contain literal double quotes; try quote_mode='literal' or 'by_delimiter'." - if self.quote == "csv" and str(exc).startswith("CSV parse error") + if self.quote_mode == "csv" and str(exc).startswith("CSV parse error") else "" ) raise IngestError(f"{self.path.name}: {exc}.{hint}") from exc diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index 6f104ca..84c2275 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -7,7 +7,7 @@ and index handling. """ import logging -from collections.abc import Collection, Mapping +from collections.abc import Mapping from contextlib import contextmanager from pathlib import Path from time import perf_counter @@ -20,6 +20,7 @@ from ..backends.resolve import resolve_backend from ..conversion import ConversionPolicy +from ..helpers import IngestError from ..loaders import FileReader, LoadReport, LoaderContext, open_reader, stage_file from .orm_table import ORMTableBase from .typing import FileTableProtocol @@ -64,7 +65,6 @@ def create_staging_table( session: so.Session, *, staging_schema_tag: str | None = None, - without: Collection[str] = (), ): """ Create a fresh staging table for ingestion. @@ -81,10 +81,6 @@ def create_staging_table( schema_translate_map tag the staging table should be created under. ``None`` means no schema qualification (backend-default behavior). - without - Names of target columns the staging table leaves out. The merge - then leaves them out of its insert. - Raises ------ RuntimeError @@ -94,7 +90,7 @@ def create_staging_table( """ _require_bind(session, cls) backend = resolve_backend(session, staging_schema_tag=staging_schema_tag, mapper=cls) - backend.create_staging_table(cls, session, without=without) + backend.create_staging_table(cls, session) @classmethod @contextmanager @@ -112,6 +108,11 @@ def manage_indices( ``"keep"``. The backend decides what ``"auto"`` means. At the moment SQLite keeps indexes by default, while PostgreSQL drops and rebuilds them. + + Raises + ------ + IngestError + An index could not be rebuilt after a merge that was committed. """ backend = resolve_backend(session, staging_schema_tag=staging_schema_tag, mapper=cls) resolved_index_strategy = backend.resolve_index_strategy(index_strategy) @@ -143,6 +144,7 @@ def manage_indices( f"{_format_elapsed(perf_counter() - commit_started)}." ) + merge_failed = False try: with backend.relaxed_foreign_keys(session, mapper=cls, no_autoflush=False): yield @@ -154,44 +156,83 @@ def manage_indices( f"{_format_elapsed(perf_counter() - commit_started)}." ) except Exception as e: + merge_failed = True session.rollback() logger.error(f"Table `{table_name}`: Merge operation failed - {e}") raise finally: if indices: - logger.info(f"Table `{table_name}`: Verifying/Rebuilding indices.") - rebuild_started = perf_counter() - connection = session.connection(bind_arguments={"mapper": cls}) - inspector = sa.inspect(connection) - schema = physical_schema_of(connection, schema_tag=cls.__table__.schema) - existing_idx_names = {idx['name'] for idx in inspector.get_indexes(table_name, schema=schema)} - - for idx in indices: - if idx.name not in existing_idx_names: - try: - logger.info(f"Table `{table_name}`: Restoring missing index: {idx.name}") - create_started = perf_counter() - session.execute(sa.schema.CreateIndex(idx)) - logger.info( - f"Table `{table_name}`: Restored missing index `{idx.name}` in " - f"{_format_elapsed(perf_counter() - create_started)}." - ) - logger.info(f"Table `{table_name}`: Committing restored index `{idx.name}`.") - commit_started = perf_counter() - session.commit() - logger.info( - f"Table `{table_name}`: Commit after restoring index `{idx.name}` " - f"completed in {_format_elapsed(perf_counter() - commit_started)}." - ) - except Exception as e: - session.rollback() - logger.error(f"Table `{table_name}`: Failed to restore {idx.name}: {e}") - else: - logger.debug(f"Table `{table_name}`: Index {idx.name} already exists on disk. Skipping.") - logger.info( - f"Table `{table_name}`: Index verification/rebuild completed in " - f"{_format_elapsed(perf_counter() - rebuild_started)}." - ) + try: + failed_indices = cls._restore_indices(session, indices) + except Exception as e: + if not merge_failed: + raise + logger.error(f"Table `{table_name}`: Index rebuild after the failed merge also failed - {e}") + else: + if failed_indices and not merge_failed: + raise IngestError( + f"Table `{table_name}`: the merge was committed but index(es) {failed_indices} could not be rebuilt." + ) + + + @classmethod + def _restore_indices( + cls: Type['FileTableProtocol'], + session: so.Session, + indices: list[sa.Index], + ) -> list[str]: + """ + Recreate the given indexes that are missing from the table. + + Parameters + ---------- + session + An active SQLAlchemy session. + indices + Indexes the table is expected to have. + + Returns + ------- + list[str] + Names of the indexes that could not be recreated. + """ + table_name = cls.__tablename__ + failed_indices: list[str] = [] + logger.info(f"Table `{table_name}`: Verifying/Rebuilding indices.") + rebuild_started = perf_counter() + connection = session.connection(bind_arguments={"mapper": cls}) + inspector = sa.inspect(connection) + schema = physical_schema_of(connection, schema_tag=cls.__table__.schema) + existing_idx_names = {idx['name'] for idx in inspector.get_indexes(table_name, schema=schema)} + + for idx in indices: + if idx.name not in existing_idx_names: + try: + logger.info(f"Table `{table_name}`: Restoring missing index: {idx.name}") + create_started = perf_counter() + session.execute(sa.schema.CreateIndex(idx)) + logger.info( + f"Table `{table_name}`: Restored missing index `{idx.name}` in " + f"{_format_elapsed(perf_counter() - create_started)}." + ) + logger.info(f"Table `{table_name}`: Committing restored index `{idx.name}`.") + commit_started = perf_counter() + session.commit() + logger.info( + f"Table `{table_name}`: Commit after restoring index `{idx.name}` " + f"completed in {_format_elapsed(perf_counter() - commit_started)}." + ) + except Exception as e: + session.rollback() + failed_indices.append(str(idx.name)) + logger.error(f"Table `{table_name}`: Failed to restore {idx.name}: {e}") + else: + logger.debug(f"Table `{table_name}`: Index {idx.name} already exists on disk. Skipping.") + logger.info( + f"Table `{table_name}`: Index verification/rebuild completed in " + f"{_format_elapsed(perf_counter() - rebuild_started)}." + ) + return failed_indices @classmethod @@ -262,38 +303,10 @@ def load_staging( """ session, tag = loader_context.session, loader_context.staging_schema_tag _require_bind(session, cls) - left_out = set() if reader.empty else cls.autoincrement_columns() - set(reader.columns) - cls.create_staging_table(session, staging_schema_tag=tag, without=left_out) + cls.create_staging_table(session, staging_schema_tag=tag) staging = cls.get_staging_table(session, staging_schema_tag=tag) return stage_file(loader_context, reader, staging) - @classmethod - def _select_reader(cls: Type[FileTableProtocol], ctx: LoaderContext) -> FileReader: - """ - Select the reader for a file. Override to read other formats. - - The default picks the reader from the file suffix: Parquet for - ``.parquet`` and delimited text for every other suffix. - - Parameters - ---------- - ctx - Context of the load. - - Returns - ------- - FileReader - Reader for the file. - """ - return open_reader( - ctx.path, - ctx.tableclass.model_columns(), - nulls=ctx.policy.nulls, - quote_mode=ctx.quote_mode, - chunksize=ctx.chunksize, - skip_malformed=ctx.policy.errors.on_error == "reject_row", - ) - @classmethod def load_file( cls: Type[FileTableProtocol], @@ -305,7 +318,7 @@ def load_file( dedupe: bool = False, chunksize: int | None = None, merge_strategy: str = "replace", - quote_mode: str = "by_delimiter", + quote_mode: str | None = None, index_strategy: str = "auto", merge_batch_size: int | None = None, staging_schema_tag: str | None = None, @@ -334,8 +347,10 @@ def load_file( path Path to the input file. reader - Reader for the file. The table's ``_select_reader`` chooses one - when omitted, and ``chunksize`` and ``quote_mode`` then configure it. + Reader for the file. ``open_reader`` chooses one from the file + suffix when omitted, and ``chunksize`` and ``quote_mode`` then + configure it. Passing either together with a reader is an error, because the + reader is already configured. normalise Whether to convert values to the column types under the policy. When false, values are written as read. @@ -346,7 +361,7 @@ def load_file( rows. merge_strategy Merge strategy to apply (e.g. ``replace``, ``upsert``, or - ``insert_if_empty``). + ``insert``). quote_mode How double quotes in delimited text are read: ``"csv"`` (field quoting), ``"literal"`` (kept as data), ``"by_delimiter"`` @@ -403,23 +418,19 @@ def load_file( f"File name '{path.name}' does not match table '{cls.__tablename__}'" ) - if merge_strategy == "insert_if_empty": - logger.info( - f"Table `{cls.__tablename__}`: Checking whether target table is empty before staging load." - ) - check_started = perf_counter() - has_rows = cls._target_has_rows(session=session) - logger.info( - f"Table `{cls.__tablename__}`: Pre-load empty-table check completed in " - f"{_format_elapsed(perf_counter() - check_started)}." + if reader is not None and (chunksize is not None or quote_mode is not None): + raise ValueError("chunksize and quote_mode configure the default reader; set them on the reader passed as `reader`") + policy = policy or ConversionPolicy() + if reader is None: + reader = open_reader( + path, + cls.model_columns(), + nulls=policy.nulls, + quote_mode=quote_mode, + chunksize=chunksize, + skip_malformed=policy.errors.on_error == "reject_row", ) - if has_rows: - raise ValueError( - f"Table `{cls.__tablename__}` is not empty; cannot use merge strategy " - f"'insert_if_empty'" - ) - if merge_strategy == "replace": backend = resolve_backend(session, staging_schema_tag=staging_schema_tag, mapper=cls) if backend.resolve_index_strategy(index_strategy) == "keep": @@ -432,12 +443,10 @@ def load_file( tableclass=cls, session=session, path=path, - chunksize=chunksize, normalise=normalise, dedupe=dedupe, - quote_mode=quote_mode, staging_schema_tag=staging_schema_tag, - policy=policy or ConversionPolicy(), + policy=policy, rejects_dir=rejects_dir, report=report or LoadReport(), ) @@ -445,7 +454,7 @@ def load_file( try: # Load to staging (Indices are already excluded via updated create_staging_table) logger.info(f"Table `{cls.__tablename__}`: Loading data into staging table") - total = cls.load_staging(loader_context, reader or cls._select_reader(loader_context)) + total = cls.load_staging(loader_context, reader) # Merge staging to target (Wrapped in our index dropper!) logger.info(f"Table `{cls.__tablename__}`: Merging staging data into target table") @@ -559,7 +568,7 @@ def merge_from_staging( An active SQLAlchemy session. merge_strategy Merge strategy to apply (for example ``replace``, - ``upsert``, or ``insert_if_empty``). + ``upsert``, or ``insert``). staging_schema_tag schema_translate_map tag the staging table lives in. ``None`` means no schema qualification (backend-default behavior). @@ -581,25 +590,6 @@ def merge_from_staging( resolved=resolved, physical_schemas=physical_schemas, ) - target_empty_confirmed = False - if merge_strategy in {"replace", "upsert"}: - logger.info( - f"Table `{target}`: Checking whether target table is empty for merge optimisation." - ) - check_started = perf_counter() - has_rows = cls._target_has_rows(session=session) - logger.info( - f"Table `{target}`: Empty-table optimisation check completed in " - f"{_format_elapsed(perf_counter() - check_started)}." - ) - if not has_rows: - logger.info( - f"Table `{target}`: Target table is empty; routing merge strategy " - f"`{merge_strategy}` to insert-if-empty fast path." - ) - target_empty_confirmed = True - merge_strategy = "insert_if_empty" - if merge_strategy == "replace": logger.info(f"Table `{target}`: Merge replace delete phase starting.") delete_started = perf_counter() @@ -620,6 +610,11 @@ def merge_from_staging( f"{_format_elapsed(perf_counter() - insert_started)}." ) elif merge_strategy == "upsert": + if not cls._target_has_rows(session=session): + logger.info( + f"Table `{target}`: The target table is empty, so merge strategy `insert` " + "with `dedupe=True` loads the same rows faster." + ) logger.info(f"Table `{target}`: Merge upsert phase starting.") upsert_started = perf_counter() backend.merge_upsert( @@ -629,29 +624,14 @@ def merge_from_staging( f"Table `{target}`: Merge upsert phase completed in " f"{_format_elapsed(perf_counter() - upsert_started)}." ) - elif merge_strategy == "insert_if_empty": - if not target_empty_confirmed: - logger.info(f"Table `{target}`: Checking whether target table is empty.") - check_started = perf_counter() - has_rows = cls._target_has_rows(session=session) - logger.info( - f"Table `{target}`: Empty-table check completed in " - f"{_format_elapsed(perf_counter() - check_started)}." - ) - - if has_rows: - raise ValueError( - f"Table `{target}` is not empty; cannot use merge strategy " - f"'insert_if_empty'" - ) - - logger.info(f"Table `{target}`: Merge insert-if-empty phase starting.") + elif merge_strategy == "insert": + logger.info(f"Table `{target}`: Merge insert phase starting.") insert_started = perf_counter() backend.merge_insert( cls, session, merge_batch_size=merge_batch_size, validate_deferred_fks=validate_deferred_fks ) logger.info( - f"Table `{target}`: Merge insert-if-empty phase completed in " + f"Table `{target}`: Merge insert phase completed in " f"{_format_elapsed(perf_counter() - insert_started)}." ) else: diff --git a/src/orm_loader/tables/orm_table.py b/src/orm_loader/tables/orm_table.py index a299cdc..0f04922 100644 --- a/src/orm_loader/tables/orm_table.py +++ b/src/orm_loader/tables/orm_table.py @@ -10,25 +10,6 @@ logger = logging.getLogger(__name__) -def _resolves_to_autoincrement(c: sa.Column) -> bool: - """Whether a column will actually autoincrement, resolving SQLAlchemy's - own default ``autoincrement="auto"`` sentinel the same way SQLAlchemy - itself does at DDL-emission time. - - SQLAlchemy restricts real autoincrement behaviour to a single-column - integer primary key (or an explicit ``autoincrement=True``). - Treating every `"auto"` column as autoincrementing would wrongly - exclude genuinely-required string PKs from `required_columns()`. - """ - if c.autoincrement is False: - return False - if c.autoincrement is True: - return True - if not c.primary_key: - return False - pk_columns = list(c.table.primary_key.columns) - return len(pk_columns) == 1 and pk_columns[0] is c and isinstance(c.type, sa.Integer) - """ ORMTableBase ============ @@ -144,8 +125,6 @@ def required_columns(cls) -> set[str]: A column is considered required if it is: - non-nullable - has no default that can be supplied (see ``column_defaults``) - - not an autoincrementing column (the database assigns its value; - it's never expected in inbound data at all) This method is intended for ingestion-time validation and does not attempt to enforce schema semantics beyond insert viability. @@ -159,21 +138,9 @@ def required_columns(cls) -> set[str]: return { c.key for c in cls.mapper_for().columns - if not c.nullable and c.key not in defaults and not _resolves_to_autoincrement(c) + if not c.nullable and c.key not in defaults } - @classmethod - def autoincrement_columns(cls) -> set[str]: - """ - Return the names of the columns whose value the database assigns. - - Returns - ------- - set[str] - A set of column names. - """ - return {c.key for c in cls.mapper_for().columns if _resolves_to_autoincrement(c)} - @classmethod def column_defaults(cls) -> dict[str, Any]: """ diff --git a/src/orm_loader/tables/typing.py b/src/orm_loader/tables/typing.py index df548d7..45d179f 100644 --- a/src/orm_loader/tables/typing.py +++ b/src/orm_loader/tables/typing.py @@ -1,7 +1,7 @@ from typing import Protocol, ClassVar, runtime_checkable, TYPE_CHECKING, Optional, Type, Dict, Any, Unpack, TypedDict import sqlalchemy.orm as so import sqlalchemy as sa -from collections.abc import Collection, Mapping +from collections.abc import Mapping from pathlib import Path from contextlib import AbstractContextManager if TYPE_CHECKING: @@ -52,7 +52,6 @@ def model_columns(cls) -> dict[str, sa.ColumnElement[Any]]: ... def required_columns(cls) -> set[str]: ... @classmethod - def autoincrement_columns(cls) -> set[str]: ... @classmethod def column_defaults(cls) -> dict[str, Any]: ... @@ -73,12 +72,9 @@ class FileTableProtocol(ORMTableProtocol, Protocol): No assumptions are made about the underlying database or schema. """ - @classmethod - def _select_reader(cls, ctx: "LoaderContext") -> "FileReader": ... - @classmethod def create_staging_table( - cls, session: so.Session, *, staging_schema_tag: str | None = None, without: Collection[str] = () + cls, session: so.Session, *, staging_schema_tag: str | None = None ) -> None: ... @classmethod @@ -95,7 +91,7 @@ def load_file( dedupe: bool = False, chunksize: int | None = None, merge_strategy: str = "replace", - quote_mode: str = "by_delimiter", + quote_mode: str | None = None, index_strategy: str = "auto", merge_batch_size: int | None = None, staging_schema_tag: str | None = None, @@ -139,6 +135,9 @@ def validate_foreign_keys( @classmethod def _target_has_rows(cls, session: so.Session) -> bool: ... + @classmethod + def _restore_indices(cls, session: so.Session, indices: list[sa.Index]) -> list[str]: ... + @classmethod def manage_indices( cls, session: so.Session, index_strategy: str = "auto", *, staging_schema_tag: str | None = None diff --git a/tests/backends/test_postgres_backend.py b/tests/backends/test_postgres_backend.py index 5823233..34447c3 100644 --- a/tests/backends/test_postgres_backend.py +++ b/tests/backends/test_postgres_backend.py @@ -683,13 +683,13 @@ def test_upsert_does_not_update_existing_rows(pg_session, tmp_path): assert _simple_rows(pg_session) == [(1, "alpha")] -def test_insert_if_empty_loads_an_empty_target_and_refuses_a_non_empty_one(pg_session, tmp_path): +def test_insert_loads_new_keys_and_fails_on_an_existing_key(pg_session, tmp_path): path = _write_simple(tmp_path, [(1, "alpha"), (2, "beta")]) - assert SimpleTable.load_file(pg_session, path, merge_strategy="insert_if_empty", staging_schema_tag=STAGING_SCHEMA) == 2 + assert SimpleTable.load_file(pg_session, path, merge_strategy="insert", staging_schema_tag=STAGING_SCHEMA) == 2 pg_session.commit() assert _simple_rows(pg_session) == [(1, "alpha"), (2, "beta")] - with pytest.raises(ValueError, match="is not empty; cannot use merge strategy 'insert_if_empty'"): - SimpleTable.load_file(pg_session, path, merge_strategy="insert_if_empty", staging_schema_tag=STAGING_SCHEMA) + with pytest.raises(sa.exc.IntegrityError): + SimpleTable.load_file(pg_session, path, merge_strategy="insert", staging_schema_tag=STAGING_SCHEMA) def test_a_native_enum_column_loads_member_names(pg_session, tmp_path): diff --git a/tests/backends/test_sinks.py b/tests/backends/test_sinks.py index ec4a7ec..847b198 100644 --- a/tests/backends/test_sinks.py +++ b/tests/backends/test_sinks.py @@ -85,9 +85,9 @@ def batches(): backend.insert_batches(ctx, staging, batches(), COLUMNS) -def _staging_for(session, table, **kwargs): +def _staging_for(session, table): tag = None if session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA - table.create_staging_table(session, staging_schema_tag=tag, **kwargs) + table.create_staging_table(session, staging_schema_tag=tag) return table.get_staging_table(session, staging_schema_tag=tag) @@ -109,11 +109,6 @@ def test_a_staging_table_has_no_key_and_no_default(dialect_session): assert all(c.server_default is None for c in staging.c if c.name != "_rownum") -def test_columns_can_be_left_out_of_a_staging_table(dialect_session): - staging = _staging_for(dialect_session, _Demo, without={"demo_label"}) - assert "demo_label" not in staging.c and "demo_id" in staging.c - - def test_a_staging_table_can_be_dropped_twice_and_created_again(dialect_session): staging = _staging_for(dialect_session, _Demo) tag = None if dialect_session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA diff --git a/tests/loaders/test_defaults.py b/tests/loaders/test_defaults.py index b615f57..b9755a5 100644 --- a/tests/loaders/test_defaults.py +++ b/tests/loaders/test_defaults.py @@ -34,7 +34,7 @@ def _rows(session) -> list[tuple]: def test_only_a_default_that_can_be_supplied_makes_a_column_optional(): - assert _Table.required_columns() == {"demo_stamp"} + assert _Table.required_columns() == {"demo_id", "demo_stamp"} assert set(_Table.column_defaults()) == {"demo_status", "demo_priority", "demo_note"} @@ -98,8 +98,9 @@ def _keyed(session, tmp_path: Path, text: str, **kwargs) -> list[tuple]: return [tuple(r) for r in session.execute(sa.select(*columns).order_by(DemoKeyedRecord.keyed_id))] -def test_a_key_the_file_lacks_is_assigned_by_the_database(dialect_session, tmp_path): - assert _keyed(dialect_session, tmp_path, "keyed_name\na\nb\n") == [(1, "a"), (2, "b")] +def test_a_primary_key_the_file_lacks_is_a_missing_required_column(dialect_session, tmp_path): + with pytest.raises(IngestError, match="keyed_id"): + _keyed(dialect_session, tmp_path, "keyed_name\na\nb\n") def test_a_key_in_the_file_is_loaded_and_may_not_be_blank(dialect_session, tmp_path): diff --git a/tests/loaders/test_loader_e2e.py b/tests/loaders/test_loader_e2e.py index 9ec7046..30b1eff 100644 --- a/tests/loaders/test_loader_e2e.py +++ b/tests/loaders/test_loader_e2e.py @@ -269,7 +269,7 @@ def test_merge_strategies(sqlite_session, tmp_path, merge_strategy, expected_row assert [(r.id, r.name) for r in rows] == expected_rows -def test_insert_if_empty_merge_strategy(sqlite_session, tmp_path): +def test_insert_merge_strategy(sqlite_session, tmp_path): csv_path = tmp_path / "test_table.csv" pd.DataFrame( @@ -283,7 +283,7 @@ def test_insert_if_empty_merge_strategy(sqlite_session, tmp_path): sqlite_session, csv_path, dedupe=False, - merge_strategy="insert_if_empty", + merge_strategy="insert", ) sqlite_session.commit() @@ -297,23 +297,22 @@ def test_insert_if_empty_merge_strategy(sqlite_session, tmp_path): ] -def test_insert_if_empty_raises_on_non_empty_target(sqlite_session, sqlite_engine, tmp_path): +def test_insert_adds_new_keys_to_a_populated_target_and_fails_on_an_existing_key(sqlite_session, sqlite_engine, tmp_path): csv_path = tmp_path / "test_table.csv" pd.DataFrame([{"id": 1, "name": "alpha"}]).to_csv(csv_path, index=False, sep="\t") - _SimpleTable.load_file(sqlite_session, csv_path, dedupe=False) sqlite_session.commit() pd.DataFrame([{"id": 2, "name": "beta"}]).to_csv(csv_path, index=False, sep="\t") + _SimpleTable.load_file(sqlite_session, csv_path, dedupe=False, merge_strategy="insert") + sqlite_session.commit() + rows = sqlite_session.execute(sa.select(SimpleTable.id, SimpleTable.name).order_by(SimpleTable.id)).all() + assert [tuple(r) for r in rows] == [(1, "alpha"), (2, "beta")] - with pytest.raises(ValueError, match="is not empty; cannot use merge strategy 'insert_if_empty'"): - _SimpleTable.load_file( - sqlite_session, - csv_path, - dedupe=False, - merge_strategy="insert_if_empty", - ) + pd.DataFrame([{"id": 2, "name": "again"}]).to_csv(csv_path, index=False, sep="\t") + with pytest.raises(sa.exc.IntegrityError): + _SimpleTable.load_file(sqlite_session, csv_path, dedupe=False, merge_strategy="insert") inspector = sa.inspect(sqlite_engine) staging_name = resolve_backend(sqlite_engine).staging_name_for_table(SimpleTable.__tablename__) @@ -321,7 +320,7 @@ def test_insert_if_empty_raises_on_non_empty_target(sqlite_session, sqlite_engin @pytest.mark.parametrize("merge_strategy", ["replace", "upsert"]) -def test_empty_target_routes_merge_to_insert_if_empty(sqlite_session, tmp_path, caplog, merge_strategy): +def test_empty_target_runs_the_requested_merge(sqlite_session, tmp_path, caplog, merge_strategy): csv_path = tmp_path / "test_table.csv" pd.DataFrame( @@ -351,16 +350,12 @@ def test_empty_target_routes_merge_to_insert_if_empty(sqlite_session, tmp_path, messages = [record.getMessage() for record in caplog.records] - assert any("Checking whether target table is empty for merge optimisation." in message for message in messages) - assert any( - f"Target table is empty; routing merge strategy `{merge_strategy}` to insert-if-empty fast path." - in message - for message in messages - ) - assert any("Merge insert-if-empty phase starting." in message for message in messages) - assert not any("Checking whether target table is empty." in message for message in messages) - assert not any("Merge replace delete phase starting." in message for message in messages) - assert not any("Merge upsert phase starting." in message for message in messages) + if merge_strategy == "replace": + assert any("Merge replace delete phase starting." in message for message in messages) + assert any("Merge insert phase starting." in message for message in messages) + else: + assert any("Merge upsert phase starting." in message for message in messages) + assert not any("Merge insert phase starting." in message for message in messages) def test_staging_table_is_created_and_dropped(sqlite_session, sqlite_engine, tmp_path): @@ -587,6 +582,30 @@ def test_explicit_drop_rebuild_on_sqlite_restores_index(sqlite_session, sqlite_e assert "ix_test_table_name" in {idx["name"] for idx in inspector.get_indexes("test_table")} +def test_a_failed_index_rebuild_raises_after_the_merge_is_committed(sqlite_session, tmp_path, monkeypatch): + csv_path = tmp_path / "test_table.csv" + pd.DataFrame([{"id": 1, "name": "alpha"}]).to_csv(csv_path, index=False, sep="\t") + + def failing_create_index(index): + raise sa.exc.OperationalError("CREATE INDEX", {}, Exception("disk full")) + + monkeypatch.setattr(sa.schema, "CreateIndex", failing_create_index) + + with pytest.raises(IngestError, match="ix_test_table_name"): + _SimpleTable.load_file(sqlite_session, csv_path, index_strategy="drop_rebuild") + + assert sqlite_session.execute(sa.select(sa.func.count()).select_from(SimpleTable)).scalar_one() == 1 + + +def test_a_failed_merge_is_not_masked_by_a_failed_index_rebuild(sqlite_session, tmp_path, monkeypatch): + csv_path = tmp_path / "test_table.csv" + pd.DataFrame([{"id": 1, "name": "a"}, {"id": 1, "name": "b"}]).to_csv(csv_path, index=False, sep="\t") + monkeypatch.setattr(sa.schema, "CreateIndex", lambda index: (_ for _ in ()).throw(RuntimeError("no index"))) + + with pytest.raises(sa.exc.IntegrityError): + _SimpleTable.load_file(sqlite_session, csv_path, index_strategy="drop_rebuild") + + def test_drop_rebuild_logging_shows_merge_phases(sqlite_session, tmp_path, caplog): csv_path = tmp_path / "test_table.csv" diff --git a/tests/loaders/test_malformed_files.py b/tests/loaders/test_malformed_files.py index 9c1cbab..e78b4c5 100644 --- a/tests/loaders/test_malformed_files.py +++ b/tests/loaders/test_malformed_files.py @@ -161,3 +161,12 @@ def test_a_column_type_without_a_handler_is_fine_when_values_are_not_interpreted with pytest.raises(PolicyError, match="demo_blob"): table.load_file(dialect_session, path, staging_schema_tag=tag) assert table.load_file(dialect_session, path, staging_schema_tag=tag, normalise=False) == 1 + + +@pytest.mark.parametrize("option", [{"quote_mode": "csv"}, {"chunksize": 10}]) +def test_reader_options_alongside_a_reader_are_refused(dialect_session, tmp_path, option): + path = tmp_path / "demo_record.csv" + path.write_text("demo_id,demo_label,demo_number\n1,a,1\n") + reader = DelimitedTextReader(path, DemoRecord.model_columns()) + with pytest.raises(ValueError, match="configure the default reader"): + _Demo.load_file(dialect_session, path, reader=reader, **option) diff --git a/tests/loaders/test_merge_semantics.py b/tests/loaders/test_merge_semantics.py new file mode 100644 index 0000000..111f25e --- /dev/null +++ b/tests/loaders/test_merge_semantics.py @@ -0,0 +1,79 @@ +"""A merge strategy gives the same result whether or not the target already has rows.""" + +from pathlib import Path +from typing import Type, cast + +import pytest +import sqlalchemy as sa + +from oa_configurator import Dialect +from orm_loader.backends import STAGING_SCHEMA +from orm_loader.tables.typing import FileTableProtocol +from tests.models import CompositeTable, SimpleTable + +_Simple = cast(Type[FileTableProtocol], SimpleTable) +_Composite = cast(Type[FileTableProtocol], CompositeTable) +DUPLICATED = "id\tname\n1\tfirst\n1\tsecond\n2\tother\n" + + +@pytest.fixture +def options(dialect_session, request): + if dialect_session.get_bind().dialect.name == "sqlite": + return {"staging_schema_tag": None} + options = {"staging_schema_tag": STAGING_SCHEMA} + if dialect_session.get_bind().dialect.name == Dialect.DUCKDB: + options["resolved"] = request.getfixturevalue("duckdb_db").resolved + return options + + +def _load(session, tmp_path: Path, text: str, options: dict, **extra) -> None: + path = tmp_path / "test_table.csv" + path.write_text(text) + _Simple.load_file(session, path, dedupe=False, **options, **extra) + session.commit() + + +def _rows(session) -> list[tuple]: + return [tuple(r) for r in session.execute(sa.select(SimpleTable.id, SimpleTable.name).order_by(SimpleTable.id))] + + +def test_upsert_keeps_the_first_duplicate_in_an_empty_target(dialect_session, tmp_path, options): + _load(dialect_session, tmp_path, DUPLICATED, options, merge_strategy="upsert") + assert _rows(dialect_session) == [(1, "first"), (2, "other")] + + +def test_upsert_keeps_the_first_duplicate_in_a_populated_target(dialect_session, tmp_path, options): + _load(dialect_session, tmp_path, "id\tname\n9\tseed\n", options, merge_strategy="upsert") + _load(dialect_session, tmp_path, DUPLICATED, options, merge_strategy="upsert") + assert _rows(dialect_session) == [(1, "first"), (2, "other"), (9, "seed")] + + +def test_upsert_into_an_empty_target_hints_at_the_faster_strategy_without_changing_it( + dialect_session, tmp_path, options, caplog +): + caplog.set_level("INFO", logger="orm_loader.tables.loadable_table") + _load(dialect_session, tmp_path, DUPLICATED, options, merge_strategy="upsert") + messages = [record.getMessage() for record in caplog.records] + assert any("`insert` with `dedupe=True` loads the same rows faster" in m for m in messages) + assert any("Merge upsert phase starting." in m for m in messages) + assert not any("Merge insert phase starting." in m for m in messages) + + +def test_replace_with_a_composite_key_overwrites_matching_rows(dialect_session, tmp_path, options): + path = tmp_path / "composite_table.csv" + for text in ("a\tb\tvalue\n1\t1\told\n1\t2\tkept\n", "a\tb\tvalue\n1\t1\tnew\n2\t2\tadded\n"): + path.write_text(text) + _Composite.load_file(dialect_session, path, dedupe=False, merge_strategy="replace", **options) + dialect_session.commit() + rows = dialect_session.execute( + sa.select(CompositeTable.a, CompositeTable.b, CompositeTable.value).order_by(CompositeTable.a, CompositeTable.b) + ) + assert [tuple(r) for r in rows] == [(1, 1, "new"), (1, 2, "kept"), (2, 2, "added")] + + +def test_insert_fails_on_an_existing_key_and_adds_new_ones(dialect_session, tmp_path, options): + _load(dialect_session, tmp_path, "id\tname\n1\ta\n", options, merge_strategy="insert") + _load(dialect_session, tmp_path, "id\tname\n2\tb\n", options, merge_strategy="insert") + assert _rows(dialect_session) == [(1, "a"), (2, "b")] + with pytest.raises(sa.exc.IntegrityError): + _load(dialect_session, tmp_path, "id\tname\n2\tagain\n", options, merge_strategy="insert") diff --git a/tests/loaders/test_pipeline.py b/tests/loaders/test_pipeline.py index 82688fb..fd3bd84 100644 --- a/tests/loaders/test_pipeline.py +++ b/tests/loaders/test_pipeline.py @@ -15,7 +15,7 @@ from orm_loader.backends import STAGING_SCHEMA from orm_loader.conversion import ConversionPolicy, ErrorPolicy from orm_loader.helpers import ConversionError, RequiredValueError -from orm_loader.loaders import DelimitedTextReader, LoaderContext, ParquetReader, stage_file +from orm_loader.loaders import DelimitedTextReader, LoaderContext, ParquetReader, open_reader, stage_file from orm_loader.tables.typing import FileTableProtocol from tests.models import DemoRecord, DemoRequiredRecord, DemoTypedRecord @@ -49,7 +49,7 @@ def _stage(session, table, path: Path, reader=None, **context): table.create_staging_table(session, staging_schema_tag=tag) staging = table.get_staging_table(session, staging_schema_tag=tag) ctx = LoaderContext(tableclass=table, session=session, path=path, staging_schema_tag=tag, **context) - loaded = stage_file(ctx, reader or table._select_reader(ctx), staging) + loaded = stage_file(ctx, reader or open_reader(path, table.model_columns(), nulls=ctx.policy.nulls), staging) return loaded, ctx, staging @@ -129,7 +129,7 @@ def test_a_file_in_many_batches_loads_every_row_in_order(dialect_session, tmp_pa def test_a_failing_reader_is_raised_and_stages_nothing(dialect_session, tmp_path, fmt): path = _write(tmp_path, DemoRecord, fmt, [("1", "a", "", "", "")]) - reader = _Demo._select_reader(LoaderContext(tableclass=_Demo, session=dialect_session, path=path)) + reader = open_reader(path, _Demo.model_columns()) def broken(): raise RuntimeError("boom") diff --git a/tests/loaders/test_schema_translate_map.py b/tests/loaders/test_schema_translate_map.py index 0ee4fa0..5ced078 100644 --- a/tests/loaders/test_schema_translate_map.py +++ b/tests/loaders/test_schema_translate_map.py @@ -59,7 +59,7 @@ def test_load_file_respects_non_default_schema_end_to_end(pg_db, tmp_path): def test_replace_merge_respects_non_default_schema_end_to_end(pg_db, tmp_path): """A second load_file() call with merge_strategy="replace" against the same non-default schema. Proves the merge path itself, not just the - initial insert-if-empty fast path, qualifies correctly.""" + initial insert, qualifies correctly.""" with schema_scoped_session(pg_db, SimpleTable.__table__, prefix="loader_replace") as (scoped, session): schema = scoped.schemas[SchemaRole.PRIMARY] diff --git a/tests/loaders/test_split_connection.py b/tests/loaders/test_split_connection.py index 1529bd4..0a78b06 100644 --- a/tests/loaders/test_split_connection.py +++ b/tests/loaders/test_split_connection.py @@ -61,7 +61,7 @@ def get_bind(self, mapper=None, clause=None, **kwargs): VocabSchemaTable.load_file( sqlite_session, csv_path, - merge_strategy="insert_if_empty", + merge_strategy="insert", dedupe=False, staging_schema_tag=staging_schema_claim().schema_tag, ) diff --git a/tests/loaders/test_stage_file.py b/tests/loaders/test_stage_file.py index cb630a3..8dd24c3 100644 --- a/tests/loaders/test_stage_file.py +++ b/tests/loaders/test_stage_file.py @@ -5,22 +5,24 @@ from orm_loader.backends import resolve_backend from orm_loader.helpers import IngestError -from orm_loader.loaders import LoaderContext, stage_file +from orm_loader.loaders import LoaderContext, open_reader, stage_file from tests.models import FileLoaderTable +def _reader(ctx, **options): + return open_reader(ctx.path, ctx.tableclass.model_columns(), nulls=ctx.policy.nulls, **options) + + def test_loader_context_fields(sqlite_session): ctx = LoaderContext( tableclass=FileLoaderTable, session=sqlite_session, path=Path("file.csv"), - chunksize=100, normalise=False, dedupe=True, ) - assert ctx.chunksize == 100 assert ctx.normalise is False assert ctx.dedupe is True @@ -41,12 +43,10 @@ def test_stage_file_case_insensitive_headers(tmp_path, sqlite_session): tableclass=FileLoaderTable, session=sqlite_session, path=csv, - chunksize=100, - normalise=True, dedupe=True, ) - n = stage_file(ctx, ctx.tableclass._select_reader(ctx), FileLoaderTable.__table__) + n = stage_file(ctx, _reader(ctx), FileLoaderTable.__table__) rows = sqlite_session.execute( sa.text('SELECT id, value FROM test_file_loader ORDER BY id') @@ -74,13 +74,10 @@ def test_stage_file_literal_mode_preserves_quotes(tmp_path, sqlite_session): tableclass=FileLoaderTable, session=sqlite_session, path=csv, - chunksize=100, - normalise=True, dedupe=True, - quote_mode="literal", ) - n = stage_file(ctx, ctx.tableclass._select_reader(ctx), FileLoaderTable.__table__) + n = stage_file(ctx, _reader(ctx, quote_mode="literal"), FileLoaderTable.__table__) rows = sqlite_session.execute( sa.text('SELECT id, value FROM test_file_loader ORDER BY id') @@ -107,13 +104,10 @@ def test_stage_file_csv_mode_strips_quotes(tmp_path, sqlite_session): tableclass=FileLoaderTable, session=sqlite_session, path=csv, - chunksize=100, - normalise=True, dedupe=True, - quote_mode="csv", ) - n = stage_file(ctx, ctx.tableclass._select_reader(ctx), FileLoaderTable.__table__) + n = stage_file(ctx, _reader(ctx, quote_mode="csv"), FileLoaderTable.__table__) rows = sqlite_session.execute( sa.text('SELECT id, value FROM test_file_loader ORDER BY id') @@ -140,13 +134,10 @@ def test_stage_file_by_delimiter_tab_preserves_quotes(tmp_path, sqlite_session): tableclass=FileLoaderTable, session=sqlite_session, path=tsv, - chunksize=100, - normalise=True, dedupe=True, - quote_mode="by_delimiter", ) - n = stage_file(ctx, ctx.tableclass._select_reader(ctx), FileLoaderTable.__table__) + n = stage_file(ctx, _reader(ctx, quote_mode="by_delimiter"), FileLoaderTable.__table__) rows = sqlite_session.execute( sa.text('SELECT id, value FROM test_file_loader ORDER BY id') @@ -165,7 +156,7 @@ def test_file_without_data_rows_loads_nothing_whatever_its_header(tmp_path, sqli session=sqlite_session, path=csv, ) - assert stage_file(ctx, ctx.tableclass._select_reader(ctx), FileLoaderTable.__table__) == 0 + assert stage_file(ctx, _reader(ctx), FileLoaderTable.__table__) == 0 assert ctx.report.rows_read == 0 @@ -187,4 +178,4 @@ def losing(self, ctx, staging, batches, columns): monkeypatch.setattr(type(backend), "insert_batches", losing) with pytest.raises(IngestError, match="read 3 rows but loaded 0"): - stage_file(ctx, ctx.tableclass._select_reader(ctx), FileLoaderTable.__table__) + stage_file(ctx, _reader(ctx), FileLoaderTable.__table__) From b20a06b1c5d479d2f87e9571667b7a970f05198e Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Fri, 30 Oct 2026 01:24:59 +0000 Subject: [PATCH 12/21] Relax foreign keys before the first staging write, verify the SQLite pragma took effect, remove unused FK and connection hooks --- docs/loaders/staging.md | 2 +- src/orm_loader/backends/base.py | 49 +++----- src/orm_loader/backends/duckdb.py | 27 ----- src/orm_loader/backends/postgres.py | 37 +----- src/orm_loader/backends/sqlite.py | 106 +++-------------- src/orm_loader/tables/loadable_table.py | 60 +++++----- tests/backends/test_base_backend.py | 15 +-- tests/backends/test_duckdb_backend.py | 16 +-- tests/backends/test_postgres_backend.py | 86 -------------- tests/backends/test_sqlite_backend.py | 146 +++++++++--------------- tests/conftest.py | 1 + tests/loaders/test_loader_e2e.py | 4 +- 12 files changed, 126 insertions(+), 423 deletions(-) diff --git a/docs/loaders/staging.md b/docs/loaders/staging.md index 77f9b54..f6f890e 100644 --- a/docs/loaders/staging.md +++ b/docs/loaders/staging.md @@ -60,7 +60,7 @@ No schema changes are performed at the loader layer. ## Referential integrity -On PostgreSQL and SQLite a load relaxes foreign keys while it merges, so the caller needs no foreign key handling and a row loaded before the row it references is accepted. To verify the result, call `Table.validate_foreign_keys(session)` after every table is loaded. It checks every foreign key of the table with an anti-join and raises `ForeignKeyViolationError` with up to five unmatched keys. Null references are not violations, and the load order does not matter. +On PostgreSQL and SQLite a load relaxes foreign keys from before it writes staging until it has merged, so the caller needs no foreign key handling and a row loaded before the row it references is accepted. SQLite ignores `PRAGMA foreign_keys` inside a transaction, so a session with uncommitted writes cannot start a load there and `load_file` raises before touching anything. To verify the result, call `Table.validate_foreign_keys(session)` after every table is loaded. It checks every foreign key of the table with an anti-join and raises `ForeignKeyViolationError` with up to five unmatched keys. Null references are not violations, and the load order does not matter. To insert ORM objects built in Python that reference each other, wrap the inserts in `backend.relaxed_foreign_keys(session)`, which relaxes the keys on PostgreSQL and SQLite and leaves them enforced on DuckDB. diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index f3ad4e0..04dffd6 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -6,7 +6,6 @@ from contextlib import contextmanager from dataclasses import dataclass from functools import wraps -from time import perf_counter from typing import TYPE_CHECKING, Any, Callable, ClassVar, ParamSpec, Type, TypeVar, cast import pyarrow as pa @@ -717,24 +716,14 @@ def deduplicate_staging( after = self.count_staging(ctx, staging) return DedupeResult(tuple(key_columns), before, before - after) - @staticmethod - @abstractmethod - def _normalize_fk_check_state(previous_state: str | int) -> str | int: - """Validate and normalise a previously-returned FK state before interpolating into SQL. + def disable_fk_check(self, session: so.Session, *, mapper: Any | None = None) -> str | int | None: + """Relax foreign key enforcement and return the state to restore. - Each backend accepts a different type (SQLite: int, Postgres: str) and must - implement this to guard restore_fk_check() against invalid or injected values. + The default relaxes nothing and returns ``None``, which is right for a + backend that enforces its keys regardless. """ + return None - @abstractmethod - def disable_fk_check(self, session: so.Session, *, mapper: Any | None = None) -> str | int: - """Disable FK checks and return the previous backend-specific state.""" - - @abstractmethod - def enable_fk_check(self, session: so.Session, *, mapper: Any | None = None) -> str | int: - """Explicitly enable FK checks and return the previous backend-specific state.""" - - @abstractmethod def restore_fk_check( self, session: so.Session, @@ -742,7 +731,7 @@ def restore_fk_check( *, mapper: Any | None = None, ) -> None: - """Restore FK checks to a previously returned backend-specific state.""" + """Restore the state returned by :meth:`disable_fk_check`.""" def check_replace_feasible( self, @@ -926,9 +915,10 @@ def relaxed_foreign_keys( """ Relax foreign key checks for the block where the backend can, and restore them on exit. - A backend without ``supports_fk_toggle`` has nothing to relax, so its - constraints stay enforced and rows must be inserted parents first. - The session is rolled back when the block raises. + A backend that cannot relax its keys leaves them enforced, so rows + must be inserted parents first. The session is rolled back when the + block raises. A backend that cannot relax the keys it is able to + relax raises before the block starts, with the session untouched. Parameters ---------- @@ -941,29 +931,22 @@ def relaxed_foreign_keys( Whether to suspend the session's autoflush inside the block. """ scope = f"Table `{mapper.__tablename__}`: " if hasattr(mapper, "__tablename__") else "" - previous_fk_state: str | int | None = None + previous_state = self.disable_fk_check(session, mapper=mapper) + if previous_state is not None: + logger.info(f"{scope}Foreign key checks disabled.") try: - if self.capabilities.supports_fk_toggle: - started = perf_counter() - raw_state = self.disable_fk_check(session, mapper=mapper) - previous_fk_state = self._normalize_fk_check_state(raw_state) - logger.info(f"{scope}Foreign key checks disabled in {perf_counter() - started:.2f}s.") - if no_autoflush: with session.no_autoflush: yield else: yield - except Exception: session.rollback() raise - finally: - if previous_fk_state is not None: - started = perf_counter() - self.restore_fk_check(session, previous_fk_state, mapper=mapper) - logger.info(f"{scope}Foreign key checks restored in {perf_counter() - started:.2f}s.") + if previous_state is not None: + self.restore_fk_check(session, previous_state, mapper=mapper) + logger.info(f"{scope}Foreign key checks restored.") @requires_capability("supports_materialized_views", "materialized views") def create_materialized_view( diff --git a/src/orm_loader/backends/duckdb.py b/src/orm_loader/backends/duckdb.py index 85c462a..6f9e56e 100644 --- a/src/orm_loader/backends/duckdb.py +++ b/src/orm_loader/backends/duckdb.py @@ -161,33 +161,6 @@ def _staging_table(self, table_cls: type["FileTableProtocol"], columns: Sequence row_key = sa.Column("_rownum", sa.BigInteger, sequence, server_default=sequence.next_value()) return sa.Table(name, sa.MetaData(), *columns, row_key, schema=self.staging_schema) - @staticmethod - def _normalize_fk_check_state(previous_state: str | int) -> str | int: - raise NotImplementedError( - "DuckDB has no foreign-key-toggle mechanism; there is no state to normalize." - ) - - def disable_fk_check(self, session: so.Session, *, mapper: Any | None = None) -> str | int: - raise NotImplementedError( - "DuckDB has no foreign-key-toggle mechanism (no PRAGMA/SET equivalent exists)." - ) - - def enable_fk_check(self, session: so.Session, *, mapper: Any | None = None) -> str | int: - raise NotImplementedError( - "DuckDB has no foreign-key-toggle mechanism (no PRAGMA/SET equivalent exists)." - ) - - def restore_fk_check( - self, - session: so.Session, - previous_state: str | int, - *, - mapper: Any | None = None, - ) -> None: - raise NotImplementedError( - "DuckDB has no foreign-key-toggle mechanism (no PRAGMA/SET equivalent exists)." - ) - def _paginate_batches( self, table_cls: type["FileTableProtocol"], diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index c108a7f..91773e1 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -2,17 +2,14 @@ import io from collections.abc import Iterable, Mapping, Sequence -from contextlib import contextmanager from typing import TYPE_CHECKING, Any import pyarrow as pa import pyarrow.csv as pacsv import sqlalchemy as sa -import sqlalchemy.event as sae import sqlalchemy.orm as so from oa_configurator import ( Dialect, - autocommit_connection, qualified, ) from sqlalchemy.dialects import postgresql @@ -154,7 +151,7 @@ def _normalize_fk_check_state(previous_state: str | int) -> str: ) return normalised - def disable_fk_check(self, session: so.Session, *, mapper: Any | None = None) -> str | int: + def disable_fk_check(self, session: so.Session, *, mapper: Any | None = None) -> str: execute_options = {"bind_arguments": {"mapper": mapper}} if mapper is not None else {} previous_state = session.execute( sa.text("SHOW session_replication_role"), **execute_options @@ -164,16 +161,6 @@ def disable_fk_check(self, session: so.Session, *, mapper: Any | None = None) -> raise RuntimeError("Expected PostgreSQL FK state to be a string") return previous_state - def enable_fk_check(self, session: so.Session, *, mapper: Any | None = None) -> str | int: - execute_options = {"bind_arguments": {"mapper": mapper}} if mapper is not None else {} - previous_state = session.execute( - sa.text("SHOW session_replication_role"), **execute_options - ).scalar() - session.execute(sa.text("SET session_replication_role = 'origin'"), **execute_options) - if not isinstance(previous_state, str): - raise RuntimeError("Expected PostgreSQL FK state to be a string") - return previous_state - def restore_fk_check( self, session: so.Session, @@ -417,25 +404,3 @@ def create_materialized_view_index( cause=error, ) ) from error - - @contextmanager - def engine_with_replica_role(self, engine: "Engine"): - def _set_replica_role( - dbapi_conn: sa.engine.interfaces.DBAPIConnection, - _, - ) -> None: - cur = dbapi_conn.cursor() - cur.execute("SET session_replication_role = 'replica'") - cur.close() - - sae.listen(engine, "connect", _set_replica_role) - - try: - yield engine - finally: - sae.remove(engine, "connect", _set_replica_role) - with autocommit_connection(engine) as autocommit_conn: - autocommit_conn.execute(sa.text("SET session_replication_role = DEFAULT")) - role = autocommit_conn.execute(sa.text("SHOW session_replication_role")).scalar() - if role != "origin": - raise RuntimeError("Failed to restore session_replication_role") diff --git a/src/orm_loader/backends/sqlite.py b/src/orm_loader/backends/sqlite.py index 017a831..a3028fc 100644 --- a/src/orm_loader/backends/sqlite.py +++ b/src/orm_loader/backends/sqlite.py @@ -2,33 +2,26 @@ from collections.abc import Iterable, Mapping, Sequence import logging -import sqlite3 -from pathlib import Path from typing import TYPE_CHECKING, Any import pyarrow as pa import pyarrow.compute as pc import sqlalchemy as sa import sqlalchemy.orm as so -from sqlalchemy import event, text +from sqlalchemy import text from sqlalchemy.dialects import sqlite as sqlite_dialect -from sqlalchemy.exc import IntegrityError from sqlalchemy.sql.compiler import IdentifierPreparer from .base import BackendCapabilities, DatabaseBackend, Dialect if TYPE_CHECKING: from oa_configurator import ResolvedDatabase - from sqlalchemy.engine import Connection, Engine from ..loaders.loader import LoaderContext from ..tables.typing import FileTableProtocol logger = logging.getLogger(__name__) -VALID_SQLITE_JOURNAL_MODES = frozenset( - {"DELETE", "TRUNCATE", "PERSIST", "MEMORY", "WAL", "OFF"} -) class SQLiteBackend(DatabaseBackend): @@ -43,9 +36,6 @@ def __init__( staging_schema: str | None = None, resolved: ResolvedDatabase | None = None, physical_schemas: Mapping[str, str | None] | None = None, - busy_timeout_ms: int = 60000, - journal_mode: str = "WAL", - defer_foreign_keys: bool = True, ) -> None: if staging_schema is not None: raise ValueError( @@ -58,19 +48,6 @@ def __init__( resolved=resolved, physical_schemas=physical_schemas, ) - self.busy_timeout_ms = busy_timeout_ms - self.journal_mode = self._validate_journal_mode(journal_mode) - self.defer_foreign_keys = defer_foreign_keys - - @staticmethod - def _validate_journal_mode(journal_mode: str) -> str: - normalised = journal_mode.strip().upper() - if normalised not in VALID_SQLITE_JOURNAL_MODES: - raise ValueError( - "Unsupported SQLite journal_mode " - f"{journal_mode!r}. Expected one of: {sorted(VALID_SQLITE_JOURNAL_MODES)}" - ) - return normalised @staticmethod def _normalize_fk_check_state(previous_state: str | int) -> str: @@ -158,20 +135,27 @@ def evaluate_default( """Evaluate without a cast: SQLite casts a date to the integer of its year.""" return session.execute(sa.select(expression), bind_arguments={"mapper": table_cls}).scalar_one() - def disable_fk_check(self, session: so.Session, *, mapper: Any | None = None) -> str | int: - execute_options = {"bind_arguments": {"mapper": mapper}} if mapper is not None else {} - previous_state = session.execute(text("PRAGMA foreign_keys"), **execute_options).scalar() - session.execute(text("PRAGMA foreign_keys = OFF"), **execute_options) - if not isinstance(previous_state, int): - raise RuntimeError("Expected SQLite FK state to be an int") - return previous_state + def disable_fk_check(self, session: so.Session, *, mapper: Any | None = None) -> str | int | None: + """Turn ``PRAGMA foreign_keys`` off and return its previous value. + + SQLite ignores the pragma inside an open transaction, so the state is + read back and a pragma that did not take effect raises. - def enable_fk_check(self, session: so.Session, *, mapper: Any | None = None) -> str | int: + Raises + ------ + RuntimeError + The connection kept enforcing foreign keys. + """ execute_options = {"bind_arguments": {"mapper": mapper}} if mapper is not None else {} previous_state = session.execute(text("PRAGMA foreign_keys"), **execute_options).scalar() - session.execute(text("PRAGMA foreign_keys = ON"), **execute_options) if not isinstance(previous_state, int): raise RuntimeError("Expected SQLite FK state to be an int") + session.execute(text("PRAGMA foreign_keys = OFF"), **execute_options) + if session.execute(text("PRAGMA foreign_keys"), **execute_options).scalar(): + raise RuntimeError( + "SQLite ignored PRAGMA foreign_keys = OFF because the connection is inside a transaction. " + "Commit the session's pending writes before loading." + ) return previous_state def restore_fk_check( @@ -290,59 +274,3 @@ def _insert(select_: Any) -> sa.Insert: for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): select_ = non_paginated_select if batch_range is None else non_paginated_select.where(batch_range) session.execute(_insert(select_)) - - def configure_dbapi_connection(self, dbapi_connection: sa.engine.interfaces.DBAPIConnection) -> None: - if dbapi_connection.__class__.__module__.startswith("sqlite3"): - cursor = dbapi_connection.cursor() - cursor.execute(f"PRAGMA busy_timeout = {self.busy_timeout_ms}") - cursor.execute(f"PRAGMA journal_mode = {self.journal_mode}") - cursor.execute("PRAGMA foreign_keys = ON;") - if self.defer_foreign_keys: - cursor.execute("PRAGMA defer_foreign_keys = ON;") - cursor.close() - - def install_engine_hooks(self, engine: "Engine") -> None: - @event.listens_for(engine, "connect") - def _enable_sqlite_foreign_keys( - dbapi_connection: sa.engine.interfaces.DBAPIConnection, - _connection_record: Any - ) -> None: - self.configure_dbapi_connection(dbapi_connection) - - def explain_fk_error( - self, - session: so.Session, - exc: IntegrityError, - *, - mapper: Any | None = None, - raise_error: bool = True, - ) -> None: - bind: Engine | Connection = session.get_bind(mapper=mapper) - if bind.dialect.name != Dialect.SQLITE: - raise exc - - with self._as_connection(bind) as conn: - rows = conn.execute(text("PRAGMA foreign_key_check")).fetchall() - - if rows: - for row in rows: - logger.error( - "FK violation: table=%s rowid=%s references=%s fk_index=%s", - row[0], row[1], row[2], row[3] - ) - - if raise_error: - raise exc - - def restore_journal_mode(self, db_path: Path) -> None: - timeout_s = max(self.busy_timeout_ms / 1000, 5) - try: - with sqlite3.connect(db_path.resolve(), timeout=timeout_s) as conn: - conn.execute("PRAGMA wal_checkpoint(TRUNCATE)") - conn.execute("PRAGMA journal_mode = DELETE") - conn.commit() - except sqlite3.OperationalError as exc: - raise RuntimeError( - "Failed to restore SQLite journal mode. " - "Close or dispose active SQLite connections before calling this helper." - ) from exc diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index 84c2275..17d2ec5 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -146,15 +146,14 @@ def manage_indices( merge_failed = False try: - with backend.relaxed_foreign_keys(session, mapper=cls, no_autoflush=False): - yield - logger.info(f"Table `{table_name}`: Committing merged rows.") - commit_started = perf_counter() - session.commit() - logger.info( - f"Table `{table_name}`: Merge commit completed in " - f"{_format_elapsed(perf_counter() - commit_started)}." - ) + yield + logger.info(f"Table `{table_name}`: Committing merged rows.") + commit_started = perf_counter() + session.commit() + logger.info( + f"Table `{table_name}`: Merge commit completed in " + f"{_format_elapsed(perf_counter() - commit_started)}." + ) except Exception as e: merge_failed = True session.rollback() @@ -431,8 +430,14 @@ def load_file( skip_malformed=policy.errors.on_error == "reject_row", ) + backend = resolve_backend( + session, + staging_schema_tag=staging_schema_tag, + mapper=cls, + resolved=resolved, + physical_schemas=physical_schemas, + ) if merge_strategy == "replace": - backend = resolve_backend(session, staging_schema_tag=staging_schema_tag, mapper=cls) if backend.resolve_index_strategy(index_strategy) == "keep": backend._require_capability( "supports_replace_with_kept_indexes", @@ -452,32 +457,25 @@ def load_file( ) try: - # Load to staging (Indices are already excluded via updated create_staging_table) - logger.info(f"Table `{cls.__tablename__}`: Loading data into staging table") - total = cls.load_staging(loader_context, reader) - - # Merge staging to target (Wrapped in our index dropper!) - logger.info(f"Table `{cls.__tablename__}`: Merging staging data into target table") - with cls.manage_indices(session, index_strategy=index_strategy, staging_schema_tag=staging_schema_tag): - cls.merge_from_staging( - session, - merge_strategy=merge_strategy, - merge_batch_size=merge_batch_size, - staging_schema_tag=staging_schema_tag, - validate_deferred_fks=validate_deferred_fks, - resolved=resolved, - physical_schemas=physical_schemas, - ) - if validate_deferred_fks: - logger.info(f"Table `{cls.__tablename__}`: Validating deferred foreign keys.") - backend = resolve_backend( + # SQLite cannot switch foreign key checks inside a transaction, so relax before staging writes. + with backend.relaxed_foreign_keys(session, mapper=cls, no_autoflush=False): + logger.info(f"Table `{cls.__tablename__}`: Loading data into staging table") + total = cls.load_staging(loader_context, reader) + + logger.info(f"Table `{cls.__tablename__}`: Merging staging data into target table") + with cls.manage_indices(session, index_strategy=index_strategy, staging_schema_tag=staging_schema_tag): + cls.merge_from_staging( session, + merge_strategy=merge_strategy, + merge_batch_size=merge_batch_size, staging_schema_tag=staging_schema_tag, - mapper=cls, + validate_deferred_fks=validate_deferred_fks, resolved=resolved, physical_schemas=physical_schemas, ) - backend.validate_deferred_fks(cls, session) + if validate_deferred_fks: + logger.info(f"Table `{cls.__tablename__}`: Validating deferred foreign keys.") + backend.validate_deferred_fks(cls, session) finally: cls.drop_staging_table(session, staging_schema_tag=staging_schema_tag) diff --git a/tests/backends/test_base_backend.py b/tests/backends/test_base_backend.py index c33069e..c7616c7 100644 --- a/tests/backends/test_base_backend.py +++ b/tests/backends/test_base_backend.py @@ -121,18 +121,10 @@ def _pagination_column( def _commits_between_pagination_batches(self) -> bool: return True - @staticmethod - def _normalize_fk_check_state(previous_state: str | int) -> str | int: - return previous_state - - def disable_fk_check(self, session: so.Session, *, mapper: Any | None = None) -> str | int: + def disable_fk_check(self, session: so.Session, *, mapper: Any | None = None) -> str | int | None: self.calls.append(("disable_fk_check", session)) return "enabled" - def enable_fk_check(self, session: so.Session, *, mapper: Any | None = None) -> str | int: - self.calls.append(("enable_fk_check", session)) - return "disabled" - def restore_fk_check( self, session: so.Session, previous_state: str | int, *, mapper: Any | None = None ) -> None: @@ -247,9 +239,8 @@ def test_relaxed_foreign_keys_toggles_fk_and_restores(sqlite_session): def test_relaxed_foreign_keys_leaves_keys_enforced_when_backend_cannot_toggle(sqlite_session): class NoFKBackend(FakeBackend): - @property - def capabilities(self) -> BackendCapabilities: - return BackendCapabilities() + def disable_fk_check(self, session: so.Session, *, mapper: Any | None = None) -> str | int | None: + return super(FakeBackend, self).disable_fk_check(session, mapper=mapper) backend = NoFKBackend() diff --git a/tests/backends/test_duckdb_backend.py b/tests/backends/test_duckdb_backend.py index 00ba145..89e747c 100644 --- a/tests/backends/test_duckdb_backend.py +++ b/tests/backends/test_duckdb_backend.py @@ -256,17 +256,11 @@ def test_duckdb_backend_rownum_populates_monotonically_on_insert(duckdb_session) assert rownums == [1, 2] -def test_duckdb_backend_fk_toggle_methods_raise(): +def test_duckdb_backend_has_nothing_to_relax(): backend = DuckDBBackend() + session = cast("so.Session", None) - with pytest.raises(NotImplementedError, match="no foreign-key-toggle mechanism"): - backend.disable_fk_check(cast("so.Session", None)) - with pytest.raises(NotImplementedError, match="no foreign-key-toggle mechanism"): - backend.enable_fk_check(cast("so.Session", None)) - with pytest.raises(NotImplementedError, match="no foreign-key-toggle mechanism"): - backend.restore_fk_check(cast("so.Session", None), "ON") - with pytest.raises(NotImplementedError, match="no state to normalize"): - DuckDBBackend._normalize_fk_check_state("ON") + assert backend.disable_fk_check(session) is None @pytest.mark.duckdb @@ -281,9 +275,9 @@ def test_duckdb_backend_relaxed_foreign_keys_runs_without_toggling_fk(duckdb_ses @pytest.mark.duckdb @pytest.mark.db_dialect -def test_a_merge_on_duckdb_does_not_log_foreign_key_toggling(duckdb_session, caplog): +def test_relaxed_foreign_keys_on_duckdb_does_not_log_toggling(duckdb_session, caplog): with caplog.at_level("INFO", logger="orm_loader.backends.base"): - with ComputedColumnTable.manage_indices(duckdb_session): # ty: ignore[invalid-argument-type] + with DuckDBBackend().relaxed_foreign_keys(duckdb_session): pass assert "foreign key" not in caplog.text.lower() diff --git a/tests/backends/test_postgres_backend.py b/tests/backends/test_postgres_backend.py index 34447c3..e7dbbfc 100644 --- a/tests/backends/test_postgres_backend.py +++ b/tests/backends/test_postgres_backend.py @@ -1,13 +1,11 @@ from __future__ import annotations -import sqlalchemy.event as sae from typing import TYPE_CHECKING, Type, cast import pytest import sqlalchemy as sa import sqlalchemy.orm as so from sqlalchemy.dialects import postgresql -from sqlalchemy.engine import Engine from oa_configurator import SCHEMA_TRANSLATE_MAP_KEY, Role, qualified from oa_configurator.testing import isolated_test_schema @@ -129,16 +127,12 @@ def test_postgres_backend_fk_methods_emit_expected_sql(): session = _FakeSession() previous = backend.disable_fk_check(_sess(session)) - enabled = backend.enable_fk_check(_sess(session)) backend.restore_fk_check(_sess(session), previous) assert previous == "origin" - assert enabled == "origin" assert session.statements == [ "SHOW session_replication_role", "SET session_replication_role = 'replica'", - "SHOW session_replication_role", - "SET session_replication_role = 'origin'", "SET session_replication_role = 'origin'", ] @@ -552,86 +546,6 @@ def test_postgres_backend_disable_fk_raises_when_show_returns_non_string(): raise AssertionError("Expected RuntimeError when SHOW returns a non-string") -def test_postgres_backend_enable_fk_raises_when_show_returns_non_string(): - backend = PostgresBackend() - session = _FakeSession(scalar_result=42) - - try: - backend.enable_fk_check(_sess(session)) - except RuntimeError as exc: - assert "Expected PostgreSQL FK state to be a string" in str(exc) - else: - 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]] = [] - statements: list[str] = [] - - class _Result: - def scalar(self): - return "origin" - - class _Conn: - def __enter__(self): - return self - - def __exit__(self, *_) -> None: - return None - - def execution_options(self, **_): - return self - - def get_isolation_level(self): - return "READ COMMITTED" - - def rollback(self) -> None: - return None - - def close(self) -> None: - return None - - def execute(self, statement): - sql = str(statement.compile(dialect=postgresql.dialect())) - statements.append(sql) - return _Result() - - class _Engine(Engine): - def __init__(self) -> None: - # only exists for autocommit_connection() to route it into its real Engine branch - pass - - def connect(self): - events.append(("connect", self, "connect")) - return _Conn() - - engine = _Engine() - - def _listen(target, name, *_) -> None: - events.append(("listen", target, name)) - - def _remove(target, name, *_) -> None: - events.append(("remove", target, name)) - - monkeypatch.setattr(sae, "listen", _listen) - monkeypatch.setattr(sae, "remove", _remove) - - with backend.engine_with_replica_role(cast(Engine, engine)): - pass - - assert events == [ - ("listen", engine, "connect"), - ("remove", engine, "connect"), - ("connect", engine, "connect"), - ] - assert statements == [ - "SET session_replication_role = DEFAULT", - "SHOW session_replication_role", - ] - - # --------------------------------------------------------------------------- # Loading files into real Postgres # --------------------------------------------------------------------------- diff --git a/tests/backends/test_sqlite_backend.py b/tests/backends/test_sqlite_backend.py index 5854b69..88384e0 100644 --- a/tests/backends/test_sqlite_backend.py +++ b/tests/backends/test_sqlite_backend.py @@ -1,7 +1,5 @@ from __future__ import annotations -import sqlite3 -from pathlib import Path from typing import TYPE_CHECKING, Type, cast import pytest @@ -10,7 +8,7 @@ from sqlalchemy.dialects import sqlite as sqlite_dialect from orm_loader.backends import Dialect, SQLiteBackend -from tests.models import ComputedColumnTable +from tests.models import ComputedColumnTable, ReferencingChildTable, SimpleTable if TYPE_CHECKING: from orm_loader.tables.typing import FileTableProtocol @@ -20,9 +18,9 @@ class _FakeSession: - def __init__(self, scalar_result: int | str = 1) -> None: + def __init__(self, *scalar_results: int | str) -> None: self.statements: list[str] = [] - self.scalar_result = scalar_result + self.scalar_results = list(scalar_results) or [1] def execute(self, statement, **kwargs): if hasattr(statement, "compile"): @@ -38,7 +36,7 @@ def __init__(self, value): def scalar(self): return self._value - return _Result(self.scalar_result) + return _Result(self.scalar_results.pop(0) if len(self.scalar_results) > 1 else self.scalar_results[0]) _ComputedTableCls = cast("Type[FileTableProtocol]", ComputedColumnTable) @@ -58,7 +56,6 @@ def test_sqlite_backend_identity_and_capabilities(): assert backend.capabilities.supports_fk_toggle is True assert backend.capabilities.supports_materialized_views is False assert backend.resolve_index_strategy("auto") == "keep" - assert backend.journal_mode == "WAL" @pytest.mark.parametrize("merge_batch_size", [0, -1]) @@ -87,30 +84,25 @@ def test_sqlite_backend_create_staging_table(sqlite_session, sqlite_engine): assert {c["name"]: c["nullable"] for c in cols} == {name: target[name] for name in ("id", "name")} -def test_sqlite_backend_disable_fk_reads_then_sets(): +def test_sqlite_backend_disable_fk_reads_sets_then_reads_back(): backend = SQLiteBackend() - sqlite_session = _FakeSession(scalar_result=1) + sqlite_session = _FakeSession(1, 0) previous = backend.disable_fk_check(_sess(sqlite_session)) assert previous == 1 assert sqlite_session.statements == [ - "PRAGMA foreign_keys", # read current state - "PRAGMA foreign_keys = OFF", # set to OFF + "PRAGMA foreign_keys", + "PRAGMA foreign_keys = OFF", + "PRAGMA foreign_keys", ] -def test_sqlite_backend_enable_fk_reads_then_sets(): +def test_sqlite_backend_disable_fk_raises_when_the_pragma_is_ignored(): backend = SQLiteBackend() - sqlite_session = _FakeSession(scalar_result=0) - - previous = backend.enable_fk_check(_sess(sqlite_session)) - assert previous == 0 - assert sqlite_session.statements == [ - "PRAGMA foreign_keys", # read current state - "PRAGMA foreign_keys = ON", # set to ON - ] + with pytest.raises(RuntimeError, match="inside a transaction"): + backend.disable_fk_check(_sess(_FakeSession(1, 1))) def test_sqlite_backend_restore_fk_normalises_int_and_emits(): @@ -189,69 +181,9 @@ def test_sqlite_backend_create_materialized_view_index_raises(sqlite_engine): backend.create_materialized_view_index(sqlite_engine, "mv_test", index) -def test_sqlite_backend_configures_bulk_load_pragmas(tmp_path: Path): - backend = SQLiteBackend() - db_path = tmp_path / "test.db" - sqlite_engine = sa.create_engine(f"sqlite:///{db_path}", future=True) - backend.install_engine_hooks(sqlite_engine) - - with sqlite_engine.connect() as conn: - busy_timeout = conn.execute(sa.text("PRAGMA busy_timeout")).scalar_one() - journal_mode = conn.execute(sa.text("PRAGMA journal_mode")).scalar_one() - foreign_keys = conn.execute(sa.text("PRAGMA foreign_keys")).scalar_one() - - assert busy_timeout == 60000 - assert str(journal_mode).lower() == "wal" - assert foreign_keys == 1 - - -def test_sqlite_backend_restore_journal_mode(tmp_path: Path): - backend = SQLiteBackend() - db_path = tmp_path / "journal.db" - sqlite_engine = sa.create_engine(f"sqlite:///{db_path}", future=True) - backend.install_engine_hooks(sqlite_engine) - - with sqlite_engine.begin() as conn: - conn.execute(sa.text("CREATE TABLE t (id INTEGER PRIMARY KEY, name TEXT)")) - conn.execute(sa.text("INSERT INTO t (name) VALUES ('x')")) - - sqlite_engine.dispose() - backend.restore_journal_mode(db_path) - - with sqlite3.connect(db_path.resolve()) as conn: - journal_mode = conn.execute("PRAGMA journal_mode").fetchone()[0] - - assert str(journal_mode).lower() == "delete" - - -def test_install_engine_hooks_sets_bulk_load_pragmas(tmp_path: Path): - db_path = tmp_path / "attached.db" - sqlite_engine = sa.create_engine(f"sqlite:///{db_path}", future=True) - - SQLiteBackend(busy_timeout_ms=45000).install_engine_hooks(sqlite_engine) - - with sqlite_engine.connect() as conn: - busy_timeout = conn.execute(sa.text("PRAGMA busy_timeout")).scalar_one() - journal_mode = conn.execute(sa.text("PRAGMA journal_mode")).scalar_one() - foreign_keys = conn.execute(sa.text("PRAGMA foreign_keys")).scalar_one() - - assert busy_timeout == 45000 - assert str(journal_mode).lower() == "wal" - assert foreign_keys == 1 - - -def test_sqlite_backend_rejects_invalid_journal_mode(): - try: - SQLiteBackend(journal_mode="wal; drop table x;") - except ValueError as exc: - assert "Unsupported SQLite journal_mode" in str(exc) - else: - raise AssertionError("Expected invalid journal_mode to raise ValueError") - - def test_sqlite_backend_disable_fk_raises_when_pragma_returns_non_int(): backend = SQLiteBackend() - sqlite_session = _FakeSession(scalar_result="not_an_int") + sqlite_session = _FakeSession("not_an_int") try: backend.disable_fk_check(_sess(sqlite_session)) @@ -261,18 +193,6 @@ def test_sqlite_backend_disable_fk_raises_when_pragma_returns_non_int(): raise AssertionError("Expected RuntimeError when PRAGMA returns a non-int") -def test_sqlite_backend_enable_fk_raises_when_pragma_returns_non_int(): - backend = SQLiteBackend() - sqlite_session = _FakeSession(scalar_result="not_an_int") - - try: - backend.enable_fk_check(_sess(sqlite_session)) - except RuntimeError as exc: - assert "Expected SQLite FK state to be an int" in str(exc) - else: - raise AssertionError("Expected RuntimeError when PRAGMA returns a non-int") - - def test_sqlite_backend_restore_fk_accepts_string_values(): backend = SQLiteBackend() sqlite_session = _FakeSession() @@ -299,10 +219,46 @@ def test_sqlite_backend_fk_toggle_round_trip(sqlite_session): assert sqlite_session.execute(sa.text("PRAGMA foreign_keys")).scalar() == 1 -def test_a_merge_on_sqlite_logs_foreign_key_toggling(sqlite_session, caplog): +def test_relaxed_foreign_keys_on_sqlite_logs_the_toggling(sqlite_session, caplog): with caplog.at_level("INFO", logger="orm_loader.backends.base"): - with ComputedColumnTable.manage_indices(sqlite_session): # ty: ignore[invalid-argument-type] + with SQLiteBackend().relaxed_foreign_keys(sqlite_session): pass assert "Foreign key checks disabled" in caplog.text assert "Foreign key checks restored" in caplog.text + + +@pytest.fixture +def enforcing_session(sqlite_engine): + """A session whose connections come up with foreign key enforcement on.""" + sa.event.listen(sqlite_engine, "connect", lambda conn, _: conn.execute("PRAGMA foreign_keys = ON")) + sqlite_engine.dispose() + with so.Session(sqlite_engine) as session: + yield session + + +def _foreign_keys(session: so.Session) -> int: + return session.execute(sa.text("PRAGMA foreign_keys")).scalar_one() + + +def test_a_load_relaxes_enforced_foreign_keys_and_restores_them(enforcing_session, tmp_path): + path = tmp_path / "referencing_child_table.csv" + path.write_text("id,parent_id\n1,1\n") + child = cast("Type[FileTableProtocol]", ReferencingChildTable) + + assert child.load_file(enforcing_session, path, staging_schema_tag=None) == 1 + + assert _foreign_keys(enforcing_session) == 1 + + +def test_a_load_on_a_session_with_pending_writes_raises_and_keeps_them(enforcing_session, tmp_path): + enforcing_session.add(SimpleTable(id=1, name="pending")) + enforcing_session.flush() + path = tmp_path / f"{SimpleTable.__tablename__}.csv" + path.write_text("id,name\n2,a\n") + simple = cast("Type[FileTableProtocol]", SimpleTable) + + with pytest.raises(RuntimeError, match="inside a transaction"): + simple.load_file(enforcing_session, path, staging_schema_tag=None) + + assert enforcing_session.get(SimpleTable, 1) is not None diff --git a/tests/conftest.py b/tests/conftest.py index e43f259..9a04436 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -80,6 +80,7 @@ def duckdb_session(duckdb_db): creates the staging schema and Base.metadata, then returns duckdb_db.session.""" conn = duckdb_db.connection Base.metadata.create_all(conn) + conn.commit() return duckdb_db.session diff --git a/tests/loaders/test_loader_e2e.py b/tests/loaders/test_loader_e2e.py index 30b1eff..cf87e63 100644 --- a/tests/loaders/test_loader_e2e.py +++ b/tests/loaders/test_loader_e2e.py @@ -643,7 +643,7 @@ def test_drop_rebuild_logging_shows_merge_phases(sqlite_session, tmp_path, caplo assert any("Finished dropping 1 active indices in " in message for message in messages) assert any("Committing after index drop." in message for message in messages) assert any("Commit after index drop completed in " in message for message in messages) - assert any("Foreign key checks disabled in " in message for message in messages) + assert any("Foreign key checks disabled." in message for message in messages) assert any("Merging staging data into target table" in message for message in messages) assert any("Merge replace delete phase starting." in message for message in messages) assert any("Merge replace delete phase completed in " in message for message in messages) @@ -652,7 +652,7 @@ def test_drop_rebuild_logging_shows_merge_phases(sqlite_session, tmp_path, caplo assert any("Merging staging data into target table" in message for message in messages) assert any("Committing merged rows." in message for message in messages) assert any("Merge commit completed in " in message for message in messages) - assert any("Foreign key checks restored in " in message for message in messages) + assert any("Foreign key checks restored." in message for message in messages) assert any("Verifying/Rebuilding indices." in message for message in messages) assert any("Restoring missing index: ix_test_table_name" in message for message in messages) assert any("Restored missing index `ix_test_table_name` in " in message for message in messages) From b69cb27e8ad40d6da915e42f70f6e67b91867de3 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Fri, 30 Oct 2026 01:47:41 +0000 Subject: [PATCH 13/21] Remove capability flags, the date parser hook and the staging wrappers; evaluate expression defaults once; move materialized views to their own backend interface --- docs/loaders/conversion.md | 6 +- src/orm_loader/backends/__init__.py | 2 - src/orm_loader/backends/base.py | 163 +----------------- src/orm_loader/backends/duckdb.py | 17 +- src/orm_loader/backends/postgres.py | 31 +--- src/orm_loader/backends/sqlite.py | 11 +- src/orm_loader/conversion/handlers/dates.py | 51 +++--- src/orm_loader/loaders/processor.py | 26 +-- src/orm_loader/mappers/__init__.py | 2 + .../mappers/materialised_view_mixin.py | 86 ++++++++- src/orm_loader/tables/loadable_table.py | 102 ++--------- src/orm_loader/tables/typing.py | 15 +- tests/backends/test_base_backend.py | 114 +----------- tests/backends/test_duckdb_backend.py | 22 +-- tests/backends/test_postgres_backend.py | 7 +- tests/backends/test_sinks.py | 6 +- tests/backends/test_sqlite_backend.py | 42 +---- tests/conversion/test_date_handlers.py | 14 +- tests/loaders/test_defaults.py | 4 +- tests/loaders/test_pipeline.py | 3 +- tests/mappers/test_materialised_view_mixin.py | 12 +- tests/tables/test_typing_protocols.py | 1 - 22 files changed, 179 insertions(+), 558 deletions(-) diff --git a/docs/loaders/conversion.md b/docs/loaders/conversion.md index 7506843..fe9c3f8 100644 --- a/docs/loaders/conversion.md +++ b/docs/loaders/conversion.md @@ -70,13 +70,11 @@ The handler for a column is the first match of: the entry in `ConversionPolicy.c | `FloatHandler` | `Float` | Decimal or exponent form; `Infinity` and values that overflow to infinity are errors | none | | `NumericHandler` | `Numeric` | Decimal form, converted to the declared precision and scale without rounding; a precision without a scale has scale 0 and a column without a precision reads as `decimal128(38, 18)` | none | | `BooleanHandler` | `Boolean` | `true`, `t`, `yes`, `y`, `1`, `on` and their opposites | `true_words`, `false_words` | -| `DateHandler` | `Date` | One of `formats`, validated against the calendar; year 0 is an error | `formats`, `parser` | -| `DateTimeHandler` | `DateTime` | ISO date and time, or a date in one of `formats`; see time zones below | `formats`, `parser` | +| `DateHandler` | `Date` | One of `formats`, validated against the calendar; year 0 is an error | `formats` | +| `DateTimeHandler` | `DateTime` | ISO date and time, or a date in one of `formats`; see time zones below | `formats` | Dates default to `yyyymmdd`, `dd-mon-yyyy` and `yyyy-mm-dd`. Ambiguous forms such as `02/01/2020` are never guessed; declare `DateHandler(formats=("dd/mm/yyyy",))` or `("mm/dd/yyyy",)`. A timestamp column reads date-only values with its own `formats`. -`DateHandler(parser=...)` replaces the date parser with a function that has the signature of `parse_dates`. - ### Time zones The column type decides, so no value is converted by a convention the schema did not ask for. diff --git a/src/orm_loader/backends/__init__.py b/src/orm_loader/backends/__init__.py index c076d13..c8e9cca 100644 --- a/src/orm_loader/backends/__init__.py +++ b/src/orm_loader/backends/__init__.py @@ -3,7 +3,6 @@ from .resolve import resolve_backend from .sqlite import SQLiteBackend from .base import ( - BackendCapabilities, DatabaseBackend, DedupeResult, Dialect, @@ -19,7 +18,6 @@ ) __all__ = [ - "BackendCapabilities", "ConcurrentRefreshNotEligibleError", "DatabaseBackend", "DedupeResult", diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 04dffd6..4f5cf24 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -5,8 +5,7 @@ from collections.abc import Generator, Iterable, Iterator, Mapping, Sequence from contextlib import contextmanager from dataclasses import dataclass -from functools import wraps -from typing import TYPE_CHECKING, Any, Callable, ClassVar, ParamSpec, Type, TypeVar, cast +from typing import TYPE_CHECKING, Any, ClassVar, Type import pyarrow as pa import sqlalchemy as sa @@ -27,7 +26,6 @@ if TYPE_CHECKING: from ..loaders.loader import LoaderContext - from ..mappers.materialised_view_contracts import MaterializedViewIndex from ..tables.typing import FileTableProtocol logger = logging.getLogger(__name__) @@ -57,21 +55,6 @@ def rows_after(self) -> int: return self.rows_before - self.rows_removed -@dataclass(frozen=True) -class BackendCapabilities: - """ - Capability flags exposed by a database backend. - - These defaults are intentionally conservative. Concrete backends should - opt into capabilities explicitly. - """ - - supports_unlogged_staging: bool = False - supports_fk_toggle: bool = False - supports_materialized_views: bool = False - supports_replace_with_kept_indexes: bool = False - - class ForeignKeyViolationError(RuntimeError): """Rows reference a row that does not exist in the referred table. @@ -130,42 +113,21 @@ def staging_schema_claim(physical_schema: str | None = None, *, reserved: bool = ) -P = ParamSpec("P") -R = TypeVar("R") - - -def requires_capability( - capability_name: str, - feature_name: str, -) -> Callable[[Callable[P, R]], Callable[P, R]]: - """Guard a concrete backend method with a capability requirement. - - Subclass overrides must apply this decorator themselves; Python does not - automatically retain decorators when a method is overridden. - """ - - def decorator(method: Callable[P, R]) -> Callable[P, R]: - @wraps(method) - def guarded(*args: P.args, **kwargs: P.kwargs) -> R: - # ParamSpec keeps the complete method signature, but type checkers - # cannot infer that args[0] is the backend instance here. - self = cast("DatabaseBackend", args[0]) - self._require_capability(capability_name, feature_name) - return method(*args, **kwargs) - - return guarded - - return decorator - - class DatabaseBackend(ABC): """ Abstract base class for database-specific loader behavior. This class defines the stable contract for future backend implementations without changing existing loader orchestration yet. + + Attributes + ---------- + supports_replace_with_kept_indexes + Whether a replace merge works while the target's indexes are kept. """ + supports_replace_with_kept_indexes: ClassVar[bool] = True + def __init__( self, *, @@ -258,15 +220,6 @@ def dialect(self) -> Dialect: def identifier_preparer(self) -> IdentifierPreparer: """The dialect-specific identifier preparer used to quote/escape SQL identifiers.""" - @property - @abstractmethod - def capabilities(self) -> BackendCapabilities: - """Capability flags supported by this backend.""" - - def supports_dialect(self, dialect: Dialect) -> bool: - """Return ``True`` when the backend handles the given dialect.""" - return self.dialect == dialect - @property def default_index_strategy(self) -> str: """Default index strategy used when callers request ``auto``.""" @@ -285,19 +238,6 @@ def resolve_index_strategy(self, index_strategy: str) -> str: return self.default_index_strategy return index_strategy - def _require_capability(self, capability_name: str, feature_name: str) -> None: - """ - Raise a clear error when a backend capability is not supported. - """ - if not hasattr(self.capabilities, capability_name): - raise AttributeError( - f"Unknown backend capability {capability_name!r} on {type(self.capabilities).__name__}" - ) - if not getattr(self.capabilities, capability_name): - raise NotImplementedError( - f"Backend '{self.name}' does not support {feature_name}" - ) - @contextmanager def _as_connection( self, @@ -948,90 +888,3 @@ def relaxed_foreign_keys( self.restore_fk_check(session, previous_state, mapper=mapper) logger.info(f"{scope}Foreign key checks restored.") - @requires_capability("supports_materialized_views", "materialized views") - def create_materialized_view( - self, - bind: "Engine | Connection", - name: str, - selectable: sa.sql.Select[Any], - *, - schema: str | None = None, - with_data: bool = True, - if_not_exists: bool = True, - ) -> None: - """Create a materialized view for the supplied selectable. - - *schema* is the view's already-resolved physical schema (or None for - the connection's own default); callers resolve which - schema_translate_map key to use before calling. - """ - raise NotImplementedError( - f"Backend '{self.name}' has not implemented create_materialized_view()" - ) - - @requires_capability("supports_materialized_views", "materialized views") - def refresh_materialized_view( - self, - bind: "Engine | Connection", - name: str, - *, - schema: str | None = None, - concurrently: bool = False, - declared_indexes: tuple["MaterializedViewIndex", ...] = (), - ) -> None: - """Refresh a materialized view. - - *schema* is the view's already-resolved physical schema, matching - ``create_materialized_view``. - - ``declared_indexes`` lets supporting backends validate a concurrent - refresh request without defining a second catalog-based eligibility - rule. Other backends may ignore it. - """ - raise NotImplementedError( - f"Backend '{self.name}' has not implemented refresh_materialized_view()" - ) - - @requires_capability("supports_materialized_views", "materialized views") - def drop_materialized_view( - self, - bind: "Engine | Connection", - name: str, - *, - schema: str | None = None, - if_exists: bool = True, - cascade: bool = False, - ) -> None: - """Drop a materialized view. - - *schema* is the view's already-resolved physical schema, matching - ``create_materialized_view``. This is deliberately non-abstract: the - default implementation requires the capability flag and then raises - ``NotImplementedError``. Older third-party backend subclasses need - no override to receive a clear error when they do not support this - operation. - """ - raise NotImplementedError( - f"Backend '{self.name}' has not implemented drop_materialized_view()" - ) - - @requires_capability("supports_materialized_views", "materialized views") - def create_materialized_view_index( - self, - bind: "Engine | Connection", - name: str, - index: "MaterializedViewIndex", - *, - schema: str | None = None, - if_not_exists: bool = True, - ) -> None: - """Create an index on a materialized view. - - *schema* is the view's already-resolved physical schema, matching - ``create_materialized_view``. This is deliberately non-abstract for - the same compatibility reason as :meth:`drop_materialized_view`. - """ - raise NotImplementedError( - f"Backend '{self.name}' has not implemented " - "create_materialized_view_index()" - ) diff --git a/src/orm_loader/backends/duckdb.py b/src/orm_loader/backends/duckdb.py index 6f9e56e..1816c13 100644 --- a/src/orm_loader/backends/duckdb.py +++ b/src/orm_loader/backends/duckdb.py @@ -13,7 +13,7 @@ ) from sqlalchemy.sql.compiler import IdentifierPreparer -from .base import BackendCapabilities, DatabaseBackend, DeferredForeignKeyViolationError, Dialect +from .base import DatabaseBackend, DeferredForeignKeyViolationError, Dialect if TYPE_CHECKING: from ..loaders.loader import LoaderContext @@ -21,6 +21,8 @@ class DuckDBBackend(DatabaseBackend): + supports_replace_with_kept_indexes = False + @staticmethod def staging_name_for_table(tablename: str) -> str: return f"_staging_{tablename}" @@ -43,15 +45,6 @@ def identifier_preparer(self) -> IdentifierPreparer: # the memory of a load. STAGING_COMMIT_EVERY: ClassVar[int] = 5 - @property - def capabilities(self) -> BackendCapabilities: - return BackendCapabilities( - supports_unlogged_staging=False, - supports_fk_toggle=False, - supports_materialized_views=False, - supports_replace_with_kept_indexes=False, - ) - def _unenforced_constraints( self, tables: tuple[sa.Table, ...] ) -> tuple[UnenforcedForeignKey, ...]: @@ -330,8 +323,8 @@ def merge_replace( DuckDB can't delete a still-referenced row, so replace is an UPDATE rather than delete-then-reinsert. Requires index_strategy to have - dropped indexes first; capabilities.supports_replace_with_kept_indexes - is False for this reason. + dropped indexes first; supports_replace_with_kept_indexes is False + for this reason. A self-referential FK can point an existing row at one this load hasn't inserted yet. Doing replace then insert as separate phases diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index 91773e1..2177340 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -16,13 +16,14 @@ from sqlalchemy.exc import OperationalError from sqlalchemy.sql.compiler import IdentifierPreparer +from ..mappers.materialised_view_mixin import MaterializedViewBackend from ..mappers.materialised_view_errors import ( ConcurrentRefreshNotEligibleError, MaterializationError, MaterializationFailure, MaterializationOperation, ) -from .base import BackendCapabilities, DatabaseBackend, requires_capability +from .base import DatabaseBackend if TYPE_CHECKING: from oa_configurator import ResolvedDatabase @@ -35,7 +36,7 @@ _VALID_PG_REPLICATION_ROLES = frozenset({"origin", "local", "replica"}) -class PostgresBackend(DatabaseBackend): +class PostgresBackend(DatabaseBackend, MaterializedViewBackend): def __init__( self, *, @@ -67,28 +68,6 @@ def dialect(self) -> Dialect: def identifier_preparer(self) -> IdentifierPreparer: return postgresql.dialect().identifier_preparer - @property - def capabilities(self) -> BackendCapabilities: - return BackendCapabilities( - supports_unlogged_staging=True, - supports_fk_toggle=True, - supports_materialized_views=True, - supports_replace_with_kept_indexes=True, - ) - - def evaluate_default( - self, - session: so.Session, - table_cls: type["FileTableProtocol"], - expression: sa.ColumnElement[Any], - column_type: Any, - ) -> Any: - """Evaluate on a separate connection, because the session's is busy with the COPY.""" - bind = session.get_bind(mapper=table_cls) - engine = bind.engine if isinstance(bind, sa.Connection) else bind - with engine.connect() as connection: - return connection.execute(sa.select(sa.cast(expression, column_type))).scalar_one() - def _staging_table(self, table_cls: type["FileTableProtocol"], columns: Sequence[sa.Column[Any]]) -> sa.Table: """Unlogged, with an identity column, which paginates the merge in O(N log N) without an index.""" return sa.Table( @@ -269,7 +248,6 @@ def _insert(select_: Any) -> sa.Insert: table_cls, session, batch_range, validate_deferred_fks=validate_deferred_fks ) - @requires_capability("supports_materialized_views", "materialized views") def create_materialized_view( self, bind: Engine | Connection, @@ -303,7 +281,6 @@ def create_materialized_view( ) ) from error - @requires_capability("supports_materialized_views", "materialized views") def refresh_materialized_view( self, bind: Engine | Connection, @@ -344,7 +321,6 @@ def refresh_materialized_view( ) ) from error - @requires_capability("supports_materialized_views", "materialized views") def drop_materialized_view( self, bind: Engine | Connection, @@ -374,7 +350,6 @@ def drop_materialized_view( ) ) from error - @requires_capability("supports_materialized_views", "materialized views") def create_materialized_view_index( self, bind: Engine | Connection, diff --git a/src/orm_loader/backends/sqlite.py b/src/orm_loader/backends/sqlite.py index a3028fc..77e6169 100644 --- a/src/orm_loader/backends/sqlite.py +++ b/src/orm_loader/backends/sqlite.py @@ -12,7 +12,7 @@ from sqlalchemy.dialects import sqlite as sqlite_dialect from sqlalchemy.sql.compiler import IdentifierPreparer -from .base import BackendCapabilities, DatabaseBackend, Dialect +from .base import DatabaseBackend, Dialect if TYPE_CHECKING: from oa_configurator import ResolvedDatabase @@ -79,15 +79,6 @@ def dialect(self) -> Dialect: def identifier_preparer(self) -> IdentifierPreparer: return sqlite_dialect.dialect().identifier_preparer - @property - def capabilities(self) -> BackendCapabilities: - return BackendCapabilities( - supports_unlogged_staging=False, - supports_fk_toggle=True, - supports_materialized_views=False, - supports_replace_with_kept_indexes=True, - ) - @property def default_index_strategy(self) -> str: return "keep" diff --git a/src/orm_loader/conversion/handlers/dates.py b/src/orm_loader/conversion/handlers/dates.py index 7e2aeee..8188152 100644 --- a/src/orm_loader/conversion/handlers/dates.py +++ b/src/orm_loader/conversion/handlers/dates.py @@ -1,9 +1,10 @@ -"""Handlers for date and timestamp columns, with the date formats and the default date parser.""" +"""Handlers for date and timestamp columns, with the date formats and the date parser.""" from __future__ import annotations -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass +from types import MappingProxyType from typing import Any, ClassVar import pyarrow as pa @@ -14,7 +15,6 @@ from .base import ConvertingHandler DateShape = Callable[[pa.Array], tuple[pa.Array, pa.Array]] -DateParser = Callable[[pa.Array, tuple[str, ...]], tuple[pa.Array, pa.Array]] class DateFormats: @@ -31,7 +31,7 @@ class DateFormats: MONTHS Lowercase month abbreviations used by ``dd-mon-yyyy``. SHAPES - Shape function by format name; the one place a format is registered. + Read-only shape function by format name. """ DEFAULT: ClassVar[tuple[str, ...]] = ("yyyymmdd", "dd-mon-yyyy", "yyyy-mm-dd") @@ -69,23 +69,15 @@ def mm_dd_yyyy(values: pa.Array) -> tuple[pa.Array, pa.Array]: cut = pc.utf8_slice_codeunits return ok, pc.binary_join_element_wise(cut(values, 6, 10), cut(values, 0, 2), cut(values, 3, 5), "-") - SHAPES: ClassVar[dict[str, DateShape]] = { - "yyyymmdd": yyyymmdd, - "yyyy-mm-dd": yyyy_mm_dd, - "dd-mon-yyyy": dd_mon_yyyy, - "dd/mm/yyyy": dd_mm_yyyy, - "mm/dd/yyyy": mm_dd_yyyy, - } - - @classmethod - def shape(cls, name: str) -> DateShape: - """Return the shape function for a format name.""" - return cls.SHAPES[name] - - @classmethod - def names(cls) -> tuple[str, ...]: - """Return every supported format name.""" - return tuple(cls.SHAPES) + SHAPES: ClassVar[Mapping[str, DateShape]] = MappingProxyType( + { + "yyyymmdd": yyyymmdd, + "yyyy-mm-dd": yyyy_mm_dd, + "dd-mon-yyyy": dd_mon_yyyy, + "dd/mm/yyyy": dd_mm_yyyy, + "mm/dd/yyyy": mm_dd_yyyy, + } + ) def parse_dates(values: pa.Array, formats: tuple[str, ...]) -> tuple[pa.Array, pa.Array]: @@ -100,7 +92,7 @@ def parse_dates(values: pa.Array, formats: tuple[str, ...]) -> tuple[pa.Array, p values String array with missing values already set to null. formats - Ordered names from ``DateFormats.names()``. + Ordered keys of ``DateFormats.SHAPES``. Returns ------- @@ -111,7 +103,7 @@ def parse_dates(values: pa.Array, formats: tuple[str, ...]) -> tuple[pa.Array, p iso = pa.nulls(len(values), pa.string()) matched = no_failures(len(values)) for name in formats: - ok, candidate = DateFormats.shape(name)(values) + ok, candidate = DateFormats.SHAPES[name](values) ok = pc.and_(pc.fill_null(ok, False), pc.is_valid(candidate)) ok = pc.and_(ok, pc.invert(matched)) iso = pc.if_else(ok, candidate, iso) @@ -132,26 +124,23 @@ class DateHandler(ConvertingHandler): Attributes ---------- formats - Ordered names from ``DateFormats.names()``. - parser - Replaceable function with the signature of ``parse_dates``. + Ordered keys of ``DateFormats.SHAPES``. """ distinct: ClassVar[bool] = True formats: tuple[str, ...] = DateFormats.DEFAULT - parser: DateParser = parse_dates def __post_init__(self) -> None: - unknown = [name for name in self.formats if name not in DateFormats.names()] + unknown = [name for name in self.formats if name not in DateFormats.SHAPES] if unknown: - raise ValueError(f"Unknown date formats {unknown}; choose from {sorted(DateFormats.names())}") + raise ValueError(f"Unknown date formats {unknown}; choose from {sorted(DateFormats.SHAPES)}") def arrow_type(self, sa_type: TypeEngine[Any]) -> pa.DataType: return pa.date32() def parse(self, cleaned: pa.Array, sa_type: TypeEngine[Any]) -> Interpretation: - return Interpretation(*self.parser(cleaned, self.formats)) + return Interpretation(*parse_dates(cleaned, self.formats)) def check(self, values: pa.Array, sa_type: TypeEngine[Any]) -> pa.Array | None: return pc.fill_null(pc.less(pc.year(values), 1), False) @@ -200,7 +189,7 @@ def parse(self, cleaned: pa.Array, sa_type: TypeEngine[Any]) -> Interpretation: if aware: return Interpretation(stamps, pc.or_(stamp_failed, pc.and_(pc.is_valid(cleaned), pc.invert(is_iso)))) rest = pc.dictionary_encode(pc.if_else(is_iso, NULL_STRING, cleaned)) - dates, invalid = self.parser(rest.dictionary, self.formats) + dates, invalid = parse_dates(rest.dictionary, self.formats) dates = pc.take(dates, rest.indices) invalid = pc.fill_null(pc.take(invalid, rest.indices), False) values = pc.coalesce(stamps, pc.cast(dates, target)) diff --git a/src/orm_loader/loaders/processor.py b/src/orm_loader/loaders/processor.py index 70331e2..4df1594 100644 --- a/src/orm_loader/loaders/processor.py +++ b/src/orm_loader/loaders/processor.py @@ -105,7 +105,8 @@ class BatchProcessor: rows read and rejected. evaluate Function that returns the value of a SQL expression default as a - column of the given type stores it. + column of the given type stores it. It is called once for each such + default when the processor is built. Attributes ---------- @@ -117,8 +118,7 @@ class BatchProcessor: ----- A missing value takes the column's default, and a column the file lacks is filled with it. The default is read by the column's handler like a - value from the file. A SQL expression default is evaluated for each batch - that needs it. A missing value left in a non-nullable column is a + value from the file. A missing value left in a non-nullable column is a required value that is missing. """ @@ -142,16 +142,16 @@ def __init__( self.normalise = normalise self.table_name = tableclass.__tablename__ model_columns = tableclass.model_columns() - self.defaults = tableclass.column_defaults() if normalise else {} - self.out_columns = [*columns, *(name for name in self.defaults if name not in columns)] + defaults = tableclass.column_defaults() if normalise else {} + self.out_columns = [*columns, *(name for name in defaults if name not in columns)] self.types = {name: underlying_type(model_columns[name].type) for name in self.out_columns} self.handlers = ( {name: policy.for_column(self.types[name], self.table_name, name) for name in self.out_columns} if normalise else {} ) + self.defaults = {name: self._evaluated(name, default, evaluate) for name, default in defaults.items()} self.required = [name for name in self.out_columns if not model_columns[name].nullable] - self.evaluate = evaluate self._literals: dict[str, pa.Scalar] = {} self.rejects = RejectsWriter(rejects_path) if policy.errors.on_error == "reject_row": @@ -222,12 +222,16 @@ def _fill_defaults(self, results: dict[str, Interpretation], count: int) -> None filled = pc.if_else(missing, self._fill_value(name, default), values) results[name] = Interpretation(filled, invalid) + def _evaluated(self, name: str, default: Any, evaluate: Callable[[Any, Any], Any] | None) -> Any: + """Return a literal default as it is and the value of a SQL expression default.""" + if not isinstance(default, sa.ClauseElement): + return default + if evaluate is None: + raise PolicyError(table_name=self.table_name, column_name=name, column_type="SQL expression default") + return evaluate(default, self.types[name]) + def _fill_value(self, name: str, default: Any) -> pa.Scalar: - """Return the scalar that fills a column: a literal is read once, an expression is evaluated for each batch.""" - if isinstance(default, sa.ClauseElement): - if self.evaluate is None: - raise PolicyError(table_name=self.table_name, column_name=name, column_type="SQL expression default") - return self._read_default(name, self.evaluate(default, self.types[name])) + """Return the scalar that fills a column, read through the column's handler once.""" if name not in self._literals: self._literals[name] = self._read_default(name, default) return self._literals[name] diff --git a/src/orm_loader/mappers/__init__.py b/src/orm_loader/mappers/__init__.py index 1fd7352..fe637e3 100644 --- a/src/orm_loader/mappers/__init__.py +++ b/src/orm_loader/mappers/__init__.py @@ -16,6 +16,7 @@ ) from .materialised_view_mixin import ( CreateMaterializedView, + MaterializedViewBackend, MaterializedViewMixin, refresh_all_mvs, resolve_mv_refresh_order, @@ -35,6 +36,7 @@ "MaterializationError", "MaterializationFailure", "MaterializationOperation", + "MaterializedViewBackend", "MaterializedViewIndex", "MaterializedViewMixin", "refresh_all_mvs", diff --git a/src/orm_loader/mappers/materialised_view_mixin.py b/src/orm_loader/mappers/materialised_view_mixin.py index 39a2447..19d5938 100644 --- a/src/orm_loader/mappers/materialised_view_mixin.py +++ b/src/orm_loader/mappers/materialised_view_mixin.py @@ -1,3 +1,4 @@ +from abc import ABC, abstractmethod from collections import defaultdict, deque from typing import Any @@ -77,6 +78,69 @@ def _create_view( population = "" if element.with_data else " WITH NO DATA" return f"CREATE MATERIALIZED VIEW {existence}{element.name} as {compiled}{population}" + +class MaterializedViewBackend(ABC): + """Interface of a database backend that supports materialized views. + + A backend without materialized views does not inherit this class. + ``schema`` arguments are the view's already-resolved physical schema, or + None for the connection's own default. + """ + + @abstractmethod + def create_materialized_view( + self, + bind: "sa.engine.Engine | sa.engine.Connection", + name: str, + selectable: sa.sql.Select[Any], + *, + schema: str | None = None, + with_data: bool = True, + if_not_exists: bool = True, + ) -> None: + """Create a materialized view for the supplied selectable.""" + + @abstractmethod + def refresh_materialized_view( + self, + bind: "sa.engine.Engine | sa.engine.Connection", + name: str, + *, + schema: str | None = None, + concurrently: bool = False, + declared_indexes: tuple[MaterializedViewIndex, ...] = (), + ) -> None: + """Refresh a materialized view. + + ``declared_indexes`` lets the backend validate a concurrent refresh + request against the view's declared unique index. + """ + + @abstractmethod + def drop_materialized_view( + self, + bind: "sa.engine.Engine | sa.engine.Connection", + name: str, + *, + schema: str | None = None, + if_exists: bool = True, + cascade: bool = False, + ) -> None: + """Drop a materialized view.""" + + @abstractmethod + def create_materialized_view_index( + self, + bind: "sa.engine.Engine | sa.engine.Connection", + name: str, + index: MaterializedViewIndex, + *, + schema: str | None = None, + if_not_exists: bool = True, + ) -> None: + """Create an index on a materialized view.""" + + class MaterializedViewMixin: """ @@ -203,6 +267,16 @@ def _resolve_schema_tag(cls, schema_tag: str | None) -> str | None: return table.schema return cls.__mv_schema_tag__ + @staticmethod + def _materialized_view_backend(bind: "sa.engine.Connection | sa.engine.Engine") -> MaterializedViewBackend: + """Return the bind's backend, or raise when it has no materialized views.""" + from ..backends.resolve import resolve_backend + + backend = resolve_backend(bind) + if not isinstance(backend, MaterializedViewBackend): + raise NotImplementedError(f"Backend '{backend.name}' does not support materialized views") + return backend + @classmethod def create_mv( cls, @@ -261,9 +335,7 @@ def create_mv( if not isinstance(bind, (sa.engine.Engine, sa.engine.Connection)): raise TypeError("bind must be a SQLAlchemy Engine or Connection") - from ..backends.resolve import resolve_backend - - backend = resolve_backend(bind) + backend = cls._materialized_view_backend(bind) tag = cls._resolve_schema_tag(schema_tag) schema = physical_schema_of(bind, schema_tag=tag) @@ -323,9 +395,7 @@ def refresh_mv( RecentObservationMV.refresh_mv(conn) ``` """ - from ..backends.resolve import resolve_backend - - backend = resolve_backend(bind) + backend = cls._materialized_view_backend(bind) tag = cls._resolve_schema_tag(schema_tag) backend.refresh_materialized_view( bind, @@ -354,9 +424,7 @@ def drop_mv( schema_translate_map key the view's own physical schema resolves through; see :meth:`_resolve_schema_tag` for the default. """ - from ..backends.resolve import resolve_backend - - backend = resolve_backend(bind) + backend = cls._materialized_view_backend(bind) tag = cls._resolve_schema_tag(schema_tag) backend.drop_materialized_view( bind, cls.__mv_name__, schema=physical_schema_of(bind, schema_tag=tag), diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index 17d2ec5..77351fd 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -59,39 +59,6 @@ class FileLoadableTableInterface(ORMTableBase): __abstract__ = True - @classmethod - def create_staging_table( - cls: Type[FileTableProtocol], - session: so.Session, - *, - staging_schema_tag: str | None = None, - ): - """ - Create a fresh staging table for ingestion. - - Any existing staging table with the same name is dropped first. - The staging table has the target's columns with their types and - nullability. - - Parameters - ---------- - session - An active SQLAlchemy session bound to an engine. - staging_schema_tag - schema_translate_map tag the staging table should be created - under. ``None`` means no schema qualification (backend-default - behavior). - Raises - ------ - RuntimeError - If the session is not bound to an engine. - NotImplementedError - If the database dialect is unsupported. - """ - _require_bind(session, cls) - backend = resolve_backend(session, staging_schema_tag=staging_schema_tag, mapper=cls) - backend.create_staging_table(cls, session) - @classmethod @contextmanager def manage_indices( @@ -279,33 +246,6 @@ def get_staging_table( schema=backend.staging_schema, ) - @classmethod - def load_staging( - cls: Type[FileTableProtocol], - loader_context: LoaderContext, - reader: FileReader, - ) -> int: - """ - Create the staging table and load the file into it. - - Parameters - ---------- - loader_context - Context object containing session, path, and load options. - reader - Reader that yields the file as Arrow batches. - - Returns - ------- - int - Number of rows loaded into the staging table. - """ - session, tag = loader_context.session, loader_context.staging_schema_tag - _require_bind(session, cls) - cls.create_staging_table(session, staging_schema_tag=tag) - staging = cls.get_staging_table(session, staging_schema_tag=tag) - return stage_file(loader_context, reader, staging) - @classmethod def load_file( cls: Type[FileTableProtocol], @@ -417,6 +357,7 @@ def load_file( f"File name '{path.name}' does not match table '{cls.__tablename__}'" ) + _require_bind(session, cls) if reader is not None and (chunksize is not None or quote_mode is not None): raise ValueError("chunksize and quote_mode configure the default reader; set them on the reader passed as `reader`") policy = policy or ConversionPolicy() @@ -437,12 +378,14 @@ def load_file( resolved=resolved, physical_schemas=physical_schemas, ) - if merge_strategy == "replace": - if backend.resolve_index_strategy(index_strategy) == "keep": - backend._require_capability( - "supports_replace_with_kept_indexes", - "merge_strategy='replace' with index_strategy='keep'", - ) + if ( + merge_strategy == "replace" + and not backend.supports_replace_with_kept_indexes + and backend.resolve_index_strategy(index_strategy) == "keep" + ): + raise NotImplementedError( + f"Backend '{backend.name}' does not support merge_strategy='replace' with index_strategy='keep'" + ) loader_context = LoaderContext( tableclass=cls, @@ -460,7 +403,9 @@ def load_file( # SQLite cannot switch foreign key checks inside a transaction, so relax before staging writes. with backend.relaxed_foreign_keys(session, mapper=cls, no_autoflush=False): logger.info(f"Table `{cls.__tablename__}`: Loading data into staging table") - total = cls.load_staging(loader_context, reader) + backend.create_staging_table(cls, session) + staging = cls.get_staging_table(session, staging_schema_tag=staging_schema_tag) + total = stage_file(loader_context, reader, staging) logger.info(f"Table `{cls.__tablename__}`: Merging staging data into target table") with cls.manage_indices(session, index_strategy=index_strategy, staging_schema_tag=staging_schema_tag): @@ -477,7 +422,7 @@ def load_file( logger.info(f"Table `{cls.__tablename__}`: Validating deferred foreign keys.") backend.validate_deferred_fks(cls, session) finally: - cls.drop_staging_table(session, staging_schema_tag=staging_schema_tag) + backend.drop_staging_table(cls, session) logger.info(f"Table `{cls.__tablename__}`: Successfully finished ingestion. Total rows: {total}") return total @@ -634,24 +579,3 @@ def merge_from_staging( ) else: raise ValueError(f"Unknown merge strategy '{merge_strategy}'") - - @classmethod - def drop_staging_table( - cls: Type[FileTableProtocol], - session: so.Session, - *, - staging_schema_tag: str | None = None, - ): - """ - Drop the staging table if it exists. - - Parameters - ---------- - session - An active SQLAlchemy session bound to an engine. - staging_schema_tag - schema_translate_map tag the staging table lives in. ``None`` - means no schema qualification (backend-default behavior). - """ - backend = resolve_backend(session, staging_schema_tag=staging_schema_tag, mapper=cls) - backend.drop_staging_table(cls, session) diff --git a/src/orm_loader/tables/typing.py b/src/orm_loader/tables/typing.py index 45d179f..e9173ab 100644 --- a/src/orm_loader/tables/typing.py +++ b/src/orm_loader/tables/typing.py @@ -1,4 +1,4 @@ -from typing import Protocol, ClassVar, runtime_checkable, TYPE_CHECKING, Optional, Type, Dict, Any, Unpack, TypedDict +from typing import Protocol, ClassVar, runtime_checkable, TYPE_CHECKING, Optional, Dict, Any, Unpack, TypedDict import sqlalchemy.orm as so import sqlalchemy as sa from collections.abc import Mapping @@ -8,7 +8,7 @@ from oa_configurator import ResolvedDatabase from ..conversion import ConversionPolicy - from ..loaders import FileReader, LoadReport, LoaderContext + from ..loaders import FileReader, LoadReport class ToDictKwargs(TypedDict, total=False): include_nulls: bool @@ -72,14 +72,6 @@ class FileTableProtocol(ORMTableProtocol, Protocol): No assumptions are made about the underlying database or schema. """ - @classmethod - def create_staging_table( - cls, session: so.Session, *, staging_schema_tag: str | None = None - ) -> None: ... - - @classmethod - def load_staging(cls: Type["FileTableProtocol"], loader_context: "LoaderContext", reader: "FileReader") -> int: ... - @classmethod def load_file( cls, @@ -119,9 +111,6 @@ def merge_from_staging( physical_schemas: Optional[Mapping[str, str | None]] = None, ) -> None: ... - @classmethod - def drop_staging_table(cls, session: so.Session, *, staging_schema_tag: str | None = None) -> None: ... - @classmethod def validate_foreign_keys( cls, diff --git a/tests/backends/test_base_backend.py b/tests/backends/test_base_backend.py index c7616c7..ed73130 100644 --- a/tests/backends/test_base_backend.py +++ b/tests/backends/test_base_backend.py @@ -10,10 +10,9 @@ import pytest import sqlalchemy as sa import sqlalchemy.orm as so -from sqlalchemy.engine import Connection, Engine +from sqlalchemy.engine import Engine from orm_loader.backends import ( - BackendCapabilities, DatabaseBackend, Dialect, resolve_backend, @@ -60,12 +59,6 @@ def identifier_preparer(self): return sqlite.dialect().identifier_preparer - @property - def capabilities(self) -> BackendCapabilities: - return BackendCapabilities( - supports_fk_toggle=True, - ) - @staticmethod def staging_name_for_table(tablename: str) -> str: return f"_staging_{tablename}" @@ -130,29 +123,6 @@ def restore_fk_check( ) -> None: self.calls.append(("restore_fk_check", previous_state)) - def create_materialized_view( - self, - bind: Engine | Connection, - name: str, - selectable: sa.sql.Select[Any], - *, - schema: str | None = None, - with_data: bool = True, - if_not_exists: bool = True, - ) -> None: - return None - - def refresh_materialized_view( - self, - bind: Engine | Connection, - name: str, - *, - schema: str | None = None, - concurrently: bool = False, - declared_indexes: tuple = (), - ) -> None: - return None - class _ComputedTable: __table__ = sa.Table( @@ -167,14 +137,6 @@ class _ComputedTable: _ComputedTableCls = cast("Type[FileTableProtocol]", _ComputedTable) -def test_backend_capabilities_defaults(): - caps = BackendCapabilities() - - assert caps.supports_unlogged_staging is False - assert caps.supports_fk_toggle is False - assert caps.supports_materialized_views is False - - def test_database_backend_is_abstract(): with pytest.raises(TypeError): DatabaseBackend() # type: ignore @@ -185,33 +147,11 @@ def test_fake_backend_can_implement_contract(): assert backend.name == "fake" assert backend.dialect == Dialect.SQLITE - assert backend.capabilities.supports_fk_toggle is True - assert backend.supports_dialect(Dialect.SQLITE) is True - assert backend.supports_dialect(Dialect.POSTGRESQL) is False + assert backend.supports_replace_with_kept_indexes is True assert backend.resolve_index_strategy("auto") == "drop_rebuild" assert backend.resolve_index_strategy("keep") == "keep" -def test_require_capability_passes_for_supported_feature(): - backend = FakeBackend() - - backend._require_capability("supports_fk_toggle", "foreign key toggling") - - -def test_require_capability_raises_for_unsupported_feature(): - backend = FakeBackend() - - with pytest.raises(NotImplementedError, match="does not support materialized views"): - backend._require_capability("supports_materialized_views", "materialized views") - - -def test_require_capability_raises_for_unknown_flag(): - backend = FakeBackend() - - with pytest.raises(AttributeError, match="Unknown backend capability"): - backend._require_capability("not_a_capability", "something") - - def test_resolve_index_strategy_raises_for_invalid_value(): backend = FakeBackend() @@ -267,7 +207,6 @@ def test_backends_package_exports(): import orm_loader.backends as backends assert backends.DatabaseBackend is DatabaseBackend - assert backends.BackendCapabilities is BackendCapabilities assert backends.Dialect is Dialect assert backends.resolve_backend is resolve_backend @@ -294,55 +233,6 @@ class dialect: with pytest.raises(NotImplementedError, match="Unsupported SQLAlchemy dialect"): resolve_backend(cast(Engine, _Unknown())) - -def test_unsupported_fake_backend_inherits_default_materialized_view_guards(): - backend = FakeBackend() - - with pytest.raises(NotImplementedError, match="does not support materialized views"): - backend.drop_materialized_view(cast(Engine, None), "mv_test") - - with pytest.raises(NotImplementedError, match="does not support materialized views"): - backend.create_materialized_view_index(cast(Engine, None), "mv_test", cast(Any, None)) - - -def test_capable_backend_inheriting_default_materialized_view_methods_fails_loudly(): - class CapableBackend(FakeBackend): - create_materialized_view = DatabaseBackend.create_materialized_view - refresh_materialized_view = DatabaseBackend.refresh_materialized_view - - @property - def capabilities(self) -> BackendCapabilities: - return BackendCapabilities( - supports_fk_toggle=True, - supports_materialized_views=True, - ) - - backend = CapableBackend() - - with pytest.raises(NotImplementedError, match="has not implemented drop_materialized_view"): - backend.drop_materialized_view(cast(Engine, None), "mv_test") - - with pytest.raises( - NotImplementedError, - match="has not implemented create_materialized_view", - ): - backend.create_materialized_view( - cast(Engine, None), "mv_test", cast(Any, None) - ) - - with pytest.raises( - NotImplementedError, - match="has not implemented refresh_materialized_view", - ): - backend.refresh_materialized_view(cast(Engine, None), "mv_test") - - with pytest.raises( - NotImplementedError, - match="has not implemented create_materialized_view_index", - ): - backend.create_materialized_view_index(cast(Engine, None), "mv_test", cast(Any, None)) - - def test_backends_import_does_not_require_psycopg(): blocker = _BlockPsycopg() original = sys.modules.pop("orm_loader.backends", None) diff --git a/tests/backends/test_duckdb_backend.py b/tests/backends/test_duckdb_backend.py index 89e747c..bf9d319 100644 --- a/tests/backends/test_duckdb_backend.py +++ b/tests/backends/test_duckdb_backend.py @@ -9,6 +9,7 @@ import sqlalchemy.orm as so from oa_configurator import CDMDatabaseConfig, ConnectionConfig, Resolver, StackConfig +from orm_loader.mappers import MaterializedViewBackend from orm_loader.backends import Dialect, DuckDBBackend, staging_schema_claim from orm_loader.backends.base import DeferredForeignKeyViolationError from orm_loader.helpers.bootstrap import create_tables @@ -159,11 +160,7 @@ def test_duckdb_backend_identity_and_capabilities(): assert backend.name == "duckdb" assert backend.dialect == Dialect.DUCKDB - assert backend.supports_dialect(Dialect.DUCKDB) is True - assert backend.capabilities.supports_unlogged_staging is False - assert backend.capabilities.supports_fk_toggle is False - assert backend.capabilities.supports_materialized_views is False - assert backend.capabilities.supports_replace_with_kept_indexes is False + assert backend.supports_replace_with_kept_indexes is False assert backend.resolve_index_strategy("auto") == "drop_rebuild" @@ -283,19 +280,8 @@ def test_relaxed_foreign_keys_on_duckdb_does_not_log_toggling(duckdb_session, ca assert "foreign key" not in caplog.text.lower() -@pytest.mark.duckdb -@pytest.mark.db_dialect -def test_duckdb_backend_materialized_view_methods_raise(duckdb_session): - backend = DuckDBBackend() - selectable = sa.select(sa.literal(1).label("n")) - engine = duckdb_session.get_bind() - - with pytest.raises(NotImplementedError, match="does not support materialized views"): - backend.create_materialized_view(engine, "mv_test", selectable) - with pytest.raises(NotImplementedError, match="does not support materialized views"): - backend.refresh_materialized_view(engine, "mv_test") - with pytest.raises(NotImplementedError, match="does not support materialized views"): - backend.drop_materialized_view(engine, "mv_test") +def test_duckdb_backend_has_no_materialized_views(): + assert not isinstance(DuckDBBackend(), MaterializedViewBackend) @pytest.mark.duckdb diff --git a/tests/backends/test_postgres_backend.py b/tests/backends/test_postgres_backend.py index e7dbbfc..a6fd6d9 100644 --- a/tests/backends/test_postgres_backend.py +++ b/tests/backends/test_postgres_backend.py @@ -82,10 +82,7 @@ def test_postgres_backend_identity_and_capabilities(): assert backend.name == "postgres" assert backend.dialect == Dialect.POSTGRESQL - assert backend.supports_dialect(Dialect.POSTGRESQL) is True - assert backend.capabilities.supports_unlogged_staging is True - assert backend.capabilities.supports_fk_toggle is True - assert backend.capabilities.supports_materialized_views is True + assert backend.supports_replace_with_kept_indexes is True @pytest.mark.parametrize("merge_batch_size", [0, -1]) @@ -562,7 +559,7 @@ def _simple_rows(session) -> list[tuple[int, str]]: def test_staging_table_has_the_target_columns_and_an_identity_row_number(pg_session): - SimpleTable.create_staging_table(pg_session, staging_schema_tag=STAGING_SCHEMA) + resolve_backend(pg_session, staging_schema_tag=STAGING_SCHEMA, mapper=SimpleTable).create_staging_table(SimpleTable, pg_session) staging_name = resolve_backend(pg_session).staging_name_for_table(SimpleTable.__tablename__) cols = pg_session.execute(sa.text( "SELECT column_name, data_type FROM information_schema.columns" diff --git a/tests/backends/test_sinks.py b/tests/backends/test_sinks.py index 847b198..0abf09c 100644 --- a/tests/backends/test_sinks.py +++ b/tests/backends/test_sinks.py @@ -18,7 +18,7 @@ def _sink(session, tmp_path): tag = None if session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA - _Demo.create_staging_table(session, staging_schema_tag=tag) + resolve_backend(session, staging_schema_tag=tag, mapper=_Demo).create_staging_table(_Demo, session) staging = _Demo.get_staging_table(session, staging_schema_tag=tag) ctx = LoaderContext(tableclass=_Demo, session=session, path=tmp_path / "demo_record.csv", staging_schema_tag=tag) return resolve_backend(session, staging_schema_tag=tag, mapper=_Demo), ctx, staging @@ -87,7 +87,7 @@ def batches(): def _staging_for(session, table): tag = None if session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA - table.create_staging_table(session, staging_schema_tag=tag) + resolve_backend(session, staging_schema_tag=tag, mapper=table).create_staging_table(table, session) return table.get_staging_table(session, staging_schema_tag=tag) @@ -113,7 +113,7 @@ def test_a_staging_table_can_be_dropped_twice_and_created_again(dialect_session) staging = _staging_for(dialect_session, _Demo) tag = None if dialect_session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA for _ in range(2): - _Demo.drop_staging_table(dialect_session, staging_schema_tag=tag) + resolve_backend(dialect_session, staging_schema_tag=tag, mapper=_Demo).drop_staging_table(_Demo, dialect_session) with pytest.raises(sa.exc.NoSuchTableError): _Demo.get_staging_table(dialect_session, staging_schema_tag=tag) assert _staging_for(dialect_session, _Demo).name == staging.name diff --git a/tests/backends/test_sqlite_backend.py b/tests/backends/test_sqlite_backend.py index 88384e0..f7a8a1f 100644 --- a/tests/backends/test_sqlite_backend.py +++ b/tests/backends/test_sqlite_backend.py @@ -8,6 +8,7 @@ from sqlalchemy.dialects import sqlite as sqlite_dialect from orm_loader.backends import Dialect, SQLiteBackend +from orm_loader.mappers import MaterializedViewBackend from tests.models import ComputedColumnTable, ReferencingChildTable, SimpleTable if TYPE_CHECKING: @@ -51,10 +52,7 @@ def test_sqlite_backend_identity_and_capabilities(): assert backend.name == "sqlite" assert backend.dialect == Dialect.SQLITE - assert backend.supports_dialect(Dialect.SQLITE) is True - assert backend.capabilities.supports_unlogged_staging is False - assert backend.capabilities.supports_fk_toggle is True - assert backend.capabilities.supports_materialized_views is False + assert backend.supports_replace_with_kept_indexes is True assert backend.resolve_index_strategy("auto") == "keep" @@ -145,40 +143,8 @@ def test_sqlite_backend_normalize_fk_check_state(): raise AssertionError("Expected ValueError for unrecognised string") -def test_sqlite_backend_materialized_view_methods_raise(sqlite_engine): - backend = SQLiteBackend() - selectable = sa.select(sa.literal(1).label("n")) - - try: - backend.create_materialized_view(sqlite_engine, "mv_test", selectable) - except NotImplementedError as exc: - assert "does not support materialized views" in str(exc) - else: - raise AssertionError("Expected create_materialized_view() to raise NotImplementedError") - - try: - backend.refresh_materialized_view(sqlite_engine, "mv_test") - except NotImplementedError as exc: - assert "does not support materialized views" in str(exc) - else: - raise AssertionError("Expected refresh_materialized_view() to raise NotImplementedError") - - -def test_sqlite_backend_drop_materialized_view_raises(sqlite_engine): - backend = SQLiteBackend() - - with pytest.raises(NotImplementedError, match="does not support materialized views"): - backend.drop_materialized_view(sqlite_engine, "mv_test") - - -def test_sqlite_backend_create_materialized_view_index_raises(sqlite_engine): - from orm_loader.mappers.materialised_view_contracts import MaterializedViewIndex - - backend = SQLiteBackend() - index = MaterializedViewIndex(name="mv_test_uq", columns=("id",), unique=True) - - with pytest.raises(NotImplementedError, match="does not support materialized views"): - backend.create_materialized_view_index(sqlite_engine, "mv_test", index) +def test_sqlite_backend_has_no_materialized_views(): + assert not isinstance(SQLiteBackend(), MaterializedViewBackend) def test_sqlite_backend_disable_fk_raises_when_pragma_returns_non_int(): diff --git a/tests/conversion/test_date_handlers.py b/tests/conversion/test_date_handlers.py index c32f427..1678b9c 100644 --- a/tests/conversion/test_date_handlers.py +++ b/tests/conversion/test_date_handlers.py @@ -49,18 +49,6 @@ def test_unknown_date_format_name_is_rejected(): DateHandler(formats=("yyyy/dd/mm",)) -def test_date_parser_is_replaceable(): - calls = [] - - def parser(values, formats): - calls.append(len(values)) - return pa.array([date(2000, 1, 1)] * len(values), pa.date32()), pa.array([False] * len(values)) - - values, _ = interpret(["a", "a", "b"], sa.Date(), ConversionPolicy(date=DateHandler(parser=parser))) - assert values == [date(2000, 1, 1)] * 3 - assert calls == [2] - - def test_dictionary_encoded_dates_parse_the_same(): column = pc.dictionary_encode(pa.array(["20200229", None, "", "20200301", "20200229"])) assert interpret(column, sa.Date())[0] == [ @@ -131,7 +119,7 @@ def test_arrow_strptime_still_rolls_impossible_dates_over(): def test_default_formats_are_mutually_distinguishable(): probes = ["20200229", "29-FEB-2020", "2020-02-29"] for name in DateFormats.DEFAULT: - matches = [bool(DateFormats.shape(name)(pa.array([p]))[0][0].as_py()) for p in probes] + matches = [bool(DateFormats.SHAPES[name](pa.array([p]))[0][0].as_py()) for p in probes] assert sum(matches) == 1 diff --git a/tests/loaders/test_defaults.py b/tests/loaders/test_defaults.py index b9755a5..f5d8054 100644 --- a/tests/loaders/test_defaults.py +++ b/tests/loaders/test_defaults.py @@ -109,7 +109,7 @@ def test_a_key_in_the_file_is_loaded_and_may_not_be_blank(dialect_session, tmp_p _keyed(dialect_session, tmp_path, "keyed_id,keyed_name\n,a\n") -def test_an_expression_default_is_evaluated_for_each_batch_that_needs_it(dialect_session, tmp_path, monkeypatch): +def test_an_expression_default_is_evaluated_once_for_the_load(dialect_session, tmp_path, monkeypatch): from orm_loader.backends import resolve_backend from orm_loader.loaders import DelimitedTextReader @@ -130,7 +130,7 @@ def counting(self, session, table_cls, expression, column_type): table.load_file(dialect_session, path, staging_schema_tag=tag, reader=reader) batches = len(list(DelimitedTextReader(path, DemoDefaultKinds.model_columns(), block_size=64 << 10).batches())) now_calls = [c for c in calls if isinstance(c, sa.sql.functions.FunctionElement)] - assert batches > 1 and len(now_calls) == batches + assert batches > 1 and len(now_calls) == 1 def test_a_default_the_handler_cannot_read_is_an_error_naming_the_column(dialect_session, tmp_path): diff --git a/tests/loaders/test_pipeline.py b/tests/loaders/test_pipeline.py index fd3bd84..e086ef7 100644 --- a/tests/loaders/test_pipeline.py +++ b/tests/loaders/test_pipeline.py @@ -12,6 +12,7 @@ import pytest import sqlalchemy as sa +from orm_loader.backends import resolve_backend from orm_loader.backends import STAGING_SCHEMA from orm_loader.conversion import ConversionPolicy, ErrorPolicy from orm_loader.helpers import ConversionError, RequiredValueError @@ -46,7 +47,7 @@ def _write(tmp_path: Path, table, fmt: str, rows: list[tuple], header: list[str] def _stage(session, table, path: Path, reader=None, **context): tag = None if session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA - table.create_staging_table(session, staging_schema_tag=tag) + resolve_backend(session, staging_schema_tag=tag, mapper=table).create_staging_table(table, session) staging = table.get_staging_table(session, staging_schema_tag=tag) ctx = LoaderContext(tableclass=table, session=session, path=path, staging_schema_tag=tag, **context) loaded = stage_file(ctx, reader or open_reader(path, table.model_columns(), nulls=ctx.policy.nulls), staging) diff --git a/tests/mappers/test_materialised_view_mixin.py b/tests/mappers/test_materialised_view_mixin.py index e4cf523..afb8df5 100644 --- a/tests/mappers/test_materialised_view_mixin.py +++ b/tests/mappers/test_materialised_view_mixin.py @@ -11,6 +11,7 @@ from oa_configurator.testing import scoped_test_schema from orm_loader.mappers.materialised_view_contracts import MaterializedViewIndex from orm_loader.mappers.materialised_view_mixin import ( + MaterializedViewBackend, MaterializedViewMixin, refresh_all_mvs, resolve_mv_refresh_order, @@ -96,7 +97,7 @@ def test_refresh_mv_concurrently_against_postgres(pg_db) -> None: _IndexedPgMv.refresh_mv(conn, concurrently=True) -class _FakeBackend: +class _FakeBackend(MaterializedViewBackend): """Capture lifecycle calls without touching a real database.""" def __init__(self) -> None: @@ -417,3 +418,12 @@ def test_refresh_all_mvs_refreshes_in_dependency_order(fake_backend: _FakeBacken refreshed_names = [call[1][1] for call in fake_backend.calls] assert refreshed_names == ["mv_a", "mv_b"] + + +def test_a_backend_without_materialized_views_is_refused(bind): + with pytest.raises(NotImplementedError, match="Backend 'sqlite' does not support materialized views"): + _NoIndexMv.create_mv(bind) + with pytest.raises(NotImplementedError, match="does not support materialized views"): + _NoIndexMv.refresh_mv(bind) + with pytest.raises(NotImplementedError, match="does not support materialized views"): + _NoIndexMv.drop_mv(bind) diff --git a/tests/tables/test_typing_protocols.py b/tests/tables/test_typing_protocols.py index 12ce537..276772f 100644 --- a/tests/tables/test_typing_protocols.py +++ b/tests/tables/test_typing_protocols.py @@ -46,5 +46,4 @@ class P(FileLoadableTableInterface, ORMTableBase, Base): assert isinstance(P, FileTableProtocol) # Required class-level API assert callable(P.load_file) - assert callable(P.create_staging_table) assert callable(P.merge_from_staging) \ No newline at end of file From f49285f1dcb4b46a5d2e8daf10770114c0eb9f14 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Fri, 30 Oct 2026 02:32:33 +0000 Subject: [PATCH 14/21] Resolve the backend once per load, pass the staging table to the merges, let stage_file roll back --- docs/loaders/context.md | 1 + src/orm_loader/backends/base.py | 11 ++- src/orm_loader/backends/duckdb.py | 6 +- src/orm_loader/backends/postgres.py | 9 +- src/orm_loader/backends/sqlite.py | 6 +- src/orm_loader/loaders/loader.py | 46 +++++----- src/orm_loader/tables/loadable_table.py | 95 +++++---------------- src/orm_loader/tables/typing.py | 11 +-- tests/backends/test_duckdb_backend.py | 107 ++++++++---------------- tests/backends/test_postgres_backend.py | 1 + tests/backends/test_shared_backend.py | 42 ++++------ tests/backends/test_sinks.py | 15 ++-- tests/backends/test_sqlite_backend.py | 1 + tests/loaders/test_loader_e2e.py | 14 ++++ tests/loaders/test_pipeline.py | 24 ++++-- tests/loaders/test_stage_file.py | 7 ++ tests/models.py | 2 +- 17 files changed, 168 insertions(+), 230 deletions(-) diff --git a/docs/loaders/context.md b/docs/loaders/context.md index 56e7752..60658a8 100644 --- a/docs/loaders/context.md +++ b/docs/loaders/context.md @@ -21,6 +21,7 @@ on globals or implicit configuration. | `tableclass` | Target ORM table class | | `session` | Active SQLAlchemy session | | `path` | Path to the input file | +| `backend` | The session's database backend, resolved once for the load | | `normalise` | Whether to interpret values under the policy; when false, values are written as read | | `dedupe` | Whether the staging table keeps only the first row for each primary key | | `policy` | `ConversionPolicy`: how values are interpreted and what happens to unreadable ones | diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 4f5cf24..0d505bf 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -482,7 +482,7 @@ def create_staging_table( self, table_cls: Type["FileTableProtocol"], session: so.Session, - ) -> None: + ) -> sa.Table: """ Create the staging table for a table class, replacing an existing one. @@ -497,6 +497,11 @@ def create_staging_table( Table the staging table is for. session Session whose connection creates the table. + + Returns + ------- + sqlalchemy.Table + The staging table as created. """ self.drop_staging_table(table_cls, session) columns = [ @@ -507,6 +512,7 @@ def create_staging_table( staging = self._staging_table(table_cls, columns) staging.create(bind=session.connection(bind_arguments={"mapper": table_cls}), checkfirst=True) session.commit() + return staging def drop_staging_table( self, @@ -694,6 +700,7 @@ def merge_replace( table_cls: Type["FileTableProtocol"], session: so.Session, pk_cols: list[str], + staging: sa.Table, *, merge_batch_size: int | None = None, validate_deferred_fks: bool = True, @@ -712,6 +719,7 @@ def merge_upsert( table_cls: Type["FileTableProtocol"], session: so.Session, pk_cols: list[str], + staging: sa.Table, *, merge_batch_size: int | None = None, validate_deferred_fks: bool = True, @@ -727,6 +735,7 @@ def merge_insert( self, table_cls: Type["FileTableProtocol"], session: so.Session, + staging: sa.Table, *, merge_batch_size: int | None = None, validate_deferred_fks: bool = True, diff --git a/src/orm_loader/backends/duckdb.py b/src/orm_loader/backends/duckdb.py index 1816c13..29c0839 100644 --- a/src/orm_loader/backends/duckdb.py +++ b/src/orm_loader/backends/duckdb.py @@ -315,6 +315,7 @@ def merge_replace( table_cls: type["FileTableProtocol"], session: so.Session, pk_cols: list[str], + staging: sa.Table, *, merge_batch_size: int | None = None, validate_deferred_fks: bool = True, @@ -341,7 +342,6 @@ def merge_replace( self._validate_merge_batch_size(merge_batch_size) self.check_replace_feasible(table_cls, pk_cols, session) target = table_cls.__table__ - staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) update_cols = [c for c in self._insertable_column_names(table_cls, staging) if c not in pk_cols] pk_join = sa.and_(*(target.c[c] == staging.c[c] for c in pk_cols)) @@ -411,6 +411,7 @@ def merge_upsert( table_cls: type["FileTableProtocol"], session: so.Session, pk_cols: list[str], + staging: sa.Table, *, merge_batch_size: int | None = None, validate_deferred_fks: bool = True, @@ -419,7 +420,6 @@ def merge_upsert( self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ - staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) insertable_cols = self._insertable_column_names(table_cls, staging) def _upsert(select_: Any) -> sa.Insert: @@ -447,13 +447,13 @@ def merge_insert( self, table_cls: type["FileTableProtocol"], session: so.Session, + staging: sa.Table, *, merge_batch_size: int | None = None, validate_deferred_fks: bool = True, ) -> None: self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ - staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) insertable_cols = self._insertable_column_names(table_cls, staging) non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index 2177340..61094ea 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -107,9 +107,6 @@ def insert_batches( pacsv.write_csv(table, buffer, options) copy.write(buffer.getvalue()) ctx.session.flush() - except Exception: - ctx.session.rollback() - raise finally: cursor.close() @@ -171,6 +168,7 @@ def merge_replace( table_cls: type["FileTableProtocol"], session: so.Session, pk_cols: list[str], + staging: sa.Table, *, merge_batch_size: int | None = None, validate_deferred_fks: bool = True, @@ -178,7 +176,6 @@ def merge_replace( """Delete every target row whose primary key appears in staging.""" self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ - staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) pk_join = sa.and_(*(target.c[c] == staging.c[c] for c in pk_cols)) for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): @@ -193,13 +190,13 @@ def merge_upsert( table_cls: type["FileTableProtocol"], session: so.Session, pk_cols: list[str], + staging: sa.Table, *, merge_batch_size: int | None = None, validate_deferred_fks: bool = True, ) -> None: self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ - staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) insertable_cols = self._insertable_column_names(table_cls, staging) def _upsert(select_: Any) -> sa.Insert: @@ -223,6 +220,7 @@ def merge_insert( self, table_cls: type["FileTableProtocol"], session: so.Session, + staging: sa.Table, *, merge_batch_size: int | None = None, validate_deferred_fks: bool = True, @@ -234,7 +232,6 @@ def merge_insert( """ self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ - staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) insertable_cols = self._insertable_column_names(table_cls, staging) non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) diff --git a/src/orm_loader/backends/sqlite.py b/src/orm_loader/backends/sqlite.py index 77e6169..6a3c2cd 100644 --- a/src/orm_loader/backends/sqlite.py +++ b/src/orm_loader/backends/sqlite.py @@ -185,6 +185,7 @@ def merge_replace( table_cls: type["FileTableProtocol"], session: so.Session, pk_cols: list[str], + staging: sa.Table, *, merge_batch_size: int | None = None, validate_deferred_fks: bool = True, @@ -199,7 +200,6 @@ def merge_replace( """ self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ - staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) # SQLite's DELETE has no USING, so staging keys go in a subquery. IN # rather than a correlated EXISTS, which rescans the unindexed staging @@ -222,13 +222,13 @@ def merge_upsert( table_cls: type["FileTableProtocol"], session: so.Session, pk_cols: list[str], + staging: sa.Table, *, merge_batch_size: int | None = None, validate_deferred_fks: bool = True, ) -> None: self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ - staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) insertable_cols = self._insertable_column_names(table_cls, staging) def _upsert(select_: Any) -> sa.Insert: @@ -249,13 +249,13 @@ def merge_insert( self, table_cls: type["FileTableProtocol"], session: so.Session, + staging: sa.Table, *, merge_batch_size: int | None = None, validate_deferred_fks: bool = True, ) -> None: self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ - staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) insertable_cols = self._insertable_column_names(table_cls, staging) non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) diff --git a/src/orm_loader/loaders/loader.py b/src/orm_loader/loaders/loader.py index 3a729d6..ad88819 100644 --- a/src/orm_loader/loaders/loader.py +++ b/src/orm_loader/loaders/loader.py @@ -18,8 +18,7 @@ import sqlalchemy as sa import sqlalchemy.orm as so -from ..backends.base import DedupeResult -from ..backends.resolve import resolve_backend +from ..backends.base import DatabaseBackend, DedupeResult from ..conversion import ConversionPolicy from ..helpers import IngestError from .processor import BatchProcessor @@ -83,14 +82,12 @@ class LoaderContext: Active SQLAlchemy session. path Path to the source data file. + backend + Backend of the session's database, resolved once for the load. normalise Whether to interpret values under the policy. dedupe Whether the staging table keeps only the first row for each primary key. - staging_schema_tag - schema_translate_map tag the staging table lives in, passed to - resolve_backend() so every backend resolution within this load - resolves the same schema. policy How values are interpreted and what happens to unreadable ones. rejects_dir @@ -103,10 +100,10 @@ class LoaderContext: tableclass: type[FileTableProtocol] session: so.Session path: Path + backend: DatabaseBackend normalise: bool = True dedupe: bool = True - staging_schema_tag: str | None = None policy: ConversionPolicy = field(default_factory=ConversionPolicy) rejects_dir: Path | None = None report: LoadReport = field(default_factory=LoadReport) @@ -145,7 +142,9 @@ def stage_file(ctx: LoaderContext, reader: FileReader, staging: sa.Table) -> int The reader streams the file in bounded batches, so memory does not grow with the row count. A file without data rows loads nothing and is not checked for required columns. With ``ctx.dedupe`` set, the staging table keeps the - first row for each primary key. + first row for each primary key. A failure after the first write rolls + the session back; DuckDB commits every few batches, so earlier batches + stay committed there. Parameters ---------- @@ -184,7 +183,7 @@ def stage_file(ctx: LoaderContext, reader: FileReader, staging: sa.Table) -> int logger.warning("File %s has columns that %s does not have, which are ignored: %s", ctx.path.name, table_name, ignored) if ctx.normalise: _require_columns_present(ctx.tableclass, reader.columns) - backend = resolve_backend(ctx.session, staging_schema_tag=ctx.staging_schema_tag, mapper=ctx.tableclass) + backend = ctx.backend rejects_dir = ctx.rejects_dir or ctx.path.parent / "rejects" processor = BatchProcessor( ctx.tableclass, @@ -196,16 +195,21 @@ def stage_file(ctx: LoaderContext, reader: FileReader, staging: sa.Table) -> int reader.skipped, lambda expression, column_type: backend.evaluate_default(ctx.session, ctx.tableclass, expression, column_type), ) - backend.insert_batches(ctx, staging, processor.stream(chain([] if first is None else [first], batches)), processor.out_columns) - total = backend.count_staging(ctx, staging) - keys = [c for c in ctx.tableclass.pk_names() if c in reader.columns] - if ctx.dedupe and keys: - ctx.report.dedupe = backend.deduplicate_staging(ctx, staging, keys) - total = ctx.report.dedupe.rows_after - ctx.report.rows_loaded += total - if not ctx.report.accounted: - raise IngestError( - f"{table_name}: read {ctx.report.rows_read} rows but loaded {ctx.report.rows_loaded}, " - f"rejected {ctx.report.rows_rejected} and dropped {ctx.report.rows_deduplicated} duplicates" - ) + try: + batches = processor.stream(chain([] if first is None else [first], batches)) + backend.insert_batches(ctx, staging, batches, processor.out_columns) + total = backend.count_staging(ctx, staging) + keys = [c for c in ctx.tableclass.pk_names() if c in reader.columns] + if ctx.dedupe and keys: + ctx.report.dedupe = backend.deduplicate_staging(ctx, staging, keys) + total = ctx.report.dedupe.rows_after + ctx.report.rows_loaded += total + if not ctx.report.accounted: + raise IngestError( + f"{table_name}: read {ctx.report.rows_read} rows but loaded {ctx.report.rows_loaded}, " + f"rejected {ctx.report.rows_rejected} and dropped {ctx.report.rows_deduplicated} duplicates" + ) + except Exception: + ctx.session.rollback() + raise return total diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index 77351fd..2d3adf4 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -18,6 +18,7 @@ from oa_configurator import ResolvedDatabase, physical_schema_of from sqlalchemy.exc import InvalidRequestError, UnboundExecutionError +from ..backends.base import DatabaseBackend from ..backends.resolve import resolve_backend from ..conversion import ConversionPolicy from ..helpers import IngestError @@ -66,7 +67,7 @@ def manage_indices( session: so.Session, index_strategy: str = "auto", *, - staging_schema_tag: str | None = None, + backend: DatabaseBackend, ) -> Generator[None, None, None]: """ Manage non-primary-key indexes around a staged merge. @@ -81,7 +82,6 @@ def manage_indices( IngestError An index could not be rebuilt after a merge that was committed. """ - backend = resolve_backend(session, staging_schema_tag=staging_schema_tag, mapper=cls) resolved_index_strategy = backend.resolve_index_strategy(index_strategy) table_name = cls.__tablename__ @@ -201,51 +201,6 @@ def _restore_indices( return failed_indices - @classmethod - def get_staging_table( - cls: Type[FileTableProtocol], - session: so.Session, - *, - staging_schema_tag: str | None = None, - ) -> sa.Table: - """ - Return the reflected staging table. - - Parameters - ---------- - session - An active SQLAlchemy session bound to an engine. - staging_schema_tag - schema_translate_map tag the staging table lives in. ``None`` - means no schema qualification (backend-default behavior). - - Returns - ------- - sqlalchemy.Table - The reflected staging table. - - Raises - ------ - sqlalchemy.exc.NoSuchTableError - The staging table has not been created. - - Notes - ----- - Reflects via ``session.connection()``, not the bare engine. On - SQLite's SingletonThreadPool a second connection opened straight - from the engine is the same underlying DBAPI connection, and closing - that second wrapper resets its perceived transaction state, silently - discarding the session's own uncommitted work. - """ - _require_bind(session, cls) - backend = resolve_backend(session, staging_schema_tag=staging_schema_tag, mapper=cls) - return sa.Table( - backend.staging_name_for_table(cls.__tablename__), - sa.MetaData(), # throwaway — keeps staging table out of Base.metadata - autoload_with=session.connection(bind_arguments={"mapper": cls}), - schema=backend.staging_schema, - ) - @classmethod def load_file( cls: Type[FileTableProtocol], @@ -391,9 +346,9 @@ def load_file( tableclass=cls, session=session, path=path, + backend=backend, normalise=normalise, dedupe=dedupe, - staging_schema_tag=staging_schema_tag, policy=policy, rejects_dir=rejects_dir, report=report or LoadReport(), @@ -403,20 +358,18 @@ def load_file( # SQLite cannot switch foreign key checks inside a transaction, so relax before staging writes. with backend.relaxed_foreign_keys(session, mapper=cls, no_autoflush=False): logger.info(f"Table `{cls.__tablename__}`: Loading data into staging table") - backend.create_staging_table(cls, session) - staging = cls.get_staging_table(session, staging_schema_tag=staging_schema_tag) + staging = backend.create_staging_table(cls, session) total = stage_file(loader_context, reader, staging) logger.info(f"Table `{cls.__tablename__}`: Merging staging data into target table") - with cls.manage_indices(session, index_strategy=index_strategy, staging_schema_tag=staging_schema_tag): + with cls.manage_indices(session, index_strategy, backend=backend): cls.merge_from_staging( session, - merge_strategy=merge_strategy, + merge_strategy, + backend=backend, + staging=staging, merge_batch_size=merge_batch_size, - staging_schema_tag=staging_schema_tag, validate_deferred_fks=validate_deferred_fks, - resolved=resolved, - physical_schemas=physical_schemas, ) if validate_deferred_fks: logger.info(f"Table `{cls.__tablename__}`: Validating deferred foreign keys.") @@ -496,11 +449,10 @@ def merge_from_staging( session: so.Session, merge_strategy: str = "replace", *, + backend: DatabaseBackend, + staging: sa.Table, merge_batch_size: int | None = None, - staging_schema_tag: str | None = None, validate_deferred_fks: bool = True, - resolved: ResolvedDatabase | None = None, - physical_schemas: Mapping[str, str | None] | None = None, ): """ Merge data from the staging table into the target table. @@ -512,32 +464,25 @@ def merge_from_staging( merge_strategy Merge strategy to apply (for example ``replace``, ``upsert``, or ``insert``). - staging_schema_tag - schema_translate_map tag the staging table lives in. ``None`` - means no schema qualification (backend-default behavior). + backend + Backend of the session's database. + staging + Staging table the rows are merged from. + merge_batch_size + Rows per merge batch, or ``None`` for one statement. validate_deferred_fks Forwarded to the backend's merge methods, which validate deferred FKs before every intermediate commit a batched merge makes. See ``load_file()``. - resolved, physical_schemas - Forwarded to ``resolve_backend()``. See ``load_file()``. """ target = cls.__tablename__ pk_cols = cls.pk_names() - _require_bind(session, cls) - backend = resolve_backend( - session, - staging_schema_tag=staging_schema_tag, - mapper=cls, - resolved=resolved, - physical_schemas=physical_schemas, - ) if merge_strategy == "replace": logger.info(f"Table `{target}`: Merge replace delete phase starting.") delete_started = perf_counter() backend.merge_replace( - cls, session, pk_cols, merge_batch_size=merge_batch_size, validate_deferred_fks=validate_deferred_fks + cls, session, pk_cols, staging, merge_batch_size=merge_batch_size, validate_deferred_fks=validate_deferred_fks ) logger.info( f"Table `{target}`: Merge replace delete phase completed in " @@ -546,7 +491,7 @@ def merge_from_staging( logger.info(f"Table `{target}`: Merge insert phase starting.") insert_started = perf_counter() backend.merge_insert( - cls, session, merge_batch_size=merge_batch_size, validate_deferred_fks=validate_deferred_fks + cls, session, staging, merge_batch_size=merge_batch_size, validate_deferred_fks=validate_deferred_fks ) logger.info( f"Table `{target}`: Merge insert phase completed in " @@ -561,7 +506,7 @@ def merge_from_staging( logger.info(f"Table `{target}`: Merge upsert phase starting.") upsert_started = perf_counter() backend.merge_upsert( - cls, session, pk_cols, merge_batch_size=merge_batch_size, validate_deferred_fks=validate_deferred_fks + cls, session, pk_cols, staging, merge_batch_size=merge_batch_size, validate_deferred_fks=validate_deferred_fks ) logger.info( f"Table `{target}`: Merge upsert phase completed in " @@ -571,7 +516,7 @@ def merge_from_staging( logger.info(f"Table `{target}`: Merge insert phase starting.") insert_started = perf_counter() backend.merge_insert( - cls, session, merge_batch_size=merge_batch_size, validate_deferred_fks=validate_deferred_fks + cls, session, staging, merge_batch_size=merge_batch_size, validate_deferred_fks=validate_deferred_fks ) logger.info( f"Table `{target}`: Merge insert phase completed in " diff --git a/src/orm_loader/tables/typing.py b/src/orm_loader/tables/typing.py index e9173ab..23160af 100644 --- a/src/orm_loader/tables/typing.py +++ b/src/orm_loader/tables/typing.py @@ -7,6 +7,7 @@ if TYPE_CHECKING: from oa_configurator import ResolvedDatabase + from ..backends.base import DatabaseBackend from ..conversion import ConversionPolicy from ..loaders import FileReader, LoadReport @@ -95,20 +96,16 @@ def load_file( report: Optional["LoadReport"] = None, ) -> int: ... - @classmethod - def get_staging_table(cls, session: so.Session, *, staging_schema_tag: str | None = None) -> sa.Table: ... - @classmethod def merge_from_staging( cls, session: so.Session, merge_strategy: str = "replace", *, + backend: "DatabaseBackend", + staging: sa.Table, merge_batch_size: int | None = None, - staging_schema_tag: str | None = None, validate_deferred_fks: bool = True, - resolved: Optional["ResolvedDatabase"] = None, - physical_schemas: Optional[Mapping[str, str | None]] = None, ) -> None: ... @classmethod @@ -129,7 +126,7 @@ def _restore_indices(cls, session: so.Session, indices: list[sa.Index]) -> list[ @classmethod def manage_indices( - cls, session: so.Session, index_strategy: str = "auto", *, staging_schema_tag: str | None = None + cls, session: so.Session, index_strategy: str = "auto", *, backend: "DatabaseBackend" ) -> AbstractContextManager[None]: ... diff --git a/tests/backends/test_duckdb_backend.py b/tests/backends/test_duckdb_backend.py index bf9d319..670a315 100644 --- a/tests/backends/test_duckdb_backend.py +++ b/tests/backends/test_duckdb_backend.py @@ -173,6 +173,7 @@ def test_duckdb_backend_rejects_non_positive_merge_batch_size(merge_batch_size, args = [_ComputedTableCls, duckdb_session] if merge_strategy in {"replace", "upsert"}: args.append(["id"]) + args.append(sa.Table("staging", sa.MetaData())) merge = getattr(backend, f"merge_{merge_strategy}") with pytest.raises(ValueError, match="merge_batch_size must be a positive integer or None"): @@ -244,8 +245,7 @@ def test_duckdb_backend_creates_staging_again_when_only_its_sequence_was_left_be @pytest.mark.db_dialect def test_duckdb_backend_rownum_populates_monotonically_on_insert(duckdb_session): backend = DuckDBBackend() - backend.create_staging_table(_ComputedTableCls, duckdb_session) - staging = _ComputedTableCls.get_staging_table(duckdb_session, staging_schema_tag=backend.staging_schema_tag) + staging = backend.create_staging_table(_ComputedTableCls, duckdb_session) duckdb_session.execute(sa.insert(staging), [{"id": 1, "name": "a"}, {"id": 2, "name": "b"}]) @@ -323,8 +323,7 @@ def test_duckdb_backend_merge_replace_updates_in_place(duckdb_session): staging rows for the follow-up merge_insert() call. """ backend = DuckDBBackend() - backend.create_staging_table(_ComputedTableCls, duckdb_session) - staging = _ComputedTableCls.get_staging_table(duckdb_session, staging_schema_tag=backend.staging_schema_tag) + staging = backend.create_staging_table(_ComputedTableCls, duckdb_session) duckdb_session.execute( sa.insert(ComputedColumnTable), @@ -334,8 +333,8 @@ def test_duckdb_backend_merge_replace_updates_in_place(duckdb_session): sa.insert(staging), [{"id": 1, "name": "alpha-staged"}, {"id": 3, "name": "gamma"}] ) - backend.merge_replace(_ComputedTableCls, duckdb_session, ["id"]) - backend.merge_insert(_ComputedTableCls, duckdb_session) + backend.merge_replace(_ComputedTableCls, duckdb_session, ["id"], staging) + backend.merge_insert(_ComputedTableCls, duckdb_session, staging) rows = {r.id: r.name for r in duckdb_session.execute(sa.select(ComputedColumnTable)).scalars().all()} assert rows == {1: "alpha-staged", 2: "beta", 3: "gamma"} @@ -345,16 +344,15 @@ def test_duckdb_backend_merge_replace_updates_in_place(duckdb_session): @pytest.mark.db_dialect def test_duckdb_backend_merge_replace_paginated_updates_in_place(duckdb_session): backend = DuckDBBackend() - backend.create_staging_table(_ComputedTableCls, duckdb_session) - staging = _ComputedTableCls.get_staging_table(duckdb_session, staging_schema_tag=backend.staging_schema_tag) + staging = backend.create_staging_table(_ComputedTableCls, duckdb_session) duckdb_session.execute( sa.insert(ComputedColumnTable), [{"id": i, "name": f"orig{i}"} for i in range(10)] ) duckdb_session.execute(sa.insert(staging), [{"id": i, "name": f"staged{i}"} for i in range(10)]) - backend.merge_replace(_ComputedTableCls, duckdb_session, ["id"], merge_batch_size=3) - backend.merge_insert(_ComputedTableCls, duckdb_session, merge_batch_size=3) + backend.merge_replace(_ComputedTableCls, duckdb_session, ["id"], staging, merge_batch_size=3) + backend.merge_insert(_ComputedTableCls, duckdb_session, staging, merge_batch_size=3) rows = {r.id: r.name for r in duckdb_session.execute(sa.select(ComputedColumnTable)).scalars().all()} assert rows == {i: f"staged{i}" for i in range(10)} @@ -366,10 +364,7 @@ def test_duckdb_backend_merge_replace_rejects_referenced_table_with_own_fk(duckd """merge_replace must reject a still-referenced row up front, before touching any data, rather than partway through a batch. """ - duckdb_backend.create_staging_table(_ReferencedParentTableCls, duckdb_session) - staging = _ReferencedParentTableCls.get_staging_table( - duckdb_session, staging_schema_tag=duckdb_backend.staging_schema_tag - ) + staging = duckdb_backend.create_staging_table(_ReferencedParentTableCls, duckdb_session) duckdb_session.execute(sa.insert(UnreferencedGrandparentTable), [{"id": 1, "name": "g1"}]) duckdb_session.execute( @@ -379,7 +374,7 @@ def test_duckdb_backend_merge_replace_rejects_referenced_table_with_own_fk(duckd duckdb_session.execute(sa.insert(staging), [{"id": 1, "grandparent_id": 1, "name": "p1-updated"}]) with pytest.raises(ValueError, match="referenced_parent_table"): - duckdb_backend.merge_replace(_ReferencedParentTableCls, duckdb_session, ["id"]) + duckdb_backend.merge_replace(_ReferencedParentTableCls, duckdb_session, ["id"], staging) @pytest.mark.duckdb @@ -392,12 +387,11 @@ def test_duckdb_backend_merge_insert_batch_validation_prevents_partial_commit( already validated and committed in their own batches, survive. """ session = duckdb_session_with_committed_schema - duckdb_backend.create_staging_table(_CycleATableCls, session) - staging = _CycleATableCls.get_staging_table(session, staging_schema_tag=duckdb_backend.staging_schema_tag) + staging = duckdb_backend.create_staging_table(_CycleATableCls, session) session.execute(sa.insert(staging), [{"id": 1, "b_id": None}, {"id": 2, "b_id": None}, {"id": 3, "b_id": 999}]) with pytest.raises(DeferredForeignKeyViolationError, match="cycle_a_table"): - duckdb_backend.merge_insert(_CycleATableCls, session, merge_batch_size=1) + duckdb_backend.merge_insert(_CycleATableCls, session, staging, merge_batch_size=1) session.rollback() remaining = session.execute(sa.select(CycleATable.id).order_by(CycleATable.id)).scalars().all() @@ -413,16 +407,13 @@ def test_duckdb_backend_merge_insert_loads_self_referential_rows_in_dependency_o row not yet committed); merge_insert must still load every row, in however many batches that takes, regardless of insertion order. """ - duckdb_backend.create_staging_table(_SelfReferentialTableCls, duckdb_session) - staging = _SelfReferentialTableCls.get_staging_table( - duckdb_session, staging_schema_tag=duckdb_backend.staging_schema_tag - ) + staging = duckdb_backend.create_staging_table(_SelfReferentialTableCls, duckdb_session) duckdb_session.execute( sa.insert(staging), [{"id": 3, "parent_id": 2}, {"id": 2, "parent_id": 1}, {"id": 1, "parent_id": None}], ) - duckdb_backend.merge_insert(_SelfReferentialTableCls, duckdb_session, merge_batch_size=1) + duckdb_backend.merge_insert(_SelfReferentialTableCls, duckdb_session, staging, merge_batch_size=1) rows = {r.id: r.parent_id for r in duckdb_session.execute(sa.select(SelfReferentialTable)).scalars().all()} assert rows == {1: None, 2: 1, 3: 2} @@ -433,13 +424,10 @@ def test_duckdb_backend_merge_insert_loads_self_referential_rows_in_dependency_o def test_duckdb_backend_merge_upsert_loads_self_referential_rows_in_dependency_order( duckdb_session, duckdb_backend, ): - duckdb_backend.create_staging_table(_SelfReferentialTableCls, duckdb_session) - staging = _SelfReferentialTableCls.get_staging_table( - duckdb_session, staging_schema_tag=duckdb_backend.staging_schema_tag - ) + staging = duckdb_backend.create_staging_table(_SelfReferentialTableCls, duckdb_session) duckdb_session.execute(sa.insert(staging), [{"id": 2, "parent_id": 1}, {"id": 1, "parent_id": None}]) - duckdb_backend.merge_upsert(_SelfReferentialTableCls, duckdb_session, ["id"], merge_batch_size=1) + duckdb_backend.merge_upsert(_SelfReferentialTableCls, duckdb_session, ["id"], staging, merge_batch_size=1) rows = {r.id: r.parent_id for r in duckdb_session.execute(sa.select(SelfReferentialTable)).scalars().all()} assert rows == {1: None, 2: 1} @@ -455,21 +443,18 @@ def test_duckdb_backend_merge_replace_updates_and_inserts_self_referential_rows_ exists; merge_replace must do both together, in dependency order. """ duckdb_session.execute(sa.insert(SelfReferentialTable), [{"id": 1, "parent_id": None}]) - duckdb_backend.create_staging_table(_SelfReferentialTableCls, duckdb_session) - staging = _SelfReferentialTableCls.get_staging_table( - duckdb_session, staging_schema_tag=duckdb_backend.staging_schema_tag - ) + staging = duckdb_backend.create_staging_table(_SelfReferentialTableCls, duckdb_session) duckdb_session.execute( sa.insert(staging), [{"id": 1, "parent_id": 2}, {"id": 2, "parent_id": None}] ) - duckdb_backend.merge_replace(_SelfReferentialTableCls, duckdb_session, ["id"], merge_batch_size=1) + duckdb_backend.merge_replace(_SelfReferentialTableCls, duckdb_session, ["id"], staging, merge_batch_size=1) rows = {r.id: r.parent_id for r in duckdb_session.execute(sa.select(SelfReferentialTable)).scalars().all()} assert rows == {1: 2, 2: None} assert duckdb_session.execute(sa.select(staging.c.id)).scalars().all() == [] - duckdb_backend.merge_insert(_SelfReferentialTableCls, duckdb_session) + duckdb_backend.merge_insert(_SelfReferentialTableCls, duckdb_session, staging) rows_after_insert = { r.id: r.parent_id for r in duckdb_session.execute(sa.select(SelfReferentialTable)).scalars().all() } @@ -488,14 +473,11 @@ def test_duckdb_backend_merge_replace_catches_a_genuine_orphan( """ session = duckdb_session_with_committed_schema session.execute(sa.insert(SelfReferentialTable), [{"id": 1, "parent_id": None}]) - duckdb_backend.create_staging_table(_SelfReferentialTableCls, session) - staging = _SelfReferentialTableCls.get_staging_table( - session, staging_schema_tag=duckdb_backend.staging_schema_tag - ) + staging = duckdb_backend.create_staging_table(_SelfReferentialTableCls, session) session.execute(sa.insert(staging), [{"id": 1, "parent_id": 999}]) with pytest.raises(DeferredForeignKeyViolationError, match="self_referential_table"): - duckdb_backend.merge_replace(_SelfReferentialTableCls, session, ["id"], merge_batch_size=1) + duckdb_backend.merge_replace(_SelfReferentialTableCls, session, ["id"], staging, merge_batch_size=1) session.rollback() assert session.execute(sa.select(SelfReferentialTable.parent_id)).scalars().all() == [None] @@ -512,14 +494,11 @@ def test_duckdb_backend_merge_insert_catches_a_genuine_orphan_in_its_own_batch( already committed and survives rolling back the failed one. """ session = duckdb_session_with_committed_schema - duckdb_backend.create_staging_table(_SelfReferentialTableCls, session) - staging = _SelfReferentialTableCls.get_staging_table( - session, staging_schema_tag=duckdb_backend.staging_schema_tag - ) + staging = duckdb_backend.create_staging_table(_SelfReferentialTableCls, session) session.execute(sa.insert(staging), [{"id": 1, "parent_id": None}, {"id": 2, "parent_id": 999}]) with pytest.raises(DeferredForeignKeyViolationError, match="self_referential_table"): - duckdb_backend.merge_insert(_SelfReferentialTableCls, session, merge_batch_size=1) + duckdb_backend.merge_insert(_SelfReferentialTableCls, session, staging, merge_batch_size=1) session.rollback() assert session.execute(sa.select(SelfReferentialTable.id)).scalars().all() == [1] @@ -538,13 +517,10 @@ def test_duckdb_backend_merge_insert_unbatched_has_no_intermediate_commit( but with no batching requested at all. """ session = duckdb_session_with_committed_schema - duckdb_backend.create_staging_table(_SelfReferentialTableCls, session) - staging = _SelfReferentialTableCls.get_staging_table( - session, staging_schema_tag=duckdb_backend.staging_schema_tag - ) + staging = duckdb_backend.create_staging_table(_SelfReferentialTableCls, session) session.execute(sa.insert(staging), [{"id": 1, "parent_id": None}, {"id": 2, "parent_id": 999}]) - duckdb_backend.merge_insert(_SelfReferentialTableCls, session, merge_batch_size=None) + duckdb_backend.merge_insert(_SelfReferentialTableCls, session, staging, merge_batch_size=None) with pytest.raises(DeferredForeignKeyViolationError, match="self_referential_table"): duckdb_backend.validate_deferred_fks(_SelfReferentialTableCls, session) @@ -558,13 +534,10 @@ def test_duckdb_backend_merge_insert_accepts_a_self_link(duckdb_session, duckdb_ """A row whose parent is itself is valid once it exists, not an orphan merely because the target was empty when it was considered. """ - duckdb_backend.create_staging_table(_SelfReferentialTableCls, duckdb_session) - staging = _SelfReferentialTableCls.get_staging_table( - duckdb_session, staging_schema_tag=duckdb_backend.staging_schema_tag - ) + staging = duckdb_backend.create_staging_table(_SelfReferentialTableCls, duckdb_session) duckdb_session.execute(sa.insert(staging), [{"id": 1, "parent_id": 1}]) - duckdb_backend.merge_insert(_SelfReferentialTableCls, duckdb_session, merge_batch_size=1) + duckdb_backend.merge_insert(_SelfReferentialTableCls, duckdb_session, staging, merge_batch_size=1) duckdb_backend.validate_deferred_fks(_SelfReferentialTableCls, duckdb_session) rows = {r.id: r.parent_id for r in duckdb_session.execute(sa.select(SelfReferentialTable)).scalars().all()} @@ -577,13 +550,10 @@ def test_duckdb_backend_merge_insert_accepts_a_two_row_cycle(duckdb_session, duc """Two rows referencing each other are valid once both exist. Neither can ever become ready on its own, which is not a reason to reject them. """ - duckdb_backend.create_staging_table(_SelfReferentialTableCls, duckdb_session) - staging = _SelfReferentialTableCls.get_staging_table( - duckdb_session, staging_schema_tag=duckdb_backend.staging_schema_tag - ) + staging = duckdb_backend.create_staging_table(_SelfReferentialTableCls, duckdb_session) duckdb_session.execute(sa.insert(staging), [{"id": 1, "parent_id": 2}, {"id": 2, "parent_id": 1}]) - duckdb_backend.merge_insert(_SelfReferentialTableCls, duckdb_session, merge_batch_size=1) + duckdb_backend.merge_insert(_SelfReferentialTableCls, duckdb_session, staging, merge_batch_size=1) duckdb_backend.validate_deferred_fks(_SelfReferentialTableCls, duckdb_session) rows = {r.id: r.parent_id for r in duckdb_session.execute(sa.select(SelfReferentialTable)).scalars().all()} @@ -600,13 +570,10 @@ def test_duckdb_backend_merge_upsert_ignores_a_duplicate_with_a_bad_reference( """ session = duckdb_session_with_committed_schema session.execute(sa.insert(SelfReferentialTable), [{"id": 1, "parent_id": None}]) - duckdb_backend.create_staging_table(_SelfReferentialTableCls, session) - staging = _SelfReferentialTableCls.get_staging_table( - session, staging_schema_tag=duckdb_backend.staging_schema_tag - ) + staging = duckdb_backend.create_staging_table(_SelfReferentialTableCls, session) session.execute(sa.insert(staging), [{"id": 1, "parent_id": 999}]) - duckdb_backend.merge_upsert(_SelfReferentialTableCls, session, ["id"], merge_batch_size=1) + duckdb_backend.merge_upsert(_SelfReferentialTableCls, session, ["id"], staging, merge_batch_size=1) duckdb_backend.validate_deferred_fks(_SelfReferentialTableCls, session) rows = {r.id: r.parent_id for r in session.execute(sa.select(SelfReferentialTable)).scalars().all()} @@ -622,14 +589,11 @@ def test_duckdb_backend_merge_insert_respects_validate_deferred_fks_false( self-referential batching path. """ session = duckdb_session_with_committed_schema - duckdb_backend.create_staging_table(_SelfReferentialTableCls, session) - staging = _SelfReferentialTableCls.get_staging_table( - session, staging_schema_tag=duckdb_backend.staging_schema_tag - ) + staging = duckdb_backend.create_staging_table(_SelfReferentialTableCls, session) session.execute(sa.insert(staging), [{"id": 1, "parent_id": 999}]) duckdb_backend.merge_insert( - _SelfReferentialTableCls, session, merge_batch_size=1, validate_deferred_fks=False + _SelfReferentialTableCls, session, staging, merge_batch_size=1, validate_deferred_fks=False ) rows = {r.id: r.parent_id for r in session.execute(sa.select(SelfReferentialTable)).scalars().all()} @@ -701,11 +665,10 @@ def test_duckdb_backend_merge_insert_batch_validation_is_a_no_op_without_deferre query on every batch commit. """ backend = DuckDBBackend() - backend.create_staging_table(_ComputedTableCls, duckdb_session) - staging = _ComputedTableCls.get_staging_table(duckdb_session, staging_schema_tag=backend.staging_schema_tag) + staging = backend.create_staging_table(_ComputedTableCls, duckdb_session) duckdb_session.execute(sa.insert(staging), [{"id": i, "name": f"n{i}"} for i in range(5)]) - backend.merge_insert(_ComputedTableCls, duckdb_session, merge_batch_size=1) + backend.merge_insert(_ComputedTableCls, duckdb_session, staging, merge_batch_size=1) rows = duckdb_session.execute(sa.select(ComputedColumnTable.id)).scalars().all() assert sorted(rows) == list(range(5)) diff --git a/tests/backends/test_postgres_backend.py b/tests/backends/test_postgres_backend.py index a6fd6d9..31dc745 100644 --- a/tests/backends/test_postgres_backend.py +++ b/tests/backends/test_postgres_backend.py @@ -95,6 +95,7 @@ def test_postgres_backend_rejects_non_positive_merge_batch_size(merge_batch_size args = [_ComputedTableCls, session] if merge_strategy in {"replace", "upsert"}: args.append(["id"]) + args.append(sa.Table("staging", sa.MetaData())) with pytest.raises(ValueError, match="merge_batch_size must be a positive integer or None"): merge(*args, merge_batch_size=merge_batch_size) diff --git a/tests/backends/test_shared_backend.py b/tests/backends/test_shared_backend.py index f86fd9e..a21a91b 100644 --- a/tests/backends/test_shared_backend.py +++ b/tests/backends/test_shared_backend.py @@ -43,8 +43,7 @@ def merge_backend(request: pytest.FixtureRequest) -> tuple[DatabaseBackend, "so. @pytest.mark.parametrize("merge_backend", _REPLACE_DELETE_SEMANTICS_PARAMS, indirect=True) def test_merge_replace_single_pk(merge_backend: tuple[DatabaseBackend, "so.Session"]) -> None: backend, session = merge_backend - backend.create_staging_table(_ComputedTableCls, session) - staging = _ComputedTableCls.get_staging_table(session, staging_schema_tag=backend.staging_schema_tag) + staging = backend.create_staging_table(_ComputedTableCls, session) session.execute( sa.insert(ComputedColumnTable), @@ -52,7 +51,7 @@ def test_merge_replace_single_pk(merge_backend: tuple[DatabaseBackend, "so.Sessi ) session.execute(sa.insert(staging), [{"id": 1, "name": "alpha-staged"}]) - backend.merge_replace(_ComputedTableCls, session, ["id"]) + backend.merge_replace(_ComputedTableCls, session, ["id"], staging) remaining = session.execute(sa.select(ComputedColumnTable.id)).scalars().all() assert remaining == [2] @@ -61,8 +60,7 @@ def test_merge_replace_single_pk(merge_backend: tuple[DatabaseBackend, "so.Sessi @pytest.mark.parametrize("merge_backend", _REPLACE_DELETE_SEMANTICS_PARAMS, indirect=True) def test_merge_replace_composite_pk(merge_backend: tuple[DatabaseBackend, "so.Session"]) -> None: backend, session = merge_backend - backend.create_staging_table(_CompositeTableCls, session) - staging = _CompositeTableCls.get_staging_table(session, staging_schema_tag=backend.staging_schema_tag) + staging = backend.create_staging_table(_CompositeTableCls, session) session.execute( sa.insert(CompositeTable), @@ -70,7 +68,7 @@ def test_merge_replace_composite_pk(merge_backend: tuple[DatabaseBackend, "so.Se ) session.execute(sa.insert(staging), [{"a": 1, "b": 1, "value": "staged"}]) - backend.merge_replace(_CompositeTableCls, session, ["a", "b"]) + backend.merge_replace(_CompositeTableCls, session, ["a", "b"], staging) remaining = session.execute(sa.select(CompositeTable.a, CompositeTable.b)).all() assert remaining == [(2, 2)] @@ -78,11 +76,10 @@ def test_merge_replace_composite_pk(merge_backend: tuple[DatabaseBackend, "so.Se def test_merge_insert_excludes_computed_columns(merge_backend: tuple[DatabaseBackend, "so.Session"]) -> None: backend, session = merge_backend - backend.create_staging_table(_ComputedTableCls, session) - staging = _ComputedTableCls.get_staging_table(session, staging_schema_tag=backend.staging_schema_tag) + staging = backend.create_staging_table(_ComputedTableCls, session) session.execute(sa.insert(staging), [{"id": 1, "name": "alpha"}]) - backend.merge_insert(_ComputedTableCls, session) + backend.merge_insert(_ComputedTableCls, session, staging) row = session.execute(sa.select(ComputedColumnTable)).scalars().one() assert (row.id, row.name, row.slug) == (1, "alpha", "alpha") @@ -90,15 +87,14 @@ def test_merge_insert_excludes_computed_columns(merge_backend: tuple[DatabaseBac def test_merge_upsert_excludes_computed_columns(merge_backend: tuple[DatabaseBackend, "so.Session"]) -> None: backend, session = merge_backend - backend.create_staging_table(_ComputedTableCls, session) - staging = _ComputedTableCls.get_staging_table(session, staging_schema_tag=backend.staging_schema_tag) + staging = backend.create_staging_table(_ComputedTableCls, session) session.execute(sa.insert(ComputedColumnTable), [{"id": 1, "name": "existing"}]) session.execute( sa.insert(staging), [{"id": 1, "name": "ignored"}, {"id": 2, "name": "new"}] ) - backend.merge_upsert(_ComputedTableCls, session, ["id"]) + backend.merge_upsert(_ComputedTableCls, session, ["id"], staging) rows = {r.id: r.name for r in session.execute(sa.select(ComputedColumnTable)).scalars().all()} assert rows == {1: "existing", 2: "new"} @@ -108,10 +104,7 @@ def test_merge_upsert_excludes_computed_columns(merge_backend: tuple[DatabaseBac @pytest.mark.db_dialect def test_postgres_merge_insert_paginates_rownums_with_initial_gap(pg_session) -> None: backend = resolve_backend(pg_session, staging_schema_tag=STAGING_SCHEMA) - backend.create_staging_table(_ComputedTableCls, pg_session) - staging = _ComputedTableCls.get_staging_table( - pg_session, staging_schema_tag=backend.staging_schema_tag - ) + staging = backend.create_staging_table(_ComputedTableCls, pg_session) pg_session.execute( sa.insert(staging), [{"id": i, "name": f"row{i}"} for i in range(10)] ) @@ -129,7 +122,7 @@ def test_postgres_merge_insert_paginates_rownums_with_initial_gap(pg_session) -> ) ) - backend.merge_insert(_ComputedTableCls, pg_session, merge_batch_size=2) + backend.merge_insert(_ComputedTableCls, pg_session, staging, merge_batch_size=2) ids = pg_session.execute( sa.select(ComputedColumnTable.id).order_by(ComputedColumnTable.id) @@ -140,15 +133,14 @@ def test_postgres_merge_insert_paginates_rownums_with_initial_gap(pg_session) -> @pytest.mark.parametrize("merge_backend", _REPLACE_DELETE_SEMANTICS_PARAMS, indirect=True) def test_merge_replace_paginated_path(merge_backend: tuple[DatabaseBackend, "so.Session"]) -> None: backend, session = merge_backend - backend.create_staging_table(_ComputedTableCls, session) - staging = _ComputedTableCls.get_staging_table(session, staging_schema_tag=backend.staging_schema_tag) + staging = backend.create_staging_table(_ComputedTableCls, session) session.execute( sa.insert(ComputedColumnTable), [{"id": i, "name": f"orig{i}"} for i in range(10)] ) session.execute(sa.insert(staging), [{"id": i, "name": f"staged{i}"} for i in range(10)]) - backend.merge_replace(_ComputedTableCls, session, ["id"], merge_batch_size=3) + backend.merge_replace(_ComputedTableCls, session, ["id"], staging, merge_batch_size=3) remaining = session.execute( sa.select(sa.func.count()).select_from(ComputedColumnTable.__table__) @@ -158,11 +150,10 @@ def test_merge_replace_paginated_path(merge_backend: tuple[DatabaseBackend, "so. def test_merge_insert_paginated_path(merge_backend: tuple[DatabaseBackend, "so.Session"]) -> None: backend, session = merge_backend - backend.create_staging_table(_ComputedTableCls, session) - staging = _ComputedTableCls.get_staging_table(session, staging_schema_tag=backend.staging_schema_tag) + staging = backend.create_staging_table(_ComputedTableCls, session) session.execute(sa.insert(staging), [{"id": i, "name": f"row{i}"} for i in range(10)]) - backend.merge_insert(_ComputedTableCls, session, merge_batch_size=3) + backend.merge_insert(_ComputedTableCls, session, staging, merge_batch_size=3) ids = sorted(session.execute(sa.select(ComputedColumnTable.id)).scalars().all()) assert ids == list(range(10)) @@ -170,8 +161,7 @@ def test_merge_insert_paginated_path(merge_backend: tuple[DatabaseBackend, "so.S def test_merge_upsert_paginated_path(merge_backend: tuple[DatabaseBackend, "so.Session"]) -> None: backend, session = merge_backend - backend.create_staging_table(_ComputedTableCls, session) - staging = _ComputedTableCls.get_staging_table(session, staging_schema_tag=backend.staging_schema_tag) + staging = backend.create_staging_table(_ComputedTableCls, session) session.execute( sa.insert(ComputedColumnTable), [{"id": i, "name": "kept"} for i in range(5)] @@ -180,7 +170,7 @@ def test_merge_upsert_paginated_path(merge_backend: tuple[DatabaseBackend, "so.S sa.insert(staging), [{"id": i, "name": "should-not-overwrite"} for i in range(10)] ) - backend.merge_upsert(_ComputedTableCls, session, ["id"], merge_batch_size=3) + backend.merge_upsert(_ComputedTableCls, session, ["id"], staging, merge_batch_size=3) rows = {r.id: r.name for r in session.execute(sa.select(ComputedColumnTable)).scalars().all()} assert rows == {**{i: "kept" for i in range(5)}, **{i: "should-not-overwrite" for i in range(5, 10)}} diff --git a/tests/backends/test_sinks.py b/tests/backends/test_sinks.py index 0abf09c..d807846 100644 --- a/tests/backends/test_sinks.py +++ b/tests/backends/test_sinks.py @@ -18,10 +18,10 @@ def _sink(session, tmp_path): tag = None if session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA - resolve_backend(session, staging_schema_tag=tag, mapper=_Demo).create_staging_table(_Demo, session) - staging = _Demo.get_staging_table(session, staging_schema_tag=tag) - ctx = LoaderContext(tableclass=_Demo, session=session, path=tmp_path / "demo_record.csv", staging_schema_tag=tag) - return resolve_backend(session, staging_schema_tag=tag, mapper=_Demo), ctx, staging + backend = resolve_backend(session, staging_schema_tag=tag, mapper=_Demo) + staging = backend.create_staging_table(_Demo, session) + ctx = LoaderContext(tableclass=_Demo, session=session, path=tmp_path / "demo_record.csv", backend=backend) + return backend, ctx, staging def _batch(i: int, label: str | None) -> pa.Table: @@ -87,8 +87,8 @@ def batches(): def _staging_for(session, table): tag = None if session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA - resolve_backend(session, staging_schema_tag=tag, mapper=table).create_staging_table(table, session) - return table.get_staging_table(session, staging_schema_tag=tag) + created = resolve_backend(session, staging_schema_tag=tag, mapper=table).create_staging_table(table, session) + return sa.Table(created.name, sa.MetaData(), autoload_with=session.connection(), schema=created.schema) def test_a_staging_table_mirrors_the_columns_and_nullability_of_the_target(dialect_session): @@ -114,6 +114,5 @@ def test_a_staging_table_can_be_dropped_twice_and_created_again(dialect_session) tag = None if dialect_session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA for _ in range(2): resolve_backend(dialect_session, staging_schema_tag=tag, mapper=_Demo).drop_staging_table(_Demo, dialect_session) - with pytest.raises(sa.exc.NoSuchTableError): - _Demo.get_staging_table(dialect_session, staging_schema_tag=tag) + assert not sa.inspect(dialect_session.connection()).has_table(staging.name, schema=staging.schema) assert _staging_for(dialect_session, _Demo).name == staging.name diff --git a/tests/backends/test_sqlite_backend.py b/tests/backends/test_sqlite_backend.py index f7a8a1f..5b4a477 100644 --- a/tests/backends/test_sqlite_backend.py +++ b/tests/backends/test_sqlite_backend.py @@ -65,6 +65,7 @@ def test_sqlite_backend_rejects_non_positive_merge_batch_size(merge_batch_size, args = [_ComputedTableCls, sqlite_session] if merge_strategy in {"replace", "upsert"}: args.append(["id"]) + args.append(sa.Table("staging", sa.MetaData())) with pytest.raises(ValueError, match="merge_batch_size must be a positive integer or None"): merge(*args, merge_batch_size=merge_batch_size) diff --git a/tests/loaders/test_loader_e2e.py b/tests/loaders/test_loader_e2e.py index cf87e63..39f79fb 100644 --- a/tests/loaders/test_loader_e2e.py +++ b/tests/loaders/test_loader_e2e.py @@ -269,6 +269,20 @@ def test_merge_strategies(sqlite_session, tmp_path, merge_strategy, expected_row assert [(r.id, r.name) for r in rows] == expected_rows +def test_a_load_resolves_its_backend_once(sqlite_session, tmp_path, monkeypatch): + import orm_loader.tables.loadable_table as loadable_table + + calls = [] + original = loadable_table.resolve_backend + monkeypatch.setattr(loadable_table, "resolve_backend", lambda *a, **k: calls.append(1) or original(*a, **k)) + csv_path = tmp_path / "test_table.csv" + pd.DataFrame([{"id": 1, "name": "alpha"}]).to_csv(csv_path, index=False, sep="\t") + + _SimpleTable.load_file(sqlite_session, csv_path, dedupe=False, merge_strategy="insert") + + assert len(calls) == 1 + + def test_insert_merge_strategy(sqlite_session, tmp_path): csv_path = tmp_path / "test_table.csv" diff --git a/tests/loaders/test_pipeline.py b/tests/loaders/test_pipeline.py index e086ef7..0b350d8 100644 --- a/tests/loaders/test_pipeline.py +++ b/tests/loaders/test_pipeline.py @@ -47,13 +47,25 @@ def _write(tmp_path: Path, table, fmt: str, rows: list[tuple], header: list[str] def _stage(session, table, path: Path, reader=None, **context): tag = None if session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA - resolve_backend(session, staging_schema_tag=tag, mapper=table).create_staging_table(table, session) - staging = table.get_staging_table(session, staging_schema_tag=tag) - ctx = LoaderContext(tableclass=table, session=session, path=path, staging_schema_tag=tag, **context) + backend = resolve_backend(session, staging_schema_tag=tag, mapper=table) + staging = backend.create_staging_table(table, session) + ctx = LoaderContext(tableclass=table, session=session, path=path, backend=backend, **context) loaded = stage_file(ctx, reader or open_reader(path, table.model_columns(), nulls=ctx.policy.nulls), staging) return loaded, ctx, staging +def _reflected(session, table) -> sa.Table: + """The staging table of a table as the database holds it.""" + tag = None if session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA + backend = resolve_backend(session, staging_schema_tag=tag, mapper=table) + return sa.Table( + backend.staging_name_for_table(table.__tablename__), + sa.MetaData(), + autoload_with=session.connection(), + schema=backend.staging_schema, + ) + + def _staged(session, staging, *columns: str) -> list[tuple]: order = staging.c._rownum if "_rownum" in staging.c else sa.literal_column("rowid") return [tuple(r) for r in session.execute(sa.select(*(staging.c[c] for c in columns)).order_by(order)).all()] @@ -89,8 +101,7 @@ def test_an_unreadable_value_raises_and_stages_nothing(dialect_session, tmp_path path = _write(tmp_path, DemoRecord, fmt, rows) with pytest.raises(ConversionError, match="demo_number"): _stage(dialect_session, _Demo, path) - staging = _Demo.get_staging_table(dialect_session, staging_schema_tag=None if dialect_session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA) - assert _staged(dialect_session, staging, "demo_id") == [] + assert _staged(dialect_session, _reflected(dialect_session, _Demo), "demo_id") == [] def test_a_missing_required_value_raises(dialect_session, tmp_path, fmt): @@ -139,8 +150,7 @@ def broken(): reader.batches = broken # type: ignore[method-assign] with pytest.raises(RuntimeError, match="boom"): _stage(dialect_session, _Demo, path, reader=reader) - tag = None if dialect_session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA - assert _staged(dialect_session, _Demo.get_staging_table(dialect_session, staging_schema_tag=tag), "demo_id") == [] + assert _staged(dialect_session, _reflected(dialect_session, _Demo), "demo_id") == [] TYPED_HEADER = [ diff --git a/tests/loaders/test_stage_file.py b/tests/loaders/test_stage_file.py index 8dd24c3..7241461 100644 --- a/tests/loaders/test_stage_file.py +++ b/tests/loaders/test_stage_file.py @@ -18,6 +18,7 @@ def test_loader_context_fields(sqlite_session): ctx = LoaderContext( tableclass=FileLoaderTable, session=sqlite_session, + backend=resolve_backend(sqlite_session, mapper=FileLoaderTable), path=Path("file.csv"), normalise=False, dedupe=True, @@ -42,6 +43,7 @@ def test_stage_file_case_insensitive_headers(tmp_path, sqlite_session): ctx = LoaderContext( tableclass=FileLoaderTable, session=sqlite_session, + backend=resolve_backend(sqlite_session, mapper=FileLoaderTable), path=csv, dedupe=True, ) @@ -73,6 +75,7 @@ def test_stage_file_literal_mode_preserves_quotes(tmp_path, sqlite_session): ctx = LoaderContext( tableclass=FileLoaderTable, session=sqlite_session, + backend=resolve_backend(sqlite_session, mapper=FileLoaderTable), path=csv, dedupe=True, ) @@ -103,6 +106,7 @@ def test_stage_file_csv_mode_strips_quotes(tmp_path, sqlite_session): ctx = LoaderContext( tableclass=FileLoaderTable, session=sqlite_session, + backend=resolve_backend(sqlite_session, mapper=FileLoaderTable), path=csv, dedupe=True, ) @@ -133,6 +137,7 @@ def test_stage_file_by_delimiter_tab_preserves_quotes(tmp_path, sqlite_session): ctx = LoaderContext( tableclass=FileLoaderTable, session=sqlite_session, + backend=resolve_backend(sqlite_session, mapper=FileLoaderTable), path=tsv, dedupe=True, ) @@ -154,6 +159,7 @@ def test_file_without_data_rows_loads_nothing_whatever_its_header(tmp_path, sqli ctx = LoaderContext( tableclass=FileLoaderTable, session=sqlite_session, + backend=resolve_backend(sqlite_session, mapper=FileLoaderTable), path=csv, ) assert stage_file(ctx, _reader(ctx), FileLoaderTable.__table__) == 0 @@ -167,6 +173,7 @@ def test_accounting_is_checked_against_the_staging_table(tmp_path, sqlite_sessio ctx = LoaderContext( tableclass=FileLoaderTable, session=sqlite_session, + backend=resolve_backend(sqlite_session, mapper=FileLoaderTable), path=csv, dedupe=False, ) diff --git a/tests/models.py b/tests/models.py index d7b15c2..a6e9425 100644 --- a/tests/models.py +++ b/tests/models.py @@ -90,7 +90,7 @@ class EnumTable(Base, FileLoadableTableInterface): class ComputedColumnTable(Base, FileLoadableTableInterface): """A real, registered table with a computed column, for merge-method - tests that need get_staging_table() to work. Unlike a bare + tests that need a staging table to be reflected. Unlike a bare __tablename__/__table__ pair, this actually implements FileLoadableTableInterface.""" From 53892bc343c988756d6b402221909ce6b69bbb82 Mon Sep 17 00:00:00 2001 From: Georgie Kennedy Date: Tue, 3 Nov 2026 11:49:46 +1100 Subject: [PATCH 15/21] Apply SQL-expression defaults in the merge, fix numeric, null-marker and column-name handling, report rejects per row - SQL-expression defaults are applied per row by COALESCE in the merge; their staging columns are nullable and evaluate_default is removed - Numeric without precision, or above 76 digits, is read as text; 39-76 digits use decimal256; invalid scale is rejected - Legacy null markers cover the pandas/pyarrow defaults and case variants - Reject counts and reasons are reported per row; the row limit is checked after processing - Primary keys and model columns are keyed by database column name; computed columns are skipped; primary keys stay required without normalisation - infer_delim is exported; staging and conversion docs updated --- docs/loaders/conversion.md | 6 +- docs/loaders/staging.md | 10 +- src/orm_loader/backends/base.py | 32 +- src/orm_loader/backends/duckdb.py | 30 +- src/orm_loader/backends/postgres.py | 15 +- src/orm_loader/backends/sqlite.py | 27 +- src/orm_loader/conversion/handlers/numbers.py | 37 ++- src/orm_loader/conversion/policy.py | 13 + src/orm_loader/helpers/errors.py | 16 +- src/orm_loader/loaders/__init__.py | 2 + src/orm_loader/loaders/loader.py | 16 +- src/orm_loader/loaders/processor.py | 103 ++++--- src/orm_loader/loaders/readers/__init__.py | 3 +- src/orm_loader/loaders/readers/base.py | 5 + src/orm_loader/loaders/readers/delimited.py | 28 +- src/orm_loader/tables/loadable_table.py | 22 +- src/orm_loader/tables/orm_table.py | 32 +- src/orm_loader/tables/serialisable_table.py | 5 +- src/orm_loader/tables/typing.py | 3 +- tests/conversion/test_number_handlers.py | 34 ++- tests/loaders/test_defaults.py | 31 +- tests/loaders/test_review_regressions.py | 285 ++++++++++++++++++ 22 files changed, 594 insertions(+), 161 deletions(-) create mode 100644 tests/loaders/test_review_regressions.py diff --git a/docs/loaders/conversion.md b/docs/loaders/conversion.md index fe9c3f8..9a91297 100644 --- a/docs/loaders/conversion.md +++ b/docs/loaders/conversion.md @@ -42,7 +42,9 @@ policy = ConversionPolicy( ## NullPolicy - Text columns keep every nonempty value verbatim. Only empty and whitespace-only text is null. -- `text_markers` adds exact, case-sensitive words that are null in every text column of the load. To treat a word as missing in some columns only, subclass `TextHandler`, override `missing()` and register it in `ConversionPolicy.columns`. +- Compared with 2.x, words such as `NULL`, `None`, `N/A` and `NaN` stay as text by default. To restore the legacy marker set, pass `NullPolicy(text_markers=NullPolicy.LEGACY_TEXT_MARKERS)`. You can add or remove markers with set operations. +- `text_markers` adds exact, case-sensitive words that are null in every text column of the load. Matching ignores surrounding whitespace in normal conversion. With `marker_quoting="unquoted_only"` in CSV quote mode, Arrow matches markers exactly because it applies them while parsing. To treat a word as missing in some columns only, subclass `TextHandler`, override `missing()` and register it in `ConversionPolicy.columns`. +- `empty_is_null=False` keeps empty text fields as empty strings in delimited files. Typed columns still treat empty fields as missing. - `marker_quoting="unquoted_only"` matches a marker only when the source did not quote it. - Typed columns accept a valid value, a missing word (`null`, `none`, `na`, `n/a`, `nan` and similar) or a real null, and fail otherwise. @@ -68,7 +70,7 @@ The handler for a column is the first match of: the entry in `ConversionPolicy.c | `EnumHandler` | `Enum` | One of the labels the column stores | `strip` | | `IntegerHandler` | `Integer`, `SmallInteger`, `BigInteger` | Optional sign and digits, with an optional all-zero fraction (`1.0`); a value outside the 16, 32 or 64 bits of the column type is an error | none | | `FloatHandler` | `Float` | Decimal or exponent form; `Infinity` and values that overflow to infinity are errors | none | -| `NumericHandler` | `Numeric` | Decimal form, converted to the declared precision and scale without rounding; a precision without a scale has scale 0 and a column without a precision reads as `decimal128(38, 18)` | none | +| `NumericHandler` | `Numeric` | Decimal form, checked against declared precision and scale without rounding; a precision without a scale has scale 0. Unconstrained and wider declarations stay as decimal text until the database stores them | none | | `BooleanHandler` | `Boolean` | `true`, `t`, `yes`, `y`, `1`, `on` and their opposites | `true_words`, `false_words` | | `DateHandler` | `Date` | One of `formats`, validated against the calendar; year 0 is an error | `formats` | | `DateTimeHandler` | `DateTime` | ISO date and time, or a date in one of `formats`; see time zones below | `formats` | diff --git a/docs/loaders/staging.md b/docs/loaders/staging.md index f6f890e..8b8b281 100644 --- a/docs/loaders/staging.md +++ b/docs/loaders/staging.md @@ -11,14 +11,14 @@ - Values are interpreted by column type, as described in [Conversion](conversion.md). - With the default policy an unreadable value or a missing required value raises and nothing is loaded. - With `ErrorPolicy(on_error="reject_row")` those rows, and delimited lines with the wrong number of fields, are written to `/.rejects.csv` and the rest load. A malformed line is listed with its row number in the file, and its text is in the reason. `rejects_dir` defaults to a `rejects` directory inside the input file's directory, and the file of an earlier load is removed. `LoadReport` accounts for every row read, checked against the rows counted in the staging table. -- With `normalise=False` raw text is passed to the staging table unchanged, without defaults, required checks or type conversion. What the database does with a value that does not fit its column then depends on the dialect. -- A blank value, and a column the file lacks, take the column's default when it has one: a server default written as text, a scalar Python default, or a SQL expression such as `current_date`. The column's handler reads the default like a value from the file, and an expression is evaluated by the database for each batch that needs it. A column with a callable default stays required. +- With `normalise=False` raw text is passed to the staging table unchanged, without literal defaults, non-key required checks or type conversion. SQL-expression defaults still apply to null staging values during merge. Primary-key columns must still be present. What the database does with a value that does not fit its column then depends on the dialect. +- A blank value, and a column the file lacks, take the column's default when it has one: a server default written as text, a scalar Python default, or a SQL expression such as `current_date`. Literal defaults are converted by the column's handler. SQL expressions are applied in the merge with `COALESCE`, so the database evaluates them once for each row that needs a default. A column with a callable default stays required. --- ## Staging table -Staging mirrors the target: its columns, their types and their nullability, with no key, index, default or constraint, and a `_rownum` column that orders the rows (SQLite uses its rowid). The processor completes every value before it is written, so a NOT NULL column never receives NULL. Every column that is required and has no default must be in the file, primary key columns included, so a file without its key fails before anything is loaded. +Staging mirrors the target’s loadable columns and types, excluding computed columns, with no key, index, default or constraint, and a `_rownum` column that orders the rows (SQLite uses its rowid). Columns with SQL-expression defaults allow NULL in staging so the merge can apply their defaults. Other columns preserve the target’s nullability. Every column that is required and has no default must be in the file, primary key columns included, so a file without its key fails before anything is loaded. --- @@ -52,7 +52,9 @@ Each strategy does what its name says, whatever the target already holds. | `upsert` | Inserts the file and skips rows whose primary key already exists, keeping the first of a repeated key | | `insert` | Inserts the file without any key check, so a key that already exists fails with the database's duplicate key error | -`insert` suits a first load into an empty table, or any file whose keys are new. With `dedupe=True` it is several times faster than `upsert`, because `upsert` checks every row against the key index. A load with `upsert` into an empty table logs this hint and runs unchanged. With `merge_batch_size` set, batches commit separately on PostgreSQL, so a duplicate key partway through leaves the earlier batches in the table. +`insert` suits a first load into an empty table, or any file whose keys are new. With `dedupe=True` it is several times faster than `upsert`, because `upsert` checks every row against the key index. A load with `upsert` into an empty table logs this hint and runs unchanged. + +With `merge_batch_size` set, PostgreSQL and DuckDB commit the session after each merge batch. If a later batch or deferred-key validation fails, earlier target batches stay committed, along with any other pending session work. A `replace` failure can leave some rows deleted or updated. SQLite keeps the batches in the caller's transaction. Use batching only when you can accept partial results and the session has no unrelated pending work. No schema changes are performed at the loader layer. diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 0d505bf..2e453c1 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -299,6 +299,15 @@ def _insertable_column_names( if c.computed is None and (staging is None or c.name in staging.c) ] + @staticmethod + def _staging_value( + table_cls: Type["FileTableProtocol"], staging: sa.Table, name: str + ) -> sa.ColumnElement[Any]: + """Apply a SQL-expression default to null staging values during merge.""" + column = staging.c[name] + default = table_cls.sql_expression_defaults().get(name) + return column if default is None else sa.func.coalesce(column, default).label(name) + def deferred_fk_constraints( self, table_cls: Type["FileTableProtocol"], @@ -488,8 +497,8 @@ def create_staging_table( The staging table has the target's columns with their types and nullability, and no key, index, default or constraint, plus the row - key column of the backend. The processor completes every value before - it is written, so the target's NOT NULL holds in staging. + key column of the backend. Columns with SQL-expression defaults are + nullable in staging so the merge can evaluate their defaults per row. Parameters ---------- @@ -504,8 +513,9 @@ def create_staging_table( The staging table as created. """ self.drop_staging_table(table_cls, session) + expression_defaults = table_cls.sql_expression_defaults() columns = [ - sa.Column(c.name, c.type, nullable=c.nullable) + sa.Column(c.name, c.type, nullable=c.nullable or c.name in expression_defaults) for c in table_cls.__table__.columns if c.computed is None ] @@ -547,22 +557,6 @@ def _staging_table(self, table_cls: Type["FileTableProtocol"], columns: Sequence schema=self.staging_schema, ) - def evaluate_default( - self, - session: so.Session, - table_cls: Type["FileTableProtocol"], - expression: sa.ColumnElement[Any], - column_type: Any, - ) -> Any: - """ - Return the value of a SQL expression default, as a column of the given type stores it. - - The database evaluates the expression and casts the result to the - column type, as it would for the default of an inserted row. - """ - select = sa.select(sa.cast(expression, column_type)) - return session.execute(select, bind_arguments={"mapper": table_cls}).scalar_one() - @abstractmethod def insert_batches( self, diff --git a/src/orm_loader/backends/duckdb.py b/src/orm_loader/backends/duckdb.py index 29c0839..ff5aeda 100644 --- a/src/orm_loader/backends/duckdb.py +++ b/src/orm_loader/backends/duckdb.py @@ -291,11 +291,12 @@ def check_replace_feasible( rejects a SET on an FK-constrained column of a referenced row. """ target = table_cls.__table__ + target_columns = {column.name: column for column in target.columns} update_cols = [c for c in self._insertable_column_names(table_cls) if c not in pk_cols] deferred_cols = {col.name for fk in self.deferred_fk_constraints(table_cls) for col in fk.columns} enforced_fk_cols = [ c for c in update_cols - if c in target.c and target.c[c].foreign_keys and c not in deferred_cols + if c in target_columns and target_columns[c].foreign_keys and c not in deferred_cols ] if not enforced_fk_cols: return @@ -342,8 +343,9 @@ def merge_replace( self._validate_merge_batch_size(merge_batch_size) self.check_replace_feasible(table_cls, pk_cols, session) target = table_cls.__table__ + target_columns = {column.name: column for column in target.columns} update_cols = [c for c in self._insertable_column_names(table_cls, staging) if c not in pk_cols] - pk_join = sa.and_(*(target.c[c] == staging.c[c] for c in pk_cols)) + pk_join = sa.and_(*(target_columns[c] == staging.c[c] for c in pk_cols)) def _batch_keys(extra: sa.ColumnElement[bool] | None = None) -> sa.Select[Any]: keys = sa.select(*(staging.c[c] for c in pk_cols)) @@ -353,11 +355,13 @@ def _batch_keys(extra: sa.ColumnElement[bool] | None = None) -> sa.Select[Any]: def _pk_in(keys: sa.Select[Any]) -> sa.ColumnElement[bool]: if len(pk_cols) == 1: - return target.c[pk_cols[0]].in_(keys) - return sa.tuple_(*(target.c[c] for c in pk_cols)).in_(keys) + return target_columns[pk_cols[0]].in_(keys) + return sa.tuple_(*(target_columns[c] for c in pk_cols)).in_(keys) def _update(extra: sa.ColumnElement[bool] | None = None) -> None: - stmt = sa.update(target).values({c: staging.c[c] for c in update_cols}).where(pk_join) + stmt = sa.update(target).values( + {target_columns[c]: self._staging_value(table_cls, staging, c) for c in update_cols} + ).where(pk_join) if extra is not None: stmt = stmt.where(_pk_in(_batch_keys(extra))) session.execute(stmt) @@ -376,10 +380,10 @@ def _drop_matched_from_staging(extra: sa.ColumnElement[bool] | None = None) -> N self_referential_fks = self._self_referential_deferred_fks(table_cls) if self_referential_fks: insertable_cols = self._insertable_column_names(table_cls, staging) - non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) + non_paginated_select = sa.select(*(self._staging_value(table_cls, staging, c) for c in insertable_cols)) def _insert(select_: Any) -> sa.Insert: - return sa.insert(target).from_select(insertable_cols, select_) + return sa.insert(target).from_select([target_columns[c] for c in insertable_cols], select_) def _process_batch(batch_range: sa.ColumnElement[bool] | None) -> None: # batch_range's rows already in target are updated; what's @@ -420,16 +424,17 @@ def merge_upsert( self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ + target_columns = {column.name: column for column in target.columns} insertable_cols = self._insertable_column_names(table_cls, staging) def _upsert(select_: Any) -> sa.Insert: return ( duckdb_sqlalchemy.insert(target) - .from_select(insertable_cols, select_) - .on_conflict_do_nothing(index_elements=pk_cols) + .from_select([target_columns[c] for c in insertable_cols], select_) + .on_conflict_do_nothing(index_elements=[target_columns[c] for c in pk_cols]) ) - non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) + non_paginated_select = sa.select(*(self._staging_value(table_cls, staging, c) for c in insertable_cols)) self_referential = bool(self._self_referential_deferred_fks(table_cls)) for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): @@ -454,11 +459,12 @@ def merge_insert( ) -> None: self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ + target_columns = {column.name: column for column in target.columns} insertable_cols = self._insertable_column_names(table_cls, staging) - non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) + non_paginated_select = sa.select(*(self._staging_value(table_cls, staging, c) for c in insertable_cols)) def _insert(select_: Any) -> sa.Insert: - return sa.insert(target).from_select(insertable_cols, select_) + return sa.insert(target).from_select([target_columns[c] for c in insertable_cols], select_) self_referential = bool(self._self_referential_deferred_fks(table_cls)) diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index 61094ea..d421514 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -176,7 +176,8 @@ def merge_replace( """Delete every target row whose primary key appears in staging.""" self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ - pk_join = sa.and_(*(target.c[c] == staging.c[c] for c in pk_cols)) + target_columns = {column.name: column for column in target.columns} + pk_join = sa.and_(*(target_columns[c] == staging.c[c] for c in pk_cols)) for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): where = pk_join if batch_range is None else sa.and_(pk_join, batch_range) @@ -197,17 +198,18 @@ def merge_upsert( ) -> None: self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ + target_columns = {column.name: column for column in target.columns} insertable_cols = self._insertable_column_names(table_cls, staging) def _upsert(select_: Any) -> sa.Insert: # sa.insert() has no .on_conflict_do_nothing() return ( postgresql.insert(target) - .from_select(insertable_cols, select_) - .on_conflict_do_nothing(index_elements=pk_cols) + .from_select([target_columns[c] for c in insertable_cols], select_) + .on_conflict_do_nothing(index_elements=[target_columns[c] for c in pk_cols]) ) - non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) + non_paginated_select = sa.select(*(self._staging_value(table_cls, staging, c) for c in insertable_cols)) for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): select_ = non_paginated_select if batch_range is None else non_paginated_select.where(batch_range) @@ -232,11 +234,12 @@ def merge_insert( """ self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ + target_columns = {column.name: column for column in target.columns} insertable_cols = self._insertable_column_names(table_cls, staging) - non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) + non_paginated_select = sa.select(*(self._staging_value(table_cls, staging, c) for c in insertable_cols)) def _insert(select_: Any) -> sa.Insert: - return sa.insert(target).from_select(insertable_cols, select_) + return sa.insert(target).from_select([target_columns[c] for c in insertable_cols], select_) for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): select_ = non_paginated_select if batch_range is None else non_paginated_select.where(batch_range) diff --git a/src/orm_loader/backends/sqlite.py b/src/orm_loader/backends/sqlite.py index 6a3c2cd..76f39dc 100644 --- a/src/orm_loader/backends/sqlite.py +++ b/src/orm_loader/backends/sqlite.py @@ -116,16 +116,6 @@ def _storable(column: pa.ChunkedArray) -> list[Any]: column = pc.cast(column, pa.string()) return column.to_pylist() - def evaluate_default( - self, - session: so.Session, - table_cls: type["FileTableProtocol"], - expression: sa.ColumnElement[Any], - column_type: Any, - ) -> Any: - """Evaluate without a cast: SQLite casts a date to the integer of its year.""" - return session.execute(sa.select(expression), bind_arguments={"mapper": table_cls}).scalar_one() - def disable_fk_check(self, session: so.Session, *, mapper: Any | None = None) -> str | int | None: """Turn ``PRAGMA foreign_keys`` off and return its previous value. @@ -200,6 +190,7 @@ def merge_replace( """ self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ + target_columns = {column.name: column for column in target.columns} # SQLite's DELETE has no USING, so staging keys go in a subquery. IN # rather than a correlated EXISTS, which rescans the unindexed staging @@ -209,9 +200,9 @@ def _delete(extra: sa.ColumnElement[bool] | None = None) -> sa.Delete: if extra is not None: keys = keys.where(extra) if len(pk_cols) == 1: - return sa.delete(target).where(target.c[pk_cols[0]].in_(keys)) + return sa.delete(target).where(target_columns[pk_cols[0]].in_(keys)) return sa.delete(target).where( - sa.tuple_(*(target.c[c] for c in pk_cols)).in_(keys) + sa.tuple_(*(target_columns[c] for c in pk_cols)).in_(keys) ) for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): @@ -229,16 +220,17 @@ def merge_upsert( ) -> None: self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ + target_columns = {column.name: column for column in target.columns} insertable_cols = self._insertable_column_names(table_cls, staging) def _upsert(select_: Any) -> sa.Insert: return ( sqlite_dialect.insert(target) - .from_select(insertable_cols, select_) - .on_conflict_do_nothing(index_elements=pk_cols) + .from_select([target_columns[c] for c in insertable_cols], select_) + .on_conflict_do_nothing(index_elements=[target_columns[c] for c in pk_cols]) ) - non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) + non_paginated_select = sa.select(*(self._staging_value(table_cls, staging, c) for c in insertable_cols)) for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): @@ -256,11 +248,12 @@ def merge_insert( ) -> None: self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ + target_columns = {column.name: column for column in target.columns} insertable_cols = self._insertable_column_names(table_cls, staging) - non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) + non_paginated_select = sa.select(*(self._staging_value(table_cls, staging, c) for c in insertable_cols)) def _insert(select_: Any) -> sa.Insert: - return sa.insert(target).from_select(insertable_cols, select_) + return sa.insert(target).from_select([target_columns[c] for c in insertable_cols], select_) for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): select_ = non_paginated_select if batch_range is None else non_paginated_select.where(batch_range) diff --git a/src/orm_loader/conversion/handlers/numbers.py b/src/orm_loader/conversion/handlers/numbers.py index 5b78166..01f16a1 100644 --- a/src/orm_loader/conversion/handlers/numbers.py +++ b/src/orm_loader/conversion/handlers/numbers.py @@ -84,21 +84,32 @@ def check(self, values: pa.Array, sa_type: TypeEngine[Any]) -> pa.Array | None: class NumericHandler(NumberHandler): """Accepts decimal forms and keeps the column's declared precision and scale. - A precision without a scale has scale 0. A column without a precision is - read as ``decimal128(38, 18)``. Trailing zeros of the fraction are - dropped, and a value with more digits than the column holds is unreadable. + A precision without a scale has scale 0. Values without a declared + precision, and supported declarations wider than Arrow's decimal types, + stay as decimal text until the database stores them. Trailing fractional + zeros are dropped, and a value outside a declared precision or scale is + unreadable. """ @staticmethod - def precision_scale(sa_type: TypeEngine[Any]) -> tuple[int, int]: - """Return the precision and scale a column is read with.""" + def precision_scale(sa_type: TypeEngine[Any]) -> tuple[int | None, int | None]: + """Return a column's declared precision and scale.""" precision = getattr(sa_type, "precision", None) - if precision is None: - return 38, 18 - return precision, getattr(sa_type, "scale", None) or 0 + scale = getattr(sa_type, "scale", None) + if precision is not None and ( + precision < 1 or (scale is not None and (scale < 0 or scale > precision)) + ): + raise ValueError(f"unsupported Numeric precision/scale declaration ({precision}, {scale})") + return precision, scale def arrow_type(self, sa_type: TypeEngine[Any]) -> pa.DataType: - return pa.decimal128(*self.precision_scale(sa_type)) + precision, scale = self.precision_scale(sa_type) + if precision is None or precision > 76: + return pa.string() + scale = scale or 0 + if precision > 38: + return pa.decimal256(precision, scale) + return pa.decimal128(precision, scale) def clean(self, raw: pa.Array, trimmed: pa.Array) -> pa.Array: return pc.replace_substring_regex(trimmed, r"(\.\d*?)0+$", r"\1") @@ -110,6 +121,14 @@ def pattern(self, sa_type: TypeEngine[Any]) -> str: than a decimal holds, so both parts are bounded here. """ precision, scale = self.precision_scale(sa_type) + if precision is None: + if scale is None: + return r"^[+-]?(?:\d+(?:\.\d*)?|\.\d+)$" + forms = [rf"\d+(\.\d{{0,{scale}}})?"] + if scale: + forms.append(rf"\.\d{{1,{scale}}}") + return rf"^[+-]?({'|'.join(forms)})$" + scale = scale or 0 forms = [rf"0*\d{{0,{precision - scale}}}(\.\d{{0,{scale}}})?"] if scale: forms.append(rf"\.\d{{1,{scale}}}") diff --git a/src/orm_loader/conversion/policy.py b/src/orm_loader/conversion/policy.py index c512bc4..29ea2c3 100644 --- a/src/orm_loader/conversion/policy.py +++ b/src/orm_loader/conversion/policy.py @@ -4,6 +4,7 @@ from collections.abc import Mapping from dataclasses import dataclass, field +from itertools import product from types import MappingProxyType from typing import Any, ClassVar, Literal @@ -36,6 +37,12 @@ from .values import no_failures, underlying_type +def _case_variants(value: str) -> frozenset[str]: + """Return every upper/lower-case spelling of a short null marker.""" + choices = [(char.lower(), char.upper()) if char.isalpha() else (char,) for char in value] + return frozenset("".join(chars) for chars in product(*choices)) + + @dataclass(frozen=True) class NullPolicy: """What counts as a missing value, shared by every column handler. @@ -61,6 +68,12 @@ class NullPolicy: DEFAULT_TYPED_MISSING: ClassVar[frozenset[str]] = frozenset( {"null", "none", "na", "n/a", "nan", "#n/a", "-nan", ""} ) + LEGACY_TEXT_MARKERS: ClassVar[frozenset[str]] = frozenset( + { + "#N/A", "#N/A N/A", "#NA", "-1.#IND", "-1.#QNAN", "-NaN", "-nan", "1.#IND", "1.#QNAN", + "", "N/A", "NA", "NULL", "NaN", "None", "n/a", "nan", "null", + } + ) | frozenset().union(*(_case_variants(word) for word in ("nan", "null", "none", "na", "n/a"))) empty_is_null: bool = True text_markers: frozenset[str] = frozenset() diff --git a/src/orm_loader/helpers/errors.py b/src/orm_loader/helpers/errors.py index 71edd51..cb5cb3e 100644 --- a/src/orm_loader/helpers/errors.py +++ b/src/orm_loader/helpers/errors.py @@ -19,16 +19,26 @@ class PolicyError(IngestError): Name of the column type that has no handler. """ - def __init__(self, *, table_name: str | None, column_name: str | None, column_type: str) -> None: + def __init__( + self, + *, + table_name: str | None, + column_name: str | None, + column_type: str, + detail: str | None = None, + ) -> None: self.table_name = table_name self.column_name = column_name self.column_type = column_type target = ".".join(part for part in (table_name, column_name) if part) or "column" - super().__init__( - f"{target}: no handler for column type {column_type}. Register one in " + message = ( + f"{target}: unsupported column type {column_type}: {detail}." + if detail + else f"{target}: no handler for column type {column_type}. Register one in " f"ConversionPolicy(handlers=...) or ConversionPolicy(columns=...), " f"or PassthroughHandler() to load the values unchanged." ) + super().__init__(message) class ConversionError(IngestError): diff --git a/src/orm_loader/loaders/__init__.py b/src/orm_loader/loaders/__init__.py index 2d62ef8..e5b244c 100644 --- a/src/orm_loader/loaders/__init__.py +++ b/src/orm_loader/loaders/__init__.py @@ -1,9 +1,11 @@ from .loader import LoaderContext, LoadReport, stage_file from .readers import DelimitedTextReader, FileReader, ParquetReader, open_reader +from .readers.detection import infer_delim __all__ = [ "DelimitedTextReader", "FileReader", + "infer_delim", "LoadReport", "LoaderContext", "ParquetReader", diff --git a/src/orm_loader/loaders/loader.py b/src/orm_loader/loaders/loader.py index ad88819..7bfaaf5 100644 --- a/src/orm_loader/loaders/loader.py +++ b/src/orm_loader/loaders/loader.py @@ -109,7 +109,9 @@ class LoaderContext: report: LoadReport = field(default_factory=LoadReport) -def _require_columns_present(tableclass: type[FileTableProtocol], present_columns: Iterable[str]) -> None: +def _require_columns_present( + tableclass: type[FileTableProtocol], present_columns: Iterable[str], *, normalise: bool = True +) -> None: """Raise before staging if a required column is entirely absent from the file. Partial-column loads are not supported, so every merge strategy @@ -122,6 +124,8 @@ def _require_columns_present(tableclass: type[FileTableProtocol], present_column Table being loaded. present_columns Model columns present in the file. + normalise + When false, require only primary-key columns. Raises ------ @@ -129,10 +133,12 @@ def _require_columns_present(tableclass: type[FileTableProtocol], present_column A non-nullable column without a default is not in the file. """ present = set(present_columns) - missing = sorted(c for c in tableclass.required_columns() if c not in present) + required = tableclass.required_columns() if normalise else set(tableclass.pk_names()) + missing = sorted(c for c in required if c not in present) if missing: raise IngestError( - f"{tableclass.__tablename__}: source data is missing required column(s) {missing} (non-nullable, no default)" + f"{tableclass.__tablename__}: source data is missing required column(s) {missing} " + "(non-nullable without a default, or primary-key columns)" ) @@ -181,8 +187,7 @@ def stage_file(ctx: LoaderContext, reader: FileReader, staging: sa.Table) -> int ignored = [name for name in reader.file_columns if name not in reader.columns] if ignored: logger.warning("File %s has columns that %s does not have, which are ignored: %s", ctx.path.name, table_name, ignored) - if ctx.normalise: - _require_columns_present(ctx.tableclass, reader.columns) + _require_columns_present(ctx.tableclass, reader.columns, normalise=ctx.normalise) backend = ctx.backend rejects_dir = ctx.rejects_dir or ctx.path.parent / "rejects" processor = BatchProcessor( @@ -193,7 +198,6 @@ def stage_file(ctx: LoaderContext, reader: FileReader, staging: sa.Table) -> int rejects_dir / f"{ctx.path.stem}.rejects.csv", ctx.normalise, reader.skipped, - lambda expression, column_type: backend.evaluate_default(ctx.session, ctx.tableclass, expression, column_type), ) try: batches = processor.stream(chain([] if first is None else [first], batches)) diff --git a/src/orm_loader/loaders/processor.py b/src/orm_loader/loaders/processor.py index 4df1594..615e4a8 100644 --- a/src/orm_loader/loaders/processor.py +++ b/src/orm_loader/loaders/processor.py @@ -5,14 +5,13 @@ import logging import os from bisect import bisect_left -from collections.abc import Callable, Iterable, Iterator +from collections.abc import Iterable, Iterator from pathlib import Path from typing import TYPE_CHECKING, Any import pyarrow as pa import pyarrow.compute as pc import pyarrow.csv as pv -import sqlalchemy as sa from ..conversion import ConversionPolicy, Interpretation from ..conversion.values import no_failures, underlying_type @@ -103,23 +102,17 @@ class BatchProcessor: Lines the reader left out as malformed, each as its line number in the file and the reason. They are written to the rejects file and count as rows read and rejected. - evaluate - Function that returns the value of a SQL expression default as a - column of the given type stores it. It is called once for each such - default when the processor is built. - Attributes ---------- out_columns - Columns of the typed batches: the file's columns, then the columns - the file lacks that have a default. + Columns of typed batches: loadable file columns, then the columns the + file lacks that have a literal or SQL-expression default. Notes ----- - A missing value takes the column's default, and a column the file lacks - is filled with it. The default is read by the column's handler like a - value from the file. A missing value left in a non-nullable column is a - required value that is missing. + Missing values take literal defaults. SQL-expression defaults remain + null in staging so the database can evaluate them per row during merge. + A missing value left in a required column is an error. """ def __init__( @@ -131,9 +124,7 @@ def __init__( rejects_path: Path, normalise: bool = True, skipped: list[tuple[int | None, str]] | None = None, - evaluate: Callable[[Any, Any], Any] | None = None, ) -> None: - self.columns = columns self.skipped = skipped if skipped is not None else [] self._parsed = 0 self._shifts: list[int] = [] @@ -141,17 +132,44 @@ def __init__( self.report = report self.normalise = normalise self.table_name = tableclass.__tablename__ - model_columns = tableclass.model_columns() - defaults = tableclass.column_defaults() if normalise else {} - self.out_columns = [*columns, *(name for name in defaults if name not in columns)] + model_columns = { + name: column + for name, column in tableclass.model_columns().items() + if column.computed is None + } + self.columns = [name for name in columns if name in model_columns] + all_defaults = tableclass.column_defaults() if normalise else {} + self.sql_expression_defaults = tableclass.sql_expression_defaults() if normalise else {} + defaults = { + name: default + for name, default in all_defaults.items() + if name not in self.sql_expression_defaults + } + self.out_columns = [*self.columns, *(name for name in all_defaults if name not in self.columns)] self.types = {name: underlying_type(model_columns[name].type) for name in self.out_columns} self.handlers = ( {name: policy.for_column(self.types[name], self.table_name, name) for name in self.out_columns} if normalise else {} ) - self.defaults = {name: self._evaluated(name, default, evaluate) for name, default in defaults.items()} - self.required = [name for name in self.out_columns if not model_columns[name].nullable] + for name, handler in self.handlers.items(): + try: + handler.arrow_type(self.types[name]) + except ValueError as exc: + raise PolicyError( + table_name=self.table_name, + column_name=name, + column_type=repr(self.types[name]), + detail=str(exc), + ) from exc + self.defaults = defaults + primary_keys = set(tableclass.pk_names()) + self.required = [ + name + for name in self.out_columns + if (not model_columns[name].nullable or name in primary_keys) + and name not in self.sql_expression_defaults + ] self._literals: dict[str, pa.Scalar] = {} self.rejects = RejectsWriter(rejects_path) if policy.errors.on_error == "reject_row": @@ -162,8 +180,9 @@ def stream(self, batches: Iterable[pa.RecordBatch]) -> Iterator[pa.Table]: """Yield typed tables for the raw batches, closing the rejects file afterwards.""" try: for raw in batches: - self._reject_skipped() + self._reject_skipped(check_limit=False) table = self.process(raw) + self._check_limit() if table.num_rows: yield table self._reject_skipped() @@ -197,7 +216,7 @@ def process(self, raw: pa.RecordBatch) -> pa.Table: self._parsed += raw.num_rows self.report.rows_read += raw.num_rows if not self.normalise: - return pa.Table.from_batches([raw]) + return pa.Table.from_batches([raw.select(self.columns)]) results = { name: self.handlers[name].interpret(raw.column(name), self.types[name], self.policy.nulls) for name in self.columns @@ -212,6 +231,12 @@ def process(self, raw: pa.RecordBatch) -> pa.Table: def _fill_defaults(self, results: dict[str, Interpretation], count: int) -> None: """Give a missing value, and a column the file lacks, the column's default.""" + for name in self.sql_expression_defaults: + if name not in results: + results[name] = Interpretation( + pa.nulls(count, self.handlers[name].arrow_type(self.types[name])), + no_failures(count), + ) for name, default in self.defaults.items(): if name not in results: results[name] = Interpretation(pa.repeat(self._fill_value(name, default), count), no_failures(count)) @@ -222,14 +247,6 @@ def _fill_defaults(self, results: dict[str, Interpretation], count: int) -> None filled = pc.if_else(missing, self._fill_value(name, default), values) results[name] = Interpretation(filled, invalid) - def _evaluated(self, name: str, default: Any, evaluate: Callable[[Any, Any], Any] | None) -> Any: - """Return a literal default as it is and the value of a SQL expression default.""" - if not isinstance(default, sa.ClauseElement): - return default - if evaluate is None: - raise PolicyError(table_name=self.table_name, column_name=name, column_type="SQL expression default") - return evaluate(default, self.types[name]) - def _fill_value(self, name: str, default: Any) -> pa.Scalar: """Return the scalar that fills a column, read through the column's handler once.""" if name not in self._literals: @@ -249,16 +266,16 @@ def _read_default(self, name: str, value: Any) -> pa.Scalar: ) return values[0] - def _failures(self, results: dict[str, Any]) -> dict[str, tuple[pa.Array, str]]: - found: dict[str, tuple[pa.Array, str]] = {} + def _failures(self, results: dict[str, Any]) -> dict[str, list[tuple[pa.Array, str]]]: + found: dict[str, list[tuple[pa.Array, str]]] = {} for name in self.columns: invalid = results[name].invalid if pc.any(invalid).as_py(): - found[name] = (invalid, f"cannot be read as {self.handlers[name].expected(self.types[name])}") + found[name] = [(invalid, f"cannot be read as {self.handlers[name].expected(self.types[name])}")] for name in self.required: missing = pc.and_(pc.is_null(results[name].values), pc.invert(results[name].invalid)) if pc.any(missing).as_py(): - found[name] = (pc.or_(found[name][0], missing) if name in found else missing, "required value is missing") + found.setdefault(name, []).append((missing, "required value is missing")) return found def _conversion_error(self, name: str, raw: pa.Array | pa.ChunkedArray, invalid: pa.Array) -> ConversionError | None: @@ -311,17 +328,18 @@ def _reject( self, raw: pa.RecordBatch, results: dict[str, Any], - failures: dict[str, tuple[pa.Array, str]], + failures: dict[str, list[tuple[pa.Array, str]]], first: int, ) -> pa.Table: rejected = no_failures(raw.num_rows) column = pa.nulls(raw.num_rows, pa.string()) reason = pa.nulls(raw.num_rows, pa.string()) - for name, (mask, why) in failures.items(): - first_hit = pc.and_(mask, pc.invert(rejected)) - column = pc.if_else(first_hit, pa.scalar(name), column) - reason = pc.if_else(first_hit, pa.scalar(why), reason) - rejected = pc.or_(rejected, mask) + for name, reasons in failures.items(): + for mask, why in reasons: + first_hit = pc.and_(mask, pc.invert(rejected)) + column = pc.if_else(first_hit, pa.scalar(name), column) + reason = pc.if_else(first_hit, pa.scalar(why), reason) + rejected = pc.or_(rejected, mask) row_numbers = self._row_numbers(first, raw.num_rows) plain = [c.dictionary_decode() if pa.types.is_dictionary(c.type) else c for c in raw.columns] listing = pa.table( @@ -348,7 +366,7 @@ def _row_numbers(self, first: int, count: int) -> pa.Array: positions = range(first + 1, first + count + 1) return pa.array([p + bisect_left(self._shifts, p) for p in positions], pa.int64()) - def _reject_skipped(self) -> None: + def _reject_skipped(self, *, check_limit: bool = True) -> None: """Write the lines the reader left out to the rejects file and remember where they were.""" count = len(self.skipped) if not count: @@ -370,7 +388,8 @@ def _reject_skipped(self) -> None: self.report.rows_read += count self.report.rows_rejected += count self.report.rejects_path = self.rejects.path - self._check_limit() + if check_limit: + self._check_limit() def _check_limit(self) -> None: limit = self.policy.errors.max_rejects diff --git a/src/orm_loader/loaders/readers/__init__.py b/src/orm_loader/loaders/readers/__init__.py index a2ad040..dfbdd8a 100644 --- a/src/orm_loader/loaders/readers/__init__.py +++ b/src/orm_loader/loaders/readers/__init__.py @@ -7,6 +7,7 @@ from ...conversion import NullPolicy from .base import FileReader from .delimited import DelimitedTextReader +from .detection import infer_delim from .parquet import ParquetReader @@ -52,4 +53,4 @@ def open_reader( -__all__ = ["DelimitedTextReader", "FileReader", "ParquetReader", "open_reader"] +__all__ = ["DelimitedTextReader", "FileReader", "ParquetReader", "infer_delim", "open_reader"] diff --git a/src/orm_loader/loaders/readers/base.py b/src/orm_loader/loaders/readers/base.py index 8bc360a..be398ee 100644 --- a/src/orm_loader/loaders/readers/base.py +++ b/src/orm_loader/loaders/readers/base.py @@ -90,8 +90,13 @@ def null_markers(batch: pa.RecordBatch, markers: frozenset[str]) -> pa.RecordBat values = pa.array(sorted(markers), pa.string()) arrays = [] for column in batch.columns: + was_dictionary = pa.types.is_dictionary(column.type) and pa.types.is_string(column.type.value_type) + if was_dictionary: + column = column.dictionary_decode() if pa.types.is_string(column.type): hit = pc.fill_null(pc.is_in(pc.utf8_trim_whitespace(column), value_set=values), False) column = pc.if_else(hit, pa.scalar(None, pa.string()), column) + if was_dictionary: + column = pc.dictionary_encode(column) arrays.append(column) return pa.RecordBatch.from_arrays(arrays, names=batch.schema.names) diff --git a/src/orm_loader/loaders/readers/delimited.py b/src/orm_loader/loaders/readers/delimited.py index f987de0..8398bec 100644 --- a/src/orm_loader/loaders/readers/delimited.py +++ b/src/orm_loader/loaders/readers/delimited.py @@ -25,10 +25,11 @@ class DelimitedTextReader(FileReader): """A delimited text file read as raw strings in bounded blocks. - Date columns are dictionary-encoded. An unquoted empty field is null and - a quoted empty field is an empty string. Header names are normalised as - described by ``normalise_headers``. File columns that are not model - columns are ignored. + Date columns are dictionary-encoded. With the default null policy, an + unquoted empty field is null and a quoted empty field is an empty string. + With ``empty_is_null=False``, unquoted empty text remains empty. Header + names are normalised as described by ``normalise_headers``. File columns + that are not model columns are ignored. Parameters ---------- @@ -42,7 +43,8 @@ class DelimitedTextReader(FileReader): ``"csv"`` reads double quotes as field quoting, ``"literal"`` keeps them as data, ``"by_delimiter"`` chooses ``"literal"`` for tab delimited files and ``"csv"`` otherwise. ``None`` uses - ``DEFAULT_QUOTE_MODE``. + ``DEFAULT_QUOTE_MODE``. In CSV mode, unquoted-only null markers must + match exactly, including surrounding whitespace. block_size Bytes per block; memory use follows this, not the row count. skip_malformed @@ -105,10 +107,23 @@ def _warn_if_everything_is_quoted(self) -> None: sorted(self.nulls.text_markers), ) + def _null_values(self, markers: frozenset[str]) -> list[str]: + """Return exact import markers; padded markers are handled after parsing.""" + empty = [""] if self.nulls.empty_is_null else [] + quoted_markers = sorted(markers) if self.quote_mode == "csv" else [] + return empty + quoted_markers + + def _apply_reader_markers(self, batch: pa.RecordBatch, markers: frozenset[str]) -> pa.RecordBatch: + """Apply whitespace-tolerant markers after parsing when quoting is not tracked.""" + if self.quote_mode == "csv" or not markers: + return batch + return self.null_markers(batch, markers) + def batches(self) -> Iterator[pa.RecordBatch]: """Yield the file's rows in blocks of ``block_size`` bytes.""" if self.empty: return + markers = self.reader_markers() string, encoded = pa.string(), pa.dictionary(pa.int32(), pa.string()) column_types = { name: encoded if isinstance(self.model_columns[name].type, sa.Date) else string @@ -134,13 +149,14 @@ def batches(self) -> Iterator[pa.RecordBatch]: convert_options=pv.ConvertOptions( column_types=column_types, include_columns=self.columns, - null_values=[""] + sorted(self.reader_markers()), + null_values=self._null_values(markers), strings_can_be_null=True, quoted_strings_can_be_null=False, ), ) with reader: for batch in reader: + batch = self._apply_reader_markers(batch, markers) if batch.num_rows: yield batch except UnicodeDecodeError as exc: diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index 2d3adf4..f8e0580 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -263,6 +263,12 @@ def load_file( index_strategy Index handling strategy during merge. Use ``"auto"`` to let the backend choose a sensible default. + merge_batch_size + Rows per merge batch, or ``None`` to merge in one statement. On + PostgreSQL and DuckDB, each merge batch commits the session. A + later failure leaves earlier target batches committed, along + with any other pending session work. SQLite keeps the batches in + the caller's transaction. staging_schema_tag schema_translate_map tag the staging table lives in. ``None`` means no schema qualification (backend-default behavior). @@ -319,12 +325,21 @@ def load_file( if reader is None: reader = open_reader( path, - cls.model_columns(), + { + name: column + for name, column in cls.model_columns().items() + if column.computed is None + }, nulls=policy.nulls, quote_mode=quote_mode, chunksize=chunksize, skip_malformed=policy.errors.on_error == "reject_row", ) + elif reader.nulls != policy.nulls: + raise ValueError( + "the reader's null policy differs from policy.nulls; " + "build the reader with nulls=policy.nulls" + ) backend = resolve_backend( session, @@ -469,7 +484,10 @@ def merge_from_staging( staging Staging table the rows are merged from. merge_batch_size - Rows per merge batch, or ``None`` for one statement. + Rows per merge batch, or ``None`` for one statement. On + PostgreSQL and DuckDB, each batch commits the session; earlier + batches stay committed if a later batch fails. SQLite keeps the + batches in the caller's transaction. validate_deferred_fks Forwarded to the backend's merge methods, which validate deferred FKs before every intermediate commit a batched merge diff --git a/src/orm_loader/tables/orm_table.py b/src/orm_loader/tables/orm_table.py index 0f04922..1e336d3 100644 --- a/src/orm_loader/tables/orm_table.py +++ b/src/orm_loader/tables/orm_table.py @@ -102,7 +102,7 @@ def pk_names(cls) -> list[str]: list[str] A list of primary key column names. """ - return [c.key for c in cls.pk_columns() if c.key is not None] + return [c.name for c in cls.pk_columns() if c.name is not None] @classmethod def model_columns(cls) -> dict[str, sa.ColumnElement[Any]]: @@ -112,19 +112,21 @@ def model_columns(cls) -> dict[str, sa.ColumnElement[Any]]: Returns ------- dict[str, sqlalchemy.ColumnElement] - A mapping of column name to column object. + A mapping of database column name to column object. This is the name + used by file headers and staging tables, even when ``Column.key`` + differs. """ mapper = cls.mapper_for() - return {c.key: c for c in mapper.columns} + return {c.name: c for c in mapper.columns} @classmethod def required_columns(cls) -> set[str]: """ Return column names that must be present in inbound data. - A column is considered required if it is: - - non-nullable - - has no default that can be supplied (see ``column_defaults``) + A non-nullable column is required when it has no default that can be + supplied (see ``column_defaults``). Primary keys are always required + by the file loader. This method is intended for ingestion-time validation and does not attempt to enforce schema semantics beyond insert viability. @@ -136,15 +138,16 @@ def required_columns(cls) -> set[str]: """ defaults = cls.column_defaults() return { - c.key + c.name for c in cls.mapper_for().columns - if not c.nullable and c.key not in defaults + if c.computed is None + and (c.primary_key or (not c.nullable and c.name not in defaults)) } @classmethod def column_defaults(cls) -> dict[str, Any]: """ - Return the default of each column that has one that can be supplied. + Return defaults the file loader can supply for non-key, non-computed columns. A server default given as text and a scalar Python default are returned as the value itself. A server default given as a SQL @@ -159,6 +162,8 @@ def column_defaults(cls) -> dict[str, Any]: """ defaults: dict[str, Any] = {} for c in cls.mapper_for().columns: + if c.computed is not None or c.primary_key: + continue server_default = getattr(c.server_default, "arg", None) if server_default is not None: defaults[c.name] = server_default @@ -166,6 +171,15 @@ def column_defaults(cls) -> dict[str, Any]: defaults[c.name] = c.default.arg # ty: ignore[unresolved-attribute] return defaults + @classmethod + def sql_expression_defaults(cls) -> dict[str, sa.ClauseElement]: + """Return server defaults that should be evaluated for each merged row.""" + return { + name: default + for name, default in cls.column_defaults().items() + if isinstance(default, sa.ClauseElement) + } + @classmethod def max_id(cls, session: so.Session) -> int: """ diff --git a/src/orm_loader/tables/serialisable_table.py b/src/orm_loader/tables/serialisable_table.py index 2310855..0ec0b05 100644 --- a/src/orm_loader/tables/serialisable_table.py +++ b/src/orm_loader/tables/serialisable_table.py @@ -84,7 +84,8 @@ def to_dict( A dictionary representation of the ORM row. """ data: dict[str, Any] = {} - for key, _ in self.model_columns().items(): + for column in self.mapper_for().columns: + key = column.key if only and key not in only: continue if exclude and key in exclude: @@ -163,4 +164,4 @@ def __json__(self) -> dict[str, Any]: dict[str, Any] A dictionary representation of the ORM row. """ - return self.to_dict() \ No newline at end of file + return self.to_dict() diff --git a/src/orm_loader/tables/typing.py b/src/orm_loader/tables/typing.py index 23160af..23e3db4 100644 --- a/src/orm_loader/tables/typing.py +++ b/src/orm_loader/tables/typing.py @@ -53,9 +53,10 @@ def model_columns(cls) -> dict[str, sa.ColumnElement[Any]]: ... def required_columns(cls) -> set[str]: ... @classmethod + def column_defaults(cls) -> dict[str, Any]: ... @classmethod - def column_defaults(cls) -> dict[str, Any]: ... + def sql_expression_defaults(cls) -> dict[str, sa.ClauseElement]: ... @runtime_checkable class FileTableProtocol(ORMTableProtocol, Protocol): diff --git a/tests/conversion/test_number_handlers.py b/tests/conversion/test_number_handlers.py index aed12bf..fe00e26 100644 --- a/tests/conversion/test_number_handlers.py +++ b/tests/conversion/test_number_handlers.py @@ -76,10 +76,26 @@ def test_numeric_precision_without_scale_has_scale_zero(): assert unreadable("1.5", sa.Numeric(30)) -def test_numeric_without_precision_is_exact_and_rejects_what_does_not_fit(): - assert read("12345678901234567890.123456789", sa.Numeric()) == Decimal("12345678901234567890.123456789000000000") - assert unreadable("1" + "0" * 30, sa.Numeric()) - assert unreadable("0.1234567890123456789", sa.Numeric()) +@pytest.mark.parametrize("raw", [ + "12345678901234567890.123456789", "1" + "0" * 80, "0.1234567890123456789", +]) +def test_numeric_without_precision_preserves_every_digit_as_text(raw): + assert read(raw, sa.Numeric()) == raw + + +def test_wide_numeric_uses_decimal256_and_enforces_its_declared_shape(): + from orm_loader.conversion.handlers import NumericHandler + + assert NumericHandler().arrow_type(sa.Numeric(50, 2)) == pa.decimal256(50, 2) + raw = "123456789012345678901234567890123456789012.25" + assert read(raw, sa.Numeric(50, 2)) == Decimal(raw) + assert unreadable("1.555", sa.Numeric(50, 2)) + assert unreadable("9" * 49, sa.Numeric(50, 2)) + + +@pytest.mark.parametrize("raw", ["abc", "1e5", ".", "--1"]) +def test_unconstrained_numeric_still_rejects_non_decimal_text(raw): + assert unreadable(raw, sa.Numeric()) def test_numeric_digits_beyond_the_precision_do_not_wrap_around(): @@ -140,3 +156,13 @@ def test_numeric_with_scale_zero_accepts_a_zero_fraction_only(): [Decimal("2"), Decimal("2"), None, None], [False, False, True, True], ) + + +def test_numeric_wider_than_decimal256_preserves_text_and_declared_bounds(): + from orm_loader.conversion.handlers import NumericHandler + + assert NumericHandler().arrow_type(sa.Numeric(80, 2)) == pa.string() + raw = "9" * 78 + ".12" + assert read(raw, sa.Numeric(80, 2)) == raw + assert unreadable("9" * 79, sa.Numeric(80, 2)) + assert unreadable("1.234", sa.Numeric(80, 2)) diff --git a/tests/loaders/test_defaults.py b/tests/loaders/test_defaults.py index f5d8054..ce6c056 100644 --- a/tests/loaders/test_defaults.py +++ b/tests/loaders/test_defaults.py @@ -109,28 +109,27 @@ def test_a_key_in_the_file_is_loaded_and_may_not_be_blank(dialect_session, tmp_p _keyed(dialect_session, tmp_path, "keyed_id,keyed_name\n,a\n") -def test_an_expression_default_is_evaluated_once_for_the_load(dialect_session, tmp_path, monkeypatch): - from orm_loader.backends import resolve_backend +def test_expression_defaults_fill_nulls_across_reader_batches(dialect_session, tmp_path): from orm_loader.loaders import DelimitedTextReader - calls: list[object] = [] - backend_class = type(resolve_backend(dialect_session, mapper=DemoDefaultKinds)) - original = backend_class.evaluate_default - - def counting(self, session, table_cls, expression, column_type): - calls.append(expression) - return original(self, session, table_cls, expression, column_type) - - monkeypatch.setattr(backend_class, "evaluate_default", counting) path = tmp_path / "demo_default_kinds.csv" - path.write_text("kind_id,kind_now\n" + "".join(f"{i},\n" for i in range(1, 30_001))) + path.write_text( + "kind_id,kind_now\n1,2020-01-01 00:00:00\n" + + "".join(f"{i},\n" for i in range(2, 201)) + ) tag = None if dialect_session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA table = cast(Type[FileTableProtocol], DemoDefaultKinds) - reader = DelimitedTextReader(path, DemoDefaultKinds.model_columns(), block_size=64 << 10) + reader = DelimitedTextReader(path, DemoDefaultKinds.model_columns(), block_size=256) + assert len(list(reader.batches())) > 1 table.load_file(dialect_session, path, staging_schema_tag=tag, reader=reader) - batches = len(list(DelimitedTextReader(path, DemoDefaultKinds.model_columns(), block_size=64 << 10).batches())) - now_calls = [c for c in calls if isinstance(c, sa.sql.functions.FunctionElement)] - assert batches > 1 and len(now_calls) == 1 + dialect_session.commit() + rows = dialect_session.execute( + sa.select(DemoDefaultKinds.kind_id, DemoDefaultKinds.kind_now) + .order_by(DemoDefaultKinds.kind_id) + ).all() + assert len(rows) == 200 + assert rows[0].kind_now == datetime(2020, 1, 1) + assert all(isinstance(row.kind_now, datetime) for row in rows) def test_a_default_the_handler_cannot_read_is_an_error_naming_the_column(dialect_session, tmp_path): diff --git a/tests/loaders/test_review_regressions.py b/tests/loaders/test_review_regressions.py new file mode 100644 index 0000000..d4dc01b --- /dev/null +++ b/tests/loaders/test_review_regressions.py @@ -0,0 +1,285 @@ +"""Regression coverage for PR #46's ingestion review findings.""" + +import csv +import inspect +from datetime import date +from typing import cast + +import pyarrow as pa +import pytest +import sqlalchemy as sa +import sqlalchemy.orm as so + +from orm_loader.backends import STAGING_SCHEMA +from orm_loader.conversion import ConversionPolicy, ErrorPolicy, NullPolicy +from orm_loader.helpers import IngestError, PolicyError +from orm_loader.loaders import DelimitedTextReader, LoadReport, infer_delim +from orm_loader.loaders.processor import BatchProcessor +from orm_loader.tables import FileLoadableTableInterface +from orm_loader.tables.typing import FileTableProtocol, ORMTableProtocol + + +class ReviewBase(so.DeclarativeBase): + pass + + +class ReviewRecord(ReviewBase, FileLoadableTableInterface): + __tablename__ = "review_record" + record_id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True, autoincrement=False) + score: so.Mapped[int] = so.mapped_column(sa.Integer, nullable=False) + note: so.Mapped[str | None] = so.mapped_column(sa.String) + event_date: so.Mapped[date | None] = so.mapped_column(sa.Date) + twice: so.Mapped[int | None] = so.mapped_column(sa.Integer, sa.Computed("score * 2")) + + +class RequiredComputedRecord(ReviewBase, FileLoadableTableInterface): + __tablename__ = "review_required_computed" + record_id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) + score: so.Mapped[int] = so.mapped_column(sa.Integer, nullable=False) + twice: so.Mapped[int] = so.mapped_column(sa.Integer, sa.Computed("score * 2"), nullable=False) + + +class KeyedRecord(ReviewBase, FileLoadableTableInterface): + __tablename__ = "review_keyed_record" + record_id: so.Mapped[int] = so.mapped_column( + "id_col", sa.Integer, key="id_key", primary_key=True, autoincrement=False, + ) + value: so.Mapped[int] = so.mapped_column("val_col", sa.Integer, key="val_key") + + +class CompositeKeyedRecord(ReviewBase, FileLoadableTableInterface): + __tablename__ = "review_composite_keyed_record" + record_id: so.Mapped[int] = so.mapped_column( + "id_col", sa.Integer, key="id_key", primary_key=True, autoincrement=False, + ) + group_id: so.Mapped[int] = so.mapped_column( + "group_col", sa.Integer, key="group_key", primary_key=True, autoincrement=False, + ) + value: so.Mapped[int] = so.mapped_column("val_col", sa.Integer, key="val_key") + + +class ExpressionDefaultRecord(ReviewBase, FileLoadableTableInterface): + __tablename__ = "review_expression_default" + record_id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True, autoincrement=False) + token: so.Mapped[float] = so.mapped_column(sa.Float, server_default=sa.func.random()) + + +class DefaultKeyRecord(ReviewBase, FileLoadableTableInterface): + __tablename__ = "review_default_key" + record_id: so.Mapped[int] = so.mapped_column( + sa.Integer, primary_key=True, nullable=True, default=99, autoincrement=False, + ) + note: so.Mapped[str | None] = so.mapped_column(sa.String) + + +class InvalidNumericRecord(ReviewBase, FileLoadableTableInterface): + __tablename__ = "review_invalid_numeric" + record_id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) + amount: so.Mapped[float] = so.mapped_column(sa.Numeric(2, 5)) + + +def _create(session, model): + model.__table__.create(session.connection()) + session.commit() + + +def _load(session, tmp_path, model, text, **kwargs): + path = tmp_path / f"{model.__tablename__}.csv" + path.write_text(text, encoding="utf-8") + tag = None if session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA + table = cast(type[FileTableProtocol], model) + loaded = table.load_file(session, path, staging_schema_tag=tag, **kwargs) + session.commit() + return loaded + + +@pytest.fixture +def review_session(dialect_session): + _create(dialect_session, ReviewRecord) + return dialect_session + + +def test_malformed_row_ratio_counts_the_current_parsed_block(review_session, tmp_path): + policy = ConversionPolicy(errors=ErrorPolicy(on_error="reject_row", max_rejects=0.01)) + report = LoadReport() + loaded = _load( + review_session, tmp_path, ReviewRecord, + "record_id,score\n1,2,3\n" + "".join(f"{i},{i}\n" for i in range(2, 201)), + policy=policy, report=report, rejects_dir=tmp_path, + ) + assert loaded == 199 + assert (report.rows_read, report.rows_loaded, report.rows_rejected) == (200, 199, 1) + assert report.accounted + + +def test_unreadable_and_missing_values_keep_separate_rejection_reasons(review_session, tmp_path): + policy = ConversionPolicy(errors=ErrorPolicy(on_error="reject_row")) + _load( + review_session, tmp_path, ReviewRecord, + "record_id,score\n1,bad\n2,\n3,3\n", policy=policy, rejects_dir=tmp_path, + ) + with (tmp_path / "review_record.rejects.csv").open(newline="") as handle: + rejected = list(csv.DictReader(handle)) + assert [(row["row"], row["column"], row["reason"]) for row in rejected] == [ + ("1", "score", "cannot be read as Integer"), + ("2", "score", "required value is missing"), + ] + assert review_session.scalars(sa.select(ReviewRecord.record_id)).all() == [3] + + +def test_empty_text_can_be_preserved(review_session, tmp_path): + policy = ConversionPolicy(nulls=NullPolicy(empty_is_null=False)) + _load( + review_session, tmp_path, ReviewRecord, + "record_id,score,note\n1,2,\n", policy=policy, + ) + assert review_session.scalar(sa.select(ReviewRecord.note)) == "" + + +@pytest.mark.parametrize("supplied_reader", [False, True]) +def test_computed_columns_are_never_loaded(review_session, tmp_path, supplied_reader): + text = "record_id,score,twice\n1,2,wrong\n" + kwargs = {} + if supplied_reader: + path = tmp_path / "review_record.csv" + path.write_text(text) + kwargs["reader"] = DelimitedTextReader(path, ReviewRecord.model_columns()) + _load(review_session, tmp_path, ReviewRecord, text, **kwargs) + assert review_session.scalar(sa.select(ReviewRecord.twice)) == 4 + + +def test_supplied_reader_must_share_the_conversion_null_policy(review_session, tmp_path): + path = tmp_path / "review_record.csv" + path.write_text("record_id,score\n1,1\n") + reader = DelimitedTextReader(path, ReviewRecord.model_columns()) + policy = ConversionPolicy(nulls=NullPolicy(empty_is_null=False)) + with pytest.raises(ValueError, match="reader's null policy differs"): + _load(review_session, tmp_path, ReviewRecord, path.read_text(), reader=reader, policy=policy) + + +def test_tab_markers_ignore_padding_without_quote_semantics(review_session, tmp_path): + path = tmp_path / "review_record.tsv" + path.write_text( + "record_id\tscore\tnote\tevent_date\n" + "1\t2\t MISSING \t MISSING \n2\t3\tMISSING\tMISSING\n" + ) + policy = ConversionPolicy(nulls=NullPolicy( + text_markers=frozenset({"MISSING"}), marker_quoting="unquoted_only", + )) + reader = DelimitedTextReader(path, ReviewRecord.model_columns(), nulls=policy.nulls) + _load(review_session, tmp_path, ReviewRecord, "unused", reader=reader, policy=policy) + rows = review_session.execute(sa.select(ReviewRecord.note, ReviewRecord.event_date)).all() + assert [tuple(row) for row in rows] == [(None, None), (None, None)] + + +def test_legacy_text_markers_cover_all_case_variants(review_session, tmp_path): + markers = sorted(NullPolicy.LEGACY_TEXT_MARKERS) + assert {"NULL", "nUlL", "None", "nan", "N/a", "#N/A"} <= set(markers) + text = "record_id,score,note\n" + "".join( + f"{i},1,{marker}\n" for i, marker in enumerate([*markers, "keep"], 1) + ) + policy = ConversionPolicy(nulls=NullPolicy(text_markers=NullPolicy.LEGACY_TEXT_MARKERS)) + _load(review_session, tmp_path, ReviewRecord, text, policy=policy) + notes = review_session.scalars(sa.select(ReviewRecord.note).order_by(ReviewRecord.record_id)).all() + assert notes == [None] * len(markers) + ["keep"] + + +@pytest.mark.parametrize("normalise", [False, True]) +def test_defaulted_nullable_primary_key_must_be_in_the_file(dialect_session, tmp_path, normalise): + _create(dialect_session, DefaultKeyRecord) + with pytest.raises(IngestError, match="record_id"): + _load(dialect_session, tmp_path, DefaultKeyRecord, "note\na\n", normalise=normalise) + + +def test_null_primary_key_is_rejected_instead_of_taking_its_default(dialect_session, tmp_path): + _create(dialect_session, DefaultKeyRecord) + policy = ConversionPolicy(errors=ErrorPolicy(on_error="reject_row")) + _load( + dialect_session, tmp_path, DefaultKeyRecord, + "record_id,note\n,a\n1,b\n", policy=policy, rejects_dir=tmp_path, + ) + assert dialect_session.scalars(sa.select(DefaultKeyRecord.record_id)).all() == [1] + + +@pytest.mark.parametrize("model", [KeyedRecord, CompositeKeyedRecord]) +@pytest.mark.parametrize("strategy", ["insert", "upsert", "replace"]) +@pytest.mark.parametrize("batch_size", [None, 1]) +def test_merge_uses_database_names_when_column_keys_differ( + dialect_session, tmp_path, model, strategy, batch_size, +): + _create(dialect_session, model) + composite = model is CompositeKeyedRecord + header = "id_col,group_col,val_col\n" if composite else "id_col,val_col\n" + def row(key, value): + return f"{key},1,{value}\n" if composite else f"{key},{value}\n" + + _load(dialect_session, tmp_path, model, header + row(1, 5), merge_strategy="insert") + data = row(2, 20) if strategy == "insert" else row(1, 10) + row(2, 20) + _load( + dialect_session, tmp_path, model, header + data, + merge_strategy=strategy, merge_batch_size=batch_size, + ) + rows = dialect_session.execute(sa.select(model.record_id, model.value).order_by(model.record_id)) + assert [tuple(row) for row in rows] == [(1, 10 if strategy == "replace" else 5), (2, 20)] + + +@pytest.mark.parametrize("strategy", ["insert", "upsert", "replace"]) +@pytest.mark.parametrize("batch_size", [None, 2]) +@pytest.mark.parametrize("normalise", [False, True]) +def test_sql_expression_defaults_are_evaluated_for_each_merged_row( + dialect_session, tmp_path, strategy, batch_size, normalise, +): + _create(dialect_session, ExpressionDefaultRecord) + table = ExpressionDefaultRecord + if strategy != "insert": + _load(dialect_session, tmp_path, table, "record_id,token\n1,42\n", merge_strategy="insert") + text = "record_id,token\n" + "".join(f"{i},\n" for i in range(1, 11)) + "11,123\n" + _load( + dialect_session, tmp_path, table, text, merge_strategy=strategy, + merge_batch_size=batch_size, normalise=normalise, + ) + tokens = dialect_session.scalars(sa.select(table.token).order_by(table.record_id)).all() + assert len(tokens) == 11 + assert all(value is not None for value in tokens) + assert len(set(tokens[:10])) == 10 + assert tokens[-1] == 123 + if strategy == "upsert": + assert tokens[0] == 42 + + +def test_unsupported_numeric_declaration_names_the_table_and_column(tmp_path): + with pytest.raises(PolicyError, match=r"review_invalid_numeric.amount.*\(2, 5\)"): + BatchProcessor( + cast(type[FileTableProtocol], InvalidNumericRecord), ["record_id", "amount"], + ConversionPolicy(), LoadReport(), tmp_path / "rejects.csv", + ) + + +def test_infer_delim_remains_importable_from_loaders(tmp_path): + path = tmp_path / "data.tsv" + path.write_text("a\tb\n") + assert infer_delim(path) == "\t" + + +@pytest.mark.parametrize("normalise", [False, True]) +def test_nonnullable_computed_column_can_be_absent_from_the_file(tmp_path, normalise): + table = cast(type[FileTableProtocol], RequiredComputedRecord) + assert table.required_columns() == {"record_id", "score"} + processor = BatchProcessor( + table, ["record_id", "score"], ConversionPolicy(), LoadReport(), + tmp_path / "rejects.csv", normalise=normalise, + ) + output = processor.process(pa.record_batch({"record_id": ["1"], "score": ["2"]})) + assert output.column_names == ["record_id", "score"] + assert output.num_rows == 1 + + +def test_defaulted_nullable_primary_key_is_always_required(): + assert DefaultKeyRecord.required_columns() == {"record_id"} + assert "record_id" not in DefaultKeyRecord.column_defaults() + + +def test_column_defaults_protocol_has_a_callable_classmethod_function(): + descriptor = vars(ORMTableProtocol)["column_defaults"] + assert isinstance(descriptor, classmethod) + assert inspect.isfunction(descriptor.__func__) From b97b592c254a7edf9cb58daa0c2ef1695c7c14e5 Mon Sep 17 00:00:00 2001 From: Georgie Kennedy Date: Tue, 3 Nov 2026 11:49:55 +1100 Subject: [PATCH 16/21] Keep PostgreSQL foreign keys relaxed across commits without leaking the role into the pool relaxed_foreign_keys sets session_replication_role with SET LOCAL and reapplies it whenever a transaction begins on the selected bind, so a commit that switches pooled connections keeps keys relaxed and no connection is returned to the pool in replica. The role found in an open transaction is restored on exit. A scoped_session relaxes only the current thread's session. Docs say to flush ORM inserts inside the block. --- docs/loaders/staging.md | 11 ++ src/orm_loader/backends/postgres.py | 68 +++++++++- tests/backends/test_postgres_fk_scope.py | 155 +++++++++++++++++++++++ 3 files changed, 231 insertions(+), 3 deletions(-) create mode 100644 tests/backends/test_postgres_fk_scope.py diff --git a/docs/loaders/staging.md b/docs/loaders/staging.md index 8b8b281..1fd033d 100644 --- a/docs/loaders/staging.md +++ b/docs/loaders/staging.md @@ -66,5 +66,16 @@ On PostgreSQL and SQLite a load relaxes foreign keys from before it writes stagi To insert ORM objects built in Python that reference each other, wrap the inserts in `backend.relaxed_foreign_keys(session)`, which relaxes the keys on PostgreSQL and SQLite and leaves them enforced on DuckDB. +Autoflush is suspended by default. Flush the pending inserts inside the block so they run while checks are relaxed: + +```python +with backend.relaxed_foreign_keys(session): + session.add_all(objects) + session.flush() +session.commit() +``` + +On PostgreSQL, relaxation is transaction-local and is reapplied when a new transaction begins, even if the session switches pooled connections. Exiting the scope restores the prior role on any transaction still active. + DuckDB enforces the keys it creates, so load parent tables first. `load_file` already validates the keys DuckDB never enforces (cross-schema, cyclic and self-referential) after each table's merge, and `validate_foreign_keys` repeats that check, so it works the same way on every backend. A DuckDB table with foreign keys needs `resolved=` for both. diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index d421514..8f294b3 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -1,8 +1,10 @@ from __future__ import annotations import io +import logging +from contextlib import contextmanager, nullcontext from collections.abc import Iterable, Mapping, Sequence -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast import pyarrow as pa import pyarrow.csv as pacsv @@ -12,6 +14,7 @@ Dialect, qualified, ) +from sqlalchemy import event from sqlalchemy.dialects import postgresql from sqlalchemy.exc import OperationalError from sqlalchemy.sql.compiler import IdentifierPreparer @@ -33,6 +36,8 @@ from ..mappers.materialised_view_contracts import MaterializedViewIndex from ..tables.typing import FileTableProtocol +logger = logging.getLogger(__name__) + _VALID_PG_REPLICATION_ROLES = frozenset({"origin", "local", "replica"}) @@ -150,6 +155,63 @@ def restore_fk_check( sa.text(f"SET session_replication_role = '{safe_state}'"), **execute_options ) + @contextmanager + def relaxed_foreign_keys( + self, + session: so.Session | so.scoped_session[so.Session], + *, + mapper: Any | None = None, + no_autoflush: bool = True, + ): + """Relax checks in each transaction on the selected bind. + + SET LOCAL resets when a transaction ends. The after_begin listener + reapplies it even when a commit or rollback changes the pooled + connection. An existing transaction is relaxed immediately, and its + original role is restored when the scope exits. + """ + if isinstance(session, so.scoped_session): + session = cast(so.Session, session()) + target_engine = session.get_bind(mapper=mapper).engine + active_connection: sa.Connection | None = None + previous_role = "origin" + + def after_begin(_session, _transaction, connection: sa.Connection) -> None: + nonlocal active_connection, previous_role + if connection.engine is not target_engine: + return + if connection is not active_connection: + previous_role = self._normalize_fk_check_state( + connection.scalar(sa.text("SHOW session_replication_role")) + ) + active_connection = connection + connection.exec_driver_sql("SET LOCAL session_replication_role = 'replica'") + + scope = f"Table `{mapper.__tablename__}`: " if hasattr(mapper, "__tablename__") else "" + event.listen(session, "after_begin", after_begin) + try: + connection = session.connection(bind_arguments={"mapper": mapper}) + if connection is not active_connection: + after_begin(session, session.get_transaction(), connection) + logger.info("%sForeign key checks disabled.", scope) + with session.no_autoflush if no_autoflush else nullcontext(): + yield + except Exception: + session.rollback() + raise + finally: + event.remove(session, "after_begin", after_begin) + if ( + active_connection is not None + and not active_connection.closed + and not active_connection.invalidated + and active_connection.in_transaction() + ): + active_connection.exec_driver_sql( + f"SET LOCAL session_replication_role = '{previous_role}'" + ) + logger.info("%sForeign key checks restored.", scope) + @property def _commits_between_pagination_batches(self) -> bool: """True: bounds the WAL growth, lock duration, and dead-tuple @@ -229,8 +291,8 @@ def merge_insert( ) -> None: """Insert all staging rows into the target table. - session_replication_role='replica' is session-level and persists - across commits, so FK checks stay disabled for every batch. + relaxed_foreign_keys reapplies transaction-local replication-role + settings after each commit, keeping FK checks disabled for every batch. """ self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ diff --git a/tests/backends/test_postgres_fk_scope.py b/tests/backends/test_postgres_fk_scope.py new file mode 100644 index 0000000..29f4bab --- /dev/null +++ b/tests/backends/test_postgres_fk_scope.py @@ -0,0 +1,155 @@ +"""Foreign-key relaxation follows transactions, without leaking into the pool.""" + +import pytest +import sqlalchemy as sa +import sqlalchemy.orm as so +from sqlalchemy.exc import IntegrityError + +from orm_loader.backends import PostgresBackend + +pytestmark = [pytest.mark.postgresql, pytest.mark.db_dialect] + + +@pytest.fixture +def pooled_pg_engine(pg_db): + engine = sa.create_engine( + pg_db.connection.engine.url, pool_size=2, max_overflow=0, pool_use_lifo=False, + ) + try: + first = engine.connect() + second = engine.connect() + first.close() + second.close() + yield engine + finally: + engine.dispose() + + +@pytest.mark.parametrize("end_transaction", ["commit", "rollback"]) +def test_relaxation_follows_connection_swaps_without_affecting_other_sessions( + pooled_pg_engine, end_transaction, +): + backend = PostgresBackend() + with so.Session(pooled_pg_engine) as session: + with backend.relaxed_foreign_keys(session): + first_pid = session.scalar(sa.text("SELECT pg_backend_pid()")) + assert session.scalar(sa.text("SHOW session_replication_role")) == "replica" + getattr(session, end_transaction)() + second_pid = session.scalar(sa.text("SELECT pg_backend_pid()")) + assert second_pid != first_pid + assert session.scalar(sa.text("SHOW session_replication_role")) == "replica" + with so.Session(pooled_pg_engine) as unrelated: + assert unrelated.scalar(sa.text("SELECT pg_backend_pid()")) == first_pid + assert unrelated.scalar(sa.text("SHOW session_replication_role")) == "origin" + # Restore the still-open transaction immediately on context exit. + assert session.scalar(sa.text("SHOW session_replication_role")) == "origin" + session.commit() + # A later transaction must not re-enable relaxation. + assert session.scalar(sa.text("SHOW session_replication_role")) == "origin" + with pooled_pg_engine.connect() as first, pooled_pg_engine.connect() as second: + assert first.scalar(sa.text("SHOW session_replication_role")) == "origin" + assert second.scalar(sa.text("SHOW session_replication_role")) == "origin" + + +def test_failed_scope_rolls_back_and_removes_transaction_listener(pooled_pg_engine): + with so.Session(pooled_pg_engine) as session: + with pytest.raises(RuntimeError, match="stop the load"): + with PostgresBackend().relaxed_foreign_keys(session): + session.execute(sa.text("SELECT 1")) + session.commit() + assert session.scalar(sa.text("SHOW session_replication_role")) == "replica" + raise RuntimeError("stop the load") + assert session.scalar(sa.text("SHOW session_replication_role")) == "origin" + with pooled_pg_engine.connect() as first, pooled_pg_engine.connect() as second: + assert first.scalar(sa.text("SHOW session_replication_role")) == "origin" + assert second.scalar(sa.text("SHOW session_replication_role")) == "origin" + + +def test_nested_scopes_restore_the_existing_transaction_role(pooled_pg_engine): + backend = PostgresBackend() + with so.Session(pooled_pg_engine) as session: + session.execute(sa.text("SET LOCAL session_replication_role = 'local'")) + with backend.relaxed_foreign_keys(session): + assert session.scalar(sa.text("SHOW session_replication_role")) == "replica" + with backend.relaxed_foreign_keys(session): + assert session.scalar(sa.text("SHOW session_replication_role")) == "replica" + assert session.scalar(sa.text("SHOW session_replication_role")) == "replica" + assert session.scalar(sa.text("SHOW session_replication_role")) == "local" + session.rollback() + assert session.scalar(sa.text("SHOW session_replication_role")) == "origin" + + +def test_relaxation_is_limited_to_the_mapper_bind(pooled_pg_engine): + from tests.models import EnumTable, SimpleTable + + other_engine = sa.create_engine(pooled_pg_engine.url) + try: + with so.Session(binds={SimpleTable: pooled_pg_engine, EnumTable: other_engine}) as session: + with PostgresBackend().relaxed_foreign_keys(session, mapper=SimpleTable): + for _ in range(2): + target = session.connection(bind_arguments={"mapper": SimpleTable}) + unrelated = session.connection(bind_arguments={"mapper": EnumTable}) + assert target.scalar(sa.text("SHOW session_replication_role")) == "replica" + assert unrelated.scalar(sa.text("SHOW session_replication_role")) == "origin" + session.commit() + finally: + other_engine.dispose() + + +def test_scoped_session_relaxation_does_not_affect_another_thread(pooled_pg_engine): + from concurrent.futures import ThreadPoolExecutor + + registry = so.scoped_session(so.sessionmaker(bind=pooled_pg_engine)) + + def other_thread_role(): + try: + return registry.scalar(sa.text("SHOW session_replication_role")) + finally: + registry.remove() + + try: + with ThreadPoolExecutor(max_workers=1) as worker: + with PostgresBackend().relaxed_foreign_keys(registry): + assert registry.scalar(sa.text("SHOW session_replication_role")) == "replica" + other_role = worker.submit(other_thread_role).result(timeout=10) + restored_role = registry.scalar(sa.text("SHOW session_replication_role")) + assert (other_role, restored_role) == ("origin", "origin") + registry.commit() + assert registry.scalar(sa.text("SHOW session_replication_role")) == "origin" + finally: + registry.remove() + + +class _ScopeBase(so.DeclarativeBase): + pass + + +class _ScopeParent(_ScopeBase): + __tablename__ = "_fk_scope_parent" + id: so.Mapped[int] = so.mapped_column(primary_key=True, autoincrement=False) + + +class _ScopeChild(_ScopeBase): + __tablename__ = "_fk_scope_child" + id: so.Mapped[int] = so.mapped_column(primary_key=True, autoincrement=False) + parent_id: so.Mapped[int] = so.mapped_column(sa.ForeignKey("_fk_scope_parent.id")) + + +@pytest.mark.parametrize("flush_inside", [False, True]) +def test_orm_inserts_must_be_flushed_inside_relaxation(pooled_pg_engine, flush_inside): + _ScopeBase.metadata.create_all(pooled_pg_engine) + try: + with so.Session(pooled_pg_engine) as session: + with PostgresBackend().relaxed_foreign_keys(session): + session.add_all([_ScopeChild(id=1, parent_id=999)]) + if flush_inside: + session.flush() + if flush_inside: + session.commit() + else: + with pytest.raises(IntegrityError): + session.commit() + with so.Session(pooled_pg_engine) as verify: + assert verify.scalar(sa.select(sa.func.count()).select_from(_ScopeChild)) == int(flush_inside) + finally: + _ScopeBase.metadata.drop_all(pooled_pg_engine) From d0d01a98948e8c02284a25140bcc468ac72985c5 Mon Sep 17 00:00:00 2001 From: Georgie Kennedy Date: Tue, 3 Nov 2026 11:51:20 +1100 Subject: [PATCH 17/21] dependabot --- uv.lock | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/uv.lock b/uv.lock index 5d15524..0811e6d 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 5 +revision = 3 requires-python = ">=3.12" resolution-markers = [ "python_full_version >= '3.14'", @@ -62,11 +62,11 @@ wheels = [ [[package]] name = "certifi" -version = "2026.1.4" +version = "2026.10.23" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/e0/2d/a891ca51311197f6ad14a7ef42e2399f36cf2f9bd44752b3dc4eab60fdc5/certifi-2026.1.4.tar.gz", hash = "sha256:ac726dd470482006e014ad384921ed6438c457018f4b3d204aea4281258b2120", size = 154268, upload-time = "2026-01-04T02:42:41.825Z" } +sdist = { url = "https://files.pythonhosted.org/packages/fd/11/b524cbd9bbdaca58608e4cf4053fd463417549e1959fa8d3aa18319f9555/certifi-2026.10.23.tar.gz", hash = "sha256:a9895e6e2758b371e0545e321a3a8fd1fcfd9829ada8824a936dfb16aaaf36cd", size = 136225, upload-time = "2026-10-23T11:36:16.478Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e6/ad/3cc14f097111b4de0040c83a525973216457bbeeb63739ef1ed275c1c021/certifi-2026.1.4-py3-none-any.whl", hash = "sha256:9943707519e4add1115f44c2bc244f782c0249876bf51b6599fee1ffbedd685c", size = 152900, upload-time = "2026-01-04T02:42:40.15Z" }, + { url = "https://files.pythonhosted.org/packages/32/d1/a3652fe2c19134d0d15922177821d863cbd7e842d53a21bbe7e8e3a1b5ac/certifi-2026.10.23-py3-none-any.whl", hash = "sha256:e9b54de76e530bb6cf0ccd036787103acb58fd42ba930ea8bf948ddbaa2f61d1", size = 135082, upload-time = "2026-10-23T11:36:14.763Z" }, ] [[package]] @@ -1151,11 +1151,11 @@ wheels = [ [[package]] name = "setuptools" -version = "84.0.0" +version = "84.2.1" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/6d/44/f5da03a8ef95d369145c5bb53050e7877c9f3d312e128605fd9504829143/setuptools-84.0.0.tar.gz", hash = "sha256:f4695c21257f0d9b537ec2692c941d02ee143b7cc1276941349a546573b2ef73", size = 1168449, upload-time = "2026-08-08T18:27:58.365Z" } +sdist = { url = "https://files.pythonhosted.org/packages/70/b3/c19426f783614d625699a4b247faf628ea9ab1db228674491767b417f9e1/setuptools-84.2.1.tar.gz", hash = "sha256:3ff8b3b5492a0aa66960f4419365a86b84e013070d4313e4d4a6894f5d64d6bc", size = 1230289, upload-time = "2026-10-20T19:01:37.742Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/95/9c/c510029fc6ef33a6275cd2c5d3cecd6613dfd6aa401d57c54f1c18852ccf/setuptools-84.0.0-py3-none-any.whl", hash = "sha256:51a52592b3b99e102b609654876bd65f19f999935166d1352678931132b0c670", size = 818216, upload-time = "2026-08-08T18:27:56.719Z" }, + { url = "https://files.pythonhosted.org/packages/ef/55/6dc6d6bb953bdfc52519a0059592c52b6e3b8ad51cf8256dffceb55b0610/setuptools-84.2.1-py3-none-any.whl", hash = "sha256:ca140b45d5325a9478f880e351484f62fef15f336462e449ca0c03c873d589c2", size = 876263, upload-time = "2026-10-20T19:01:35.678Z" }, ] [[package]] From f24b7cb1ea9c5b450f7c0a2d6e76bc5435945dd2 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 4 Nov 2026 22:53:24 +0000 Subject: [PATCH 18/21] Apply expression defaults per row in the merge, check fractional reject limits at the end, keep computed columns out of loads - model_columns() excludes computed columns, so the reader, processor and required_columns agree - a primary key is always required in the file and never takes a default - SQL-expression defaults are applied by COALESCE in the merge and are not written by the processor - a count limit stops a load at once, a fraction is checked against all rows read, and a load that rejected rows logs a warning - Numeric is cast up to 38 digits and kept as validated text beyond that or without a precision - handlers validate the declared column type through validate(), which PassthroughHandler no longer trips over - the review regression tests move into topical test files --- src/orm_loader/conversion/handlers/base.py | 3 + src/orm_loader/conversion/handlers/numbers.py | 45 ++-- src/orm_loader/loaders/processor.py | 77 ++++--- src/orm_loader/tables/loadable_table.py | 12 +- src/orm_loader/tables/orm_table.py | 49 ++-- tests/conversion/test_number_handlers.py | 22 +- tests/loaders/test_processor.py | 63 ++++- ...regressions.py => test_special_columns.py} | 216 +++++------------- tests/tables/test_typing_protocols.py | 8 + 9 files changed, 233 insertions(+), 262 deletions(-) rename tests/loaders/{test_review_regressions.py => test_special_columns.py} (50%) diff --git a/src/orm_loader/conversion/handlers/base.py b/src/orm_loader/conversion/handlers/base.py index 2dc36a7..e785135 100644 --- a/src/orm_loader/conversion/handlers/base.py +++ b/src/orm_loader/conversion/handlers/base.py @@ -47,6 +47,9 @@ def expected(self, sa_type: TypeEngine[Any]) -> str: """Return the name reported in an error for an unreadable value.""" return type(sa_type).__name__ + def validate(self, sa_type: TypeEngine[Any]) -> None: + """Raise ``ValueError`` when the handler cannot read the column type as declared.""" + class PassthroughHandler(ColumnHandler): """Returns a column unchanged, for column types that are not converted.""" diff --git a/src/orm_loader/conversion/handlers/numbers.py b/src/orm_loader/conversion/handlers/numbers.py index 01f16a1..cfdb9f6 100644 --- a/src/orm_loader/conversion/handlers/numbers.py +++ b/src/orm_loader/conversion/handlers/numbers.py @@ -84,32 +84,28 @@ def check(self, values: pa.Array, sa_type: TypeEngine[Any]) -> pa.Array | None: class NumericHandler(NumberHandler): """Accepts decimal forms and keeps the column's declared precision and scale. - A precision without a scale has scale 0. Values without a declared - precision, and supported declarations wider than Arrow's decimal types, - stay as decimal text until the database stores them. Trailing fractional - zeros are dropped, and a value outside a declared precision or scale is - unreadable. + A precision without a scale has scale 0. A value is checked against the + declared digits and is unreadable when it does not fit, and trailing + fractional zeros are dropped. Up to 38 digits the value is cast to an + Arrow decimal. A column without a precision, or with more than 38 digits, + stays decimal text until the database stores it. """ @staticmethod def precision_scale(sa_type: TypeEngine[Any]) -> tuple[int | None, int | None]: """Return a column's declared precision and scale.""" - precision = getattr(sa_type, "precision", None) - scale = getattr(sa_type, "scale", None) - if precision is not None and ( - precision < 1 or (scale is not None and (scale < 0 or scale > precision)) - ): + return getattr(sa_type, "precision", None), getattr(sa_type, "scale", None) + + def validate(self, sa_type: TypeEngine[Any]) -> None: + precision, scale = self.precision_scale(sa_type) + if precision is not None and (precision < 1 or (scale is not None and not 0 <= scale <= precision)): raise ValueError(f"unsupported Numeric precision/scale declaration ({precision}, {scale})") - return precision, scale def arrow_type(self, sa_type: TypeEngine[Any]) -> pa.DataType: precision, scale = self.precision_scale(sa_type) - if precision is None or precision > 76: + if precision is None or precision > 38: return pa.string() - scale = scale or 0 - if precision > 38: - return pa.decimal256(precision, scale) - return pa.decimal128(precision, scale) + return pa.decimal128(precision, scale or 0) def clean(self, raw: pa.Array, trimmed: pa.Array) -> pa.Array: return pc.replace_substring_regex(trimmed, r"(\.\d*?)0+$", r"\1") @@ -121,15 +117,10 @@ def pattern(self, sa_type: TypeEngine[Any]) -> str: than a decimal holds, so both parts are bounded here. """ precision, scale = self.precision_scale(sa_type) - if precision is None: - if scale is None: - return r"^[+-]?(?:\d+(?:\.\d*)?|\.\d+)$" - forms = [rf"\d+(\.\d{{0,{scale}}})?"] - if scale: - forms.append(rf"\.\d{{1,{scale}}}") - return rf"^[+-]?({'|'.join(forms)})$" - scale = scale or 0 - forms = [rf"0*\d{{0,{precision - scale}}}(\.\d{{0,{scale}}})?"] - if scale: - forms.append(rf"\.\d{{1,{scale}}}") + digits = None if precision is None and scale is None else scale or 0 + whole = r"\d+" if precision is None else rf"0*\d{{0,{precision - (digits or 0)}}}" + fraction = r"\d*" if digits is None else rf"\d{{0,{digits}}}" + forms = [rf"{whole}(\.{fraction})?"] + if digits != 0: + forms.append(r"\.\d+" if digits is None else rf"\.\d{{1,{digits}}}") return rf"^[+-]?({'|'.join(forms)})$" diff --git a/src/orm_loader/loaders/processor.py b/src/orm_loader/loaders/processor.py index 615e4a8..62fa9d2 100644 --- a/src/orm_loader/loaders/processor.py +++ b/src/orm_loader/loaders/processor.py @@ -102,17 +102,21 @@ class BatchProcessor: Lines the reader left out as malformed, each as its line number in the file and the reason. They are written to the rejects file and count as rows read and rejected. + Attributes ---------- out_columns - Columns of typed batches: loadable file columns, then the columns the - file lacks that have a literal or SQL-expression default. + Columns of the typed batches: the file's columns, then the columns + the file lacks that have a literal default. Notes ----- - Missing values take literal defaults. SQL-expression defaults remain - null in staging so the database can evaluate them per row during merge. - A missing value left in a required column is an error. + A missing value takes the column's literal default, and a column the file + lacks is filled with it. The default is read by the column's handler like + a value from the file. A SQL-expression default is not applied here: the + value stays null in staging and the merge applies the default for each + row. A missing value left in any other non-nullable column is a required + value that is missing. """ def __init__( @@ -132,20 +136,15 @@ def __init__( self.report = report self.normalise = normalise self.table_name = tableclass.__tablename__ - model_columns = { - name: column - for name, column in tableclass.model_columns().items() - if column.computed is None - } - self.columns = [name for name in columns if name in model_columns] - all_defaults = tableclass.column_defaults() if normalise else {} - self.sql_expression_defaults = tableclass.sql_expression_defaults() if normalise else {} + model_columns = tableclass.model_columns() + self.columns = columns + expressions = tableclass.sql_expression_defaults() if normalise else {} defaults = { name: default - for name, default in all_defaults.items() - if name not in self.sql_expression_defaults + for name, default in (tableclass.column_defaults() if normalise else {}).items() + if name not in expressions } - self.out_columns = [*self.columns, *(name for name in all_defaults if name not in self.columns)] + self.out_columns = [*columns, *(name for name in defaults if name not in columns)] self.types = {name: underlying_type(model_columns[name].type) for name in self.out_columns} self.handlers = ( {name: policy.for_column(self.types[name], self.table_name, name) for name in self.out_columns} @@ -154,7 +153,7 @@ def __init__( ) for name, handler in self.handlers.items(): try: - handler.arrow_type(self.types[name]) + handler.validate(self.types[name]) except ValueError as exc: raise PolicyError( table_name=self.table_name, @@ -167,8 +166,7 @@ def __init__( self.required = [ name for name in self.out_columns - if (not model_columns[name].nullable or name in primary_keys) - and name not in self.sql_expression_defaults + if (not model_columns[name].nullable or name in primary_keys) and name not in expressions ] self._literals: dict[str, pa.Scalar] = {} self.rejects = RejectsWriter(rejects_path) @@ -180,12 +178,13 @@ def stream(self, batches: Iterable[pa.RecordBatch]) -> Iterator[pa.Table]: """Yield typed tables for the raw batches, closing the rejects file afterwards.""" try: for raw in batches: - self._reject_skipped(check_limit=False) + self._reject_skipped() table = self.process(raw) - self._check_limit() if table.num_rows: yield table self._reject_skipped() + self._check_limit(final=True) + self._warn_if_rejected() finally: self.rejects.close() @@ -216,7 +215,7 @@ def process(self, raw: pa.RecordBatch) -> pa.Table: self._parsed += raw.num_rows self.report.rows_read += raw.num_rows if not self.normalise: - return pa.Table.from_batches([raw.select(self.columns)]) + return pa.Table.from_batches([raw]) results = { name: self.handlers[name].interpret(raw.column(name), self.types[name], self.policy.nulls) for name in self.columns @@ -231,12 +230,6 @@ def process(self, raw: pa.RecordBatch) -> pa.Table: def _fill_defaults(self, results: dict[str, Interpretation], count: int) -> None: """Give a missing value, and a column the file lacks, the column's default.""" - for name in self.sql_expression_defaults: - if name not in results: - results[name] = Interpretation( - pa.nulls(count, self.handlers[name].arrow_type(self.types[name])), - no_failures(count), - ) for name, default in self.defaults.items(): if name not in results: results[name] = Interpretation(pa.repeat(self._fill_value(name, default), count), no_failures(count)) @@ -366,7 +359,7 @@ def _row_numbers(self, first: int, count: int) -> pa.Array: positions = range(first + 1, first + count + 1) return pa.array([p + bisect_left(self._shifts, p) for p in positions], pa.int64()) - def _reject_skipped(self, *, check_limit: bool = True) -> None: + def _reject_skipped(self) -> None: """Write the lines the reader left out to the rejects file and remember where they were.""" count = len(self.skipped) if not count: @@ -388,14 +381,32 @@ def _reject_skipped(self, *, check_limit: bool = True) -> None: self.report.rows_read += count self.report.rows_rejected += count self.report.rejects_path = self.rejects.path - if check_limit: - self._check_limit() + self._check_limit() - def _check_limit(self) -> None: + def _check_limit(self, final: bool = False) -> None: + """Raise when the rejected rows pass ``max_rejects``. + + A count is checked whenever rows are rejected. A fraction needs the + total number of rows, so it is checked once, when ``final`` is set at + the end of the stream. + """ limit = self.policy.errors.max_rejects if limit is None: return rejected, read = self.report.rows_rejected, self.report.rows_read - over = rejected / read > limit if isinstance(limit, float) and limit < 1 else rejected > limit + if isinstance(limit, float) and limit < 1: + over = final and read > 0 and rejected / read > limit + else: + over = rejected > limit if over: raise RejectLimitError(rejected=rejected, read=read, limit=limit) + + def _warn_if_rejected(self) -> None: + if self.report.rows_rejected: + logger.warning( + "%s: %d of %d rows rejected, see %s", + self.table_name, + self.report.rows_rejected, + self.report.rows_read, + self.report.rejects_path, + ) diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index f8e0580..3f18d25 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -325,21 +325,11 @@ def load_file( if reader is None: reader = open_reader( path, - { - name: column - for name, column in cls.model_columns().items() - if column.computed is None - }, - nulls=policy.nulls, + cls.model_columns(), quote_mode=quote_mode, chunksize=chunksize, skip_malformed=policy.errors.on_error == "reject_row", ) - elif reader.nulls != policy.nulls: - raise ValueError( - "the reader's null policy differs from policy.nulls; " - "build the reader with nulls=policy.nulls" - ) backend = resolve_backend( session, diff --git a/src/orm_loader/tables/orm_table.py b/src/orm_loader/tables/orm_table.py index 1e336d3..6f72954 100644 --- a/src/orm_loader/tables/orm_table.py +++ b/src/orm_loader/tables/orm_table.py @@ -107,26 +107,26 @@ def pk_names(cls) -> list[str]: @classmethod def model_columns(cls) -> dict[str, sa.ColumnElement[Any]]: """ - Return all mapped columns for the table. + Return the mapped columns a file can supply. + + Columns are keyed by their database name, which is the name that file + headers and staging tables use, even when ``Column.key`` differs. + Computed columns are generated by the database and are not included. Returns ------- dict[str, sqlalchemy.ColumnElement] - A mapping of database column name to column object. This is the name - used by file headers and staging tables, even when ``Column.key`` - differs. + A mapping of database column name to column object. """ - mapper = cls.mapper_for() - return {c.name: c for c in mapper.columns} - + return {c.name: c for c in cls.mapper_for().columns if c.computed is None} + @classmethod def required_columns(cls) -> set[str]: """ Return column names that must be present in inbound data. - A non-nullable column is required when it has no default that can be - supplied (see ``column_defaults``). Primary keys are always required - by the file loader. + A column is required if it is a primary key, or if it is non-nullable + and has no default that can be supplied (see ``column_defaults``). This method is intended for ingestion-time validation and does not attempt to enforce schema semantics beyond insert viability. @@ -138,22 +138,22 @@ def required_columns(cls) -> set[str]: """ defaults = cls.column_defaults() return { - c.name - for c in cls.mapper_for().columns - if c.computed is None - and (c.primary_key or (not c.nullable and c.name not in defaults)) + name + for name, c in cls.model_columns().items() + if c.primary_key or (not c.nullable and name not in defaults) } @classmethod def column_defaults(cls) -> dict[str, Any]: """ - Return defaults the file loader can supply for non-key, non-computed columns. + Return the default of each non-key column that has one that can be supplied. A server default given as text and a scalar Python default are returned as the value itself. A server default given as a SQL expression is returned as that expression, for the database to evaluate. A callable or context-sensitive default cannot be supplied - for a whole file, so its column is not listed. + for a whole file, so its column is not listed. A primary key never + takes a default, because a file has to name it. Returns ------- @@ -161,19 +161,26 @@ def column_defaults(cls) -> dict[str, Any]: A mapping of column name to the value, or SQL expression, of its default. """ defaults: dict[str, Any] = {} - for c in cls.mapper_for().columns: - if c.computed is not None or c.primary_key: + for name, c in cls.model_columns().items(): + if c.primary_key: continue server_default = getattr(c.server_default, "arg", None) if server_default is not None: - defaults[c.name] = server_default + defaults[name] = server_default elif c.default is not None and c.default.is_scalar: - defaults[c.name] = c.default.arg # ty: ignore[unresolved-attribute] + defaults[name] = c.default.arg return defaults @classmethod def sql_expression_defaults(cls) -> dict[str, sa.ClauseElement]: - """Return server defaults that should be evaluated for each merged row.""" + """ + Return the defaults that are SQL expressions, which the database evaluates for each row. + + Returns + ------- + dict[str, sqlalchemy.ClauseElement] + A mapping of column name to its SQL expression default. + """ return { name: default for name, default in cls.column_defaults().items() diff --git a/tests/conversion/test_number_handlers.py b/tests/conversion/test_number_handlers.py index fe00e26..2bcd1c6 100644 --- a/tests/conversion/test_number_handlers.py +++ b/tests/conversion/test_number_handlers.py @@ -83,14 +83,6 @@ def test_numeric_without_precision_preserves_every_digit_as_text(raw): assert read(raw, sa.Numeric()) == raw -def test_wide_numeric_uses_decimal256_and_enforces_its_declared_shape(): - from orm_loader.conversion.handlers import NumericHandler - - assert NumericHandler().arrow_type(sa.Numeric(50, 2)) == pa.decimal256(50, 2) - raw = "123456789012345678901234567890123456789012.25" - assert read(raw, sa.Numeric(50, 2)) == Decimal(raw) - assert unreadable("1.555", sa.Numeric(50, 2)) - assert unreadable("9" * 49, sa.Numeric(50, 2)) @pytest.mark.parametrize("raw", ["abc", "1e5", ".", "--1"]) @@ -158,11 +150,13 @@ def test_numeric_with_scale_zero_accepts_a_zero_fraction_only(): ) -def test_numeric_wider_than_decimal256_preserves_text_and_declared_bounds(): +@pytest.mark.parametrize("precision", [50, 80]) +def test_numeric_wider_than_an_arrow_decimal_preserves_text_and_declared_bounds(precision): from orm_loader.conversion.handlers import NumericHandler - assert NumericHandler().arrow_type(sa.Numeric(80, 2)) == pa.string() - raw = "9" * 78 + ".12" - assert read(raw, sa.Numeric(80, 2)) == raw - assert unreadable("9" * 79, sa.Numeric(80, 2)) - assert unreadable("1.234", sa.Numeric(80, 2)) + column = sa.Numeric(precision, 2) + assert NumericHandler().arrow_type(column) == pa.string() + raw = "9" * (precision - 2) + ".12" + assert read(raw, column) == raw + assert unreadable("9" * (precision - 1), column) + assert unreadable("1.234", column) diff --git a/tests/loaders/test_processor.py b/tests/loaders/test_processor.py index 97108fe..74fc878 100644 --- a/tests/loaders/test_processor.py +++ b/tests/loaders/test_processor.py @@ -1,6 +1,7 @@ """BatchProcessor: raw batches in, typed rows, rejects and accounting out, no database.""" import csv +import logging from datetime import date from pathlib import Path @@ -118,9 +119,67 @@ def test_max_rejects_stops_a_load_over_a_count_or_a_fraction(tmp_path, limit, ro rows = [(str(i), "a", "x" if i <= rows_bad else "1", "") for i in range(1, 5)] if should_stop: with pytest.raises(RejectLimitError): - processor.process(_batch(*rows)) + list(processor.stream([_batch(*rows)])) else: - assert processor.process(_batch(*rows)).num_rows == 4 - rows_bad + assert sum(t.num_rows for t in processor.stream([_batch(*rows)])) == 4 - rows_bad + + +def test_a_count_limit_stops_the_stream_before_the_next_batch(tmp_path): + policy = ConversionPolicy(errors=ErrorPolicy(on_error="reject_row", max_rejects=1)) + processor, report = _processor(tmp_path, policy) + batches = [_batch(("1", "a", "x", ""), ("2", "a", "y", "")), _batch(("3", "a", "1", ""))] + with pytest.raises(RejectLimitError): + list(processor.stream(batches)) + assert report.rows_read == 2 + + +def test_a_fraction_is_checked_against_all_rows_read_not_the_rows_so_far(tmp_path): + policy = ConversionPolicy(errors=ErrorPolicy(on_error="reject_row", max_rejects=0.01)) + processor, report = _processor(tmp_path, policy) + early_rejects = _batch(*[(str(i), "a", "x" if i <= 5 else "1", "") for i in range(1, 51)]) + clean = _batch(*[(str(i), "a", "1", "") for i in range(51, 5051)]) + assert sum(t.num_rows for t in processor.stream([early_rejects, clean])) == 5045 + assert (report.rows_read, report.rows_rejected) == (5050, 5) + + +def test_a_fraction_over_the_limit_stops_the_load_when_the_stream_ends(tmp_path): + policy = ConversionPolicy(errors=ErrorPolicy(on_error="reject_row", max_rejects=0.01)) + processor, _ = _processor(tmp_path, policy) + batches = [_batch(("1", "a", "x", ""), *[(str(i), "a", "1", "") for i in range(2, 51)])] + with pytest.raises(RejectLimitError, match="1 of 50"): + list(processor.stream(batches)) + + +def test_a_malformed_line_is_counted_with_the_rows_of_its_own_block(tmp_path): + policy = ConversionPolicy(errors=ErrorPolicy(on_error="reject_row", max_rejects=0.01)) + processor, report = _processor(tmp_path, policy, skipped=[(2, "expected 4 columns, got 5")]) + clean = _batch(*[(str(i), "a", "1", "") for i in range(2, 201)]) + assert sum(t.num_rows for t in processor.stream([clean])) == 199 + assert (report.rows_read, report.rows_rejected) == (200, 1) + + +def test_rejected_rows_are_reported_in_a_warning(tmp_path, caplog): + processor, _ = _processor(tmp_path, REJECT) + with caplog.at_level(logging.WARNING): + list(processor.stream([_batch(("1", "a", "x", ""), ("2", "a", "1", ""))])) + assert "1 of 2 rows rejected" in caplog.text + + +def test_a_clean_stream_logs_no_warning(tmp_path, caplog): + processor, _ = _processor(tmp_path, REJECT) + with caplog.at_level(logging.WARNING): + list(processor.stream([_batch(("1", "a", "1", ""))])) + assert "rejected" not in caplog.text + + +def test_an_unreadable_value_and_a_missing_value_keep_their_own_reasons(tmp_path): + processor, _ = _processor(tmp_path, REJECT, table=DemoRequiredRecord, columns=["demo_id", "demo_label"]) + processor.process(_batch(("1", "a"), ("x", "b"), ("2", ""), columns=["demo_id", "demo_label"])) + processor.rejects.close() + assert [(r["row"], r["column"], r["reason"]) for r in _rejects(tmp_path)] == [ + ("2", "demo_id", "cannot be read as Integer"), + ("3", "demo_label", "required value is missing"), + ] def test_stream_skips_empty_results_and_closes_the_rejects_file(tmp_path): diff --git a/tests/loaders/test_review_regressions.py b/tests/loaders/test_special_columns.py similarity index 50% rename from tests/loaders/test_review_regressions.py rename to tests/loaders/test_special_columns.py index d4dc01b..58d1c70 100644 --- a/tests/loaders/test_review_regressions.py +++ b/tests/loaders/test_special_columns.py @@ -1,8 +1,5 @@ -"""Regression coverage for PR #46's ingestion review findings.""" +"""Columns a file load treats specially: computed, renamed, defaulted and primary-key columns.""" -import csv -import inspect -from datetime import date from typing import cast import pyarrow as pa @@ -11,44 +8,42 @@ import sqlalchemy.orm as so from orm_loader.backends import STAGING_SCHEMA -from orm_loader.conversion import ConversionPolicy, ErrorPolicy, NullPolicy +from orm_loader.conversion import ConversionPolicy, ErrorPolicy from orm_loader.helpers import IngestError, PolicyError -from orm_loader.loaders import DelimitedTextReader, LoadReport, infer_delim +from orm_loader.loaders import DelimitedTextReader, LoadReport from orm_loader.loaders.processor import BatchProcessor from orm_loader.tables import FileLoadableTableInterface -from orm_loader.tables.typing import FileTableProtocol, ORMTableProtocol +from orm_loader.tables.typing import FileTableProtocol -class ReviewBase(so.DeclarativeBase): +class SpecialBase(so.DeclarativeBase): pass -class ReviewRecord(ReviewBase, FileLoadableTableInterface): - __tablename__ = "review_record" +class ComputedRecord(SpecialBase, FileLoadableTableInterface): + __tablename__ = "special_computed" record_id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True, autoincrement=False) score: so.Mapped[int] = so.mapped_column(sa.Integer, nullable=False) - note: so.Mapped[str | None] = so.mapped_column(sa.String) - event_date: so.Mapped[date | None] = so.mapped_column(sa.Date) twice: so.Mapped[int | None] = so.mapped_column(sa.Integer, sa.Computed("score * 2")) -class RequiredComputedRecord(ReviewBase, FileLoadableTableInterface): - __tablename__ = "review_required_computed" +class RequiredComputedRecord(SpecialBase, FileLoadableTableInterface): + __tablename__ = "special_required_computed" record_id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) score: so.Mapped[int] = so.mapped_column(sa.Integer, nullable=False) twice: so.Mapped[int] = so.mapped_column(sa.Integer, sa.Computed("score * 2"), nullable=False) -class KeyedRecord(ReviewBase, FileLoadableTableInterface): - __tablename__ = "review_keyed_record" +class KeyedRecord(SpecialBase, FileLoadableTableInterface): + __tablename__ = "special_keyed" record_id: so.Mapped[int] = so.mapped_column( "id_col", sa.Integer, key="id_key", primary_key=True, autoincrement=False, ) value: so.Mapped[int] = so.mapped_column("val_col", sa.Integer, key="val_key") -class CompositeKeyedRecord(ReviewBase, FileLoadableTableInterface): - __tablename__ = "review_composite_keyed_record" +class CompositeKeyedRecord(SpecialBase, FileLoadableTableInterface): + __tablename__ = "special_composite_keyed" record_id: so.Mapped[int] = so.mapped_column( "id_col", sa.Integer, key="id_key", primary_key=True, autoincrement=False, ) @@ -58,22 +53,22 @@ class CompositeKeyedRecord(ReviewBase, FileLoadableTableInterface): value: so.Mapped[int] = so.mapped_column("val_col", sa.Integer, key="val_key") -class ExpressionDefaultRecord(ReviewBase, FileLoadableTableInterface): - __tablename__ = "review_expression_default" +class ExpressionDefaultRecord(SpecialBase, FileLoadableTableInterface): + __tablename__ = "special_expression_default" record_id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True, autoincrement=False) token: so.Mapped[float] = so.mapped_column(sa.Float, server_default=sa.func.random()) -class DefaultKeyRecord(ReviewBase, FileLoadableTableInterface): - __tablename__ = "review_default_key" +class DefaultKeyRecord(SpecialBase, FileLoadableTableInterface): + __tablename__ = "special_default_key" record_id: so.Mapped[int] = so.mapped_column( sa.Integer, primary_key=True, nullable=True, default=99, autoincrement=False, ) note: so.Mapped[str | None] = so.mapped_column(sa.String) -class InvalidNumericRecord(ReviewBase, FileLoadableTableInterface): - __tablename__ = "review_invalid_numeric" +class InvalidNumericRecord(SpecialBase, FileLoadableTableInterface): + __tablename__ = "special_invalid_numeric" record_id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) amount: so.Mapped[float] = so.mapped_column(sa.Numeric(2, 5)) @@ -93,112 +88,30 @@ def _load(session, tmp_path, model, text, **kwargs): return loaded -@pytest.fixture -def review_session(dialect_session): - _create(dialect_session, ReviewRecord) - return dialect_session - - -def test_malformed_row_ratio_counts_the_current_parsed_block(review_session, tmp_path): - policy = ConversionPolicy(errors=ErrorPolicy(on_error="reject_row", max_rejects=0.01)) - report = LoadReport() - loaded = _load( - review_session, tmp_path, ReviewRecord, - "record_id,score\n1,2,3\n" + "".join(f"{i},{i}\n" for i in range(2, 201)), - policy=policy, report=report, rejects_dir=tmp_path, - ) - assert loaded == 199 - assert (report.rows_read, report.rows_loaded, report.rows_rejected) == (200, 199, 1) - assert report.accounted - - -def test_unreadable_and_missing_values_keep_separate_rejection_reasons(review_session, tmp_path): - policy = ConversionPolicy(errors=ErrorPolicy(on_error="reject_row")) - _load( - review_session, tmp_path, ReviewRecord, - "record_id,score\n1,bad\n2,\n3,3\n", policy=policy, rejects_dir=tmp_path, - ) - with (tmp_path / "review_record.rejects.csv").open(newline="") as handle: - rejected = list(csv.DictReader(handle)) - assert [(row["row"], row["column"], row["reason"]) for row in rejected] == [ - ("1", "score", "cannot be read as Integer"), - ("2", "score", "required value is missing"), - ] - assert review_session.scalars(sa.select(ReviewRecord.record_id)).all() == [3] - - -def test_empty_text_can_be_preserved(review_session, tmp_path): - policy = ConversionPolicy(nulls=NullPolicy(empty_is_null=False)) - _load( - review_session, tmp_path, ReviewRecord, - "record_id,score,note\n1,2,\n", policy=policy, - ) - assert review_session.scalar(sa.select(ReviewRecord.note)) == "" - - @pytest.mark.parametrize("supplied_reader", [False, True]) -def test_computed_columns_are_never_loaded(review_session, tmp_path, supplied_reader): +def test_computed_columns_are_never_loaded(dialect_session, tmp_path, supplied_reader): + _create(dialect_session, ComputedRecord) text = "record_id,score,twice\n1,2,wrong\n" kwargs = {} if supplied_reader: - path = tmp_path / "review_record.csv" + path = tmp_path / "special_computed.csv" path.write_text(text) - kwargs["reader"] = DelimitedTextReader(path, ReviewRecord.model_columns()) - _load(review_session, tmp_path, ReviewRecord, text, **kwargs) - assert review_session.scalar(sa.select(ReviewRecord.twice)) == 4 - - -def test_supplied_reader_must_share_the_conversion_null_policy(review_session, tmp_path): - path = tmp_path / "review_record.csv" - path.write_text("record_id,score\n1,1\n") - reader = DelimitedTextReader(path, ReviewRecord.model_columns()) - policy = ConversionPolicy(nulls=NullPolicy(empty_is_null=False)) - with pytest.raises(ValueError, match="reader's null policy differs"): - _load(review_session, tmp_path, ReviewRecord, path.read_text(), reader=reader, policy=policy) - - -def test_tab_markers_ignore_padding_without_quote_semantics(review_session, tmp_path): - path = tmp_path / "review_record.tsv" - path.write_text( - "record_id\tscore\tnote\tevent_date\n" - "1\t2\t MISSING \t MISSING \n2\t3\tMISSING\tMISSING\n" - ) - policy = ConversionPolicy(nulls=NullPolicy( - text_markers=frozenset({"MISSING"}), marker_quoting="unquoted_only", - )) - reader = DelimitedTextReader(path, ReviewRecord.model_columns(), nulls=policy.nulls) - _load(review_session, tmp_path, ReviewRecord, "unused", reader=reader, policy=policy) - rows = review_session.execute(sa.select(ReviewRecord.note, ReviewRecord.event_date)).all() - assert [tuple(row) for row in rows] == [(None, None), (None, None)] - - -def test_legacy_text_markers_cover_all_case_variants(review_session, tmp_path): - markers = sorted(NullPolicy.LEGACY_TEXT_MARKERS) - assert {"NULL", "nUlL", "None", "nan", "N/a", "#N/A"} <= set(markers) - text = "record_id,score,note\n" + "".join( - f"{i},1,{marker}\n" for i, marker in enumerate([*markers, "keep"], 1) - ) - policy = ConversionPolicy(nulls=NullPolicy(text_markers=NullPolicy.LEGACY_TEXT_MARKERS)) - _load(review_session, tmp_path, ReviewRecord, text, policy=policy) - notes = review_session.scalars(sa.select(ReviewRecord.note).order_by(ReviewRecord.record_id)).all() - assert notes == [None] * len(markers) + ["keep"] + kwargs["reader"] = DelimitedTextReader(path, ComputedRecord.model_columns()) + _load(dialect_session, tmp_path, ComputedRecord, text, **kwargs) + assert dialect_session.scalar(sa.select(ComputedRecord.twice)) == 4 @pytest.mark.parametrize("normalise", [False, True]) -def test_defaulted_nullable_primary_key_must_be_in_the_file(dialect_session, tmp_path, normalise): - _create(dialect_session, DefaultKeyRecord) - with pytest.raises(IngestError, match="record_id"): - _load(dialect_session, tmp_path, DefaultKeyRecord, "note\na\n", normalise=normalise) - - -def test_null_primary_key_is_rejected_instead_of_taking_its_default(dialect_session, tmp_path): - _create(dialect_session, DefaultKeyRecord) - policy = ConversionPolicy(errors=ErrorPolicy(on_error="reject_row")) - _load( - dialect_session, tmp_path, DefaultKeyRecord, - "record_id,note\n,a\n1,b\n", policy=policy, rejects_dir=tmp_path, +def test_nonnullable_computed_column_can_be_absent_from_the_file(tmp_path, normalise): + table = cast(type[FileTableProtocol], RequiredComputedRecord) + assert table.required_columns() == {"record_id", "score"} + processor = BatchProcessor( + table, ["record_id", "score"], ConversionPolicy(), LoadReport(), + tmp_path / "rejects.csv", normalise=normalise, ) - assert dialect_session.scalars(sa.select(DefaultKeyRecord.record_id)).all() == [1] + output = processor.process(pa.record_batch({"record_id": ["1"], "score": ["2"]})) + assert output.column_names == ["record_id", "score"] + assert output.num_rows == 1 @pytest.mark.parametrize("model", [KeyedRecord, CompositeKeyedRecord]) @@ -225,19 +138,13 @@ def row(key, value): @pytest.mark.parametrize("strategy", ["insert", "upsert", "replace"]) @pytest.mark.parametrize("batch_size", [None, 2]) -@pytest.mark.parametrize("normalise", [False, True]) -def test_sql_expression_defaults_are_evaluated_for_each_merged_row( - dialect_session, tmp_path, strategy, batch_size, normalise, -): +def test_sql_expression_defaults_are_evaluated_for_each_merged_row(dialect_session, tmp_path, strategy, batch_size): _create(dialect_session, ExpressionDefaultRecord) table = ExpressionDefaultRecord if strategy != "insert": _load(dialect_session, tmp_path, table, "record_id,token\n1,42\n", merge_strategy="insert") text = "record_id,token\n" + "".join(f"{i},\n" for i in range(1, 11)) + "11,123\n" - _load( - dialect_session, tmp_path, table, text, merge_strategy=strategy, - merge_batch_size=batch_size, normalise=normalise, - ) + _load(dialect_session, tmp_path, table, text, merge_strategy=strategy, merge_batch_size=batch_size) tokens = dialect_session.scalars(sa.select(table.token).order_by(table.record_id)).all() assert len(tokens) == 11 assert all(value is not None for value in tokens) @@ -247,39 +154,40 @@ def test_sql_expression_defaults_are_evaluated_for_each_merged_row( assert tokens[0] == 42 -def test_unsupported_numeric_declaration_names_the_table_and_column(tmp_path): - with pytest.raises(PolicyError, match=r"review_invalid_numeric.amount.*\(2, 5\)"): - BatchProcessor( - cast(type[FileTableProtocol], InvalidNumericRecord), ["record_id", "amount"], - ConversionPolicy(), LoadReport(), tmp_path / "rejects.csv", - ) +@pytest.mark.parametrize("normalise", [False, True]) +def test_a_sql_expression_default_is_evaluated_for_each_row_of_a_missing_column(dialect_session, tmp_path, normalise): + _create(dialect_session, ExpressionDefaultRecord) + text = "record_id\n" + "".join(f"{i}\n" for i in range(1, 11)) + _load(dialect_session, tmp_path, ExpressionDefaultRecord, text, normalise=normalise) + tokens = dialect_session.scalars(sa.select(ExpressionDefaultRecord.token)).all() + assert len(set(tokens)) == 10 -def test_infer_delim_remains_importable_from_loaders(tmp_path): - path = tmp_path / "data.tsv" - path.write_text("a\tb\n") - assert infer_delim(path) == "\t" +@pytest.mark.parametrize("normalise", [False, True]) +def test_a_primary_key_with_a_default_must_still_be_in_the_file(dialect_session, tmp_path, normalise): + _create(dialect_session, DefaultKeyRecord) + with pytest.raises(IngestError, match="record_id"): + _load(dialect_session, tmp_path, DefaultKeyRecord, "note\na\n", normalise=normalise) -@pytest.mark.parametrize("normalise", [False, True]) -def test_nonnullable_computed_column_can_be_absent_from_the_file(tmp_path, normalise): - table = cast(type[FileTableProtocol], RequiredComputedRecord) - assert table.required_columns() == {"record_id", "score"} - processor = BatchProcessor( - table, ["record_id", "score"], ConversionPolicy(), LoadReport(), - tmp_path / "rejects.csv", normalise=normalise, +def test_a_blank_primary_key_is_rejected_instead_of_taking_its_default(dialect_session, tmp_path): + _create(dialect_session, DefaultKeyRecord) + policy = ConversionPolicy(errors=ErrorPolicy(on_error="reject_row")) + _load( + dialect_session, tmp_path, DefaultKeyRecord, + "record_id,note\n,a\n1,b\n", policy=policy, rejects_dir=tmp_path, ) - output = processor.process(pa.record_batch({"record_id": ["1"], "score": ["2"]})) - assert output.column_names == ["record_id", "score"] - assert output.num_rows == 1 + assert dialect_session.scalars(sa.select(DefaultKeyRecord.record_id)).all() == [1] -def test_defaulted_nullable_primary_key_is_always_required(): +def test_a_primary_key_is_required_and_takes_no_default(): assert DefaultKeyRecord.required_columns() == {"record_id"} assert "record_id" not in DefaultKeyRecord.column_defaults() -def test_column_defaults_protocol_has_a_callable_classmethod_function(): - descriptor = vars(ORMTableProtocol)["column_defaults"] - assert isinstance(descriptor, classmethod) - assert inspect.isfunction(descriptor.__func__) +def test_unsupported_numeric_declaration_names_the_table_and_column(tmp_path): + with pytest.raises(PolicyError, match=r"special_invalid_numeric.amount.*\(2, 5\)"): + BatchProcessor( + cast(type[FileTableProtocol], InvalidNumericRecord), ["record_id", "amount"], + ConversionPolicy(), LoadReport(), tmp_path / "rejects.csv", + ) diff --git a/tests/tables/test_typing_protocols.py b/tests/tables/test_typing_protocols.py index 276772f..49a6a5f 100644 --- a/tests/tables/test_typing_protocols.py +++ b/tests/tables/test_typing_protocols.py @@ -1,3 +1,5 @@ +import inspect + import sqlalchemy as sa import sqlalchemy.orm as so from orm_loader.tables.orm_table import ORMTableBase @@ -22,6 +24,12 @@ class Q(ORMTableBase, Base): assert callable(Q.pk_names) assert callable(Q.model_columns) +def test_every_protocol_member_is_a_plain_classmethod(): + for name in ("column_defaults", "sql_expression_defaults"): + descriptor = vars(ORMTableProtocol)[name] + assert isinstance(descriptor, classmethod) + assert inspect.isfunction(descriptor.__func__) + def test_serialized_table_protocol_conformance(): class G(SerialisableTableInterface, ORMTableBase, Base): __tablename__ = "g" From 79c017fe0674d3a37b64194bad64abb002456d1e Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 4 Nov 2026 22:53:31 +0000 Subject: [PATCH 19/21] Let the processor alone decide what is missing: readers null nothing, marker_quoting and LEGACY_TEXT_MARKERS are removed - readers take no NullPolicy and hand over text as read, so empty_is_null=False keeps empty strings in every format - text markers match after trimming for every format, and a reader can no longer disagree with the policy - infer_delim and infer_encoding are no longer exported from orm_loader.loaders - docs describe how to restore the 2.x words with text_markers and the raw pass-through of normalise=False --- docs/loaders/conversion.md | 9 ++-- docs/loaders/readers.md | 4 +- docs/loaders/staging.md | 2 +- src/orm_loader/conversion/policy.py | 25 +--------- src/orm_loader/loaders/__init__.py | 2 - src/orm_loader/loaders/readers/__init__.py | 13 ++--- src/orm_loader/loaders/readers/base.py | 32 +----------- src/orm_loader/loaders/readers/delimited.py | 54 +++------------------ src/orm_loader/loaders/readers/parquet.py | 11 +---- tests/conversion/test_policy.py | 6 --- tests/loaders/test_load_policy.py | 30 ++++++------ tests/loaders/test_pipeline.py | 2 +- tests/loaders/test_readers.py | 23 +++++---- tests/loaders/test_stage_file.py | 2 +- 14 files changed, 48 insertions(+), 167 deletions(-) diff --git a/docs/loaders/conversion.md b/docs/loaders/conversion.md index 9a91297..cbe23d4 100644 --- a/docs/loaders/conversion.md +++ b/docs/loaders/conversion.md @@ -42,10 +42,9 @@ policy = ConversionPolicy( ## NullPolicy - Text columns keep every nonempty value verbatim. Only empty and whitespace-only text is null. -- Compared with 2.x, words such as `NULL`, `None`, `N/A` and `NaN` stay as text by default. To restore the legacy marker set, pass `NullPolicy(text_markers=NullPolicy.LEGACY_TEXT_MARKERS)`. You can add or remove markers with set operations. -- `text_markers` adds exact, case-sensitive words that are null in every text column of the load. Matching ignores surrounding whitespace in normal conversion. With `marker_quoting="unquoted_only"` in CSV quote mode, Arrow matches markers exactly because it applies them while parsing. To treat a word as missing in some columns only, subclass `TextHandler`, override `missing()` and register it in `ConversionPolicy.columns`. -- `empty_is_null=False` keeps empty text fields as empty strings in delimited files. Typed columns still treat empty fields as missing. -- `marker_quoting="unquoted_only"` matches a marker only when the source did not quote it. +- Words such as `NULL`, `None`, `N/A` and `NaN` stay as text in text columns. Version 2 turned them into NULL, which dropped concepts with those names. To treat words as missing again, list them in `text_markers`, for example `NullPolicy(text_markers=frozenset({"NULL", "None", "N/A", "NaN", "nan", "null", "NA", "n/a"}))`. Set operations add or remove words from a set you keep. +- `text_markers` are exact, case-sensitive words that are null in every text column of the load. Matching ignores surrounding whitespace and applies the same way to every file format. To treat a word as missing in some columns only, subclass `TextHandler`, override `missing()` and register it in `ConversionPolicy.columns`. +- `empty_is_null=False` keeps empty text fields as empty strings. Typed columns still treat empty fields as missing. - Typed columns accept a valid value, a missing word (`null`, `none`, `na`, `n/a`, `nan` and similar) or a real null, and fail otherwise. ::: orm_loader.conversion.NullPolicy @@ -54,7 +53,7 @@ policy = ConversionPolicy( ## ErrorPolicy -`on_error` is `raise` (default) or `reject_row`. `max_rejects` stops a `reject_row` load once rejected rows pass a count or a fraction of the rows read. +`on_error` is `raise` (default) or `reject_row`. `max_rejects` stops a `reject_row` load before the merge. A count stops the load as soon as it is passed, and a fraction is checked once all rows have been read, against the total number of rows read. A load that rejected rows logs a warning with their number and the rejects file. ::: orm_loader.conversion.ErrorPolicy diff --git a/docs/loaders/readers.md b/docs/loaders/readers.md index 0f275df..2eb3772 100644 --- a/docs/loaders/readers.md +++ b/docs/loaders/readers.md @@ -1,6 +1,6 @@ # Readers -A reader yields a file as Arrow record batches holding the model columns present in the file. Text formats yield raw strings, so the conversion policy sees exactly what the file contains. Typed formats yield their stored types. +A reader yields a file as Arrow record batches holding the model columns present in the file. Text formats yield raw strings, so the conversion policy sees exactly what the file contains. Typed formats yield their stored types. A reader decides nothing about missing values: an empty field is an empty string, and the null policy is applied afterwards. `open_reader` picks the reader from the file suffix: `.parquet` is read as Parquet, any other suffix as delimited text. `load_file` opens the reader with it, taking `quote_mode` and `chunksize` as arguments, and `load_file(reader=...)` passes a configured reader explicitly. @@ -21,7 +21,7 @@ A new format is a new `FileReader` subclass that implements `batches()` and sets ## DelimitedTextReader -Reads raw strings in 4 MB blocks, so memory does not grow with the row count. Header names are lowercased and stripped of quotes, padding and a trailing `_hash`, and file columns that are not model columns are ignored with a warning. Parquet column names are normalised the same way. An unquoted empty field is null and a quoted empty field is an empty string. +Reads raw strings in 4 MB blocks, so memory does not grow with the row count. Header names are lowercased and stripped of quotes, padding and a trailing `_hash`, and file columns that are not model columns are ignored with a warning. Parquet column names are normalised the same way. `quote_mode` selects how double quotes are read: diff --git a/docs/loaders/staging.md b/docs/loaders/staging.md index 1fd033d..c8db0bd 100644 --- a/docs/loaders/staging.md +++ b/docs/loaders/staging.md @@ -11,7 +11,7 @@ - Values are interpreted by column type, as described in [Conversion](conversion.md). - With the default policy an unreadable value or a missing required value raises and nothing is loaded. - With `ErrorPolicy(on_error="reject_row")` those rows, and delimited lines with the wrong number of fields, are written to `/.rejects.csv` and the rest load. A malformed line is listed with its row number in the file, and its text is in the reason. `rejects_dir` defaults to a `rejects` directory inside the input file's directory, and the file of an earlier load is removed. `LoadReport` accounts for every row read, checked against the rows counted in the staging table. -- With `normalise=False` raw text is passed to the staging table unchanged, without literal defaults, non-key required checks or type conversion. SQL-expression defaults still apply to null staging values during merge. Primary-key columns must still be present. What the database does with a value that does not fit its column then depends on the dialect. +- With `normalise=False` raw text is passed to the staging table unchanged, empty fields included, without literal defaults, non-key required checks or type conversion. SQL-expression defaults still apply to null staging values during merge. Primary-key columns must still be present. What the database does with a value that does not fit its column then depends on the dialect. - A blank value, and a column the file lacks, take the column's default when it has one: a server default written as text, a scalar Python default, or a SQL expression such as `current_date`. Literal defaults are converted by the column's handler. SQL expressions are applied in the merge with `COALESCE`, so the database evaluates them once for each row that needs a default. A column with a callable default stays required. --- diff --git a/src/orm_loader/conversion/policy.py b/src/orm_loader/conversion/policy.py index 29ea2c3..3f5560b 100644 --- a/src/orm_loader/conversion/policy.py +++ b/src/orm_loader/conversion/policy.py @@ -4,7 +4,6 @@ from collections.abc import Mapping from dataclasses import dataclass, field -from itertools import product from types import MappingProxyType from typing import Any, ClassVar, Literal @@ -37,12 +36,6 @@ from .values import no_failures, underlying_type -def _case_variants(value: str) -> frozenset[str]: - """Return every upper/lower-case spelling of a short null marker.""" - choices = [(char.lower(), char.upper()) if char.isalpha() else (char,) for char in value] - return frozenset("".join(chars) for chars in product(*choices)) - - @dataclass(frozen=True) class NullPolicy: """What counts as a missing value, shared by every column handler. @@ -57,10 +50,6 @@ class NullPolicy: text_markers Exact, case-sensitive text values treated as missing in text columns. Matching ignores surrounding whitespace. - marker_quoting - ``"ignore"`` matches markers whether or not the source quoted them. - ``"unquoted_only"`` matches only markers the source did not quote, - which needs a source that records quoting and import-level markers. typed_missing_words Lowercase words treated as missing in non-text columns. """ @@ -68,22 +57,11 @@ class NullPolicy: DEFAULT_TYPED_MISSING: ClassVar[frozenset[str]] = frozenset( {"null", "none", "na", "n/a", "nan", "#n/a", "-nan", ""} ) - LEGACY_TEXT_MARKERS: ClassVar[frozenset[str]] = frozenset( - { - "#N/A", "#N/A N/A", "#NA", "-1.#IND", "-1.#QNAN", "-NaN", "-nan", "1.#IND", "1.#QNAN", - "", "N/A", "NA", "NULL", "NaN", "None", "n/a", "nan", "null", - } - ) | frozenset().union(*(_case_variants(word) for word in ("nan", "null", "none", "na", "n/a"))) empty_is_null: bool = True text_markers: frozenset[str] = frozenset() - marker_quoting: Literal["ignore", "unquoted_only"] = "ignore" typed_missing_words: frozenset[str] = DEFAULT_TYPED_MISSING - def __post_init__(self) -> None: - if self.marker_quoting not in ("ignore", "unquoted_only"): - raise ValueError(f"marker_quoting must be 'ignore' or 'unquoted_only', not {self.marker_quoting!r}") - def text_missing(self, trimmed: pa.Array) -> pa.Array: """Return a mask of trimmed text values that are missing in a text column. @@ -97,8 +75,7 @@ def text_missing(self, trimmed: pa.Array) -> pa.Array: pa.Array Boolean array that is true where the value is missing. """ - words = self.text_markers if self.marker_quoting == "ignore" else frozenset() - return self._mask(trimmed, words, empty=self.empty_is_null, casefold=False) + return self._mask(trimmed, self.text_markers, empty=self.empty_is_null, casefold=False) def typed_missing(self, trimmed: pa.Array) -> pa.Array: """Return a mask of trimmed text values that are missing in a typed column.""" diff --git a/src/orm_loader/loaders/__init__.py b/src/orm_loader/loaders/__init__.py index e5b244c..2d62ef8 100644 --- a/src/orm_loader/loaders/__init__.py +++ b/src/orm_loader/loaders/__init__.py @@ -1,11 +1,9 @@ from .loader import LoaderContext, LoadReport, stage_file from .readers import DelimitedTextReader, FileReader, ParquetReader, open_reader -from .readers.detection import infer_delim __all__ = [ "DelimitedTextReader", "FileReader", - "infer_delim", "LoadReport", "LoaderContext", "ParquetReader", diff --git a/src/orm_loader/loaders/readers/__init__.py b/src/orm_loader/loaders/readers/__init__.py index dfbdd8a..139ef6c 100644 --- a/src/orm_loader/loaders/readers/__init__.py +++ b/src/orm_loader/loaders/readers/__init__.py @@ -4,10 +4,8 @@ import sqlalchemy as sa -from ...conversion import NullPolicy from .base import FileReader from .delimited import DelimitedTextReader -from .detection import infer_delim from .parquet import ParquetReader @@ -15,7 +13,6 @@ def open_reader( path: Path, model_columns: Mapping[str, sa.ColumnElement[Any]], *, - nulls: NullPolicy | None = None, quote_mode: str | None = None, chunksize: int | None = None, skip_malformed: bool = False, @@ -31,8 +28,6 @@ def open_reader( File to read. model_columns Model columns by name; only these are read. - nulls - Null policy of the load. quote_mode Quote handling for delimited text. ``None`` uses the reader's default. chunksize @@ -47,10 +42,8 @@ def open_reader( The opened reader. """ if path.suffix.lower() == ".parquet": - return ParquetReader(path, model_columns, nulls, chunksize) - return DelimitedTextReader(path, model_columns, nulls, quote_mode, skip_malformed=skip_malformed) + return ParquetReader(path, model_columns, chunksize) + return DelimitedTextReader(path, model_columns, quote_mode, skip_malformed=skip_malformed) - - -__all__ = ["DelimitedTextReader", "FileReader", "ParquetReader", "infer_delim", "open_reader"] +__all__ = ["DelimitedTextReader", "FileReader", "ParquetReader", "open_reader"] diff --git a/src/orm_loader/loaders/readers/base.py b/src/orm_loader/loaders/readers/base.py index be398ee..95e991a 100644 --- a/src/orm_loader/loaders/readers/base.py +++ b/src/orm_loader/loaders/readers/base.py @@ -8,10 +8,8 @@ from typing import Any import pyarrow as pa -import pyarrow.compute as pc import sqlalchemy as sa -from ...conversion import NullPolicy from ...helpers import IngestError @@ -21,8 +19,7 @@ class FileReader(ABC): A reader yields batches holding only the model columns present in the file, named as the model names them. Text formats yield raw strings, so the null policy sees exactly what the file contains. Typed formats yield - typed columns. Markers that are quote-sensitive are applied while - reading, because quoting is not visible afterwards. + typed columns. A reader decides nothing about missing values. Parameters ---------- @@ -30,8 +27,6 @@ class FileReader(ABC): File to read. model_columns Model columns by name; only these are read. - nulls - Null policy whose quote-sensitive markers are applied while reading. Attributes ---------- @@ -51,11 +46,9 @@ def __init__( self, path: Path, model_columns: Mapping[str, sa.ColumnElement[Any]], - nulls: NullPolicy | None = None, ) -> None: self.path = path self.model_columns = model_columns - self.nulls = nulls or NullPolicy() self.columns: list[str] = [] self.file_columns: list[str] = [] self.empty = False @@ -77,26 +70,3 @@ def normalise_headers(self, names: Iterable[str]) -> list[str]: @abstractmethod def batches(self) -> Iterator[pa.RecordBatch]: """Yield the file's rows as non-empty record batches.""" - - def reader_markers(self) -> frozenset[str]: - """Return the markers that must be applied while reading.""" - if self.nulls.marker_quoting == "unquoted_only": - return self.nulls.text_markers - return frozenset() - - @staticmethod - def null_markers(batch: pa.RecordBatch, markers: frozenset[str]) -> pa.RecordBatch: - """Set string values equal to a marker, ignoring padding, to null.""" - values = pa.array(sorted(markers), pa.string()) - arrays = [] - for column in batch.columns: - was_dictionary = pa.types.is_dictionary(column.type) and pa.types.is_string(column.type.value_type) - if was_dictionary: - column = column.dictionary_decode() - if pa.types.is_string(column.type): - hit = pc.fill_null(pc.is_in(pc.utf8_trim_whitespace(column), value_set=values), False) - column = pc.if_else(hit, pa.scalar(None, pa.string()), column) - if was_dictionary: - column = pc.dictionary_encode(column) - arrays.append(column) - return pa.RecordBatch.from_arrays(arrays, names=batch.schema.names) diff --git a/src/orm_loader/loaders/readers/delimited.py b/src/orm_loader/loaders/readers/delimited.py index 8398bec..ac3b8bf 100644 --- a/src/orm_loader/loaders/readers/delimited.py +++ b/src/orm_loader/loaders/readers/delimited.py @@ -3,7 +3,6 @@ from __future__ import annotations import codecs -import logging from collections.abc import Iterator, Mapping from pathlib import Path from typing import Any, Literal @@ -12,24 +11,18 @@ import pyarrow.csv as pv import sqlalchemy as sa -from ...conversion import NullPolicy from ...helpers import IngestError from .base import FileReader from .detection import infer_delim, infer_encoding, resolve_quote_mode -logger = logging.getLogger(__name__) - -QUOTE_SAMPLE_LINES = 50 - class DelimitedTextReader(FileReader): """A delimited text file read as raw strings in bounded blocks. - Date columns are dictionary-encoded. With the default null policy, an - unquoted empty field is null and a quoted empty field is an empty string. - With ``empty_is_null=False``, unquoted empty text remains empty. Header - names are normalised as described by ``normalise_headers``. File columns - that are not model columns are ignored. + Every value is kept as read, so an empty field is an empty string and + nothing is null. Date columns are dictionary-encoded. Header names are + normalised as described by ``normalise_headers``. File columns that are + not model columns are ignored. Parameters ---------- @@ -37,14 +30,11 @@ class DelimitedTextReader(FileReader): File to read. model_columns Model columns by name; only these are read. - nulls - Null policy whose quote-sensitive markers are applied while reading. quote_mode ``"csv"`` reads double quotes as field quoting, ``"literal"`` keeps them as data, ``"by_delimiter"`` chooses ``"literal"`` for tab delimited files and ``"csv"`` otherwise. ``None`` uses - ``DEFAULT_QUOTE_MODE``. In CSV mode, unquoted-only null markers must - match exactly, including surrounding whitespace. + ``DEFAULT_QUOTE_MODE``. block_size Bytes per block; memory use follows this, not the row count. skip_malformed @@ -63,12 +53,11 @@ def __init__( self, path: Path, model_columns: Mapping[str, sa.ColumnElement[Any]], - nulls: NullPolicy | None = None, quote_mode: str | None = None, block_size: int = 4 << 20, skip_malformed: bool = False, ) -> None: - super().__init__(path, model_columns, nulls) + super().__init__(path, model_columns) self.block_size = block_size self.skip_malformed = skip_malformed if path.stat().st_size == 0: @@ -83,8 +72,6 @@ def __init__( self.quote_mode = resolve_quote_mode(quote_mode or self.DEFAULT_QUOTE_MODE, self.delimiter) self.file_columns = self._header_names() self.columns = [name for name in self.file_columns if name in model_columns] - if self.nulls.marker_quoting == "unquoted_only" and self.nulls.text_markers: - self._warn_if_everything_is_quoted() def _header_names(self) -> list[str]: with open(self.path, "rb") as handle: @@ -95,35 +82,10 @@ def _skip_row(self, row: pv.InvalidRow) -> Literal["skip"]: self.skipped.append((row.number, f"expected {row.expected_columns} columns, got {row.actual_columns}: {row.text.rstrip()}")) return "skip" - def _warn_if_everything_is_quoted(self) -> None: - with open(self.path, "rb") as handle: - handle.readline() - lines = [handle.readline().decode(self.encoding, errors="replace") for _ in range(QUOTE_SAMPLE_LINES)] - fields = [f for line in lines if line.strip() for f in line.rstrip("\r\n").split(self.delimiter)] - if fields and all(f.startswith('"') for f in fields): - logger.warning( - "%s quotes every field, so quote-sensitive text markers %s can never match", - self.path.name, - sorted(self.nulls.text_markers), - ) - - def _null_values(self, markers: frozenset[str]) -> list[str]: - """Return exact import markers; padded markers are handled after parsing.""" - empty = [""] if self.nulls.empty_is_null else [] - quoted_markers = sorted(markers) if self.quote_mode == "csv" else [] - return empty + quoted_markers - - def _apply_reader_markers(self, batch: pa.RecordBatch, markers: frozenset[str]) -> pa.RecordBatch: - """Apply whitespace-tolerant markers after parsing when quoting is not tracked.""" - if self.quote_mode == "csv" or not markers: - return batch - return self.null_markers(batch, markers) - def batches(self) -> Iterator[pa.RecordBatch]: """Yield the file's rows in blocks of ``block_size`` bytes.""" if self.empty: return - markers = self.reader_markers() string, encoded = pa.string(), pa.dictionary(pa.int32(), pa.string()) column_types = { name: encoded if isinstance(self.model_columns[name].type, sa.Date) else string @@ -149,14 +111,10 @@ def batches(self) -> Iterator[pa.RecordBatch]: convert_options=pv.ConvertOptions( column_types=column_types, include_columns=self.columns, - null_values=self._null_values(markers), - strings_can_be_null=True, - quoted_strings_can_be_null=False, ), ) with reader: for batch in reader: - batch = self._apply_reader_markers(batch, markers) if batch.num_rows: yield batch except UnicodeDecodeError as exc: diff --git a/src/orm_loader/loaders/readers/parquet.py b/src/orm_loader/loaders/readers/parquet.py index 7e0ba50..ee9e744 100644 --- a/src/orm_loader/loaders/readers/parquet.py +++ b/src/orm_loader/loaders/readers/parquet.py @@ -10,7 +10,6 @@ import pyarrow.dataset as ds import sqlalchemy as sa -from ...conversion import NullPolicy from .base import FileReader @@ -25,10 +24,6 @@ class ParquetReader(FileReader): File to read. model_columns Model columns by name; only these are read. - nulls - Null policy whose quote-sensitive markers are applied while reading. A - Parquet string was never quoted, so every marker match counts as - unquoted. chunksize Rows per batch. """ @@ -37,10 +32,9 @@ def __init__( self, path: Path, model_columns: Mapping[str, sa.ColumnElement[Any]], - nulls: NullPolicy | None = None, chunksize: int | None = None, ) -> None: - super().__init__(path, model_columns, nulls) + super().__init__(path, model_columns) self.chunksize = chunksize or 64_000 self._dataset = ds.dataset(path, format="parquet") self.file_columns = self.normalise_headers(self._dataset.schema.names) @@ -50,10 +44,7 @@ def __init__( def batches(self) -> Iterator[pa.RecordBatch]: """Yield the stored rows in batches of ``chunksize`` rows.""" - markers = self.reader_markers() for batch in self._dataset.to_batches(columns=self._stored, batch_size=self.chunksize): batch = batch.rename_columns(self.columns) - if markers: - batch = self.null_markers(batch, markers) if batch.num_rows: yield batch diff --git a/tests/conversion/test_policy.py b/tests/conversion/test_policy.py index 4b97414..ccdf64b 100644 --- a/tests/conversion/test_policy.py +++ b/tests/conversion/test_policy.py @@ -13,7 +13,6 @@ FloatHandler, IntegerHandler, Interpretation, - NullPolicy, NumericHandler, PassthroughHandler, TextHandler, @@ -115,8 +114,3 @@ def test_the_handler_and_column_mappings_cannot_be_changed_afterwards(): def test_unknown_error_mode_is_rejected(): with pytest.raises(ValueError, match="on_error"): ErrorPolicy(on_error="ignore") # type: ignore[arg-type] - - -def test_unknown_marker_quoting_is_rejected(): - with pytest.raises(ValueError, match="marker_quoting"): - NullPolicy(marker_quoting="sometimes") # type: ignore[arg-type] diff --git a/tests/loaders/test_load_policy.py b/tests/loaders/test_load_policy.py index 486ebd2..715ed1d 100644 --- a/tests/loaders/test_load_policy.py +++ b/tests/loaders/test_load_policy.py @@ -1,6 +1,5 @@ """Loading under a null policy: markers, rejected rows, row accounting and sources.""" -import logging import os from datetime import date from pathlib import Path @@ -62,21 +61,24 @@ def test_marker_in_a_required_column_raises(dialect_session, tmp_path): _load(dialect_session, tmp_path, _DemoRequired, text, policy=MARKER) -def test_quote_sensitive_marker_matches_only_unquoted(dialect_session, tmp_path): - policy = ConversionPolicy(nulls=NullPolicy(text_markers=frozenset({"NULL"}), marker_quoting="unquoted_only")) - _load(dialect_session, tmp_path, _Demo, SEVEN_ROWS, policy=policy) - assert _labels(dialect_session, DemoRecord) == { - 1: None, 2: "NULL", 3: '"NULL"', 4: None, 5: None, 6: "alpha,beta", 7: "alpha\nbeta", - } +@pytest.mark.parametrize("suffix, delimiter", [("csv", ","), ("tsv", "\t")]) +def test_empty_text_is_kept_when_the_policy_says_so(dialect_session, tmp_path, suffix, delimiter): + text = f"demo_id{delimiter}demo_label{delimiter}demo_number\n1{delimiter}{delimiter}\n2{delimiter}x{delimiter}5\n" + path = tmp_path / f"demo_record.{suffix}" + path.write_text(text, encoding="utf-8") + tag = None if dialect_session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA + policy = ConversionPolicy(nulls=NullPolicy(empty_is_null=False)) + _Demo.load_file(dialect_session, path, staging_schema_tag=tag, policy=policy) + rows = dialect_session.execute(sa.select(DemoRecord.demo_label, DemoRecord.demo_number).order_by(DemoRecord.demo_id)).all() + assert [tuple(r) for r in rows] == [("", None), ("x", 5)] -def test_quote_sensitive_marker_warns_when_every_field_is_quoted(dialect_session, tmp_path, caplog): - policy = ConversionPolicy(nulls=NullPolicy(text_markers=frozenset({"NULL"}), marker_quoting="unquoted_only")) - text = '"demo_id","demo_label"\n"1","NULL"\n"2","x"\n' - with caplog.at_level(logging.WARNING): - _load(dialect_session, tmp_path, _Demo, text, policy=policy) - assert "quotes every field" in caplog.text - assert _labels(dialect_session, DemoRecord) == {1: "NULL", 2: "x"} +def test_markers_ignore_padding_in_a_tab_delimited_file(dialect_session, tmp_path): + path = tmp_path / "demo_record.tsv" + path.write_text("demo_id\tdemo_label\n1\t NULL \n2\tNULL\n3\tkept\n", encoding="utf-8") + tag = None if dialect_session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA + _Demo.load_file(dialect_session, path, staging_schema_tag=tag, policy=MARKER) + assert _labels(dialect_session, DemoRecord) == {1: None, 2: None, 3: "kept"} def test_both_athena_date_formats_and_iso_load_in_one_file(dialect_session, tmp_path): diff --git a/tests/loaders/test_pipeline.py b/tests/loaders/test_pipeline.py index 0b350d8..b161470 100644 --- a/tests/loaders/test_pipeline.py +++ b/tests/loaders/test_pipeline.py @@ -50,7 +50,7 @@ def _stage(session, table, path: Path, reader=None, **context): backend = resolve_backend(session, staging_schema_tag=tag, mapper=table) staging = backend.create_staging_table(table, session) ctx = LoaderContext(tableclass=table, session=session, path=path, backend=backend, **context) - loaded = stage_file(ctx, reader or open_reader(path, table.model_columns(), nulls=ctx.policy.nulls), staging) + loaded = stage_file(ctx, reader or open_reader(path, table.model_columns()), staging) return loaded, ctx, staging diff --git a/tests/loaders/test_readers.py b/tests/loaders/test_readers.py index 25eb2b3..ead44fb 100644 --- a/tests/loaders/test_readers.py +++ b/tests/loaders/test_readers.py @@ -6,7 +6,6 @@ import pyarrow.parquet as pq import pytest -from orm_loader.conversion import NullPolicy from orm_loader.loaders import DelimitedTextReader, ParquetReader, open_reader from tests.models import DemoRecord @@ -50,9 +49,9 @@ def test_headers_are_lowercased_and_stripped_of_quotes_and_hash_suffix(tmp_path) assert DelimitedTextReader(path, COLUMNS).columns == ["demo_id", "demo_label"] -def test_an_unquoted_empty_field_is_null_and_a_quoted_one_is_an_empty_string(tmp_path): +def test_a_quoted_and_an_unquoted_empty_field_are_both_empty_strings(tmp_path): path = _write(tmp_path, 'demo_id,demo_label\n1,\n2,""\n') - assert [r["demo_label"] for r in _rows(DelimitedTextReader(path, COLUMNS, quote_mode="csv"))] == [None, ""] + assert [r["demo_label"] for r in _rows(DelimitedTextReader(path, COLUMNS, quote_mode="csv"))] == ["", ""] def test_embedded_newlines_and_tabs_survive(tmp_path): @@ -91,11 +90,12 @@ def test_an_empty_file_has_no_content_and_a_header_only_file_has_no_rows(tmp_pat assert not header_only.empty and list(header_only.batches()) == [] -def test_unquoted_markers_are_nulled_while_reading_but_quoted_ones_are_kept(tmp_path): - path = _write(tmp_path, 'demo_id,demo_label\n1,NULL\n2,"NULL"\n') - nulls = NullPolicy(text_markers=frozenset({"NULL"}), marker_quoting="unquoted_only") - labels = [r["demo_label"] for r in _rows(DelimitedTextReader(path, COLUMNS, nulls, quote_mode="csv"))] - assert labels == [None, "NULL"] +@pytest.mark.parametrize("name, quote_mode", [("demo_record.csv", "csv"), ("demo_record.tsv", "literal")]) +def test_a_delimited_reader_nulls_nothing(tmp_path, name, quote_mode): + delimiter = "," if name.endswith("csv") else "\t" + path = _write(tmp_path, f"demo_id{delimiter}demo_label\n1{delimiter}\n2{delimiter}NULL\n3{delimiter} NA \n", name) + labels = [r["demo_label"] for r in _rows(DelimitedTextReader(path, COLUMNS, quote_mode=quote_mode))] + assert labels == ["", "NULL", " NA "] def test_parquet_keeps_stored_types_and_ignores_other_columns(tmp_path): @@ -115,8 +115,7 @@ def test_parquet_is_read_in_batches_of_chunksize_rows(tmp_path): assert sizes == [100, 100, 50] -def test_parquet_markers_are_nulled_while_reading(tmp_path): +def test_parquet_strings_are_not_nulled_by_the_reader(tmp_path): path = tmp_path / "demo_record.parquet" - pq.write_table(pa.table({"demo_id": [1, 2], "demo_label": ["NULL", "kept"]}), path) - nulls = NullPolicy(text_markers=frozenset({"NULL"}), marker_quoting="unquoted_only") - assert [r["demo_label"] for r in _rows(ParquetReader(path, COLUMNS, nulls))] == [None, "kept"] + pq.write_table(pa.table({"demo_id": [1, 2], "demo_label": ["NULL", ""]}), path) + assert [r["demo_label"] for r in _rows(ParquetReader(path, COLUMNS))] == ["NULL", ""] diff --git a/tests/loaders/test_stage_file.py b/tests/loaders/test_stage_file.py index 7241461..3c6ae0d 100644 --- a/tests/loaders/test_stage_file.py +++ b/tests/loaders/test_stage_file.py @@ -10,7 +10,7 @@ def _reader(ctx, **options): - return open_reader(ctx.path, ctx.tableclass.model_columns(), nulls=ctx.policy.nulls, **options) + return open_reader(ctx.path, ctx.tableclass.model_columns(), **options) def test_loader_context_fields(sqlite_session): From d4edf2f5e4b241d05a72640e874a788c7eda045c Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 4 Nov 2026 22:53:31 +0000 Subject: [PATCH 20/21] Remove the PostgreSQL foreign key hooks that relaxed_foreign_keys no longer uses - relaxed_foreign_keys sets the replication role per transaction, so disable_fk_check and restore_fk_check are gone from PostgresBackend - the role check accepts only the role names Postgres reports --- src/orm_loader/backends/postgres.py | 37 +++---------------------- tests/backends/test_postgres_backend.py | 34 ----------------------- 2 files changed, 4 insertions(+), 67 deletions(-) diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index 8f294b3..e555009 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -116,45 +116,16 @@ def insert_batches( cursor.close() @staticmethod - def _normalize_fk_check_state(previous_state: str | int) -> str: - if isinstance(previous_state, int): - raise ValueError( - f"Invalid PostgreSQL session_replication_role {previous_state!r}: " - "Postgres uses string roles ('origin', 'local', 'replica'), not integers. " - "The value passed here should always come from this backend's own " - "disable_fk_check(), which returns a string." - ) - normalised = previous_state.strip().lower() + def _normalize_fk_check_state(role: str) -> str: + """Return a ``session_replication_role`` value in lowercase, or raise for an unknown role.""" + normalised = role.strip().lower() if normalised not in _VALID_PG_REPLICATION_ROLES: raise ValueError( - f"Invalid PostgreSQL session_replication_role {previous_state!r}. " + f"Invalid PostgreSQL session_replication_role {role!r}. " f"Expected one of: {sorted(_VALID_PG_REPLICATION_ROLES)}" ) return normalised - def disable_fk_check(self, session: so.Session, *, mapper: Any | None = None) -> str: - execute_options = {"bind_arguments": {"mapper": mapper}} if mapper is not None else {} - previous_state = session.execute( - sa.text("SHOW session_replication_role"), **execute_options - ).scalar() - session.execute(sa.text("SET session_replication_role = 'replica'"), **execute_options) - if not isinstance(previous_state, str): - raise RuntimeError("Expected PostgreSQL FK state to be a string") - return previous_state - - def restore_fk_check( - self, - session: so.Session, - previous_state: str | int, - *, - mapper: Any | None = None, - ) -> None: - safe_state = self._normalize_fk_check_state(previous_state) - execute_options = {"bind_arguments": {"mapper": mapper}} if mapper is not None else {} - session.execute( - sa.text(f"SET session_replication_role = '{safe_state}'"), **execute_options - ) - @contextmanager def relaxed_foreign_keys( self, diff --git a/tests/backends/test_postgres_backend.py b/tests/backends/test_postgres_backend.py index 31dc745..3d745f0 100644 --- a/tests/backends/test_postgres_backend.py +++ b/tests/backends/test_postgres_backend.py @@ -120,21 +120,6 @@ def test_postgres_backend_create_staging_table_drops_computed_columns(pg_session assert cols == {"id", "name", "_rownum"} # slug is computed, excluded -def test_postgres_backend_fk_methods_emit_expected_sql(): - backend = PostgresBackend() - session = _FakeSession() - - previous = backend.disable_fk_check(_sess(session)) - backend.restore_fk_check(_sess(session), previous) - - assert previous == "origin" - assert session.statements == [ - "SHOW session_replication_role", - "SET session_replication_role = 'replica'", - "SET session_replication_role = 'origin'", - ] - - def test_postgres_backend_materialized_view_methods_work_end_to_end(pg_db): """Real create + refresh + query, not just checking emitted SQL text. The whole point is proving this DDL actually round-trips correctly.""" @@ -524,25 +509,6 @@ def test_postgres_backend_normalize_fk_check_state(): else: raise AssertionError("Expected ValueError for unrecognised role") - try: - normalize(1) - except ValueError as exc: - assert "Postgres uses string roles" in str(exc) - else: - raise AssertionError("Expected ValueError for integer input") - - -def test_postgres_backend_disable_fk_raises_when_show_returns_non_string(): - backend = PostgresBackend() - session = _FakeSession(scalar_result=42) - - try: - backend.disable_fk_check(_sess(session)) - except RuntimeError as exc: - assert "Expected PostgreSQL FK state to be a string" in str(exc) - else: - raise AssertionError("Expected RuntimeError when SHOW returns a non-string") - # --------------------------------------------------------------------------- # Loading files into real Postgres From fb2bb867c5d7136402553a3ad6500ff73676e3d6 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Thu, 5 Nov 2026 02:47:31 +0000 Subject: [PATCH 21/21] Read Numeric columns exactly with Arrow and refuse those it cannot hold --- docs/loaders/conversion.md | 8 +- src/orm_loader/conversion/handlers/numbers.py | 127 ++++++++++++---- tests/conversion/test_number_handlers.py | 106 +++++++++++--- tests/conversion/test_typed_input.py | 17 +++ tests/loaders/test_numeric_columns.py | 137 ++++++++++++++++++ tests/loaders/test_special_columns.py | 21 ++- 6 files changed, 362 insertions(+), 54 deletions(-) create mode 100644 tests/loaders/test_numeric_columns.py diff --git a/docs/loaders/conversion.md b/docs/loaders/conversion.md index cbe23d4..ad8b628 100644 --- a/docs/loaders/conversion.md +++ b/docs/loaders/conversion.md @@ -69,7 +69,7 @@ The handler for a column is the first match of: the entry in `ConversionPolicy.c | `EnumHandler` | `Enum` | One of the labels the column stores | `strip` | | `IntegerHandler` | `Integer`, `SmallInteger`, `BigInteger` | Optional sign and digits, with an optional all-zero fraction (`1.0`); a value outside the 16, 32 or 64 bits of the column type is an error | none | | `FloatHandler` | `Float` | Decimal or exponent form; `Infinity` and values that overflow to infinity are errors | none | -| `NumericHandler` | `Numeric` | Decimal form, checked against declared precision and scale without rounding; a precision without a scale has scale 0. Unconstrained and wider declarations stay as decimal text until the database stores them | none | +| `NumericHandler` | `Numeric` | Plain and exponent decimal forms such as `1.5`, `.05` and `1.5E3`, converted to an Arrow decimal of the declared precision and scale without rounding; a precision without a scale has scale 0. A value that does not fit the declared digits is an error. The declaration must have a precision of at most 38 digits, see below | none | | `BooleanHandler` | `Boolean` | `true`, `t`, `yes`, `y`, `1`, `on` and their opposites | `true_words`, `false_words` | | `DateHandler` | `Date` | One of `formats`, validated against the calendar; year 0 is an error | `formats` | | `DateTimeHandler` | `DateTime` | ISO date and time, or a date in one of `formats`; see time zones below | `formats` | @@ -87,6 +87,12 @@ The column type decides, so no value is converted by a convention the schema did Input that already has a type, such as a Parquet column, is cast to the column type and checked like parsed text. A value that does not come back unchanged when cast back to its input type is an error, so an integer `2` in a `Boolean` column, a timestamp with a time in a `Date` column and a decimal that loses digits are all rejected. Float `NaN` is missing. +In a `Numeric` column an integer is read exactly and a float is read as its shortest decimal text, so a Parquet float column loads the same values as a CSV of the same numbers. Infinity is an error. + +### Numeric columns + +The loader converts only what an Arrow decimal can hold, which is a precision of at most 38 digits. That is also the widest decimal DuckDB stores, and SQLite stores `Numeric` as a float. A `Numeric` column with a wider precision or no precision raises a `PolicyError` before any row is read. To load it anyway, leave the value to the database with `columns={("table", "column"): PassthroughHandler()}`: the database then parses the text and rounds a value that exceeds the scale, and a value it cannot parse aborts the load. + ### Adding a column type Subclass `ConvertingHandler`, implement `arrow_type` and `parse`, and pass the handler in `ConversionPolicy.handlers` for the type class, or in `ConversionPolicy.columns` for a single column. `clean`, `missing` and `check` can be overridden. diff --git a/src/orm_loader/conversion/handlers/numbers.py b/src/orm_loader/conversion/handlers/numbers.py index cfdb9f6..39f5006 100644 --- a/src/orm_loader/conversion/handlers/numbers.py +++ b/src/orm_loader/conversion/handlers/numbers.py @@ -3,7 +3,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any, ClassVar +from typing import Any, ClassVar, cast import pyarrow as pa import pyarrow.compute as pc @@ -82,15 +82,28 @@ def check(self, values: pa.Array, sa_type: TypeEngine[Any]) -> pa.Array | None: @dataclass(frozen=True) class NumericHandler(NumberHandler): - """Accepts decimal forms and keeps the column's declared precision and scale. + """Reads plain and exponent decimals into an Arrow decimal with the column's precision and scale. - A precision without a scale has scale 0. A value is checked against the - declared digits and is unreadable when it does not fit, and trailing - fractional zeros are dropped. Up to 38 digits the value is cast to an - Arrow decimal. A column without a precision, or with more than 38 digits, - stays decimal text until the database stores it. + A precision without a scale has scale 0. The column must declare a + precision of at most 38 digits, the width of an Arrow decimal. Text is + handed to Arrow's decimal parser only after ``normalised`` has checked + that the column holds it exactly, because that parser silently wraps or + underflows on very long digit strings and very small exponents. Every + other value is unreadable. + + Attributes + ---------- + MAX_PRECISION + Widest precision an Arrow decimal holds, and the most digits its + parser accumulates without wrapping. + SHAPE + Regular expression splitting a value into sign, whole digits, + fraction digits and an exponent of at most three digits. """ + MAX_PRECISION: ClassVar[int] = 38 + SHAPE: ClassVar[str] = r"^(?P[+-]?)(?P\d*)(?:\.(?P\d*))?(?:[eE](?P[+-]?\d{1,3}))?$" + @staticmethod def precision_scale(sa_type: TypeEngine[Any]) -> tuple[int | None, int | None]: """Return a column's declared precision and scale.""" @@ -98,29 +111,91 @@ def precision_scale(sa_type: TypeEngine[Any]) -> tuple[int | None, int | None]: def validate(self, sa_type: TypeEngine[Any]) -> None: precision, scale = self.precision_scale(sa_type) - if precision is not None and (precision < 1 or (scale is not None and not 0 <= scale <= precision)): - raise ValueError(f"unsupported Numeric precision/scale declaration ({precision}, {scale})") + if precision is None or not 1 <= precision <= self.MAX_PRECISION or not 0 <= (scale or 0) <= precision: + raise ValueError( + f"Numeric needs a precision from 1 to {self.MAX_PRECISION} and a scale from 0 to the precision, " + f"declared ({precision}, {scale}). A PassthroughHandler for the column leaves the value to the database" + ) def arrow_type(self, sa_type: TypeEngine[Any]) -> pa.DataType: precision, scale = self.precision_scale(sa_type) - if precision is None or precision > 38: - return pa.string() - return pa.decimal128(precision, scale or 0) + return pa.decimal128(cast(int, precision), scale or 0) + + def normalised(self, cleaned: pa.Array, precision: int, scale: int) -> pa.Array: + """Return the text Arrow converts, null where the column cannot hold the value exactly. + + A value is kept when it has a digit and fits the column without + rounding or overflow. It must also stay inside what Arrow's parser + reads exactly: its digits, scaled to the column, and its implied + scale are at most ``MAX_PRECISION``. A zero is kept as ``0``, and + a leading plus sign is dropped. + + Parameters + ---------- + cleaned + Trimmed text with missing values set to null. + precision + Precision of the column. + scale + Scale of the column. + + Returns + ------- + pa.Array + String array, null for missing and for unreadable values. + """ + shaped, digits, fraction, exponent = self._split(cleaned) + count = pc.utf8_length(digits) + significant = pc.utf8_length(pc.utf8_ltrim(digits, "0")) + trailing = pc.subtract(count, pc.utf8_length(pc.utf8_rtrim(digits, "0"))) + implied = pc.subtract(fraction, exponent) + floor = pa.scalar(0, pa.int32()) + needed = pc.max_element_wise(pc.subtract(implied, trailing), floor) + growth = pc.max_element_wise(pc.subtract(pc.add(pa.scalar(scale, pa.int32()), exponent), fraction), floor) + holds = pc.and_( + pc.less_equal(needed, scale), pc.less_equal(pc.subtract(significant, implied), precision - scale) + ) + parsable = pc.and_( + pc.less_equal(pc.add(count, growth), self.MAX_PRECISION), pc.less_equal(implied, self.MAX_PRECISION) + ) + readable = pc.and_( + pc.and_(shaped, pc.greater_equal(count, 1)), + pc.or_(pc.equal(significant, 0), pc.and_(holds, parsable)), + ) + text = pc.if_else(pc.equal(significant, 0), "0", pc.replace_substring_regex(cleaned, r"^\+", "")) + return pc.if_else(pc.fill_null(readable, False), text, NULL_STRING) + + def _split(self, cleaned: pa.Array) -> tuple[pa.Array, pa.Array, pa.Array, pa.Array]: + """Split each value into its shape mask, digits, fraction length and exponent.""" + parts = pc.extract_regex(cleaned, self.SHAPE) + + def group(name: str) -> pa.Array: + return pc.fill_null(pc.struct_field(parts, name), "") + + exponent = pc.replace_substring_regex(group("exp"), r"^\+", "") + exponent = pc.if_else(pc.equal(exponent, ""), NULL_STRING, exponent).cast(pa.int32()) + digits = pc.binary_join_element_wise(group("whole"), group("frac"), "") + return pc.is_valid(parts), digits, pc.utf8_length(group("frac")), pc.fill_null(exponent, 0) - def clean(self, raw: pa.Array, trimmed: pa.Array) -> pa.Array: - return pc.replace_substring_regex(trimmed, r"(\.\d*?)0+$", r"\1") + def parse(self, cleaned: pa.Array, sa_type: TypeEngine[Any]) -> Interpretation: + precision, scale = self.precision_scale(sa_type) + candidate = self.normalised(cleaned, cast(int, precision), scale or 0) + values, cast_failed = safe_cast(candidate, self.arrow_type(sa_type)) + invalid = pc.and_(pc.is_valid(cleaned), pc.is_null(candidate)) + return Interpretation(values, pc.or_(invalid, cast_failed)) - def pattern(self, sa_type: TypeEngine[Any]) -> str: - """Return the decimal forms that fit the column's integer and fraction digits. + def _interpret_typed(self, arr: pa.Array, sa_type: TypeEngine[Any]) -> Interpretation: + """Read typed input as exactly as the text a writer would print for it. - Arrow wraps around instead of failing when a string has more digits - than a decimal holds, so both parts are bounded here. + Integers are widened first because Arrow casts only to a wide enough + decimal. A float is read as its shortest decimal text, so a Parquet + float column loads the same values as a CSV of the same numbers. + Not-a-number is missing and infinity is unreadable. """ - precision, scale = self.precision_scale(sa_type) - digits = None if precision is None and scale is None else scale or 0 - whole = r"\d+" if precision is None else rf"0*\d{{0,{precision - (digits or 0)}}}" - fraction = r"\d*" if digits is None else rf"\d{{0,{digits}}}" - forms = [rf"{whole}(\.{fraction})?"] - if digits != 0: - forms.append(r"\.\d+" if digits is None else rf"\.\d{{1,{digits}}}") - return rf"^[+-]?({'|'.join(forms)})$" + if pa.types.is_integer(arr.type): + arr = pc.cast(arr, pa.decimal128(self.MAX_PRECISION, 0)) + elif pa.types.is_floating(arr.type): + arr = arr.cast(pa.float32()) if pa.types.is_float16(arr.type) else arr + arr = pc.if_else(pc.fill_null(pc.is_nan(arr), False), pa.scalar(None, arr.type), arr) + return self.parse(pc.cast(arr, pa.string()), sa_type) + return super()._interpret_typed(arr, sa_type) diff --git a/tests/conversion/test_number_handlers.py b/tests/conversion/test_number_handlers.py index 2bcd1c6..1f0386e 100644 --- a/tests/conversion/test_number_handlers.py +++ b/tests/conversion/test_number_handlers.py @@ -1,5 +1,8 @@ +import itertools +import re import time from decimal import Decimal +from fractions import Fraction import pyarrow as pa import pytest @@ -76,20 +79,6 @@ def test_numeric_precision_without_scale_has_scale_zero(): assert unreadable("1.5", sa.Numeric(30)) -@pytest.mark.parametrize("raw", [ - "12345678901234567890.123456789", "1" + "0" * 80, "0.1234567890123456789", -]) -def test_numeric_without_precision_preserves_every_digit_as_text(raw): - assert read(raw, sa.Numeric()) == raw - - - - -@pytest.mark.parametrize("raw", ["abc", "1e5", ".", "--1"]) -def test_unconstrained_numeric_still_rejects_non_decimal_text(raw): - assert unreadable(raw, sa.Numeric()) - - def test_numeric_digits_beyond_the_precision_do_not_wrap_around(): assert unreadable("9" * 39, sa.Numeric(38, 0)) assert unreadable("9" * 45, sa.Numeric(38, 0)) @@ -138,8 +127,8 @@ def test_one_bad_value_in_a_large_batch_does_not_slow_the_whole_batch(): def test_numeric_fractions_longer_than_the_decimal_do_not_wrap_around(): values = ["1.5" + "0" * 45, "1.50" + "0" * 40, "0." + "1" * 50, "1.234", "1.2300"] assert interpret(values, sa.Numeric(10, 2)) == ( - [Decimal("1.50"), Decimal("1.50"), None, None, Decimal("1.23")], - [False, False, True, True, False], + [None, None, None, None, Decimal("1.23")], + [True, True, True, True, False], ) @@ -150,13 +139,82 @@ def test_numeric_with_scale_zero_accepts_a_zero_fraction_only(): ) -@pytest.mark.parametrize("precision", [50, 80]) -def test_numeric_wider_than_an_arrow_decimal_preserves_text_and_declared_bounds(precision): +@pytest.mark.parametrize("raw, expected", [ + ("1e3", "1000.00"), ("1.5E3", "1500.00"), ("1.5E+003", "1500.00"), ("5e005", "500000.00"), ("2.50e1", "25.00"), ("1e-2", "0.01"), ("-2.5E+2", "-250.00"), + (".05", "0.05"), ("5.", "5.00"), (".0", "0.00"), ("+7", "7.00"), ("1.500", "1.50"), ("007.50", "7.50"), +]) +def test_numeric_reads_plain_and_exponent_forms(raw, expected): + assert read(raw, sa.Numeric(12, 2)) == Decimal(expected) + + +@pytest.mark.parametrize("raw", [ + ".", "+", "-", "+.", "-.", "e5", "1e", "1e+", "1e+-5", "--1", "1.2.3", "0x1F", "1,5", "abc", + "1.555", "1e-3", "1.5e-3", "1e10", "123456789012.00", +]) +def test_numeric_rejects_malformed_and_inexact_values(raw): + assert unreadable(raw, sa.Numeric(12, 2)) + + +@pytest.mark.parametrize("raw", ["5e-50", "1e-40", "1.0e-39", "1" + "0" * 39, "9" * 39, "1e99999", "5e-9999"]) +def test_numeric_values_beyond_what_arrow_parses_exactly_are_unreadable(raw): + assert unreadable(raw, sa.Numeric(38, 2)) + assert unreadable(raw, sa.Numeric(38, 0)) + + +@pytest.mark.parametrize("raw", ["0", "-0", "0.000", ".0", "0e50", "0.0e-39", "000", "0" * 60]) +def test_numeric_reads_every_spelling_of_zero(raw): + assert read(raw, sa.Numeric(38, 2)) == Decimal("0.00") + + +def test_numeric_fills_the_declared_precision_exactly(): + assert read("9" * 36 + ".99", sa.Numeric(38, 2)) == Decimal("9" * 36 + ".99") + assert unreadable("9" * 37 + ".99", sa.Numeric(38, 2)) + + +def _reference(raw: str, precision: int, scale: int) -> Fraction | None: + """Exact value of a decimal or exponent literal, or None when it does not fit the column.""" + match = re.fullmatch(r"([+-]?)(\d*)\.?(\d*)(?:[eE]([+-]?\d+))?", raw) + if not match or not (match.group(2) or match.group(3)) or not re.fullmatch(r"[+-]?(\d+\.?\d*|\.\d+)([eE][+-]?\d+)?", raw): + return None + sign, whole, fraction, exponent = match.group(1), match.group(2), match.group(3), int(match.group(4) or 0) + value = Fraction(int((whole + fraction) or "0"), 10 ** len(fraction)) * Fraction(10) ** exponent + scaled = (-value if sign == "-" else value) * 10 ** scale + if scaled.denominator != 1 or abs(scaled) >= 10 ** precision: + return None + return Fraction(scaled.numerator, 10 ** scale) + + +@pytest.mark.parametrize("precision, scale", [(10, 2), (12, 0), (38, 2), (6, 6)]) +def test_numeric_reads_every_string_exactly_or_refuses_it(precision, scale): + """A string is read to its exact value or refused, and no string over this alphabet is refused though valid.""" + raws = ["".join(chars) for length in range(1, 6) for chars in itertools.product("+-.05e9", repeat=length)] + raws += [ + f"{mantissa}e{sign}{exponent}" + for mantissa in ("1", "1.5", "9.99", ".5") + for sign in ("+", "-") + for exponent in range(100) + ] + values, invalid = interpret(raws, sa.Numeric(precision, scale)) + refused_valid = [] + for raw, value, bad in zip(raws, values, invalid): + expected = _reference(raw, precision, scale) + if not bad: + assert expected is not None and Fraction(value) == expected, raw + elif expected is not None: + refused_valid.append(raw) + assert refused_valid == [], refused_valid[:10] + + +@pytest.mark.parametrize("declared", [sa.Numeric(), sa.Numeric(39, 2), sa.Numeric(50, 2), sa.Numeric(0, 0), sa.Numeric(2, 5)]) +def test_numeric_declarations_an_arrow_decimal_cannot_hold_are_refused(declared): + from orm_loader.conversion.handlers import NumericHandler + + with pytest.raises(ValueError, match="precision from 1 to 38"): + NumericHandler().validate(declared) + + +@pytest.mark.parametrize("declared", [sa.Numeric(1, 0), sa.Numeric(38, 38), sa.Numeric(18), sa.Numeric(10, 2)]) +def test_numeric_declarations_an_arrow_decimal_holds_are_accepted(declared): from orm_loader.conversion.handlers import NumericHandler - column = sa.Numeric(precision, 2) - assert NumericHandler().arrow_type(column) == pa.string() - raw = "9" * (precision - 2) + ".12" - assert read(raw, column) == raw - assert unreadable("9" * (precision - 1), column) - assert unreadable("1.234", column) + NumericHandler().validate(declared) diff --git a/tests/conversion/test_typed_input.py b/tests/conversion/test_typed_input.py index fce9678..8aca5b7 100644 --- a/tests/conversion/test_typed_input.py +++ b/tests/conversion/test_typed_input.py @@ -4,6 +4,7 @@ from decimal import Decimal import pyarrow as pa +import pytest import sqlalchemy as sa from tests.support import interpret @@ -81,3 +82,19 @@ def test_the_time_zone_rule_of_the_column_applies_to_typed_timestamps(): assert _result(naive, sa.DateTime(timezone=True))[1] == [True, False] assert _result(pa.array([date(2020, 1, 1)], pa.date32()), sa.DateTime(timezone=True))[1] == [True] assert _result(pa.array([date(2020, 1, 1)], pa.date32()), sa.DateTime())[1] == [False] + + +@pytest.mark.parametrize("source", [pa.decimal128(10, 3), pa.decimal256(55, 3)]) +def test_decimals_of_any_width_are_read_exactly_or_refused(source): + values, invalid = _result(pa.array([Decimal("1.250"), Decimal("1.555"), None], source), sa.Numeric(12, 2)) + assert values == [Decimal("1.25"), None, None] + assert invalid == [False, True, False] + + +def test_floats_and_integers_in_a_numeric_column_must_fit_exactly(): + values, invalid = _result(pa.array([0.1, 1.555, float("inf"), float("nan"), 2.5]), sa.Numeric(12, 2)) + assert values == [Decimal("0.10"), None, None, None, Decimal("2.50")] + assert invalid == [False, True, True, False, False] + values, invalid = _result(pa.array([5, 10**11]), sa.Numeric(12, 2)) + assert values == [Decimal("5.00"), None] + assert invalid == [False, True] diff --git a/tests/loaders/test_numeric_columns.py b/tests/loaders/test_numeric_columns.py new file mode 100644 index 0000000..78a17af --- /dev/null +++ b/tests/loaders/test_numeric_columns.py @@ -0,0 +1,137 @@ +"""Numeric columns load exactly the values their file holds, in text and in typed Parquet columns.""" + +from decimal import Decimal +from pathlib import Path +from typing import cast + +import numpy as np +import pyarrow as pa +import pyarrow.csv as pv +import pyarrow.parquet as pq +import pytest +import sqlalchemy as sa +import sqlalchemy.orm as so + +from orm_loader.backends import STAGING_SCHEMA +from orm_loader.conversion import ConversionPolicy, ErrorPolicy +from orm_loader.conversion.handlers import PassthroughHandler +from orm_loader.helpers import ConversionError, PolicyError +from orm_loader.loaders import LoadReport +from orm_loader.tables import FileLoadableTableInterface +from orm_loader.tables.typing import FileTableProtocol + + +class AmountBase(so.DeclarativeBase): + pass + + +class Amount(AmountBase, FileLoadableTableInterface): + __tablename__ = "numeric_amount" + record_id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True, autoincrement=False) + amount: so.Mapped[Decimal | None] = so.mapped_column(sa.Numeric(12, 2)) + + +class WideAmount(AmountBase, FileLoadableTableInterface): + __tablename__ = "numeric_wide_amount" + record_id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True, autoincrement=False) + amount: so.Mapped[Decimal | None] = so.mapped_column(sa.Numeric(50, 2)) + + +class UnboundedAmount(AmountBase, FileLoadableTableInterface): + __tablename__ = "numeric_unbounded_amount" + record_id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True, autoincrement=False) + amount: so.Mapped[Decimal | None] = so.mapped_column(sa.Numeric()) + + +REJECT = ConversionPolicy(errors=ErrorPolicy(on_error="reject_row")) + + +def _load(session, path: Path, model, **kwargs) -> int: + model.__table__.create(session.connection(), checkfirst=True) + session.commit() + tag = None if session.get_bind().dialect.name == "sqlite" else STAGING_SCHEMA + loaded = cast(type[FileTableProtocol], model).load_file(session, path, staging_schema_tag=tag, **kwargs) + session.commit() + return loaded + + +def _amounts(session, model=Amount) -> dict[int, Decimal | None]: + rows = session.execute(sa.select(model.record_id, model.amount).order_by(model.record_id)).all() + return {key: None if value is None else Decimal(str(value)) for key, value in rows} + + +def _rejected(path: Path) -> list[tuple[str, str]]: + header = [h.strip('"') for h in path.read_text().splitlines()[0].split(",")] + table = pv.read_csv(path, convert_options=pv.ConvertOptions(column_types={name: pa.string() for name in header})) + return [(row["row"], row["amount"]) for row in table.to_pylist()] + + +def test_text_amounts_load_exactly_and_the_rest_is_rejected_with_its_row(dialect_session, tmp_path): + cells = ["1.50", " 2.5 ", ".05", "1.5E3", "1e-2", "007.50", "+7", "3.100", "5.", ".0", "1.5E+003", + ".", "abc", "1.555", "1e10", "5e-50", "0e+-0", "1" + "0" * 40] + path = tmp_path / "numeric_amount.csv" + path.write_text("record_id,amount\n" + "".join(f"{i},{cell}\n" for i, cell in enumerate(cells, 1))) + report = LoadReport() + assert _load(dialect_session, path, Amount, policy=REJECT, report=report) == 11 + assert _amounts(dialect_session) == { + 1: Decimal("1.50"), 2: Decimal("2.50"), 3: Decimal("0.05"), 4: Decimal("1500.00"), 5: Decimal("0.01"), + 6: Decimal("7.50"), 7: Decimal("7.00"), 8: Decimal("3.10"), 9: Decimal("5.00"), 10: Decimal("0.00"), + 11: Decimal("1500.00"), + } + assert (report.rows_read, report.rows_loaded, report.rows_rejected) == (18, 11, 7) + assert report.accounted + assert _rejected(report.rejects_path) == [ + ("12", "."), ("13", "abc"), ("14", "1.555"), ("15", "1e10"), ("16", "5e-50"), ("17", "0e+-0"), ("18", "1" + "0" * 40), + ] + + +def test_a_lone_point_raises_a_conversion_error_before_it_reaches_the_database(dialect_session, tmp_path): + path = tmp_path / "numeric_amount.csv" + path.write_text("record_id,amount\n1,.\n2,2.50\n") + with pytest.raises(ConversionError, match="amount"): + _load(dialect_session, path, Amount) + assert _amounts(dialect_session) == {} + + +@pytest.mark.parametrize("column, loaded, rejected", [ + (pa.array([5, 10**11, 7], pa.int64()), {1: "5.00", 3: "7.00"}, ["2"]), + (pa.array([5, 10**11], pa.uint64()), {1: "5.00"}, ["2"]), + (pa.array([Decimal("1.250"), Decimal("1.555"), Decimal("2.5")], pa.decimal256(55, 3)), {1: "1.25", 3: "2.50"}, ["2"]), + (pa.array([Decimal("1.25"), Decimal("1.5")], pa.decimal128(10, 2)), {1: "1.25", 2: "1.50"}, []), + (pa.array([708693.7, 0.1 + 0.2, 2.5, float("inf"), float("nan")]), {1: "708693.70", 3: "2.50", 5: None}, ["2", "4"]), + (pa.array(np.array([-549424.4, 0.1], np.float32)), {1: "-549424.40", 2: "0.10"}, []), +], ids=["int64", "uint64", "decimal256", "decimal128", "float64", "float32"]) +def test_parquet_columns_load_exactly_or_are_rejected(dialect_session, tmp_path, column, loaded, rejected): + path = tmp_path / "numeric_amount.parquet" + pq.write_table(pa.table({"record_id": pa.array(range(1, len(column) + 1), pa.int32()), "amount": column}), path) + report = LoadReport() + _load(dialect_session, path, Amount, policy=REJECT, report=report) + assert _amounts(dialect_session) == {key: None if value is None else Decimal(value) for key, value in loaded.items()} + assert report.rows_rejected == len(rejected) + if rejected: + assert [row for row, _ in _rejected(report.rejects_path)] == rejected + assert report.accounted + + +@pytest.mark.parametrize("model, declared", [(WideAmount, r"\(50, 2\)"), (UnboundedAmount, r"\(None, None\)")]) +def test_a_numeric_column_an_arrow_decimal_cannot_hold_is_refused_before_any_row_is_read( + sqlite_session, tmp_path, model, declared +): + path = tmp_path / f"{model.__tablename__}.csv" + path.write_text("record_id,amount\n1,2.50\n") + with pytest.raises(PolicyError, match=rf"{model.__tablename__}.amount.*{declared}.*PassthroughHandler"): + _load(sqlite_session, path, model) + assert sqlite_session.scalar(sa.select(sa.func.count()).select_from(model)) == 0 + + +@pytest.mark.parametrize("model, expected", [ + (WideAmount, ["1.50", "0.05", "1500.00", "1.56", "1" + "0" * 40 + ".25"]), + (UnboundedAmount, ["1.50", "0.05", "1500", "1.555", "1" + "0" * 40 + ".25"]), +]) +def test_a_column_the_loader_cannot_convert_loads_when_left_to_the_database(pg_session, tmp_path, model, expected): + cells = ["1.50", ".05", "1.5E3", "1.555", "1" + "0" * 40 + ".25"] + path = tmp_path / f"{model.__tablename__}.csv" + path.write_text("record_id,amount\n" + "".join(f"{i},{cell}\n" for i, cell in enumerate(cells, 1))) + policy = ConversionPolicy(columns={(model.__tablename__, "amount"): PassthroughHandler()}) + assert _load(pg_session, path, model, policy=policy) == len(cells) + assert list(_amounts(pg_session, model).values()) == [Decimal(value) for value in expected] diff --git a/tests/loaders/test_special_columns.py b/tests/loaders/test_special_columns.py index 58d1c70..915dc5e 100644 --- a/tests/loaders/test_special_columns.py +++ b/tests/loaders/test_special_columns.py @@ -73,6 +73,18 @@ class InvalidNumericRecord(SpecialBase, FileLoadableTableInterface): amount: so.Mapped[float] = so.mapped_column(sa.Numeric(2, 5)) +class UnboundedNumericRecord(SpecialBase, FileLoadableTableInterface): + __tablename__ = "special_unbounded_numeric" + record_id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) + amount: so.Mapped[float] = so.mapped_column(sa.Numeric()) + + +class WideNumericRecord(SpecialBase, FileLoadableTableInterface): + __tablename__ = "special_wide_numeric" + record_id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) + amount: so.Mapped[float] = so.mapped_column(sa.Numeric(50, 2)) + + def _create(session, model): model.__table__.create(session.connection()) session.commit() @@ -185,9 +197,12 @@ def test_a_primary_key_is_required_and_takes_no_default(): assert "record_id" not in DefaultKeyRecord.column_defaults() -def test_unsupported_numeric_declaration_names_the_table_and_column(tmp_path): - with pytest.raises(PolicyError, match=r"special_invalid_numeric.amount.*\(2, 5\)"): +@pytest.mark.parametrize("model, declared", [ + (InvalidNumericRecord, r"\(2, 5\)"), (UnboundedNumericRecord, r"\(None, None\)"), (WideNumericRecord, r"\(50, 2\)"), +]) +def test_unsupported_numeric_declaration_names_the_table_and_column(tmp_path, model, declared): + with pytest.raises(PolicyError, match=rf"{model.__tablename__}.amount.*{declared}"): BatchProcessor( - cast(type[FileTableProtocol], InvalidNumericRecord), ["record_id", "amount"], + cast(type[FileTableProtocol], model), ["record_id", "amount"], ConversionPolicy(), LoadReport(), tmp_path / "rejects.csv", )