diff --git a/.github/workflows/integration-tests-mysql.yml b/.github/workflows/integration-tests-mysql.yml index f1a8d89..d3ba942 100644 --- a/.github/workflows/integration-tests-mysql.yml +++ b/.github/workflows/integration-tests-mysql.yml @@ -13,6 +13,7 @@ jobs: services: mysql: image: mysql:8.0 + command: --local-infile=1 env: MYSQL_DATABASE: test_db MYSQL_USER: user @@ -21,13 +22,18 @@ jobs: TZ: Europe/Paris ports: - 3306:3306 - options: --health-cmd="mysqladmin ping" --health-interval=10s --health-timeout=5s --health-retries=3 + options: >- + --health-cmd="mysqladmin ping" + --health-interval=10s + --health-timeout=5s + --health-retries=3 steps: - name: Verify MySQL connection run: | mysql --version sudo apt-get install -y mysql-client mysql --host 127.0.0.1 --port 3306 -uroot -psecretroot -e "SHOW DATABASES" + mysql --host 127.0.0.1 --port 3306 -uroot -psecretroot -e "SHOW VARIABLES LIKE 'local_infile'" - name: Check out repository code uses: actions/checkout@v4 diff --git a/README.md b/README.md index c1ab22f..c649b75 100644 --- a/README.md +++ b/README.md @@ -37,8 +37,8 @@ SQL views or stored procedures aimed at extracting, correcting and formatting th This package uses `sqlalchemy` to interact with the database, so it should work with different database backends. Automated integration tests run against PostgreSQL, MySQL, MS SQL Server and DuckDB. You may have to install additional -packages to connect to your database (e.g. `psycopg2` for PostgreSQL, `pymysql` for MySQL, `pyodbc` for MS SQL Server or -`duckdb_engine` for DuckDB). +packages to connect to your database (e.g. `psycopg2` or `psycopg` for PostgreSQL, `pymysql` or `mysqlclient` for +MySQL, `pyodbc` for MS SQL Server, or `duckdb-engine` for DuckDB). **Please read the [package documentation website](https://cre-dev.github.io/xml2db) for all the details!** diff --git a/docs/getting_started.md b/docs/getting_started.md index d438170..b69958b 100644 --- a/docs/getting_started.md +++ b/docs/getting_started.md @@ -33,7 +33,9 @@ data_model = DataModel( At this stage, it is not required to provide a connection string and have an actual database set up, but it will be necessary if you want to use this [`DataModel`](api/data_model.md) to actually import data. You may need to install the -Python package of the connector you use in your `sqlalchemy` connection string (`psycopg2` in the example above). +Python package of the connector you use in your `sqlalchemy` connection string (e.g. `psycopg2` or `psycopg` for +PostgreSQL, `pymysql` or `mysqlclient` for MySQL, `pyodbc` for SQL Server, `duckdb-engine` for DuckDB). See +[How it works](how_it_works.md#bulk-loading) for which drivers enable native bulk loading. You can provide an optional model configuration, which will allow forcing or preventing some schema simplification, set columns types manually, etc. By default, some simplifications will be applied when possible, in order to limit the diff --git a/docs/how_it_works.md b/docs/how_it_works.md index c21dfc7..fc86a4a 100644 --- a/docs/how_it_works.md +++ b/docs/how_it_works.md @@ -184,6 +184,34 @@ We keep the primary keys from the flat data model created at the previous stage, random one, which can be useful if you want to decompose the process of loading data and merging it with the target tables later, for instance to gain a finer control over concurrency. +### Bulk loading + +For each supported backend, `xml2db` uses a native bulk-loading mechanism to fill the temporary tables, which is +significantly faster than row-by-row inserts for large datasets. When the native path is unavailable (wrong driver, +missing tool, or server configuration), it falls back silently to SQLAlchemy's standard `executemany`. + +| Backend | Mechanism | Required driver / tool | Default threshold | +|---|---|---|---| +| PostgreSQL | `COPY FROM STDIN` | `psycopg2` or `psycopg` | 0 (always used) | +| MySQL / MariaDB | `LOAD DATA LOCAL INFILE` | `pymysql` or `mysqlclient`; server `local_infile=ON` | 100 rows | +| MS SQL Server | `bcp` utility | `bcp` on PATH; SQL or Windows/Kerberos auth | 100 rows | +| DuckDB | `read_csv()` | built-in | 100 rows | + +The threshold column means that batches smaller than that number always use `executemany` (avoiding temp-file overhead +for small inserts). PostgreSQL's `COPY` is in-protocol and has no file overhead, so there is no threshold. + +For MySQL, when a connection string is passed to `DataModel`, `local_infile=True` is injected automatically into the +connection arguments. The MySQL server must also have `local_infile=ON` (e.g. launched with `--local-infile=1`); +if not, bulk loading is silently skipped. + +You can control this behaviour via the `bulk_load` and `bulk_load_threshold` arguments of +[`Document.insert_into_target_tables`](api/document.md#xml2db.document.Document.insert_into_target_tables): + +- `bulk_load=None` (default): use the native path when available, fall back to `executemany` silently. +- `bulk_load=False`: always use `executemany`, regardless of what is available. +- `bulk_load=True`: require the native path; raise a `RuntimeError` with an actionable message if the required + driver, tool, or server setting is missing. + ### Merging the data The last step is to merge the temporary tables data into the target tables, while enforcing deduplication, keeping diff --git a/docs/index.md b/docs/index.md index d3fccb6..5254f4a 100644 --- a/docs/index.md +++ b/docs/index.md @@ -43,8 +43,9 @@ SQL views or stored procedures aimed at extracting, correcting and formatting th This package uses `sqlalchemy` to interact with the database, so it should work with different database backends. Automated integration tests run against PostgreSQL, MySQL, MS SQL Server and DuckDB. You may have to install additional -packages to connect to your database (e.g. `psycopg2` for PostgreSQL, `pymysql` for MySQL, `pyodbc` for MS SQL Server or -`duckdb_engine` for DuckDB). +packages to connect to your database (e.g. `psycopg2` or `psycopg` for PostgreSQL, `pymysql` or `mysqlclient` for +MySQL, `pyodbc` for MS SQL Server, or `duckdb-engine` for DuckDB). See [How it works](how_it_works.md#bulk-loading) +for which drivers enable native bulk loading. ## How to visualize your data model diff --git a/src/xml2db/dialect/base.py b/src/xml2db/dialect/base.py index b9c566d..27f5f94 100644 --- a/src/xml2db/dialect/base.py +++ b/src/xml2db/dialect/base.py @@ -14,6 +14,7 @@ SmallInteger, BigInteger, LargeBinary, + create_engine as _sa_create_engine, ) from sqlalchemy import inspect as sqlalchemy_inspect import sqlalchemy.schema @@ -317,11 +318,45 @@ def validate_model_config(self, config: dict) -> dict: ) return config + # ------------------------------------------------------------------ + # Engine creation + # ------------------------------------------------------------------ + + def create_engine(self, connection_string: str, **kwargs: Any) -> Any: + """Create a SQLAlchemy engine for this backend. + + The base implementation is a plain pass-through to + :func:`sqlalchemy.create_engine`. Subclasses override this to inject + backend-specific engine options (e.g. ``fast_executemany`` for MSSQL, + ``local_infile`` for MySQL). + + This method is only called when ``DataModel`` is given a + ``connection_string``; callers that supply a pre-built ``db_engine`` + bypass it entirely. + + Args: + connection_string: A SQLAlchemy database URL. + **kwargs: Additional keyword arguments forwarded to + :func:`sqlalchemy.create_engine`. + + Returns: + A :class:`sqlalchemy.Engine` instance. + """ + return _sa_create_engine(connection_string, **kwargs) + # ------------------------------------------------------------------ # Data loading # ------------------------------------------------------------------ - def bulk_insert(self, conn: Any, table: Any, records: list) -> None: + def bulk_insert( + self, + conn: Any, + table: Any, + records: list, + *, + bulk_load: bool | None = None, + bulk_load_threshold: int | None = None, + ) -> None: """Insert records into a staging table. The base implementation uses SQLAlchemy's parameterised executemany, @@ -332,6 +367,11 @@ def bulk_insert(self, conn: Any, table: Any, records: list) -> None: conn: A SQLAlchemy ``Connection`` already within a transaction. table: The SQLAlchemy ``Table`` object to insert into. records: A list of dicts mapping column keys to Python values. + bulk_load: ``True`` to require bulk loading (raise if unavailable), + ``False`` to always use executemany, or ``None`` (default) to + use bulk loading when available and fall back silently otherwise. + bulk_load_threshold: Minimum number of records to trigger bulk + loading. ``None`` delegates the choice to the subclass. """ if records: conn.execute(table.insert(), records) diff --git a/src/xml2db/dialect/duckdb.py b/src/xml2db/dialect/duckdb.py index af750b0..3c7f2ad 100644 --- a/src/xml2db/dialect/duckdb.py +++ b/src/xml2db/dialect/duckdb.py @@ -20,6 +20,9 @@ from .base import DatabaseDialect +# Records below this count skip read_csv (temp-file overhead). +_READ_CSV_THRESHOLD = 100 + class DuckDBDialect(DatabaseDialect): """Dialect for DuckDB. @@ -87,21 +90,42 @@ def _select_expr(self, key: str, col: Any) -> str: return f'CAST("{key}" AS {duckdb_type})' return f'"{key}"' # String / unknown: keep as VARCHAR - def bulk_insert(self, conn: Any, table: Any, records: list) -> None: + def bulk_insert( + self, + conn: Any, + table: Any, + records: list, + *, + bulk_load: bool | None = None, + bulk_load_threshold: int | None = None, + ) -> None: """Bulk-insert records via a temporary CSV file and DuckDB's ``read_csv``. All CSV columns are read as VARCHAR (``all_varchar=true``) and then explicitly cast to their target types in the ``SELECT`` clause. Binary columns are hex-encoded in the CSV and decoded with ``unhex()``. + Falls back to the base-class parameterised executemany when + ``bulk_load=False`` or the batch is below the effective threshold. + Args: conn: A SQLAlchemy ``Connection`` already within a transaction. table: The SQLAlchemy ``Table`` object to insert into. records: A list of dicts mapping column keys to Python values. + bulk_load: ``True`` or ``None`` (default) to use read_csv for + batches at or above the threshold; ``False`` to always use + executemany. + bulk_load_threshold: Override the minimum batch size. Defaults to + :data:`_READ_CSV_THRESHOLD` (100). """ if not records: return + threshold = bulk_load_threshold if bulk_load_threshold is not None else _READ_CSV_THRESHOLD + if bulk_load is False or len(records) < threshold: + super().bulk_insert(conn, table, records) + return + # Map column key -> SQLAlchemy Column object col_by_key = {col.key: col for col in table.columns} diff --git a/src/xml2db/dialect/mssql.py b/src/xml2db/dialect/mssql.py index 936d481..e0cb769 100644 --- a/src/xml2db/dialect/mssql.py +++ b/src/xml2db/dialect/mssql.py @@ -14,7 +14,7 @@ # Records below this count go through fast_executemany; at or above, BCP is used. -_BCP_THRESHOLD = 1000 +_BCP_THRESHOLD = 100 class MSSQLDialect(DatabaseDialect): @@ -25,19 +25,14 @@ class MSSQLDialect(DatabaseDialect): handled in this class. When the ``bcp`` utility is available on PATH and the connection uses SQL - authentication, :meth:`bulk_insert` switches to BCP for batches of - :data:`_BCP_THRESHOLD` rows or more. Smaller batches always use - ``fast_executemany`` (enabled at engine level) to avoid BCP's subprocess - overhead. Pass ``use_bcp=False`` to :class:`~xml2db.DataModel` to disable - BCP entirely. + authentication or a trusted connection, :meth:`bulk_insert` switches to BCP + for batches of :data:`_BCP_THRESHOLD` rows or more. Smaller batches always + use ``fast_executemany`` (enabled at engine level) to avoid BCP's subprocess + overhead. """ MAX_IDENTIFIER_LENGTH: int = 128 - def __init__(self, use_bcp: bool = True, **kwargs): - super().__init__(**kwargs) - self.bcp_path = shutil.which("bcp") if use_bcp else None - def validate_table_config(self, config: dict) -> dict: """Allow ``as_columnstore`` through unchanged for MSSQL.""" return config @@ -93,6 +88,12 @@ def relation_extra_indexes( ), ) + def create_engine(self, connection_string: str, **kwargs: Any) -> Any: + """Create a MSSQL engine with ``fast_executemany`` and ``SERIALIZABLE`` isolation.""" + kwargs.setdefault("fast_executemany", True) + kwargs.setdefault("isolation_level", "SERIALIZABLE") + return super().create_engine(connection_string, **kwargs) + @staticmethod def _format_bcp_value(v: Any) -> str: """Format a Python value as a BCP character-mode field (tab-separated).""" @@ -108,32 +109,71 @@ def _format_bcp_value(v: Any) -> str: # Tab is the field delimiter; replace any occurrence in string values. return s.replace("\t", " ") - def bulk_insert(self, conn: Any, table: Any, records: list) -> None: + def bulk_insert( + self, + conn: Any, + table: Any, + records: list, + *, + bulk_load: bool | None = None, + bulk_load_threshold: int | None = None, + ) -> None: """Bulk-insert records, using BCP for large batches when available. - Batches smaller than :data:`_BCP_THRESHOLD` rows, or any batch when - BCP is unavailable or the connection has neither SQL credentials nor - ``Trusted_Connection=yes``, fall back to ``fast_executemany``. - SQL auth uses ``-U``/``-P``; Kerberos/Windows auth uses ``-T``. + Batches smaller than the effective threshold, or any batch when + ``bulk_load=False``, use ``fast_executemany`` unconditionally. + When ``bulk_load=True`` and the batch meets the threshold but BCP + prerequisites are not satisfied (binary not on PATH, unsupported auth), + a :class:`RuntimeError` is raised with an actionable message. + + SQL auth uses ``-U``/``-P``; Kerberos/Windows auth uses ``-T``. BCP does not participate in the caller's SQLAlchemy transaction. Args: conn: A SQLAlchemy ``Connection`` already within a transaction. table: The SQLAlchemy ``Table`` object to insert into. records: A list of dicts mapping column keys to Python values. + bulk_load: ``True`` to require BCP (raise if unavailable for this + batch), ``False`` to always use fast_executemany, or ``None`` + (default) to use BCP when available and fall back silently. + bulk_load_threshold: Override the minimum batch size for BCP. + Defaults to :data:`_BCP_THRESHOLD` (100). """ if not records: return + threshold = bulk_load_threshold if bulk_load_threshold is not None else _BCP_THRESHOLD + + # bulk_load=False or batch too small → always use fast_executemany. + if bulk_load is False or len(records) < threshold: + super().bulk_insert(conn, table, records) + return + + # Check BCP prerequisites. url = conn.engine.url trusted = str(url.query.get("Trusted_Connection", "")).lower() == "yes" has_sql_auth = bool(url.username and url.password) - if ( - self.bcp_path is None - or len(records) < _BCP_THRESHOLD - or (not has_sql_auth and not trusted) - ): + bcp_path = shutil.which("bcp") + + if bcp_path is None: + if bulk_load is True: + raise RuntimeError( + "bulk_load=True requires the bcp utility on PATH. " + "Install mssql-tools (Linux/macOS) or SQL Server Command Line Utilities " + "(Windows), or set bulk_load=False to use fast_executemany instead." + ) + super().bulk_insert(conn, table, records) + return + + if not has_sql_auth and not trusted: + if bulk_load is True: + raise RuntimeError( + "bulk_load=True requires SQL Server authentication (username and password " + "in the connection string) or a Windows/Kerberos trusted connection " + "(Trusted_Connection=yes in the connection string query parameters). " + "Set bulk_load=False to use fast_executemany instead." + ) super().bulk_insert(conn, table, records) return @@ -167,7 +207,7 @@ def bulk_insert(self, conn: Any, table: Any, records: list) -> None: f.write("\t".join(row) + "\n") cmd = [ - self.bcp_path, full_name, "in", data_path, + bcp_path, full_name, "in", data_path, "-S", f"{url.host},{url.port or 1433}", "-d", url.database, "-c", # character (text) mode diff --git a/src/xml2db/dialect/mysql.py b/src/xml2db/dialect/mysql.py index c5d20fe..e3e72a5 100644 --- a/src/xml2db/dialect/mysql.py +++ b/src/xml2db/dialect/mysql.py @@ -1,23 +1,50 @@ +import os +import tempfile from typing import Any, TYPE_CHECKING -from sqlalchemy import String +from sqlalchemy import LargeBinary, String, text from sqlalchemy.dialects import mysql as mysql_dialect +from sqlalchemy.exc import OperationalError from .base import DatabaseDialect if TYPE_CHECKING: from ..table.column import DataModelColumn +# MySQL error codes for "LOAD DATA LOCAL INFILE not available". +_LOCAL_INFILE_ERRORS = (1148, 3948) + +# Records below this count skip LOAD DATA LOCAL INFILE (temp-file overhead). +_LOAD_DATA_THRESHOLD = 100 + class MySQLDialect(DatabaseDialect): """Dialect for MySQL / MariaDB. MySQL enforces a 64-character limit on identifiers. + + :meth:`bulk_insert` uses ``LOAD DATA LOCAL INFILE`` for the ``pymysql`` + and ``mysqldb`` drivers when the batch meets :data:`_LOAD_DATA_THRESHOLD`. + The engine must be created via :class:`~xml2db.model.DataModel` (which + injects ``local_infile=True`` automatically) or with + ``connect_args={"local_infile": True}``, and the MySQL server must have + ``local_infile=ON``. """ # further reducing the max length because SQL Alchemy adds suffixes to foreign key names MAX_IDENTIFIER_LENGTH: int = 56 + def __init__(self, **kwargs: Any): + super().__init__(**kwargs) + # Tri-state cache: None = not yet tested, True = works, False = unavailable. + self._local_infile_ok: bool | None = None + + def create_engine(self, connection_string: str, **kwargs: Any) -> Any: + """Create a MySQL engine with ``local_infile=True`` in connect_args.""" + connect_args = kwargs.pop("connect_args", {}) + connect_args.setdefault("local_infile", True) + return super().create_engine(connection_string, connect_args=connect_args, **kwargs) + def column_type(self, col: "DataModelColumn", temp: bool) -> Any: if col.occurs[1] != 1: return String(4000) @@ -29,3 +56,178 @@ def column_type(self, col: "DataModelColumn", temp: bool) -> Any: return mysql_dialect.BINARY(col.max_length) return mysql_dialect.VARBINARY(col.max_length) return super().column_type(col, temp) + + @staticmethod + def _format_value(v: Any, col: Any) -> str: + """Format a Python value for MySQL LOAD DATA LOCAL INFILE (tab-separated). + + NULL → ``\\N``. Binary columns are hex-encoded (decoded server-side + with UNHEX()). Booleans become ``1``/``0``. Strings have backslashes + and the field/line delimiters escaped so MySQL's default ESCAPED BY + handler reconstructs the original text. + """ + if v is None: + return "\\N" + if isinstance(col.type, LargeBinary): + # Decoded server-side via UNHEX(); transmit as plain hex string. + return bytes(v).hex() if isinstance(v, (bytes, bytearray, memoryview)) else "" + if isinstance(v, bool): + # Must precede int: bool subclasses int. + return "1" if v else "0" + s = str(v) + # Escape backslash first, then the characters MySQL's LOAD DATA + # interprets as special under ESCAPED BY '\\'. + s = s.replace("\\", "\\\\") + s = s.replace("\n", "\\n") + s = s.replace("\r", "\\r") + s = s.replace("\t", "\\t") + return s + + def bulk_insert( + self, + conn: Any, + table: Any, + records: list, + *, + bulk_load: bool | None = None, + bulk_load_threshold: int | None = None, + ) -> None: + """Bulk-insert records via MySQL's ``LOAD DATA LOCAL INFILE``. + + Builds a tab-separated temp file and streams it to the server using + the driver's ``LOAD DATA LOCAL INFILE`` protocol. Supported drivers: + + - **pymysql**: the engine must have ``local_infile=True`` in + ``connect_args`` (set automatically by :meth:`create_engine`). + - **mysqldb** (mysqlclient): same requirement. + + Falls back to the base-class parameterised executemany for other + drivers or when ``LOAD DATA LOCAL INFILE`` is unavailable. + + Binary columns are hex-encoded in the file and decoded server-side + with ``UNHEX()``. + + Args: + conn: A SQLAlchemy ``Connection`` already within a transaction. + table: The SQLAlchemy ``Table`` object to insert into. + records: A list of dicts mapping column keys to Python values. + bulk_load: ``True`` to require LOAD DATA LOCAL INFILE (raise if + unavailable), ``False`` to always use executemany, or ``None`` + (default) to use LOAD DATA when available and fall back silently. + bulk_load_threshold: Override the minimum batch size to trigger + LOAD DATA LOCAL INFILE. Defaults to + :data:`_LOAD_DATA_THRESHOLD` (100). + """ + if not records: + return + + threshold = bulk_load_threshold if bulk_load_threshold is not None else _LOAD_DATA_THRESHOLD + driver = conn.dialect.driver + + # Unsupported driver. + if driver not in ("pymysql", "mysqldb"): + if bulk_load is True: + raise RuntimeError( + f"bulk_load=True requires the pymysql or mysqldb driver, got '{driver}'. " + f"Use a mysql+pymysql:// or mysql+mysqldb:// connection string." + ) + super().bulk_insert(conn, table, records) + return + + # bulk_load=False or batch too small → use executemany. + if bulk_load is False or len(records) < threshold: + super().bulk_insert(conn, table, records) + return + + # Cached failure: LOAD DATA LOCAL INFILE is known to be unavailable. + if self._local_infile_ok is False: + if bulk_load is True: + raise RuntimeError( + "bulk_load=True requires LOAD DATA LOCAL INFILE, but it is unavailable on " + "this connection. Enable local_infile on the MySQL server and ensure the " + "engine is created via DataModel (which sets local_infile=True automatically) " + "or with connect_args={'local_infile': True}." + ) + super().bulk_insert(conn, table, records) + return + + col_by_key = {col.key: col for col in table.columns} + col_keys = [k for k in records[0] if k in col_by_key] + + # Python-side scalar defaults absent from records (same as other dialects). + extra_defaults: dict = {} + for col in table.columns: + if col.key not in records[0] and col.key in col_by_key: + d = col.default + if d is not None and d.is_scalar: + extra_defaults[col.key] = d.arg + + all_col_keys = col_keys + list(extra_defaults.keys()) + + fd, tsv_path = tempfile.mkstemp(suffix=".tsv") + _fallback = False + try: + with os.fdopen(fd, "w", encoding="utf-8", newline="") as f: + for record in records: + row = [] + for k in all_col_keys: + v = record.get(k) if k in col_keys else extra_defaults[k] + row.append(self._format_value(v, col_by_key[k])) + f.write("\t".join(row) + "\n") + + full_name = ( + f"`{table.schema}`.`{table.name}`" + if table.schema + else f"`{table.name}`" + ) + + # Binary columns use a user variable so UNHEX() can be applied in + # the SET clause; all other columns are addressed by name directly. + col_list_parts = [] + set_parts = [] + for k in all_col_keys: + col = col_by_key[k] + if isinstance(col.type, LargeBinary): + var = f"@__hex_{col.name}" + col_list_parts.append(var) + set_parts.append(f"`{col.name}` = UNHEX({var})") + else: + col_list_parts.append(f"`{col.name}`") + + col_clause = "(" + ", ".join(col_list_parts) + ")" + set_clause = (" SET " + ", ".join(set_parts)) if set_parts else "" + + sql = ( + f"LOAD DATA LOCAL INFILE '{tsv_path}' " + f"INTO TABLE {full_name} " + f"CHARACTER SET utf8mb4 " + f"FIELDS TERMINATED BY '\\t' ESCAPED BY '\\\\' " + f"LINES TERMINATED BY '\\n' " + f"{col_clause}{set_clause}" + ) + + try: + conn.execute(text(sql)) + self._local_infile_ok = True + except OperationalError as exc: + orig = getattr(exc, "orig", None) + code = orig.args[0] if orig and orig.args else None + if code in _LOCAL_INFILE_ERRORS: + self._local_infile_ok = False + if bulk_load is True: + raise RuntimeError( + "bulk_load=True requires LOAD DATA LOCAL INFILE, but it is " + "unavailable on this connection (MySQL error " + f"{code}). Enable local_infile on the MySQL server and ensure " + "the engine is created via DataModel (which sets local_infile=True " + "automatically) or with connect_args={'local_infile': True}." + ) from exc + _fallback = True + else: + raise + finally: + if os.path.exists(tsv_path): + os.unlink(tsv_path) + + if _fallback: + super().bulk_insert(conn, table, records) diff --git a/src/xml2db/dialect/postgresql.py b/src/xml2db/dialect/postgresql.py index d909e30..5244981 100644 --- a/src/xml2db/dialect/postgresql.py +++ b/src/xml2db/dialect/postgresql.py @@ -4,6 +4,10 @@ from .base import DatabaseDialect +# PostgreSQL COPY is in-protocol (no temp file), so the default threshold is 0 +# (COPY is always used for supported drivers regardless of batch size). +_COPY_THRESHOLD = 0 + class PostgreSQLDialect(DatabaseDialect): """Dialect for PostgreSQL. @@ -16,28 +20,53 @@ class PostgreSQLDialect(DatabaseDialect): MAX_IDENTIFIER_LENGTH: int = 63 - def bulk_insert(self, conn: Any, table: Any, records: list) -> None: + def bulk_insert( + self, + conn: Any, + table: Any, + records: list, + *, + bulk_load: bool | None = None, + bulk_load_threshold: int | None = None, + ) -> None: """Bulk-insert records via PostgreSQL's ``COPY FROM STDIN``. Builds an in-memory CSV payload and streams it to the server using the driver's native COPY protocol. Supported drivers: - - **psycopg2** — uses ``cursor.copy_expert()``. - - **psycopg** (psycopg3) — uses ``cursor.copy()``. + - **psycopg2**: uses ``cursor.copy_expert()``. + - **psycopg** (psycopg3): uses ``cursor.copy()``. Falls back to the base-class parameterised executemany for any other - driver. + driver (or when ``bulk_load=False``). Args: conn: A SQLAlchemy ``Connection`` already within a transaction. table: The SQLAlchemy ``Table`` object to insert into. records: A list of dicts mapping column keys to Python values. + bulk_load: ``True`` to require COPY (raise if driver unsupported), + ``False`` to always use executemany, or ``None`` (default) to + use COPY when available and fall back silently. + bulk_load_threshold: Minimum number of records to trigger COPY. + Defaults to :data:`_COPY_THRESHOLD` (0, meaning COPY is always + used for supported drivers). """ if not records: return + threshold = bulk_load_threshold if bulk_load_threshold is not None else _COPY_THRESHOLD driver = conn.dialect.driver + if driver not in ("psycopg2", "psycopg"): + if bulk_load is True: + raise RuntimeError( + f"bulk_load=True requires the psycopg2 or psycopg driver, got '{driver}'. " + f"Use a postgresql+psycopg2:// or postgresql+psycopg:// connection string." + ) + super().bulk_insert(conn, table, records) + return + + if bulk_load is False or len(records) < threshold: super().bulk_insert(conn, table, records) return diff --git a/src/xml2db/document.py b/src/xml2db/document.py index 237bf9c..6aa08da 100644 --- a/src/xml2db/document.py +++ b/src/xml2db/document.py @@ -368,13 +368,23 @@ def _build_node(node_type: str, node_pk: int) -> tuple: int(list(data_index[self.model.root_table]["records"].keys())[0]), ) - def insert_into_temp_tables(self, max_lines: int = -1) -> None: + def insert_into_temp_tables( + self, + max_lines: int = -1, + bulk_load: bool | None = None, + bulk_load_threshold: int | None = None, + ) -> None: """Insert data into temporary tables (Re)creates temp tables before inserting data. Args: max_lines: The maximum number of lines to insert in a single statement + bulk_load: ``True`` to require bulk loading (raise if unavailable), + ``False`` to always use executemany, or ``None`` (default) to + use bulk loading when available and fall back silently. + bulk_load_threshold: Minimum number of records to trigger bulk + loading. ``None`` delegates the choice to the dialect. """ logger.info(f"Dropping temp tables if exist for {self.xml_file_path}") self.model.drop_all_temp_tables() @@ -396,6 +406,8 @@ def insert_into_temp_tables(self, max_lines: int = -1) -> None: conn, query.table, data[start_idx : (start_idx + batch_size)], + bulk_load=bulk_load, + bulk_load_threshold=bulk_load_threshold, ) start_idx = start_idx + batch_size @@ -434,6 +446,8 @@ def insert_into_target_tables( self, single_transaction: bool = True, max_lines: int = -1, + bulk_load: bool | None = None, + bulk_load_threshold: int | None = None, ) -> int: """Insert and merge data into the database @@ -444,6 +458,11 @@ def insert_into_target_tables( scope required to ensure database consistency? max_lines: The maximum number of lines to insert in a single statement when loading data to the temporary tables + bulk_load: ``True`` to require bulk loading (raise if unavailable), + ``False`` to always use executemany, or ``None`` (default) to + use bulk loading when available and fall back silently. + bulk_load_threshold: Minimum number of records to trigger bulk + loading. ``None`` delegates the choice to the dialect. Returns: The number of inserted rows @@ -457,7 +476,7 @@ def insert_into_target_tables( logger.error(e) raise try: - self.insert_into_temp_tables(max_lines) + self.insert_into_temp_tables(max_lines, bulk_load, bulk_load_threshold) except Exception as e: logger.error( f"Error while importing into temporary tables from {self.xml_file_path}" diff --git a/src/xml2db/model.py b/src/xml2db/model.py index 0e9f6f3..89bc1ab 100644 --- a/src/xml2db/model.py +++ b/src/xml2db/model.py @@ -9,7 +9,8 @@ import xmlschema import sqlalchemy from lxml import etree -from sqlalchemy import MetaData, create_engine +from sqlalchemy import MetaData +from sqlalchemy.engine import make_url from sqlalchemy.sql.ddl import CreateIndex, CreateTable from graphlib import TopologicalSorter @@ -80,7 +81,6 @@ def __init__( db_type: str = None, db_schema: str = None, temp_prefix: str = None, - use_bcp: bool = True, ): self.model_config = self._validate_config(model_config) self.tables_config = model_config.get("tables", {}) if model_config else {} @@ -103,23 +103,15 @@ def __init__( ) self.engine = None self.db_type = db_type - else: - if db_engine: - self.engine = db_engine - else: - engine_options = {} - if "mssql" in connection_string: - engine_options = { - "fast_executemany": True, - "isolation_level": "SERIALIZABLE", - } - self.engine = create_engine( - connection_string, - **engine_options, - ) + self.dialect = get_dialect(self.db_type) + elif db_engine: + self.engine = db_engine self.db_type = self.engine.dialect.name - - self.dialect = get_dialect(self.db_type, use_bcp=use_bcp) + self.dialect = get_dialect(self.db_type) + else: + self.db_type = make_url(connection_string).drivername.split("+")[0] + self.dialect = get_dialect(self.db_type) + self.engine = self.dialect.create_engine(connection_string) self.model_config = self.dialect.validate_model_config(self.model_config) self.db_schema = db_schema self.temp_prefix = str(uuid4())[:8] if temp_prefix is None else temp_prefix diff --git a/tests/test_bulk_insert_mssql.py b/tests/test_bulk_insert_mssql.py index 4c85bdb..c595eff 100644 --- a/tests/test_bulk_insert_mssql.py +++ b/tests/test_bulk_insert_mssql.py @@ -51,10 +51,10 @@ def _make_table(engine, name, *extra_cols): return table, meta -def _insert_and_read(engine, table, records, use_bcp=False): - dialect = MSSQLDialect(use_bcp=use_bcp) +def _insert_and_read(engine, table, records, **bulk_insert_kwargs): + dialect = MSSQLDialect() with engine.begin() as conn: - dialect.bulk_insert(conn, table, records) + dialect.bulk_insert(conn, table, records, **bulk_insert_kwargs) with engine.connect() as conn: return conn.execute(select(table).order_by(table.c.id)).mappings().all() @@ -160,7 +160,7 @@ def test_mssql_bcp_basic(mssql_engine): try: records = [{"id": i, "label": f"row{i}"} for i in range(_BCP_THRESHOLD)] records.append({"id": _BCP_THRESHOLD, "label": None}) - rows = _insert_and_read(mssql_engine, table, records, use_bcp=True) + rows = _insert_and_read(mssql_engine, table, records) assert len(rows) == _BCP_THRESHOLD + 1 assert rows[0]["label"] == "row0" assert rows[-1]["label"] is None @@ -182,7 +182,7 @@ def test_mssql_bcp_numeric(mssql_engine): meta.create_all(mssql_engine) try: base = [{"i": j, "bi": 10**15 + j, "si": j % 32767} for j in range(_BCP_THRESHOLD)] - dialect = MSSQLDialect(use_bcp=True) + dialect = MSSQLDialect() with mssql_engine.begin() as conn: dialect.bulk_insert(conn, table, base) with mssql_engine.connect() as conn: @@ -208,7 +208,7 @@ def test_mssql_bcp_boolean(mssql_engine): [{"id": 0, "flag": True}, {"id": 1, "flag": False}, {"id": 2, "flag": None}] + [{"id": i + 3, "flag": i % 2 == 0} for i in range(_BCP_THRESHOLD)] ) - dialect = MSSQLDialect(use_bcp=True) + dialect = MSSQLDialect() with mssql_engine.begin() as conn: dialect.bulk_insert(conn, table, records) with mssql_engine.connect() as conn: @@ -237,7 +237,7 @@ def test_mssql_bcp_binary(mssql_engine): [{"id": 0, "data": payload}, {"id": 1, "data": None}] + [{"id": i + 2, "data": payload} for i in range(_BCP_THRESHOLD)] ) - dialect = MSSQLDialect(use_bcp=True) + dialect = MSSQLDialect() with mssql_engine.begin() as conn: dialect.bulk_insert(conn, table, records) with mssql_engine.connect() as conn: @@ -261,7 +261,7 @@ def test_mssql_bcp_scalar_default(mssql_engine): meta.create_all(mssql_engine) try: records = [{"id": i} for i in range(_BCP_THRESHOLD)] - dialect = MSSQLDialect(use_bcp=True) + dialect = MSSQLDialect() with mssql_engine.begin() as conn: dialect.bulk_insert(conn, table, records) with mssql_engine.connect() as conn: @@ -272,22 +272,53 @@ def test_mssql_bcp_scalar_default(mssql_engine): # --------------------------------------------------------------------------- -# BCP fallback when bcp binary is not found +# bulk_load=False forces fast_executemany even for large batches # --------------------------------------------------------------------------- @pytest.mark.dbtest -def test_mssql_bcp_fallback_when_unavailable(mssql_engine): - """With bcp_path=None, bulk_insert must still work via fast_executemany.""" - table, meta = _make_table(mssql_engine, "mssql_bi_fallback") +def test_mssql_bulk_load_false_bypasses_bcp(mssql_engine): + """bulk_load=False must use fast_executemany regardless of batch size.""" + table, meta = _make_table(mssql_engine, "mssql_bi_bulk_load_false") try: records = [{"id": i, "label": f"r{i}"} for i in range(_BCP_THRESHOLD)] - dialect = MSSQLDialect(use_bcp=True) - dialect.bcp_path = None # simulate bcp not on PATH - with mssql_engine.begin() as conn: - dialect.bulk_insert(conn, table, records) - with mssql_engine.connect() as conn: - count = len(conn.execute(select(table)).fetchall()) - assert count == _BCP_THRESHOLD + rows = _insert_and_read(mssql_engine, table, records, bulk_load=False) + assert len(rows) == _BCP_THRESHOLD finally: _drop(meta, mssql_engine) + + +# --------------------------------------------------------------------------- +# bulk_load=True raises when BCP is unavailable +# --------------------------------------------------------------------------- + + +@pytest.mark.dbtest +def test_mssql_bulk_load_true_raises_without_bcp(mssql_engine): + """bulk_load=True must raise RuntimeError when bcp is not on PATH.""" + import shutil + import unittest.mock as mock + + if shutil.which("bcp") is None: + # BCP already absent — just confirm the error is raised directly. + table, meta = _make_table(mssql_engine, "mssql_bi_bulk_load_true_no_bcp") + try: + records = [{"id": i, "label": f"r{i}"} for i in range(_BCP_THRESHOLD)] + dialect = MSSQLDialect() + with mssql_engine.begin() as conn: + with pytest.raises(RuntimeError, match="bcp utility"): + dialect.bulk_insert(conn, table, records, bulk_load=True) + finally: + _drop(meta, mssql_engine) + else: + # BCP present — patch shutil.which to simulate absence. + table, meta = _make_table(mssql_engine, "mssql_bi_bulk_load_true_no_bcp") + try: + records = [{"id": i, "label": f"r{i}"} for i in range(_BCP_THRESHOLD)] + dialect = MSSQLDialect() + with mock.patch("xml2db.dialect.mssql.shutil.which", return_value=None): + with mssql_engine.begin() as conn: + with pytest.raises(RuntimeError, match="bcp utility"): + dialect.bulk_insert(conn, table, records, bulk_load=True) + finally: + _drop(meta, mssql_engine) diff --git a/tests/test_bulk_insert_mysql.py b/tests/test_bulk_insert_mysql.py new file mode 100644 index 0000000..08eb090 --- /dev/null +++ b/tests/test_bulk_insert_mysql.py @@ -0,0 +1,310 @@ +"""Integration tests for MySQLDialect.bulk_insert (LOAD DATA LOCAL INFILE).""" +import datetime +import os + +import pytest +from sqlalchemy import ( + BigInteger, + Boolean, + Column, + DateTime, + Double, + Integer, + LargeBinary, + MetaData, + SmallInteger, + String, + Table, + create_engine, + select, +) +from sqlalchemy.engine import make_url + +from xml2db.dialect.mysql import MySQLDialect + + +def _make_mysql_engine(): + db_string = os.getenv("DB_STRING", "") + if not db_string: + pytest.skip("DB_STRING not set") + url = make_url(db_string) + if url.get_dialect().name not in ("mysql", "mariadb"): + pytest.skip("DB_STRING is not a MySQL/MariaDB connection") + driver = url.get_dialect().driver + if driver == "pymysql": + pytest.importorskip("pymysql") + elif driver == "mysqldb": + pytest.importorskip("MySQLdb") + return create_engine(url, connect_args={"local_infile": True}) + + +@pytest.fixture(scope="module") +def mysql_engine(): + engine = _make_mysql_engine() + yield engine + engine.dispose() + + +def _make_table(engine, name, *extra_cols): + meta = MetaData() + table = Table( + name, + meta, + Column("id", Integer, key="id"), + Column("label", String(200), key="label"), + *extra_cols, + ) + meta.create_all(engine) + return table, meta + + +def _roundtrip(engine, table, records, **bulk_insert_kwargs): + dialect = MySQLDialect() + with engine.begin() as conn: + dialect.bulk_insert(conn, table, records, **bulk_insert_kwargs) + with engine.connect() as conn: + q = select(table) + if "id" in table.c: + q = q.order_by(table.c.id) + return conn.execute(q).mappings().all() + + +def _drop(meta, engine): + meta.drop_all(engine) + + +# --------------------------------------------------------------------------- +# Basic types +# --------------------------------------------------------------------------- + + +@pytest.mark.dbtest +def test_mysql_bulk_insert_basic(mysql_engine): + table, meta = _make_table(mysql_engine, "mysql_bi_basic") + try: + records = [{"id": 1, "label": "hello"}, {"id": 2, "label": None}] + rows = _roundtrip(mysql_engine, table, records) + assert len(rows) == 2 + assert rows[0]["label"] == "hello" + assert rows[1]["label"] is None + finally: + _drop(meta, mysql_engine) + + +@pytest.mark.dbtest +def test_mysql_bulk_insert_empty(mysql_engine): + table, meta = _make_table(mysql_engine, "mysql_bi_empty") + try: + dialect = MySQLDialect() + with mysql_engine.begin() as conn: + dialect.bulk_insert(conn, table, []) + with mysql_engine.connect() as conn: + count = len(conn.execute(select(table)).fetchall()) + assert count == 0 + finally: + _drop(meta, mysql_engine) + + +@pytest.mark.dbtest +def test_mysql_bulk_insert_string_special_chars(mysql_engine): + """Strings with backslashes, tabs, and newlines survive the round-trip.""" + table, meta = _make_table(mysql_engine, "mysql_bi_special") + try: + records = [ + {"id": 1, "label": "back\\slash"}, + {"id": 2, "label": "tab\there"}, + {"id": 3, "label": "new\nline"}, + {"id": 4, "label": "null-like \\N text"}, + ] + rows = _roundtrip(mysql_engine, table, records) + assert rows[0]["label"] == "back\\slash" + assert rows[1]["label"] == "tab\there" + assert rows[2]["label"] == "new\nline" + assert rows[3]["label"] == "null-like \\N text" + finally: + _drop(meta, mysql_engine) + + +# --------------------------------------------------------------------------- +# Numeric types +# --------------------------------------------------------------------------- + + +@pytest.mark.dbtest +def test_mysql_bulk_insert_numeric_types(mysql_engine): + meta = MetaData() + table = Table( + "mysql_bi_numeric", + meta, + Column("i", Integer, key="i"), + Column("bi", BigInteger, key="bi"), + Column("si", SmallInteger, key="si"), + Column("d", Double, key="d"), + ) + meta.create_all(mysql_engine) + try: + records = [{"i": 1, "bi": 10**15, "si": 32767, "d": 3.14}] + rows = _roundtrip(mysql_engine, table, records) + assert rows[0]["i"] == 1 + assert rows[0]["bi"] == 10**15 + assert rows[0]["si"] == 32767 + assert abs(rows[0]["d"] - 3.14) < 1e-9 + finally: + meta.drop_all(mysql_engine) + + +# --------------------------------------------------------------------------- +# Boolean +# --------------------------------------------------------------------------- + + +@pytest.mark.dbtest +def test_mysql_bulk_insert_boolean(mysql_engine): + meta = MetaData() + table = Table( + "mysql_bi_bool", + meta, + Column("id", Integer, key="id"), + Column("flag", Boolean, key="flag"), + ) + meta.create_all(mysql_engine) + try: + records = [ + {"id": 1, "flag": True}, + {"id": 2, "flag": False}, + {"id": 3, "flag": None}, + ] + rows = _roundtrip(mysql_engine, table, records) + assert rows[0]["flag"] is True + assert rows[1]["flag"] is False + assert rows[2]["flag"] is None + finally: + meta.drop_all(mysql_engine) + + +# --------------------------------------------------------------------------- +# DateTime +# --------------------------------------------------------------------------- + + +@pytest.mark.dbtest +def test_mysql_bulk_insert_datetime(mysql_engine): + meta = MetaData() + table = Table( + "mysql_bi_dt", + meta, + Column("id", Integer, key="id"), + Column("ts", DateTime, key="ts"), + ) + meta.create_all(mysql_engine) + try: + dt = datetime.datetime(2023, 9, 27, 14, 35, 54) + records = [{"id": 1, "ts": dt}, {"id": 2, "ts": None}] + rows = _roundtrip(mysql_engine, table, records) + assert rows[0]["ts"] is not None + assert rows[1]["ts"] is None + finally: + meta.drop_all(mysql_engine) + + +# --------------------------------------------------------------------------- +# Binary (BLOB) +# --------------------------------------------------------------------------- + + +@pytest.mark.dbtest +def test_mysql_bulk_insert_binary(mysql_engine): + meta = MetaData() + table = Table( + "mysql_bi_binary", + meta, + Column("id", Integer, key="id"), + Column("data", LargeBinary, key="data"), + ) + meta.create_all(mysql_engine) + try: + payload = b"\xde\xad\xbe\xef" * 8 + records = [{"id": 1, "data": payload}, {"id": 2, "data": None}] + rows = _roundtrip(mysql_engine, table, records) + assert bytes(rows[0]["data"]) == payload + assert rows[1]["data"] is None + finally: + meta.drop_all(mysql_engine) + + +# --------------------------------------------------------------------------- +# Python-side scalar column defaults +# --------------------------------------------------------------------------- + + +@pytest.mark.dbtest +def test_mysql_bulk_insert_scalar_default(mysql_engine): + meta = MetaData() + table = Table( + "mysql_bi_default", + meta, + Column("id", Integer, key="id"), + Column("flag", Boolean, default=False, key="flag"), + ) + meta.create_all(mysql_engine) + try: + records = [{"id": 1}, {"id": 2}] + rows = _roundtrip(mysql_engine, table, records) + assert rows[0]["flag"] is False + assert rows[1]["flag"] is False + finally: + meta.drop_all(mysql_engine) + + +# --------------------------------------------------------------------------- +# bulk_load flag +# --------------------------------------------------------------------------- + + +@pytest.mark.dbtest +def test_mysql_bulk_load_true(mysql_engine): + """bulk_load=True succeeds when local_infile is enabled on the engine.""" + from xml2db.dialect.mysql import _LOAD_DATA_THRESHOLD + + table, meta = _make_table(mysql_engine, "mysql_bi_bulk_load_true") + try: + records = [{"id": i, "label": f"row{i}"} for i in range(_LOAD_DATA_THRESHOLD)] + rows = _roundtrip(mysql_engine, table, records, bulk_load=True) + assert len(rows) == _LOAD_DATA_THRESHOLD + assert rows[0]["label"] == "row0" + finally: + _drop(meta, mysql_engine) + + +@pytest.mark.dbtest +def test_mysql_bulk_load_false(mysql_engine): + """bulk_load=False uses executemany even for large batches.""" + from xml2db.dialect.mysql import _LOAD_DATA_THRESHOLD + + table, meta = _make_table(mysql_engine, "mysql_bi_bulk_load_false") + try: + records = [{"id": i, "label": f"row{i}"} for i in range(_LOAD_DATA_THRESHOLD)] + rows = _roundtrip(mysql_engine, table, records, bulk_load=False) + assert len(rows) == _LOAD_DATA_THRESHOLD + finally: + _drop(meta, mysql_engine) + + +@pytest.mark.dbtest +def test_mysql_bulk_load_true_raises_without_local_infile(mysql_engine): + """bulk_load=True raises RuntimeError when local_infile is not enabled.""" + from xml2db.dialect.mysql import _LOAD_DATA_THRESHOLD + + url = mysql_engine.url + # Create an engine without local_infile so LOAD DATA LOCAL INFILE is rejected. + engine_no_infile = create_engine(url) + table, meta = _make_table(engine_no_infile, "mysql_bi_bulk_load_no_infile") + try: + records = [{"id": i, "label": f"row{i}"} for i in range(_LOAD_DATA_THRESHOLD)] + dialect = MySQLDialect() + with engine_no_infile.begin() as conn: + with pytest.raises(RuntimeError, match="local_infile"): + dialect.bulk_insert(conn, table, records, bulk_load=True) + finally: + _drop(meta, engine_no_infile) + engine_no_infile.dispose()