diff --git a/.github/workflows/integration-tests-mssql.yml b/.github/workflows/integration-tests-mssql.yml index 6780a82..1546179 100644 --- a/.github/workflows/integration-tests-mssql.yml +++ b/.github/workflows/integration-tests-mssql.yml @@ -35,7 +35,7 @@ jobs: sudo dpkg -i packages-microsoft-prod.deb rm packages-microsoft-prod.deb sudo apt-get update - sudo ACCEPT_EULA=Y apt-get install -y msodbcsql18 + sudo ACCEPT_EULA=Y apt-get install -y msodbcsql18 mssql-tools18 python -m pip install --upgrade pip python -m pip install flake8 .[tests] pyodbc @@ -44,3 +44,4 @@ jobs: pytest tests -x env: DB_STRING: mssql+pyodbc://sa:MyTestPassword1@localhost:1433/master?driver=ODBC+Driver+18+for+SQL+Server&TrustServerCertificate=yes + PATH: /opt/mssql-tools18/bin:${{ env.PATH }} diff --git a/src/xml2db/dialect/__init__.py b/src/xml2db/dialect/__init__.py index c783f3a..b881586 100644 --- a/src/xml2db/dialect/__init__.py +++ b/src/xml2db/dialect/__init__.py @@ -49,7 +49,7 @@ } -def get_dialect(db_type: str | None) -> DatabaseDialect: +def get_dialect(db_type: str | None, **kwargs) -> DatabaseDialect: """Return a :class:`DatabaseDialect` instance for the given backend name. Args: @@ -57,9 +57,12 @@ def get_dialect(db_type: str | None) -> DatabaseDialect: ``"mssql"``, ``"mysql"``, ``"duckdb"``. ``None`` or any unrecognised string falls back to the base :class:`DatabaseDialect`, which uses safe generic defaults. + **kwargs: Extra keyword arguments forwarded to the dialect constructor. + Unknown kwargs are silently ignored by subclasses that do not + declare them. Returns: An instantiated :class:`DatabaseDialect` (or subclass) ready for use. """ cls = DIALECT_REGISTRY.get(db_type, DatabaseDialect) - return cls() + return cls(**kwargs) diff --git a/src/xml2db/dialect/base.py b/src/xml2db/dialect/base.py index 3ff9e42..b9c566d 100644 --- a/src/xml2db/dialect/base.py +++ b/src/xml2db/dialect/base.py @@ -39,6 +39,9 @@ class DatabaseDialect: MAX_IDENTIFIER_LENGTH: int = 63 # conservative default; matches PostgreSQL + def __init__(self, **kwargs): + pass + # ------------------------------------------------------------------ # Identifier handling # ------------------------------------------------------------------ diff --git a/src/xml2db/dialect/mssql.py b/src/xml2db/dialect/mssql.py index d5657c2..936d481 100644 --- a/src/xml2db/dialect/mssql.py +++ b/src/xml2db/dialect/mssql.py @@ -1,3 +1,7 @@ +import os +import shutil +import subprocess +import tempfile from typing import Any, List, TYPE_CHECKING from sqlalchemy import Index @@ -9,16 +13,31 @@ from ..table.column import DataModelColumn +# Records below this count go through fast_executemany; at or above, BCP is used. +_BCP_THRESHOLD = 1000 + + class MSSQLDialect(DatabaseDialect): """Dialect for Microsoft SQL Server. MSSQL supports identifiers up to 128 characters, so no truncation is needed. Columnstore index support and MSSQL-specific type mappings are 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. """ 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 @@ -73,3 +92,103 @@ def relation_extra_indexes( mssql_clustered=True, ), ) + + @staticmethod + def _format_bcp_value(v: Any) -> str: + """Format a Python value as a BCP character-mode field (tab-separated).""" + if v is None: + return "" + if isinstance(v, bool): + # Must precede int: bool is a subclass of int. + return "1" if v else "0" + if isinstance(v, (bytes, bytearray)): + # BCP character mode accepts hex without 0x prefix for binary columns. + return v.hex() + s = str(v) + # 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: + """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``. + + 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. + """ + if not records: + return + + 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) + ): + 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()) + + full_name = ( + f"[{table.schema}].[{table.name}]" + if table.schema + else f"[{table.name}]" + ) + + fd, data_path = tempfile.mkstemp(suffix=".bcp") + try: + with os.fdopen(fd, "w", encoding="utf-8", newline="") as f: + for record in records: + row = [] + for key in all_col_keys: + v = record.get(key) if key in col_keys else extra_defaults[key] + row.append(self._format_bcp_value(v)) + f.write("\t".join(row) + "\n") + + cmd = [ + self.bcp_path, full_name, "in", data_path, + "-S", f"{url.host},{url.port or 1433}", + "-d", url.database, + "-c", # character (text) mode + "-t", "\t", # tab field separator + "-r", "\n", # newline row terminator + "-k", # empty field → NULL (not column default) + "-b", "10000", + "-C", "65001", # UTF-8 + ] + if has_sql_auth: + cmd += ["-U", url.username, "-P", url.password] + else: + cmd.append("-T") # Kerberos / Windows trusted connection + if str(url.query.get("TrustServerCertificate", "")).lower() == "yes": + cmd.append("-u") # trust server certificate (mssql-tools18) + result = subprocess.run(cmd, capture_output=True, text=True) + if result.returncode != 0: + raise RuntimeError( + f"BCP failed (exit {result.returncode}):\n" + f"{result.stdout}\n{result.stderr}" + ) + finally: + if os.path.exists(data_path): + os.unlink(data_path) diff --git a/src/xml2db/model.py b/src/xml2db/model.py index de0edd7..0e9f6f3 100644 --- a/src/xml2db/model.py +++ b/src/xml2db/model.py @@ -80,6 +80,7 @@ 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 {} @@ -118,7 +119,7 @@ def __init__( ) self.db_type = self.engine.dialect.name - self.dialect = get_dialect(self.db_type) + self.dialect = get_dialect(self.db_type, use_bcp=use_bcp) 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 new file mode 100644 index 0000000..4c85bdb --- /dev/null +++ b/tests/test_bulk_insert_mssql.py @@ -0,0 +1,293 @@ +"""Integration tests for MSSQLDialect.bulk_insert (fast_executemany and BCP).""" +import os + +import pytest +from sqlalchemy import ( + BigInteger, + Boolean, + Column, + Integer, + LargeBinary, + MetaData, + SmallInteger, + String, + Table, + create_engine, + select, +) +from sqlalchemy.engine import make_url + +from xml2db.dialect.mssql import MSSQLDialect, _BCP_THRESHOLD + + +def _make_mssql_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 != "mssql": + pytest.skip("DB_STRING is not an MSSQL connection") + pytest.importorskip("pyodbc") + return create_engine(url, fast_executemany=True) + + +@pytest.fixture(scope="module") +def mssql_engine(): + engine = _make_mssql_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 _insert_and_read(engine, table, records, use_bcp=False): + dialect = MSSQLDialect(use_bcp=use_bcp) + with engine.begin() as conn: + dialect.bulk_insert(conn, table, records) + with engine.connect() as conn: + return conn.execute(select(table).order_by(table.c.id)).mappings().all() + + +def _drop(meta, engine): + meta.drop_all(engine) + + +# --------------------------------------------------------------------------- +# fast_executemany path (< _BCP_THRESHOLD rows) +# --------------------------------------------------------------------------- + + +@pytest.mark.dbtest +def test_mssql_fast_executemany_basic(mssql_engine): + table, meta = _make_table(mssql_engine, "mssql_bi_basic_fem") + try: + records = [{"id": 1, "label": "hello"}, {"id": 2, "label": None}] + rows = _insert_and_read(mssql_engine, table, records) + assert len(rows) == 2 + assert rows[0]["label"] == "hello" + assert rows[1]["label"] is None + finally: + _drop(meta, mssql_engine) + + +@pytest.mark.dbtest +def test_mssql_fast_executemany_numeric(mssql_engine): + meta = MetaData() + table = Table( + "mssql_bi_numeric_fem", + meta, + Column("id", Integer, key="id"), + Column("i", Integer, key="i"), + Column("bi", BigInteger, key="bi"), + Column("si", SmallInteger, key="si"), + ) + meta.create_all(mssql_engine) + try: + records = [{"id": 1, "i": 1, "bi": 10**15, "si": 32767}] + rows = _insert_and_read(mssql_engine, table, records) + assert rows[0]["i"] == 1 + assert rows[0]["bi"] == 10**15 + assert rows[0]["si"] == 32767 + finally: + meta.drop_all(mssql_engine) + + +@pytest.mark.dbtest +def test_mssql_fast_executemany_boolean(mssql_engine): + meta = MetaData() + table = Table( + "mssql_bi_bool_fem", + meta, + Column("id", Integer, key="id"), + Column("flag", Boolean, key="flag"), + ) + meta.create_all(mssql_engine) + try: + records = [{"id": 1, "flag": True}, {"id": 2, "flag": False}, {"id": 3, "flag": None}] + rows = _insert_and_read(mssql_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(mssql_engine) + + +@pytest.mark.dbtest +def test_mssql_fast_executemany_scalar_default(mssql_engine): + meta = MetaData() + table = Table( + "mssql_bi_default_fem", + meta, + Column("id", Integer, key="id"), + Column("flag", Boolean, default=False, key="flag"), + ) + meta.create_all(mssql_engine) + try: + records = [{"id": 1}, {"id": 2}] + rows = _insert_and_read(mssql_engine, table, records) + assert rows[0]["flag"] is False + assert rows[1]["flag"] is False + finally: + meta.drop_all(mssql_engine) + + +# --------------------------------------------------------------------------- +# BCP path (>= _BCP_THRESHOLD rows) +# --------------------------------------------------------------------------- + + +def _require_bcp(): + import shutil + if shutil.which("bcp") is None: + pytest.skip("bcp not found on PATH") + + +@pytest.mark.dbtest +def test_mssql_bcp_basic(mssql_engine): + _require_bcp() + table, meta = _make_table(mssql_engine, "mssql_bi_basic_bcp") + 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) + assert len(rows) == _BCP_THRESHOLD + 1 + assert rows[0]["label"] == "row0" + assert rows[-1]["label"] is None + finally: + _drop(meta, mssql_engine) + + +@pytest.mark.dbtest +def test_mssql_bcp_numeric(mssql_engine): + _require_bcp() + meta = MetaData() + table = Table( + "mssql_bi_numeric_bcp", + meta, + Column("i", Integer, key="i"), + Column("bi", BigInteger, key="bi"), + Column("si", SmallInteger, key="si"), + ) + 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) + with mssql_engine.begin() as conn: + dialect.bulk_insert(conn, table, base) + with mssql_engine.connect() as conn: + rows = conn.execute(select(table).order_by(table.c.i)).mappings().all() + assert rows[0]["bi"] == 10**15 + finally: + meta.drop_all(mssql_engine) + + +@pytest.mark.dbtest +def test_mssql_bcp_boolean(mssql_engine): + _require_bcp() + meta = MetaData() + table = Table( + "mssql_bi_bool_bcp", + meta, + Column("id", Integer, key="id"), + Column("flag", Boolean, key="flag"), + ) + meta.create_all(mssql_engine) + try: + records = ( + [{"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) + with mssql_engine.begin() as conn: + dialect.bulk_insert(conn, table, records) + with mssql_engine.connect() as conn: + rows = conn.execute(select(table).order_by(table.c.id)).mappings().all() + assert rows[0]["flag"] is True + assert rows[1]["flag"] is False + assert rows[2]["flag"] is None + finally: + meta.drop_all(mssql_engine) + + +@pytest.mark.dbtest +def test_mssql_bcp_binary(mssql_engine): + _require_bcp() + meta = MetaData() + table = Table( + "mssql_bi_binary_bcp", + meta, + Column("id", Integer, key="id"), + Column("data", LargeBinary, key="data"), + ) + meta.create_all(mssql_engine) + try: + payload = b"\xde\xad\xbe\xef" * 8 + records = ( + [{"id": 0, "data": payload}, {"id": 1, "data": None}] + + [{"id": i + 2, "data": payload} for i in range(_BCP_THRESHOLD)] + ) + dialect = MSSQLDialect(use_bcp=True) + with mssql_engine.begin() as conn: + dialect.bulk_insert(conn, table, records) + with mssql_engine.connect() as conn: + rows = conn.execute(select(table).order_by(table.c.id)).mappings().all() + assert bytes(rows[0]["data"]) == payload + assert rows[1]["data"] is None + finally: + meta.drop_all(mssql_engine) + + +@pytest.mark.dbtest +def test_mssql_bcp_scalar_default(mssql_engine): + _require_bcp() + meta = MetaData() + table = Table( + "mssql_bi_default_bcp", + meta, + Column("id", Integer, key="id"), + Column("flag", Boolean, default=False, key="flag"), + ) + meta.create_all(mssql_engine) + try: + records = [{"id": i} for i in range(_BCP_THRESHOLD)] + dialect = MSSQLDialect(use_bcp=True) + with mssql_engine.begin() as conn: + dialect.bulk_insert(conn, table, records) + with mssql_engine.connect() as conn: + rows = conn.execute(select(table).order_by(table.c.id)).mappings().all() + assert all(r["flag"] is False for r in rows) + finally: + meta.drop_all(mssql_engine) + + +# --------------------------------------------------------------------------- +# BCP fallback when bcp binary is not found +# --------------------------------------------------------------------------- + + +@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") + 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 + finally: + _drop(meta, mssql_engine)