From 5919288348b6391ea498355684c65310c2a73cde Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 8 Jul 2026 12:57:25 +0530 Subject: [PATCH] GH-570: Add Cursor.bulkcopy_arrow for native Arrow bulk copy Adds cursor.bulkcopy_arrow(table_name, source) to bulk-load directly from Apache Arrow data (pyarrow Table/RecordBatch/RecordBatchReader, __arrow_c_stream__/__arrow_c_array__ producers, or an iterable of batches) via the mssql_py_core Rust core, streaming through the Arrow C Data Interface with the GIL released during transfer. bulkcopy() now raises TypeError to steer Arrow inputs here. Auth setup is refactored into a shared _build_pycore_context() helper (preserves SQL, pre-acquired-token, and ServicePrincipal factory paths). Adds tests/test_024_bulkcopy_arrow.py (36 unit tests at 100% coverage of the new code + 16 live round-trip/type-matrix tests) and benchmarks/bench_bulkcopy_arrow.py. Requires mssql-py-core 0.1.5+. --- CHANGELOG.md | 12 + benchmarks/bench_bulkcopy_arrow.py | 222 ++++++++++ mssql_python/cursor.py | 366 ++++++++++++---- mssql_python/mssql_python.pyi | 33 +- tests/test_024_bulkcopy_arrow.py | 658 +++++++++++++++++++++++++++++ 5 files changed, 1212 insertions(+), 79 deletions(-) create mode 100644 benchmarks/bench_bulkcopy_arrow.py create mode 100644 tests/test_024_bulkcopy_arrow.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 61b6e571..bcd4c6e6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,18 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), ### Added - New feature: Support for macOS and Linux. - Documentation: Added API documentation in the Wiki. +- **GH-570:** New `Cursor.bulkcopy_arrow(table_name, source)` method for + high-performance bulk loading directly from Apache Arrow data. Accepts a + `pyarrow.Table`, `RecordBatch`, or `RecordBatchReader`, any object exposing + the Arrow C Data Interface (`__arrow_c_stream__` / `__arrow_c_array__` — e.g. + polars, pandas 2.2+, DuckDB, ADBC results), or an iterable of record batches. + Data is streamed to the server through the Arrow C Data Interface without + materializing intermediate Python row objects, and the GIL is released for + the duration of the network transfer. When the source data already originates + as Arrow, this avoids the Arrow→tuple conversion the classic `bulkcopy()` + path requires (measured ~1.4x–2.7x faster end-to-end for such sources). + `bulkcopy()` now raises `TypeError` steering Arrow inputs to this method. + Requires `mssql-py-core` 0.1.5+. - Bulk copy now supports `Authentication=ActiveDirectoryServicePrincipal` via an `entra_id_token_factory` callback registered on the mssql-py-core connection. The callback is invoked by mssql-tds mid-handshake (FedAuth diff --git a/benchmarks/bench_bulkcopy_arrow.py b/benchmarks/bench_bulkcopy_arrow.py new file mode 100644 index 00000000..a8328971 --- /dev/null +++ b/benchmarks/bench_bulkcopy_arrow.py @@ -0,0 +1,222 @@ +"""Benchmark: cursor.bulkcopy() (row tuples) vs cursor.bulkcopy_arrow() (Arrow). + +Feeds *identical* data to both code paths so the comparison is apples-to-apples: + - bulkcopy() receives a list of row tuples + - bulkcopy_arrow() receives the equivalent pyarrow.Table + +For each (profile, row_count) it recreates the target table, runs each method +`REPEATS` times (after a warmup), and reports the median wall-clock time and +throughput plus the arrow-vs-tuples speedup. + +Usage (PowerShell): + $env:DB_CONNECTION_STRING = "Server=...;DATABASE=...;UID=...;PWD=...;TrustServerCertificate=yes;" + python benchmarks\bench_bulkcopy_arrow.py + python benchmarks\bench_bulkcopy_arrow.py --rows 1000 10000 100000 --repeats 5 +""" + +import argparse +import datetime as dt +import os +import statistics +import sys +import time + +import pyarrow as pa + +import mssql_python + + +TABLE = "bulkcopy_arrow_bench" + + +# --------------------------------------------------------------------------- # +# Column profiles. Each column: (name, sql_type, arrow_type, value_fn) +# value_fn(i) -> python value for row i (also used to build the Arrow column). +# --------------------------------------------------------------------------- # +def _profiles(): + base_dt = dt.datetime(2024, 1, 1) + names = ["alpha", "bravo", "charlie", "delta", "echo", "foxtrot"] + + narrow = [ + ("id", "INT", pa.int32(), lambda i: i), + ("big", "BIGINT", pa.int64(), lambda i: i * 1_000_003), + ("score", "FLOAT", pa.float64(), lambda i: i * 1.5), + ] + + mixed = [ + ("id", "INT", pa.int32(), lambda i: i), + ("name", "NVARCHAR(50)", pa.string(), lambda i: names[i % len(names)]), + ("amount", "FLOAT", pa.float64(), lambda i: i * 3.14159), + ("flag", "BIT", pa.bool_(), lambda i: bool(i % 2)), + ("created", "DATETIME2", pa.timestamp("us"), lambda i: base_dt + dt.timedelta(seconds=i)), + ] + + wide = [(f"c{n}", "INT", pa.int32(), (lambda n: (lambda i: i + n))(n)) for n in range(20)] + + return {"narrow": narrow, "mixed": mixed, "wide": wide} + + +def _build_data(columns, row_count): + """Return (rows_as_tuples, arrow_table) carrying identical data.""" + per_col = [[fn(i) for i in range(row_count)] for (_, _, _, fn) in columns] + rows = list(zip(*per_col)) + schema = pa.schema([(name, atype) for (name, _, atype, _) in columns]) + table = pa.table({name: col for (name, _, _, _), col in zip(columns, per_col)}, schema=schema) + return rows, table + + +def _arrow_to_tuples(table): + """Materialize an Arrow table into row tuples the way a typical caller would. + + This is the tax the tuple path must pay when the data *originates* as Arrow + (Parquet / polars / DuckDB / pandas / ADBC): every cell becomes a boxed + Python object. bulkcopy_arrow() skips this entirely. + """ + return list(zip(*(col.to_pylist() for col in table.columns))) + + +def _recreate_table(cursor, columns): + cols_ddl = ", ".join(f"[{name}] {sql}" for (name, sql, _, _) in columns) + cursor.execute(f"IF OBJECT_ID('{TABLE}', 'U') IS NOT NULL DROP TABLE [{TABLE}];") + cursor.execute(f"CREATE TABLE [{TABLE}] ({cols_ddl});") + cursor.connection.commit() + + +def _truncate(cursor): + cursor.execute(f"TRUNCATE TABLE [{TABLE}];") + cursor.connection.commit() + + +def _time_call(fn): + start = time.perf_counter() + result = fn() + return time.perf_counter() - start, result + + +def _median_run(cursor, columns, rows, table, use_arrow, repeats): + times = [] + reported = None + for _ in range(repeats): + _truncate(cursor) + if use_arrow: + elapsed, res = _time_call(lambda: cursor.bulkcopy_arrow(TABLE, table)) + else: + elapsed, res = _time_call(lambda: cursor.bulkcopy(TABLE, rows)) + times.append(elapsed) + reported = res + return statistics.median(times), reported + + +def _median_run_from_arrow(cursor, table, repeats): + """Tuple path when the source is an Arrow table: time convert + bulkcopy. + + Returns (median_convert_s, median_insert_s, median_total_s, last_result). + """ + convert_times, insert_times, total_times = [], [], [] + reported = None + for _ in range(repeats): + _truncate(cursor) + t_conv, rows = _time_call(lambda: _arrow_to_tuples(table)) + t_ins, res = _time_call(lambda: cursor.bulkcopy(TABLE, rows)) + convert_times.append(t_conv) + insert_times.append(t_ins) + total_times.append(t_conv + t_ins) + reported = res + return ( + statistics.median(convert_times), + statistics.median(insert_times), + statistics.median(total_times), + reported, + ) + + +def run(row_counts, repeats): + conn_str = os.environ.get("DB_CONNECTION_STRING") + if not conn_str: + sys.exit("Set DB_CONNECTION_STRING to run this benchmark.") + + profiles = _profiles() + conn = mssql_python.connect(conn_str) + cursor = conn.cursor() + + # Collect measurements once; print two views afterwards. + rows_out = [] + try: + for pname, columns in profiles.items(): + for row_count in row_counts: + rows, table = _build_data(columns, row_count) + _recreate_table(cursor, columns) + + # Warmup each path once (JIT paths, connection buffers, etc.). + _truncate(cursor) + cursor.bulkcopy(TABLE, rows) + _truncate(cursor) + cursor.bulkcopy_arrow(TABLE, table) + + # Scenario A: insert-only (inputs already materialized for free). + t_tup, r_tup = _median_run(cursor, columns, rows, table, False, repeats) + t_arw, r_arw = _median_run(cursor, columns, rows, table, True, repeats) + # Scenario B: source originates as Arrow -> tuple path must convert. + t_conv, t_ins, t_total, r_fa = _median_run_from_arrow(cursor, table, repeats) + + assert r_tup["rows_copied"] == row_count, r_tup + assert r_arw["rows_copied"] == row_count, r_arw + assert r_fa["rows_copied"] == row_count, r_fa + + rows_out.append( + { + "profile": pname, + "cols": len(columns), + "rows": row_count, + "t_tup": t_tup, + "t_arw": t_arw, + "t_conv": t_conv, + "t_ins": t_ins, + "t_total": t_total, + } + ) + finally: + cursor.execute(f"IF OBJECT_ID('{TABLE}', 'U') IS NOT NULL DROP TABLE [{TABLE}];") + cursor.connection.commit() + cursor.close() + conn.close() + + # Scenario A: apples-to-apples insert, both inputs pre-materialized. + print("\n=== Scenario A: insert-only (data already in tuples / Arrow) ===") + hdr = f"{'profile':<8} {'cols':>4} {'rows':>9} {'tuples(s)':>11} {'arrow(s)':>10} {'tup r/s':>12} {'arw r/s':>12} {'speedup':>8}" + print(hdr) + print("-" * len(hdr)) + for r in rows_out: + rps_tup = r["rows"] / r["t_tup"] if r["t_tup"] else 0 + rps_arw = r["rows"] / r["t_arw"] if r["t_arw"] else 0 + speedup = r["t_tup"] / r["t_arw"] if r["t_arw"] else 0 + print( + f"{r['profile']:<8} {r['cols']:>4} {r['rows']:>9,} " + f"{r['t_tup']:>11.4f} {r['t_arw']:>10.4f} " + f"{rps_tup:>12,.0f} {rps_arw:>12,.0f} {speedup:>7.2f}x" + ) + + # Scenario B: source is Arrow. Tuple path pays convert + insert; arrow is direct. + print("\n=== Scenario B: source originates as Arrow (tuple path must convert) ===") + hdr2 = f"{'profile':<8} {'cols':>4} {'rows':>9} {'convert(s)':>11} {'+insert(s)':>11} {'total(s)':>10} {'arrow(s)':>10} {'speedup':>8}" + print(hdr2) + print("-" * len(hdr2)) + for r in rows_out: + speedup = r["t_total"] / r["t_arw"] if r["t_arw"] else 0 + print( + f"{r['profile']:<8} {r['cols']:>4} {r['rows']:>9,} " + f"{r['t_conv']:>11.4f} {r['t_ins']:>11.4f} {r['t_total']:>10.4f} " + f"{r['t_arw']:>10.4f} {speedup:>7.2f}x" + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rows", type=int, nargs="+", default=[1_000, 10_000, 100_000]) + parser.add_argument("--repeats", type=int, default=5) + args = parser.parse_args() + run(args.rows, args.repeats) + + +if __name__ == "__main__": + main() diff --git a/mssql_python/cursor.py b/mssql_python/cursor.py index 85701a40..d3966cfd 100644 --- a/mssql_python/cursor.py +++ b/mssql_python/cursor.py @@ -2815,6 +2815,117 @@ def nextset(self) -> Optional[bool]: return True # ── Mapping from ODBC connection-string keywords (lowercase, as _parse returns) + def _build_pycore_context(self) -> dict: + """Build the connection context dict expected by mssql_py_core. + + Parses the underlying ODBC connection string, validates the SERVER + parameter, and (when Azure AD auth is in use) either registers a + ServicePrincipal token factory or acquires a fresh access token, + replacing credential fields as py-core requires. Returns a dict that + can be handed to ``PyCoreConnection``. Shared by :meth:`bulkcopy` + and :meth:`bulkcopy_arrow` so both paths get identical auth handling. + + Raises: + RuntimeError: connection string unavailable or token acquisition failed. + ValueError: SERVER parameter missing. + """ + # Get and parse connection string + if not hasattr(self.connection, "connection_str"): + logger.error("_build_pycore_context: Connection string not available") + raise RuntimeError("Connection string not available for bulk copy") + + # Use the proper connection string parser that handles braced values + from mssql_python.connection_string_parser import _ConnectionStringParser + + parser = _ConnectionStringParser(validate_keywords=False) + params = parser._parse(self.connection.connection_str) + + # Check for server parameter (accepts synonyms: server, addr, address) + if not (params.get("server") or params.get("addr") or params.get("address")): + raise ValueError("SERVER parameter is required in connection string") + + # Translate parsed connection string into the dict py-core expects. + pycore_context = connstr_to_pycore_params(params) + + # Token acquisition — only thing cursor must handle (needs azure-identity SDK) + if self.connection._auth_type: + # Fresh token acquisition for mssql-py-core connection + from mssql_python.auth import AADAuth, ServicePrincipalAuth + from mssql_python.constants import _AuthInternal + + if self.connection._auth_type == _AuthInternal.SERVICE_PRINCIPAL: + # Callback-based path: tenant_id is only known from the STS URL + # the server returns mid-handshake, so we register a factory + # that py-core invokes during FedAuth (workflow 0x02). + # Cert-based ServicePrincipal is not supported on this path. + client_id = params.get("uid", "") + client_secret = params.get("pwd", "") + if not client_id or not client_secret: + raise RuntimeError( + "Bulk copy with Authentication=ActiveDirectoryServicePrincipal " + "currently supports client-secret only. " + "Provide UID (client_id) and PWD (client_secret) in the " + "connection string." + ) + try: + factory = ServicePrincipalAuth.make_token_factory(client_id, client_secret) + except (RuntimeError, ValueError) as e: + raise RuntimeError( + f"Bulk copy failed: unable to build ServicePrincipal token factory: {e}" + ) from e + pycore_context["entra_id_token_factory"] = factory + # Keep authentication/user_name/password in pycore_context — + # py-core's auth validator + transformer need them to resolve + # the auth method to ActiveDirectoryServicePrincipal before + # the factory is dispatched at handshake time. + logger.debug("Bulk copy: registered ServicePrincipal token factory") + else: + # Pre-acquired token path. Used for Default, DeviceCode, + # Interactive (non-Windows), MSI (system- or user-assigned), + # and any other AD method whose tenant_id is discoverable + # client-side via Azure Identity SDK. credential kwargs + # (e.g. user-assigned MSI client_id) were captured by + # Connection.__init__ before remove_sensitive_params stripped + # UID from connection_str. re-parsing here would miss them. + try: + raw_token = AADAuth.get_raw_token( + self.connection._auth_type, + self.connection._credential_kwargs, + ) + except (RuntimeError, ValueError) as e: + raise RuntimeError( + f"Bulk copy failed: unable to acquire Azure AD token " + f"for auth_type '{self.connection._auth_type}': {e}" + ) from e + pycore_context["access_token"] = raw_token + # Token replaces credential fields — py-core's validator rejects + # access_token combined with authentication/user_name/password. + for key in ("authentication", "user_name", "password"): + pycore_context.pop(key, None) + logger.debug( + "Bulk copy: acquired fresh Azure AD token for auth_type=%s", + self.connection._auth_type, + ) + + return pycore_context + + def _looks_like_arrow_source(self, data) -> bool: + """Return True if ``data`` should be routed to :meth:`bulkcopy_arrow`. + + Soft check: never imports pyarrow eagerly and never raises. Anything + exposing the Arrow C-stream PyCapsule protocol counts, as do the + concrete pyarrow container types when pyarrow is importable. + """ + if data is None: + return False + if hasattr(data, "__arrow_c_stream__") or hasattr(data, "__arrow_c_array__"): + return True + try: + import pyarrow as pa + except ImportError: + return False + return isinstance(data, (pa.Table, pa.RecordBatch, pa.RecordBatchReader)) + def bulkcopy( self, table_name: str, @@ -2889,13 +3000,23 @@ def bulkcopy( Raises: ImportError: If mssql_py_core library is not installed - TypeError: If data is None, not iterable, or is a string/bytes + TypeError: If data is None, not iterable, is a string/bytes, or is an + Arrow source (use :meth:`bulkcopy_arrow` for those) ValueError: If table_name is empty or parameters are invalid RuntimeError: If connection string is not available """ # Fast check if logging is enabled to avoid overhead is_logging_enabled = logger.is_debug_enabled + # Steer Arrow-shaped sources to the dedicated method instead of + # silently re-routing or failing deep inside the tuple validator. + if self._looks_like_arrow_source(data): + raise TypeError( + "bulkcopy() expects an iterable of row tuples/lists. " + "For pyarrow.Table / RecordBatch / RecordBatchReader / objects " + "implementing __arrow_c_stream__, call cursor.bulkcopy_arrow() instead." + ) + try: import mssql_py_core except ImportError as exc: @@ -2937,83 +3058,7 @@ def bulkcopy( if timeout <= 0: raise ValueError(f"timeout must be positive, got {timeout}") - # Get and parse connection string - if not hasattr(self.connection, "connection_str"): - logger.error("_bulkcopy: Connection string not available") - raise RuntimeError("Connection string not available for bulk copy") - - # Use the proper connection string parser that handles braced values - from mssql_python.connection_string_parser import _ConnectionStringParser - - parser = _ConnectionStringParser(validate_keywords=False) - params = parser._parse(self.connection.connection_str) - - # Check for server parameter (accepts synonyms: server, addr, address) - if not (params.get("server") or params.get("addr") or params.get("address")): - raise ValueError("SERVER parameter is required in connection string") - - # Translate parsed connection string into the dict py-core expects. - pycore_context = connstr_to_pycore_params(params) - - # Token acquisition — only thing cursor must handle (needs azure-identity SDK) - if self.connection._auth_type: - # Fresh token acquisition for mssql-py-core connection - from mssql_python.auth import AADAuth, ServicePrincipalAuth - from mssql_python.constants import _AuthInternal - - if self.connection._auth_type == _AuthInternal.SERVICE_PRINCIPAL: - # Callback-based path: tenant_id is only known from the STS URL - # the server returns mid-handshake, so we register a factory - # that py-core invokes during FedAuth (workflow 0x02). - # Cert-based ServicePrincipal is not supported on this path. - client_id = params.get("uid", "") - client_secret = params.get("pwd", "") - if not client_id or not client_secret: - raise RuntimeError( - "Bulk copy with Authentication=ActiveDirectoryServicePrincipal " - "currently supports client-secret only. " - "Provide UID (client_id) and PWD (client_secret) in the " - "connection string." - ) - try: - factory = ServicePrincipalAuth.make_token_factory(client_id, client_secret) - except (RuntimeError, ValueError) as e: - raise RuntimeError( - f"Bulk copy failed: unable to build ServicePrincipal token factory: {e}" - ) from e - pycore_context["entra_id_token_factory"] = factory - # Keep authentication/user_name/password in pycore_context — - # py-core's auth validator + transformer need them to resolve - # the auth method to ActiveDirectoryServicePrincipal before - # the factory is dispatched at handshake time. - logger.debug("Bulk copy: registered ServicePrincipal token factory") - else: - # Pre-acquired token path. Used for Default, DeviceCode, - # Interactive (non-Windows), MSI (system- or user-assigned), - # and any other AD method whose tenant_id is discoverable - # client-side via Azure Identity SDK. credential kwargs - # (e.g. user-assigned MSI client_id) were captured by - # Connection.__init__ before remove_sensitive_params stripped - # UID from connection_str. re-parsing here would miss them. - try: - raw_token = AADAuth.get_raw_token( - self.connection._auth_type, - self.connection._credential_kwargs, - ) - except (RuntimeError, ValueError) as e: - raise RuntimeError( - f"Bulk copy failed: unable to acquire Azure AD token " - f"for auth_type '{self.connection._auth_type}': {e}" - ) from e - pycore_context["access_token"] = raw_token - # Token replaces credential fields — py-core's validator rejects - # access_token combined with authentication/user_name/password. - for key in ("authentication", "user_name", "password"): - pycore_context.pop(key, None) - logger.debug( - "Bulk copy: acquired fresh Azure AD token for auth_type=%s", - self.connection._auth_type, - ) + pycore_context = self._build_pycore_context() pycore_connection = None pycore_cursor = None @@ -3112,6 +3157,171 @@ def _prepare_row_iterator(iterable): cleanup_error, ) + def bulkcopy_arrow( + self, + table_name: str, + source, + *, + batch_size: int = 0, + timeout: int = 30, + column_mappings: Optional[Union[List[str], List[Tuple[int, str]]]] = None, + keep_identity: bool = False, + check_constraints: bool = False, + table_lock: bool = False, + keep_nulls: bool = False, + fire_triggers: bool = False, + use_internal_transaction: bool = False, + ): + """Bulk-copy from an Apache Arrow source straight into TDS. + + ``source`` may be any of: + + * ``pyarrow.Table`` + * ``pyarrow.RecordBatch`` + * ``pyarrow.RecordBatchReader`` + * any object exposing ``__arrow_c_stream__`` (Arrow PyCapsule stream + interface — e.g. polars/pandas DataFrames ≥ 2.2, duckdb results) + * any object exposing ``__arrow_c_array__`` (single-batch PyCapsule + array interface) + * an iterable of ``pyarrow.RecordBatch`` (all batches must share the + same schema) + + Compared to :meth:`bulkcopy`, this path skips per-cell Python + round-trips: each batch's typed Arrow buffers are read directly by + the Rust core and streamed into the TDS bulk-load packets. Schema and + column-mapping semantics, the options, and the returned dictionary are + identical to :meth:`bulkcopy`. + + Args: + table_name: Target table name (may include schema, e.g. 'dbo.MyTable'). + source: Arrow source (see above). + batch_size: Rows per TDS commit. Default 0 uses server optimal. + timeout: Operation timeout in seconds. Default 30. + column_mappings: Same two formats as :meth:`bulkcopy`. When omitted, + Arrow fields map to destination columns by ordinal position. + keep_identity: Preserve identity values from the source. + check_constraints: Check constraints during bulk copy. + table_lock: Use a table-level lock instead of row-level locks. + keep_nulls: Preserve null values instead of using column defaults. + fire_triggers: Fire insert triggers on the target table. + use_internal_transaction: Use an internal transaction per batch. + + Returns: + Dictionary with ``rows_copied``, ``batch_count``, ``elapsed_time`` + and ``rows_per_second``. + + Raises: + ImportError: If the mssql_py_core library is not installed. + TypeError: If ``source`` is None. + ValueError: If ``table_name`` is empty or parameters are invalid. + RuntimeError: If the connection string is not available. + + Example: + >>> import pyarrow as pa + >>> table = pa.table({"id": [1, 2, 3], "name": ["a", "b", "c"]}) + >>> result = cursor.bulkcopy_arrow("dbo.MyTable", table) + >>> result["rows_copied"] + 3 + """ + is_logging_enabled = logger.is_debug_enabled + + try: + import mssql_py_core + except ImportError as exc: + logger.error("bulkcopy_arrow: Failed to import mssql_py_core module") + raise ImportError( + "Bulk copy requires the mssql_py_core library which is not available. " + "This is an unexpected error. " + ) from exc + + if not table_name or not isinstance(table_name, str): + logger.error("bulkcopy_arrow: Invalid table_name parameter") + raise ValueError("table_name must be a non-empty string") + + if source is None: + raise TypeError( + "source must be a pyarrow Table/RecordBatch/RecordBatchReader, " + "an iterable of RecordBatch, or an object implementing " + "__arrow_c_stream__/__arrow_c_array__, got None" + ) + + if not isinstance(batch_size, int): + raise TypeError( + f"batch_size must be a non-negative integer, got {type(batch_size).__name__}" + ) + if batch_size < 0: + raise ValueError(f"batch_size must be non-negative, got {batch_size}") + + if not isinstance(timeout, int): + raise TypeError(f"timeout must be a positive integer, got {type(timeout).__name__}") + if timeout <= 0: + raise ValueError(f"timeout must be positive, got {timeout}") + + pycore_context = self._build_pycore_context() + + pycore_connection = None + pycore_cursor = None + try: + pycore_connection = mssql_py_core.PyCoreConnection( + pycore_context, python_logger=logger if is_logging_enabled else None + ) + pycore_cursor = pycore_connection.cursor() + + result = pycore_cursor.bulkcopy_arrow( + table_name, + source, + batch_size=batch_size, + timeout=timeout, + column_mappings=column_mappings, + keep_identity=keep_identity, + check_constraints=check_constraints, + table_lock=table_lock, + keep_nulls=keep_nulls, + fire_triggers=fire_triggers, + use_internal_transaction=use_internal_transaction, + python_logger=logger if is_logging_enabled else None, + ) + + logger.info( + "bulkcopy_arrow: completed - rows_copied=%s, batch_count=%s, elapsed_time=%s", + result.get("rows_copied", "N/A"), + result.get("batch_count", "N/A"), + result.get("elapsed_time", "N/A"), + ) + return result + + except Exception as e: + logger.debug( + "bulkcopy_arrow failed for table '%s': %s: %s", + table_name, + type(e).__name__, + str(e), + ) + raise type(e)(str(e)) from None + + finally: + # Clear sensitive data to minimize memory exposure. The + # entra_id_token_factory closure captures client_secret, so drop + # our dict reference to it as well. + if pycore_context: + for key in ( + "password", + "user_name", + "access_token", + "entra_id_token_factory", + ): + pycore_context.pop(key, None) + for resource in (pycore_cursor, pycore_connection): + if resource and hasattr(resource, "close"): + try: + resource.close() + except Exception as cleanup_error: + logger.debug( + "Failed to close bulk copy resource %s: %s", + type(resource).__name__, + cleanup_error, + ) + def __enter__(self): """ Enter the runtime context for the cursor. diff --git a/mssql_python/mssql_python.pyi b/mssql_python/mssql_python.pyi index ad18756e..6f97403d 100644 --- a/mssql_python/mssql_python.pyi +++ b/mssql_python/mssql_python.pyi @@ -4,7 +4,7 @@ Licensed under the MIT license. Type stubs for mssql_python package - based on actual public API """ -from typing import Any, Dict, List, Mapping, Optional, Union, Tuple, Sequence, Callable, Iterator +from typing import Any, Dict, List, Mapping, Optional, Union, Tuple, Sequence, Callable, Iterator, Iterable import datetime import logging import pyarrow @@ -209,6 +209,37 @@ class Cursor: def arrow(self, batch_size: int = 8192) -> pyarrow.Table: ... def arrow_reader(self, batch_size: int = 8192) -> pyarrow.RecordBatchReader: ... + # Bulk Copy + def bulkcopy( + self, + table_name: str, + data: Iterable[Union[Tuple[Any, ...], Row]], + batch_size: int = 0, + timeout: int = 30, + column_mappings: Optional[Union[List[str], List[Tuple[int, str]]]] = None, + keep_identity: bool = False, + check_constraints: bool = False, + table_lock: bool = False, + keep_nulls: bool = False, + fire_triggers: bool = False, + use_internal_transaction: bool = False, + ) -> Dict[str, Any]: ... + def bulkcopy_arrow( + self, + table_name: str, + source: Any, + *, + batch_size: int = 0, + timeout: int = 30, + column_mappings: Optional[Union[List[str], List[Tuple[int, str]]]] = None, + keep_identity: bool = False, + check_constraints: bool = False, + table_lock: bool = False, + keep_nulls: bool = False, + fire_triggers: bool = False, + use_internal_transaction: bool = False, + ) -> Dict[str, Any]: ... + # DB-API 2.0 Connection Object # https://www.python.org/dev/peps/pep-0249/#connection-objects class Connection: diff --git a/tests/test_024_bulkcopy_arrow.py b/tests/test_024_bulkcopy_arrow.py new file mode 100644 index 00000000..05402e61 --- /dev/null +++ b/tests/test_024_bulkcopy_arrow.py @@ -0,0 +1,658 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Comprehensive tests for the Arrow bulk-copy feature (Cursor.bulkcopy_arrow). + +Layout: + * Unit tests (no live DB) exercise every branch of the new Python code: + - Cursor._looks_like_arrow_source (source detection) + - Cursor.bulkcopy (Arrow-steering guard) + - Cursor.bulkcopy_arrow (argument validation + dispatch + cleanup) + - Cursor._build_pycore_context (conn-string parse + AAD/ServicePrincipal auth) + * Live-DB integration tests (gated on DB_CONNECTION_STRING) round-trip real + Arrow data through TDS across the supported type matrix, the accepted source + shapes, column mappings, and the locked type-conversion decisions + (int narrowing, uint64->BIGINT, tz handling, __arrow_c_array__). +""" + +import os +import secrets +from unittest.mock import MagicMock, patch + +import pytest + +# Skip the whole module if the native core or pyarrow can't load (build +# containers without the extension, or pyarrow absent). exc_type=ImportError +# so a module that imports but fails its native extension is skipped, not +# collected as an error (pytest >= 9.1 default only skips ModuleNotFoundError). +mssql_py_core = pytest.importorskip("mssql_py_core", exc_type=ImportError) +pa = pytest.importorskip("pyarrow") + +LIVE_DB = pytest.mark.skipif( + not os.getenv("DB_CONNECTION_STRING"), + reason="DB_CONNECTION_STRING not set", +) + +SAMPLE_TOKEN = secrets.token_hex(44) + + +# ── helpers ────────────────────────────────────────────────────────────────── + + +def _bare_cursor(): + """A Cursor instance with no wiring — enough for pure validation/detection.""" + from mssql_python.cursor import Cursor + + cursor = Cursor.__new__(Cursor) + cursor.closed = True # keep __del__ a no-op during GC teardown + return cursor + + +def _cursor_with_conn(connection_str, auth_type=None, credential_kwargs=None): + """A Cursor whose .connection is a mock carrying just the auth wiring that + _build_pycore_context reads.""" + from mssql_python.cursor import Cursor + + mock_conn = MagicMock() + mock_conn.connection_str = connection_str + mock_conn._auth_type = auth_type + mock_conn._credential_kwargs = credential_kwargs + mock_conn._is_connected = True + + cursor = Cursor.__new__(Cursor) + cursor._connection = mock_conn + cursor.closed = True # keep __del__ a no-op during GC teardown + cursor.hstmt = None + return cursor + + +class _CStreamOnly: + """Object exposing only the Arrow C-stream PyCapsule protocol.""" + + def __init__(self, table): + self._t = table + + def __arrow_c_stream__(self, requested_schema=None): + return self._t.__arrow_c_stream__(requested_schema) + + +class _CArrayOnly: + """Object exposing only the Arrow C-array PyCapsule protocol (single batch).""" + + def __init__(self, batch): + self._b = batch + + def __arrow_c_array__(self, requested_schema=None): + return self._b.__arrow_c_array__(requested_schema) + + +# ── unit: _looks_like_arrow_source ─────────────────────────────────────────── + + +class TestArrowDetection: + def test_table_is_arrow(self): + assert _bare_cursor()._looks_like_arrow_source(pa.table({"a": [1, 2, 3]})) is True + + def test_record_batch_is_arrow(self): + batch = pa.record_batch([pa.array([1, 2])], names=["a"]) + assert _bare_cursor()._looks_like_arrow_source(batch) is True + + def test_record_batch_reader_is_arrow(self): + reader = pa.RecordBatchReader.from_batches( + pa.schema([("a", pa.int32())]), + [pa.record_batch([pa.array([1], type=pa.int32())], names=["a"])], + ) + assert _bare_cursor()._looks_like_arrow_source(reader) is True + + def test_c_stream_producer_is_arrow(self): + obj = _CStreamOnly(pa.table({"a": [1, 2]})) + assert _bare_cursor()._looks_like_arrow_source(obj) is True + + def test_c_array_producer_is_arrow(self): + obj = _CArrayOnly(pa.record_batch([pa.array([1, 2])], names=["a"])) + assert _bare_cursor()._looks_like_arrow_source(obj) is True + + def test_tuples_are_not_arrow(self): + assert _bare_cursor()._looks_like_arrow_source([(1,), (2,)]) is False + + def test_none_is_not_arrow(self): + assert _bare_cursor()._looks_like_arrow_source(None) is False + + def test_string_is_not_arrow(self): + assert _bare_cursor()._looks_like_arrow_source("abc") is False + + def test_bytes_are_not_arrow(self): + assert _bare_cursor()._looks_like_arrow_source(b"abc") is False + + def test_plain_object_is_not_arrow(self): + assert _bare_cursor()._looks_like_arrow_source(object()) is False + + def test_detection_survives_missing_pyarrow(self): + """A non-capsule object with pyarrow unimportable must return False, not raise.""" + cur = _bare_cursor() + with patch.dict("sys.modules", {"pyarrow": None}): + assert cur._looks_like_arrow_source(object()) is False + + def test_capsule_detection_without_pyarrow(self): + """C-stream producers are detected even when pyarrow itself is absent.""" + cur = _bare_cursor() + obj = _CStreamOnly(pa.table({"a": [1]})) + with patch.dict("sys.modules", {"pyarrow": None}): + assert cur._looks_like_arrow_source(obj) is True + + +# ── unit: bulkcopy() Arrow-steering guard ──────────────────────────────────── + + +class TestBulkcopyGuard: + def test_bulkcopy_rejects_table(self): + from mssql_python.cursor import Cursor + + with pytest.raises(TypeError, match="bulkcopy_arrow"): + Cursor.bulkcopy(_bare_cursor(), "t", pa.table({"a": [1, 2]})) + + def test_bulkcopy_rejects_record_batch(self): + from mssql_python.cursor import Cursor + + batch = pa.record_batch([pa.array([1, 2])], names=["a"]) + with pytest.raises(TypeError, match="bulkcopy_arrow"): + Cursor.bulkcopy(_bare_cursor(), "t", batch) + + def test_bulkcopy_rejects_c_stream_producer(self): + from mssql_python.cursor import Cursor + + obj = _CStreamOnly(pa.table({"a": [1, 2]})) + with pytest.raises(TypeError, match="bulkcopy_arrow"): + Cursor.bulkcopy(_bare_cursor(), "t", obj) + + +# ── unit: bulkcopy_arrow() argument validation ─────────────────────────────── + + +class TestBulkcopyArrowValidation: + def test_empty_table_name(self): + with pytest.raises(ValueError, match="table_name"): + _bare_cursor().bulkcopy_arrow("", pa.table({"a": [1]})) + + def test_non_string_table_name(self): + with pytest.raises(ValueError, match="table_name"): + _bare_cursor().bulkcopy_arrow(123, pa.table({"a": [1]})) + + def test_none_source(self): + with pytest.raises(TypeError, match="source"): + _bare_cursor().bulkcopy_arrow("t", None) + + def test_batch_size_wrong_type(self): + with pytest.raises(TypeError, match="batch_size"): + _bare_cursor().bulkcopy_arrow("t", pa.table({"a": [1]}), batch_size="5") + + def test_batch_size_negative(self): + with pytest.raises(ValueError, match="batch_size"): + _bare_cursor().bulkcopy_arrow("t", pa.table({"a": [1]}), batch_size=-1) + + def test_timeout_wrong_type(self): + with pytest.raises(TypeError, match="timeout"): + _bare_cursor().bulkcopy_arrow("t", pa.table({"a": [1]}), timeout="30") + + def test_timeout_non_positive(self): + with pytest.raises(ValueError, match="timeout"): + _bare_cursor().bulkcopy_arrow("t", pa.table({"a": [1]}), timeout=0) + + def test_missing_pycore_raises_importerror(self): + cur = _bare_cursor() + with patch.dict("sys.modules", {"mssql_py_core": None}): + with pytest.raises(ImportError, match="mssql_py_core"): + cur.bulkcopy_arrow("t", pa.table({"a": [1]})) + + +# ── unit: _build_pycore_context() ──────────────────────────────────────────── + + +class TestBuildPycoreContext: + def test_missing_connection_str_raises(self): + from mssql_python.cursor import Cursor + + cur = Cursor.__new__(Cursor) + cur._connection = object() # no connection_str attribute + cur.closed = True + with pytest.raises(RuntimeError, match="Connection string not available"): + cur._build_pycore_context() + + def test_missing_server_raises(self): + cur = _cursor_with_conn("Database=testdb;UID=sa;PWD=pwd", auth_type=None) + with pytest.raises(ValueError, match="SERVER"): + cur._build_pycore_context() + + def test_sql_auth_keeps_credentials(self): + cur = _cursor_with_conn( + "Server=localhost;Database=testdb;UID=sa;PWD=mypwd", auth_type=None + ) + ctx = cur._build_pycore_context() + assert ctx.get("user_name") == "sa" + assert ctx.get("password") == "mypwd" + assert "access_token" not in ctx + assert "entra_id_token_factory" not in ctx + + @patch("mssql_python.cursor.logger") + def test_token_path_replaces_credentials(self, mock_logger): + mock_logger.is_debug_enabled = False + cur = _cursor_with_conn( + "Server=tcp:x.database.windows.net;Database=d;" + "Authentication=ActiveDirectoryDefault;UID=u@x.com;PWD=s", + auth_type="activedirectorydefault", + ) + with patch("mssql_python.auth.AADAuth.get_raw_token", return_value=SAMPLE_TOKEN): + ctx = cur._build_pycore_context() + assert ctx.get("access_token") == SAMPLE_TOKEN + assert "authentication" not in ctx + assert "user_name" not in ctx + assert "password" not in ctx + + @patch("mssql_python.cursor.logger") + def test_token_acquisition_failure_raises(self, mock_logger): + mock_logger.is_debug_enabled = False + cur = _cursor_with_conn( + "Server=tcp:x.database.windows.net;Database=d;" + "Authentication=ActiveDirectoryDefault;UID=u@x.com;PWD=s", + auth_type="activedirectorydefault", + ) + with patch( + "mssql_python.auth.AADAuth.get_raw_token", + side_effect=RuntimeError("no token"), + ): + with pytest.raises(RuntimeError, match="unable to acquire Azure AD token"): + cur._build_pycore_context() + + @patch("mssql_python.cursor.logger") + def test_service_principal_registers_factory(self, mock_logger): + mock_logger.is_debug_enabled = False + from mssql_python.constants import _AuthInternal + + cur = _cursor_with_conn( + "Server=tcp:x.database.windows.net;Database=d;" + "Authentication=ActiveDirectoryServicePrincipal;UID=client-id;PWD=client-secret", + auth_type=_AuthInternal.SERVICE_PRINCIPAL, + ) + sentinel = object() + with patch( + "mssql_python.auth.ServicePrincipalAuth.make_token_factory", + return_value=sentinel, + ) as mk: + ctx = cur._build_pycore_context() + mk.assert_called_once_with("client-id", "client-secret") + assert ctx.get("entra_id_token_factory") is sentinel + # ServicePrincipal keeps auth/user_name/password so py-core can resolve + # the method before dispatching the factory at handshake time. + assert "access_token" not in ctx + + @patch("mssql_python.cursor.logger") + def test_service_principal_missing_secret_raises(self, mock_logger): + mock_logger.is_debug_enabled = False + from mssql_python.constants import _AuthInternal + + cur = _cursor_with_conn( + "Server=tcp:x.database.windows.net;Database=d;" + "Authentication=ActiveDirectoryServicePrincipal", + auth_type=_AuthInternal.SERVICE_PRINCIPAL, + ) + with pytest.raises(RuntimeError, match="client-secret only"): + cur._build_pycore_context() + + @patch("mssql_python.cursor.logger") + def test_service_principal_factory_build_failure_raises(self, mock_logger): + mock_logger.is_debug_enabled = False + from mssql_python.constants import _AuthInternal + + cur = _cursor_with_conn( + "Server=tcp:x.database.windows.net;Database=d;" + "Authentication=ActiveDirectoryServicePrincipal;UID=c;PWD=s", + auth_type=_AuthInternal.SERVICE_PRINCIPAL, + ) + with patch( + "mssql_python.auth.ServicePrincipalAuth.make_token_factory", + side_effect=ValueError("bad cert"), + ): + with pytest.raises(RuntimeError, match="ServicePrincipal token factory"): + cur._build_pycore_context() + + +# ── unit: bulkcopy_arrow() dispatch + cleanup ──────────────────────────────── + + +def _mock_pycore(result=None, raise_exc=None): + """Build a mock mssql_py_core module and return (module, pycore_cursor, + pycore_conn, captured) where captured['ctx'] holds the exact context dict + handed to PyCoreConnection.""" + pycore_cursor = MagicMock() + if raise_exc is not None: + pycore_cursor.bulkcopy_arrow.side_effect = raise_exc + else: + pycore_cursor.bulkcopy_arrow.return_value = result or { + "rows_copied": 2, + "batch_count": 1, + "elapsed_time": 0.01, + "rows_per_second": 200.0, + } + pycore_conn = MagicMock() + pycore_conn.cursor.return_value = pycore_cursor + + captured = {} + + def make_conn(ctx, **kwargs): + captured["ctx"] = ctx # keep the reference so we can inspect post-cleanup + return pycore_conn + + module = MagicMock() + module.PyCoreConnection = make_conn + return module, pycore_cursor, pycore_conn, captured + + +class TestBulkcopyArrowDispatch: + @patch("mssql_python.cursor.logger") + def test_success_returns_result_and_forwards_args(self, mock_logger): + mock_logger.is_debug_enabled = False + cur = _cursor_with_conn("Server=localhost;Database=d;UID=sa;PWD=p") + module, pyc_cursor, _, _ = _mock_pycore() + src = pa.table({"a": [1, 2]}) + + with patch.dict("sys.modules", {"mssql_py_core": module}): + result = cur.bulkcopy_arrow( + "dbo.t", src, batch_size=7, timeout=15, column_mappings=["a"] + ) + + assert result["rows_copied"] == 2 + pyc_cursor.bulkcopy_arrow.assert_called_once() + args, kwargs = pyc_cursor.bulkcopy_arrow.call_args + assert args[0] == "dbo.t" + assert args[1] is src + assert kwargs["batch_size"] == 7 + assert kwargs["timeout"] == 15 + assert kwargs["column_mappings"] == ["a"] + + @patch("mssql_python.cursor.logger") + def test_sensitive_fields_cleared_after_success(self, mock_logger): + mock_logger.is_debug_enabled = False + cur = _cursor_with_conn("Server=localhost;Database=d;UID=sa;PWD=secret") + module, _, _, captured = _mock_pycore() + + with patch.dict("sys.modules", {"mssql_py_core": module}): + cur.bulkcopy_arrow("t", pa.table({"a": [1]})) + + ctx = captured["ctx"] + for key in ("password", "user_name", "access_token", "entra_id_token_factory"): + assert key not in ctx + + @patch("mssql_python.cursor.logger") + def test_resources_closed_on_success(self, mock_logger): + mock_logger.is_debug_enabled = False + cur = _cursor_with_conn("Server=localhost;Database=d;UID=sa;PWD=p") + module, pyc_cursor, pyc_conn, _ = _mock_pycore() + + with patch.dict("sys.modules", {"mssql_py_core": module}): + cur.bulkcopy_arrow("t", pa.table({"a": [1]})) + + pyc_cursor.close.assert_called_once() + pyc_conn.close.assert_called_once() + + @patch("mssql_python.cursor.logger") + def test_core_exception_is_reraised_and_cleaned_up(self, mock_logger): + mock_logger.is_debug_enabled = False + cur = _cursor_with_conn("Server=localhost;Database=d;UID=sa;PWD=p") + module, pyc_cursor, pyc_conn, captured = _mock_pycore( + raise_exc=ValueError("boom") + ) + + with patch.dict("sys.modules", {"mssql_py_core": module}): + with pytest.raises(ValueError, match="boom"): + cur.bulkcopy_arrow("t", pa.table({"a": [1]})) + + # cleanup still ran despite the failure + assert "password" not in captured["ctx"] + pyc_cursor.close.assert_called_once() + pyc_conn.close.assert_called_once() + + @patch("mssql_python.cursor.logger") + def test_cleanup_swallows_close_errors(self, mock_logger): + """A failing resource.close() during teardown must not mask the result.""" + mock_logger.is_debug_enabled = False + cur = _cursor_with_conn("Server=localhost;Database=d;UID=sa;PWD=p") + module, pyc_cursor, pyc_conn, _ = _mock_pycore() + pyc_cursor.close.side_effect = RuntimeError("close failed") + + with patch.dict("sys.modules", {"mssql_py_core": module}): + result = cur.bulkcopy_arrow("t", pa.table({"a": [1]})) + + # result is still returned; the close error was swallowed and logged + assert result["rows_copied"] == 2 + pyc_cursor.close.assert_called_once() + pyc_conn.close.assert_called_once() + + +# ── live DB integration ────────────────────────────────────────────────────── + + +def _make_table(cursor, name, cols): + cursor.execute(f"IF OBJECT_ID('{name}','U') IS NOT NULL DROP TABLE {name}") + cursor.execute(f"CREATE TABLE {name} ({cols})") + cursor.connection.commit() + + +@LIVE_DB +class TestBulkcopyArrowLive: + def test_table_round_trip_with_nulls(self, cursor): + t = "mssql_python_arrow_tbl" + _make_table(cursor, t, "id INT NOT NULL, name NVARCHAR(50) NULL, score FLOAT NULL") + tbl = pa.table( + { + "id": pa.array([1, 2, 3], type=pa.int32()), + "name": pa.array(["Alice", None, "Charlie"]), + "score": pa.array([1.5, 2.5, None], type=pa.float64()), + } + ) + result = cursor.bulkcopy_arrow(t, tbl) + assert result["rows_copied"] == 3 + # returned dict shape + for key in ("rows_copied", "batch_count", "elapsed_time", "rows_per_second"): + assert key in result + cursor.execute(f"SELECT id, name, score FROM {t} ORDER BY id") + rows = cursor.fetchall() + assert [r[0] for r in rows] == [1, 2, 3] + assert rows[1][1] is None + assert rows[2][2] is None + cursor.execute(f"DROP TABLE {t}") + + def test_record_batch_source(self, cursor): + t = "mssql_python_arrow_batch" + _make_table(cursor, t, "id INT NOT NULL") + batch = pa.record_batch([pa.array([10, 20], type=pa.int32())], names=["id"]) + result = cursor.bulkcopy_arrow(t, batch) + assert result["rows_copied"] == 2 + cursor.execute(f"SELECT id FROM {t} ORDER BY id") + assert [r[0] for r in cursor.fetchall()] == [10, 20] + cursor.execute(f"DROP TABLE {t}") + + def test_record_batch_reader_multibatch(self, cursor): + t = "mssql_python_arrow_reader" + _make_table(cursor, t, "id INT NOT NULL, txt NVARCHAR(20) NULL") + schema = pa.schema([("id", pa.int32()), ("txt", pa.string())]) + batches = [ + pa.record_batch( + [pa.array([1, 2], type=pa.int32()), pa.array(["a", "b"])], schema=schema + ), + pa.record_batch( + [pa.array([3], type=pa.int32()), pa.array([None])], schema=schema + ), + ] + reader = pa.RecordBatchReader.from_batches(schema, batches) + result = cursor.bulkcopy_arrow(t, reader) + assert result["rows_copied"] == 3 + cursor.execute(f"SELECT id FROM {t} ORDER BY id") + assert [r[0] for r in cursor.fetchall()] == [1, 2, 3] + cursor.execute(f"DROP TABLE {t}") + + def test_iterable_of_record_batches(self, cursor): + t = "mssql_python_arrow_iter" + _make_table(cursor, t, "id INT NOT NULL") + schema = pa.schema([("id", pa.int32())]) + batches = [ + pa.record_batch([pa.array([1, 2], type=pa.int32())], schema=schema), + pa.record_batch([pa.array([3, 4], type=pa.int32())], schema=schema), + ] + result = cursor.bulkcopy_arrow(t, batches) + assert result["rows_copied"] == 4 + cursor.execute(f"DROP TABLE {t}") + + def test_c_stream_producer(self, cursor): + t = "mssql_python_arrow_cstream" + _make_table(cursor, t, "id INT NOT NULL") + obj = _CStreamOnly(pa.table({"id": pa.array([5, 6], type=pa.int32())})) + result = cursor.bulkcopy_arrow(t, obj) + assert result["rows_copied"] == 2 + cursor.execute(f"DROP TABLE {t}") + + def test_c_array_producer_single_batch(self, cursor): + t = "mssql_python_arrow_carray" + _make_table(cursor, t, "id INT NOT NULL, name NVARCHAR(20) NULL") + batch = pa.record_batch( + [pa.array([7, 8], type=pa.int32()), pa.array(["x", "y"])], names=["id", "name"] + ) + result = cursor.bulkcopy_arrow(t, _CArrayOnly(batch)) + assert result["rows_copied"] == 2 + cursor.execute(f"SELECT id FROM {t} ORDER BY id") + assert [r[0] for r in cursor.fetchall()] == [7, 8] + cursor.execute(f"DROP TABLE {t}") + + def test_empty_source_copies_zero_rows(self, cursor): + t = "mssql_python_arrow_empty" + _make_table(cursor, t, "id INT NOT NULL") + tbl = pa.table({"id": pa.array([], type=pa.int32())}) + result = cursor.bulkcopy_arrow(t, tbl) + assert result["rows_copied"] == 0 + cursor.execute(f"DROP TABLE {t}") + + # ── locked type-conversion decisions ──────────────────────────────────── + + def test_int64_narrows_to_int(self, cursor): + """A1: Arrow int64 loads into a SQL INT column when values fit.""" + t = "mssql_python_arrow_a1" + _make_table(cursor, t, "n INT NOT NULL") + tbl = pa.table({"n": pa.array([1, 2, 2_000_000_000], type=pa.int64())}) + result = cursor.bulkcopy_arrow(t, tbl) + assert result["rows_copied"] == 3 + cursor.execute(f"SELECT n FROM {t} ORDER BY n") + assert [r[0] for r in cursor.fetchall()] == [1, 2, 2_000_000_000] + cursor.execute(f"DROP TABLE {t}") + + def test_int64_overflow_to_int_raises(self, cursor): + """A1: out-of-range int64 -> INT must raise, not silently truncate.""" + t = "mssql_python_arrow_a1_of" + _make_table(cursor, t, "n INT NOT NULL") + tbl = pa.table({"n": pa.array([5_000_000_000], type=pa.int64())}) + with pytest.raises(Exception): + cursor.bulkcopy_arrow(t, tbl) + cursor.execute(f"DROP TABLE {t}") + + def test_uint64_to_bigint(self, cursor): + """A2: Arrow uint64 loads into BIGINT (values within i64 range).""" + t = "mssql_python_arrow_a2" + _make_table(cursor, t, "n BIGINT NOT NULL") + tbl = pa.table({"n": pa.array([1, 42, 9_000_000_000], type=pa.uint64())}) + result = cursor.bulkcopy_arrow(t, tbl) + assert result["rows_copied"] == 3 + cursor.execute(f"DROP TABLE {t}") + + def test_tz_aware_timestamp_to_datetime2_raises(self, cursor): + """C1: tz-aware timestamp -> DATETIME2 must raise and point at datetimeoffset.""" + import datetime as dt + + t = "mssql_python_arrow_c1" + _make_table(cursor, t, "ts DATETIME2 NULL") + ts = pa.array([dt.datetime(2020, 1, 1, 12)], type=pa.timestamp("us", tz="UTC")) + with pytest.raises(Exception, match="(?i)datetimeoffset"): + cursor.bulkcopy_arrow(t, pa.table({"ts": ts})) + cursor.execute(f"DROP TABLE {t}") + + # ── column mappings ───────────────────────────────────────────────────── + + def test_column_mappings_by_name(self, cursor): + t = "mssql_python_arrow_map_name" + _make_table(cursor, t, "a INT NOT NULL, b NVARCHAR(20) NULL") + src = pa.table({"x": pa.array([1, 2], type=pa.int32()), "y": pa.array(["p", "q"])}) + result = cursor.bulkcopy_arrow(t, src, column_mappings=["a", "b"]) + assert result["rows_copied"] == 2 + cursor.execute(f"SELECT a, b FROM {t} ORDER BY a") + rows = cursor.fetchall() + assert rows[0][1] == "p" and rows[1][1] == "q" + cursor.execute(f"DROP TABLE {t}") + + def test_column_mappings_by_ordinal(self, cursor): + t = "mssql_python_arrow_map_ord" + _make_table(cursor, t, "a INT NOT NULL, b NVARCHAR(20) NULL") + src = pa.table({"x": pa.array([1, 2], type=pa.int32()), "y": pa.array(["p", "q"])}) + result = cursor.bulkcopy_arrow(t, src, column_mappings=[(0, "a"), (1, "b")]) + assert result["rows_copied"] == 2 + cursor.execute(f"DROP TABLE {t}") + + # ── representative type matrix ────────────────────────────────────────── + + def test_type_matrix_numeric_and_string(self, cursor): + t = "mssql_python_arrow_types1" + _make_table( + cursor, + t, + "b BIT, i8 TINYINT, i16 SMALLINT, i32 INT, i64 BIGINT, " + "f32 REAL, f64 FLOAT, s NVARCHAR(50), v VARCHAR(50)", + ) + tbl = pa.table( + { + "b": pa.array([True, False], type=pa.bool_()), + "i8": pa.array([1, 255], type=pa.uint8()), + "i16": pa.array([1, 30000], type=pa.int16()), + "i32": pa.array([1, 2_000_000_000], type=pa.int32()), + "i64": pa.array([1, 9_000_000_000], type=pa.int64()), + "f32": pa.array([1.5, 2.5], type=pa.float32()), + "f64": pa.array([1.25, 2.75], type=pa.float64()), + "s": pa.array(["hello", "wörld"]), + "v": pa.array(["ascii", "text"]), + } + ) + result = cursor.bulkcopy_arrow(t, tbl) + assert result["rows_copied"] == 2 + cursor.execute(f"SELECT i32 FROM {t} ORDER BY i32") + assert [r[0] for r in cursor.fetchall()] == [1, 2_000_000_000] + cursor.execute(f"DROP TABLE {t}") + + def test_type_matrix_temporal_and_decimal(self, cursor): + import datetime as dt + from decimal import Decimal + + t = "mssql_python_arrow_types2" + _make_table( + cursor, + t, + "d DATE, ts DATETIME2, tm TIME, amt DECIMAL(10,2), tzo DATETIMEOFFSET", + ) + tbl = pa.table( + { + "d": pa.array([dt.date(2020, 1, 2)], type=pa.date32()), + "ts": pa.array([dt.datetime(2020, 1, 2, 3, 4, 5)], type=pa.timestamp("us")), + "tm": pa.array([dt.time(3, 4, 5)], type=pa.time64("us")), + "amt": pa.array([Decimal("123.45")], type=pa.decimal128(10, 2)), + "tzo": pa.array( + [dt.datetime(2020, 1, 2, 3, 4, 5)], type=pa.timestamp("us", tz="UTC") + ), + } + ) + result = cursor.bulkcopy_arrow(t, tbl) + assert result["rows_copied"] == 1 + cursor.execute(f"SELECT amt FROM {t}") + assert cursor.fetchone()[0] == Decimal("123.45") + cursor.execute(f"DROP TABLE {t}") + + def test_type_matrix_binary(self, cursor): + t = "mssql_python_arrow_types3" + _make_table(cursor, t, "vb VARBINARY(16)") + tbl = pa.table({"vb": pa.array([b"\x01\x02\x03", b"\xff\xee"], type=pa.binary())}) + result = cursor.bulkcopy_arrow(t, tbl) + assert result["rows_copied"] == 2 + cursor.execute(f"DROP TABLE {t}")