From 71116ed4a529423e0353f3d1b5077e5c4f63a906 Mon Sep 17 00:00:00 2001 From: Vamsi-klu Date: Sat, 14 Nov 2026 01:00:00 +0000 Subject: [PATCH] feat(postgres): bulk-load DataFrames with COPY Load pandas DataFrames from Python models and seeds into Postgres with COPY FROM STDIN through a staging table, instead of batched INSERT ... VALUES. On 500k rows locally this took 3.1s instead of 56.2s. - Type the staging table from the DataFrame dtypes and cast into the target with _casted_columns, matching the VALUES path's results. Floats are staged as NUMERIC so rounding matches VALUES literals. - Stream CSV chunks of batch_size rows through cursor.copy_expert, with a random per-load NULL marker. Only _copy_df touches the driver API. - Fall back to VALUES for unsupported dtypes or values, a single-column frame containing `\.`, drivers without copy_expert, and RisingWave. - Drop the staging table after the load, and also when the COPY fails inside query_factory, which replace_query calls outside the SourceQuery context. - Make _get_temp_table respect MAX_IDENTIFIER_LENGTH by truncating the middle of the name, so long physical names get valid staging names. Fixes #5830 Co-Authored-By: Claude Opus 5.5 Signed-off-by: Vamsi-klu --- docs/integrations/engines/postgres.md | 12 + sqlmesh/core/engine_adapter/base.py | 12 +- sqlmesh/core/engine_adapter/postgres.py | 194 ++++- sqlmesh/core/engine_adapter/risingwave.py | 4 + .../integration/test_integration_postgres.py | 788 +++++++++++++++++- tests/core/engine_adapter/test_base.py | 34 + tests/core/engine_adapter/test_postgres.py | 292 +++++++ tests/core/engine_adapter/test_risingwave.py | 15 + 8 files changed, 1345 insertions(+), 6 deletions(-) diff --git a/docs/integrations/engines/postgres.md b/docs/integrations/engines/postgres.md index cf1d3e4ce8..be2d3c6b70 100644 --- a/docs/integrations/engines/postgres.md +++ b/docs/integrations/engines/postgres.md @@ -23,3 +23,15 @@ pip install "sqlmesh[postgres]" | `role` | The role to use for authentication with the Postgres server | string | N | | `sslmode` | The security of the connection to the Postgres server | string | N | | `application_name` | The name of the application to use for the connection | string | N | + +### DataFrame loading + +When using the `psycopg2` driver, SQLMesh loads pandas DataFrames (such as those returned by Python models and used for seeds) with `COPY ... FROM STDIN (FORMAT csv)` instead of batched `INSERT ... VALUES` statements, which is much faster for large DataFrames. The data is copied into a temporary staging table, which is created in the same schema as the target table and then cast into the target table's column types. The staging table is dropped afterwards, including when the load fails. Because it is created in the target's schema, the connecting user needs the `CREATE` privilege on that schema, as for other SQLMesh DDL. + +No configuration is required and there is no option to disable this behavior. SQLMesh automatically falls back to `INSERT ... VALUES` when: + +- A column has a type or values that cannot be loaded with `COPY` identically. Supported columns are integers (up to 64 bits signed, up to 32 bits unsigned), `float64`, booleans, timestamps (with or without time zone), and strings. Other types, such as `float32`, 64-bit unsigned integers, `Decimal`, dates, bytes, dicts, lists, and categoricals, trigger the fallback. +- The DataFrame has a single column and one of its values is `\.`. +- The driver does not support `copy_expert` (for example, a driver other than `psycopg2`). + +The RisingWave engine always uses `INSERT ... VALUES`. diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py index 2dd234f70a..2e3160b7e1 100644 --- a/sqlmesh/core/engine_adapter/base.py +++ b/sqlmesh/core/engine_adapter/base.py @@ -2921,9 +2921,15 @@ def _get_temp_table( Returns the name of the temp table that should be used for the given table name. """ table = t.cast(exp.Table, exp.to_table(table).copy()) - table.set( - "this", exp.to_identifier(f"__temp_{table.name}_{random_id(short=True)}", quoted=quoted) - ) + prefix, suffix = "__temp_", f"_{random_id(short=True)}" + name = table.name + if self.MAX_IDENTIFIER_LENGTH is not None: + # Cut the middle of the name so that the temp table name is a valid identifier + keep = max(self.MAX_IDENTIFIER_LENGTH - len(prefix) - len(suffix), 0) + if len(name) > keep: + tail = keep // 2 + name = name[: keep - tail] + (name[-tail:] if tail else "") + table.set("this", exp.to_identifier(f"{prefix}{name}{suffix}", quoted=quoted)) if table_only: table.set("db", None) diff --git a/sqlmesh/core/engine_adapter/postgres.py b/sqlmesh/core/engine_adapter/postgres.py index 6794169322..82c9be4afb 100644 --- a/sqlmesh/core/engine_adapter/postgres.py +++ b/sqlmesh/core/engine_adapter/postgres.py @@ -1,5 +1,7 @@ from __future__ import annotations +import csv +import io import logging import re import typing as t @@ -14,14 +16,59 @@ logical_merge, GrantsFromInfoSchemaMixin, ) -from sqlmesh.core.engine_adapter.shared import set_catalog +from sqlmesh.core.engine_adapter.shared import SourceQuery, set_catalog +from sqlmesh.utils import random_id if t.TYPE_CHECKING: + import pandas as pd + from sqlmesh.core._typing import TableName - from sqlmesh.core.engine_adapter._typing import DF, QueryOrDF + from sqlmesh.core.engine_adapter._typing import DF, Query, QueryOrDF logger = logging.getLogger(__name__) +# pandas dtypes that can be bulk loaded with COPY, mapped to the type of the staging table column. +# The staging column types are derived from the DataFrame (not from the target table) and the final +# `SELECT` casts them to the target types, which mirrors how the generic `VALUES` implementation types +# literals by their Python value and then casts. In particular floats are staged as an unbounded +# NUMERIC, because the literals emitted by the `VALUES` implementation are NUMERIC and e.g. casting +# 2.5 to an integer rounds differently for NUMERIC (3) and DOUBLE PRECISION (2). +_COPY_DTYPE_TO_STAGING_TYPE: t.Dict[str, str] = { + **{ + dtype: "BIGINT" + for dtype in ( + "int8", + "int16", + "int32", + "int64", + "uint8", + "uint16", + "uint32", + "Int8", + "Int16", + "Int32", + "Int64", + "UInt8", + "UInt16", + "UInt32", + ) + }, + "float64": "NUMERIC", + "Float64": "NUMERIC", + "bool": "BOOLEAN", + "boolean": "BOOLEAN", +} +# What pandas' `infer_dtype` reports for object / string columns that can be copied. Python numbers in an +# object column (e.g. seeds, where NaN is replaced with None) are staged as NUMERIC, which is lossless. +_COPY_INFERRED_TO_STAGING_TYPE: t.Dict[str, str] = { + "string": "TEXT", + "empty": "TEXT", # only nulls + "boolean": "BOOLEAN", + "integer": "NUMERIC", + "floating": "NUMERIC", + "mixed-integer-float": "NUMERIC", +} + @set_catalog() class PostgresEngineAdapter( @@ -85,6 +132,149 @@ def _fetch_native_df( self._connection_pool.commit() return df + def _df_to_source_queries( + self, + df: DF, + target_columns_to_types: t.Dict[str, exp.DataType], + batch_size: int, + target_table: TableName, + source_columns: t.Optional[t.List[str]] = None, + ) -> t.List[SourceQuery]: + """Bulk loads the DataFrame into a staging table with COPY instead of using `VALUES` batches. + + The generic implementation is used when the DataFrame cannot be loaded safely with COPY. + """ + import pandas as pd + + assert isinstance(df, pd.DataFrame) + + # pandas truncates `na_rep` to 32 characters when writing float columns, so keep the marker short + null_marker = f"__sqlmesh_null_{random_id(short=True)}__" + # Same column selection and order as the generic implementation + ordered_df = df[list(source_columns or target_columns_to_types)] + staging_columns_to_types = ( + None + if ordered_df.empty or not hasattr(self.cursor, "copy_expert") + else self._get_copy_staging_columns_to_types(ordered_df, null_marker) + ) + if staging_columns_to_types is None: + return super()._df_to_source_queries( + df, + target_columns_to_types, + batch_size, + target_table, + source_columns=source_columns, + ) + + staging_table = self._get_temp_table(target_table or "pandas") + + def query_factory() -> Query: + # The factory can be called multiple times. The staging table is only loaded once, which + # assumes that every call would have produced the same result. + if not self.table_exists(staging_table): + try: + self.create_table(staging_table, staging_columns_to_types) + self._copy_df(staging_table, ordered_df, null_marker, batch_size) + except Exception: + # Callers such as `replace_query` invoke the factory outside of `with source_query`, + # so `cleanup_func` won't run and a committed staging table would leak. + try: + self._drop_copy_staging_table(staging_table) + except Exception: + logger.warning("Failed to drop staging table %s", staging_table.sql()) + raise + return exp.select( + *self._casted_columns(target_columns_to_types, source_columns=source_columns) + ).from_(staging_table) + + return [ + SourceQuery( + query_factory=query_factory, + cleanup_func=lambda: self._drop_copy_staging_table(staging_table), + ) + ] + + def _drop_copy_staging_table(self, staging_table: exp.Table) -> None: + try: + self.drop_table(staging_table) + except Exception as ex: + # The cleanup runs while an exception is propagating. If it left the transaction aborted + # the DROP fails with "current transaction is aborted", which would hide the original + # error. The staging table was created in that transaction, so its rollback removes it. + sqlstate = getattr(ex, "pgcode", None) or getattr(ex, "sqlstate", None) + if sqlstate != "25P02" or not self._connection_pool.is_transaction_active: + raise + logger.debug("Skipped dropping %s from an aborted transaction", staging_table.sql()) + + @staticmethod + def _get_copy_staging_columns_to_types( + df: pd.DataFrame, null_marker: str + ) -> t.Optional[t.Dict[str, exp.DataType]]: + """Returns the staging table column types for the DataFrame, or None if it can't be loaded with COPY. + + None is returned when a column has a dtype or values that wouldn't be read back by COPY exactly + as the generic `VALUES` implementation would write them, e.g. dicts, bytes, Decimals, dates or + categoricals, or when a string would be misread by the CSV parser. + """ + import pandas as pd + + staging_columns_to_types: t.Dict[str, exp.DataType] = {} + for column, series in df.items(): + dtype = series.dtype + staging_type: t.Optional[str] + if isinstance(dtype, pd.DatetimeTZDtype): + staging_type = "TIMESTAMPTZ" + elif pd.api.types.is_datetime64_dtype(dtype): + staging_type = "TIMESTAMP" + elif str(dtype) in _COPY_DTYPE_TO_STAGING_TYPE: + staging_type = _COPY_DTYPE_TO_STAGING_TYPE[str(dtype)] + elif dtype == object or isinstance(dtype, pd.StringDtype): + staging_type = _COPY_INFERRED_TO_STAGING_TYPE.get( + pd.api.types.infer_dtype(series, skipna=True) + ) + if staging_type == "TEXT": + if series.eq(null_marker).any(): + return None + # On PostgreSQL <= 17 an unquoted `\.` line ends the data, which happens if it is the only column + if len(df.columns) == 1 and series.eq("\\.").any(): + return None + else: + staging_type = None + if staging_type is None: + return None + staging_columns_to_types[str(column)] = exp.DataType.build( + staging_type, dialect=PostgresEngineAdapter.DIALECT + ) + return staging_columns_to_types + + def _copy_df( + self, table: exp.Table, df: pd.DataFrame, null_marker: str, batch_size: int + ) -> None: + """Loads the DataFrame into the table with `COPY ... FROM STDIN` in chunks of `batch_size` rows. + + This is the only place that calls the driver's COPY API (`psycopg2`'s `copy_expert`). + The copy runs on the adapter's cursor so it is part of the current transaction. + """ + columns = ", ".join(self._to_sql(exp.column(c, quoted=True)) for c in df.columns) + copy_sql = ( + f"COPY {self._to_sql(table)} ({columns}) " + f"FROM STDIN WITH (FORMAT csv, NULL '{null_marker}')" + ) + chunk_size = batch_size or max(len(df.index), 1) + for start in range(0, len(df.index), chunk_size): + buffer = io.StringIO() + df.iloc[start : start + chunk_size].to_csv( + buffer, + header=False, + index=False, + na_rep=null_marker, + quoting=csv.QUOTE_MINIMAL, + lineterminator="\n", + ) + buffer.seek(0) + # A text buffer is encoded by the driver with the connection's client_encoding + self.cursor.copy_expert(copy_sql, buffer) + def _create_table_like( self, target_table_name: TableName, diff --git a/sqlmesh/core/engine_adapter/risingwave.py b/sqlmesh/core/engine_adapter/risingwave.py index 61b44f5bbb..6d330ca4af 100644 --- a/sqlmesh/core/engine_adapter/risingwave.py +++ b/sqlmesh/core/engine_adapter/risingwave.py @@ -6,6 +6,7 @@ from sqlglot import exp +from sqlmesh.core.engine_adapter.base import EngineAdapter from sqlmesh.core.engine_adapter.postgres import PostgresEngineAdapter from sqlmesh.core.engine_adapter.shared import ( set_catalog, @@ -34,6 +35,9 @@ class RisingwaveEngineAdapter(PostgresEngineAdapter): MAX_IDENTIFIER_LENGTH = None SUPPORTS_GRANTS = False + # COPY into a staging table hasn't been verified on RisingWave, so DataFrames keep using `VALUES` + _df_to_source_queries = EngineAdapter._df_to_source_queries + def columns( self, table_name: TableName, include_pseudo_columns: bool = False ) -> t.Dict[str, exp.DataType]: diff --git a/tests/core/engine_adapter/integration/test_integration_postgres.py b/tests/core/engine_adapter/integration/test_integration_postgres.py index 2bc1dcc2e7..db3ed7a917 100644 --- a/tests/core/engine_adapter/integration/test_integration_postgres.py +++ b/tests/core/engine_adapter/integration/test_integration_postgres.py @@ -1,10 +1,16 @@ import typing as t import uuid +from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager +from unittest.mock import patch + +import numpy as np # noqa: TID253 +import pandas as pd # noqa: TID253 import pytest from pytest import FixtureRequest +from pytest_mock import MockerFixture from pathlib import Path -from sqlmesh.core.engine_adapter import PostgresEngineAdapter +from sqlmesh.core.engine_adapter import EngineAdapter, PostgresEngineAdapter from sqlmesh.core.config import Config, DuckDBConnectionConfig from sqlmesh.core.config.common import VirtualEnvironmentMode from tests.core.engine_adapter.integration import TestContext @@ -1230,3 +1236,783 @@ def test_grants_plan_scd_type_2_with_vde_dev_only( ) assert roles["analyst"]["username"] in dev_physical_grants.get("SELECT", []) assert roles["etl_user"]["username"] in dev_physical_grants.get("INSERT", []) + + +# -------------------------------------------------------------------------------------------- +# DataFrame bulk loading via COPY +# -------------------------------------------------------------------------------------------- + + +def _ts(values: t.Any) -> pd.DatetimeIndex: + return pd.to_datetime(values, format="ISO8601") + + +def _types(**columns: str) -> t.Dict[str, exp.DataType]: + return {name: exp.DataType.build(kind, dialect="postgres") for name, kind in columns.items()} + + +def _leftover_staging_tables( + engine_adapter: PostgresEngineAdapter, schema_like: str +) -> t.List[str]: + rows = engine_adapter.fetchall( + "SELECT table_schema || '.' || table_name FROM information_schema.tables " + f"WHERE table_name LIKE '\\_\\_temp\\_%' AND table_schema LIKE '%{schema_like}'" + ) + return [r[0] for r in rows] + + +@contextmanager +def _values_path(adapter: PostgresEngineAdapter) -> t.Iterator[None]: + """Forces DataFrames through the generic `VALUES (...)` implementation.""" + with patch.object( + PostgresEngineAdapter, "_df_to_source_queries", EngineAdapter._df_to_source_queries + ): + yield + + +def _assert_same_rows(adapter: PostgresEngineAdapter, a: exp.Table, b: exp.Table) -> None: + a_sql, b_sql = a.sql("postgres", identify=True), b.sql("postgres", identify=True) + (a_count,) = adapter.fetchone(f"SELECT COUNT(*) FROM {a_sql}") # type: ignore + (b_count,) = adapter.fetchone(f"SELECT COUNT(*) FROM {b_sql}") # type: ignore + assert a_count == b_count > 0 + for left, right in ((a_sql, b_sql), (b_sql, a_sql)): + extra = adapter.fetchall(f"SELECT * FROM {left} EXCEPT ALL SELECT * FROM {right}") + assert extra == [], f"rows only in {left}: {extra}" + + +_ALL_THE_STRINGS = [ + "", + None, + "\\N", + "NULL", + "null", + "\\", + "\\\\", + "\\x41", + 'a"b', + '"', + "a,b", + "a\nb", + "a\rb", + "a\r\nb", + "a\tb", + " padded ", + "h\u00e9llo w\u00f6rld", + "\u65e5\u672c\u8a9e", + "\U0001f600 emoji", + "it's", + "\\.", + "end\\.", + "\\.\nafter", + "a" * 5000, +] + +_TS_NAIVE = _ts( + [ + "2020-01-01 00:00:00", + "2020-06-15 12:34:56.123456", + "1969-12-31 23:59:59.999999", + "1900-01-01 00:00:00", + "1677-09-22 00:00:00", + "2262-04-11 23:47:16.854775", + "2024-02-29 00:00:00.000001", + ] +) + +# the extreme values can't be converted to every time zone +_TS_TZ = _TS_NAIVE[[0, 1, 2, 3, 6]] + +# (id, DataFrame, target types, source_columns, expect the COPY path to be used) +_PARITY_CASES: t.List[t.Any] = [ + pytest.param( + pd.DataFrame({"s": _ALL_THE_STRINGS, "i": range(len(_ALL_THE_STRINGS))}), + _types(s="text", i="int"), + None, + True, + id="strings-multi-column", + ), + pytest.param( + pd.DataFrame({"s": [x for x in _ALL_THE_STRINGS if x != "\\."]}), + _types(s="text"), + None, + True, + id="strings-single-column", + ), + pytest.param( + pd.DataFrame({"s": ["a", "\\.", "b"]}), + _types(s="text"), + None, + False, + id="strings-end-of-copy-marker-falls-back", + ), + pytest.param( + pd.DataFrame({"s": ["", "x", None, ""]}), + _types(s="varchar(10)"), + None, + True, + id="strings-single-column-empty-and-null", + ), + pytest.param( + pd.DataFrame({"i": [1, None, 3, None], "j": [5, 6, 7, 8]}), + _types(i="int", j="smallint"), + None, + True, + id="int-with-nulls-is-float-in-pandas", + ), + pytest.param( + pd.DataFrame( + { + "i": [2**63 - 1, -(2**63), 9007199254740993, 0, -1], + "j": [1, 2, 3, 4, 5], + } + ), + _types(i="bigint", j="int"), + None, + True, + id="int64-bounds-and-2^53+1", + ), + pytest.param( + pd.DataFrame( + {"i": np.array([1, -2, 3], dtype="int8"), "j": np.array([4, 5, 6], dtype="uint32")} + ), + _types(i="smallint", j="bigint"), + None, + True, + id="small-int-dtypes", + ), + pytest.param( + pd.DataFrame({"f": [1.5, 2.5, 0.5, -0.5, -1.5, -2.5, 3.4999999999, 1e10]}), + _types(f="bigint"), + None, + True, + id="float-into-bigint-rounds-like-values", + ), + pytest.param( + pd.DataFrame( + { + "f": [ + 0.1, + 1e-300, + 1e300, + 0.30000000000000004, + -0.0, + 5e-324, + 1.7976931348623157e308, + 123456789.123456789, + np.nan, + ] + } + ), + _types(f="double precision"), + None, + True, + id="float64-into-double", + ), + pytest.param( + pd.DataFrame( + {"f": [0.1, 0.30000000000000004, 16777217.0, 1e-20, 123456789.123456789, -2.5, np.nan]} + ), + _types(f="real"), + None, + True, + id="float64-into-real", + ), + pytest.param( + pd.DataFrame( + {"f": [0.1, 0.30000000000000004, 1e-12, 123456789.123456789, -2.5, 1e10, np.nan]} + ), + _types(f="numeric(38,20)"), + None, + True, + id="float64-into-numeric-high-scale", + ), + pytest.param( + pd.DataFrame({"b": [True, False, True], "c": [True, None, False]}), + _types(b="boolean", c="boolean"), + None, + True, + id="bool-and-bool-with-nulls", + ), + pytest.param( + pd.DataFrame({"ts": _TS_NAIVE, "tstz": _TS_NAIVE}), + _types(ts="timestamp", tstz="timestamptz"), + None, + True, + id="naive-timestamps", + ), + pytest.param( + pd.DataFrame( + { + "utc": _TS_TZ.tz_localize("UTC"), + "east": _TS_TZ.tz_localize("UTC").tz_convert("US/Eastern"), + "kolkata": _TS_TZ.tz_localize("UTC").tz_convert("Asia/Kolkata"), + } + ), + _types(utc="timestamptz", east="timestamp", kolkata="timestamptz"), + None, + True, + id="tz-aware-timestamps", + ), + pytest.param( + pd.DataFrame({"ts": _ts(["2020-01-01", "2020-01-02", "2020-01-03"])}), + _types(ts="timestamp"), + None, + True, + id="midnight-only-timestamps", + ), + pytest.param( + pd.DataFrame({"d": ["2020-01-01", "1969-12-31", "2024-02-29", None]}), + _types(d="date"), + None, + True, + id="date-from-strings", + ), + pytest.param( + pd.DataFrame({"n": ["1.5", "-0.000000001", "12345678901234567890.123456789", "0", None]}), + _types(n="numeric(38,9)"), + None, + True, + id="numeric-from-strings", + ), + pytest.param( + pd.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"], "ignored": [0, 0, 0]}), + _types(a="int", b="text", c="int"), + ["a", "b"], + True, + id="source-columns-subset-with-missing-target-column", + ), + pytest.param( + pd.DataFrame( + { + "Order": [1, 2], + "ID": ["a", "b"], + "select": [0.5, 1.5], + "with space": [True, False], + 'we"ird': ["x", "y"], + "\u00fcber": [1, 2], + } + ), + _types( + **{ + "Order": "int", + "ID": "text", + "select": "double precision", + "with space": "boolean", + 'we"ird': "text", + "\u00fcber": "int", + } + ), + None, + True, + id="reserved-mixed-case-and-unicode-column-names", + ), + pytest.param( + pd.DataFrame( + { + "f": pd.Series([1.5, None, 0.1, 2.5], dtype=object), + "i": pd.Series([1, None, 3, 2**70], dtype=object), + "m": pd.Series([1, 2.5, None, 3], dtype=object), + "b": pd.Series([True, None, False, True], dtype=object), + } + ), + _types(f="double precision", i="numeric(30,0)", m="numeric(10,2)", b="boolean"), + None, + True, + id="object-columns-with-python-numbers-and-none", + ), + pytest.param( + pd.DataFrame({"s": [None, None], "i": [None, None]}), + _types(s="text", i="int"), + None, + True, + id="all-null-columns", + ), +] + + +@pytest.mark.parametrize("df, types, source_columns, expect_copy", _PARITY_CASES) +@pytest.mark.parametrize("chunk_size", [10000, 3]) +def test_df_copy_matches_values_path( + engine_adapter: PostgresEngineAdapter, + ctx: TestContext, + mocker: MockerFixture, + df: pd.DataFrame, + types: t.Dict[str, exp.DataType], + source_columns: t.Optional[t.List[str]], + expect_copy: bool, + chunk_size: int, +): + """Differential test: COPY and the generic VALUES implementation must produce identical rows.""" + copy_table = ctx.table("parity_copy") + values_table = ctx.table("parity_values") + for table in (copy_table, values_table): + engine_adapter.create_table(table, types) + + engine_adapter.DEFAULT_BATCH_SIZE = chunk_size + spy = mocker.spy(PostgresEngineAdapter, "_copy_df") + engine_adapter.insert_append( + copy_table, df, target_columns_to_types=types, source_columns=source_columns + ) + assert (spy.call_count > 0) is expect_copy + + with _values_path(engine_adapter): + engine_adapter.insert_append( + values_table, df, target_columns_to_types=types, source_columns=source_columns + ) + + _assert_same_rows(engine_adapter, copy_table, values_table) + assert ( + _leftover_staging_tables(engine_adapter, exp.to_table(ctx.schema(TEST_SCHEMA)).name) == [] + ) + + +def test_df_copy_exact_values(engine_adapter: PostgresEngineAdapter, ctx: TestContext): + """Guards against the parity test passing because both paths are equally wrong.""" + table = ctx.table("exact") + types = _types( + id="int", + s="text", + i="int", + f="double precision", + b="boolean", + ts="timestamptz", + n="numeric(38,9)", + ) + engine_adapter.create_table(table, types) + df = pd.DataFrame( + { + "id": [1, 2, 3, 4], + "s": ["", None, "\\N", 'a"b,c\nd'], + "i": [1, None, 3, 4], + "f": [0.1, np.nan, 1e-300, 2.5], + "b": [True, None, False, True], + "ts": _ts(["2020-01-01 00:00:00.123456", None, "1960-01-01", "2020-01-02"]).tz_localize( + "UTC" + ), + "n": ["1.5", None, "2", "-0.000000001"], + } + ) + engine_adapter.insert_append(table, df, target_columns_to_types=types) + + rows = engine_adapter.fetchall( + f"SELECT id, s, i, f, b, to_char(ts AT TIME ZONE 'UTC', 'YYYY-MM-DD HH24:MI:SS.US'), n::text " + f"FROM {table.sql('postgres', identify=True)} ORDER BY id" + ) + assert rows == [ + (1, "", 1, 0.1, True, "2020-01-01 00:00:00.123456", "1.500000000"), + (2, None, None, None, None, None, None), + (3, "\\N", 3, 1e-300, False, "1960-01-01 00:00:00.000000", "2.000000000"), + (4, 'a"b,c\nd', 4, 2.5, True, "2020-01-02 00:00:00.000000", "-0.000000001"), + ] + + +def test_df_copy_nullable_pandas_dtypes(engine_adapter: PostgresEngineAdapter, ctx: TestContext): + """The generic implementation can't load pd.NA at all, COPY handles the nullable dtypes.""" + table = ctx.table("nullable") + types = _types(i="int", j="bigint", f="double precision", b="boolean", s="text") + engine_adapter.create_table(table, types) + df = pd.DataFrame( + { + "i": pd.array([1, None, 3], dtype="Int64"), + "j": pd.array([2**62, -(2**62), None], dtype="Int64"), + "f": pd.array([0.1, None, 2.5], dtype="Float64"), + "b": pd.array([True, None, False], dtype="boolean"), + "s": pd.array(["x", None, ""], dtype="string"), + } + ) + engine_adapter.insert_append(table, df, target_columns_to_types=types) + rows = engine_adapter.fetchall( + f"SELECT * FROM {table.sql('postgres', identify=True)} ORDER BY i NULLS LAST" + ) + assert rows == [ + (1, 2**62, 0.1, True, "x"), + (3, None, 2.5, False, ""), + (None, -(2**62), None, None, None), + ] + + +def test_df_copy_source_columns_in_a_different_order( + engine_adapter: PostgresEngineAdapter, ctx: TestContext +): + table = ctx.table("reordered") + types = _types(a="int", b="text", c="int") + engine_adapter.create_table(table, types) + df = pd.DataFrame({"b": ["x", "y"], "a": [1, 2], "ignored": [0, 0]}) + engine_adapter.insert_append( + table, df, target_columns_to_types=types, source_columns=["b", "a"] + ) + assert engine_adapter.fetchall( + f"SELECT a, b, c FROM {table.sql('postgres', identify=True)} ORDER BY a" + ) == [(1, "x", None), (2, "y", None)] + + +def test_df_copy_infinity(engine_adapter: PostgresEngineAdapter, ctx: TestContext): + if engine_adapter.server_version < (14, 0): + pytest.skip("NUMERIC infinity requires PostgreSQL 14+") + table = ctx.table("inf") + types = _types(f="double precision") + engine_adapter.create_table(table, types) + df = pd.DataFrame({"f": [float("inf"), float("-inf"), 1.0]}) + engine_adapter.insert_append(table, df, target_columns_to_types=types) + rows = engine_adapter.fetchall( + f"SELECT f FROM {table.sql('postgres', identify=True)} ORDER BY f" + ) + assert rows == [(float("-inf"),), (1.0,), (float("inf"),)] + + +def test_df_copy_null_timestamps(engine_adapter: PostgresEngineAdapter, ctx: TestContext): + table = ctx.table("nat") + types = _types(ts="timestamp", tstz="timestamptz") + engine_adapter.create_table(table, types) + df = pd.DataFrame( + { + "ts": _ts(["2020-01-01", None]), + "tstz": _ts(["2020-01-01", None]).tz_localize("UTC"), + } + ) + engine_adapter.insert_append(table, df, target_columns_to_types=types) + rows = engine_adapter.fetchall( + f"SELECT ts, tstz FROM {table.sql('postgres', identify=True)} ORDER BY ts NULLS LAST" + ) + assert rows[1] == (None, None) + assert rows[0][0] is not None and rows[0][1] is not None + + +def test_df_copy_in_non_utf8_client_encoding( + engine_adapter: PostgresEngineAdapter, ctx: TestContext +): + connection = engine_adapter._connection_pool.get() + connection.set_client_encoding("LATIN1") + try: + table = ctx.table("latin1") + types = _types(s="text") + engine_adapter.create_table(table, types) + values = ["h\u00e9llo", "caf\u00e9 \u00fcber", "plain"] + engine_adapter.insert_append( + table, pd.DataFrame({"s": values}), target_columns_to_types=types + ) + rows = engine_adapter.fetchall( + f"SELECT s FROM {table.sql('postgres', identify=True)} ORDER BY s" + ) + assert sorted(r[0] for r in rows) == sorted(values) + finally: + connection.set_client_encoding("UTF8") + + +def test_df_copy_cleans_up_staging_table_on_failed_insert( + engine_adapter: PostgresEngineAdapter, ctx: TestContext, mocker: MockerFixture +): + table = ctx.table("not_null_target") + engine_adapter.execute( + f"CREATE TABLE {table.sql('postgres', identify=True)} (a INT NOT NULL, b TEXT)" + ) + spy = mocker.spy(PostgresEngineAdapter, "_copy_df") + df = pd.DataFrame({"a": [1, None], "b": ["x", "y"]}) + + with pytest.raises(Exception, match='null value in column "a"'): + engine_adapter.insert_append(table, df, target_columns_to_types=_types(a="int", b="text")) + + assert spy.call_count == 1 + assert ( + _leftover_staging_tables(engine_adapter, exp.to_table(ctx.schema(TEST_SCHEMA)).name) == [] + ) + assert engine_adapter.fetchall(f"SELECT * FROM {table.sql('postgres', identify=True)}") == [] + # The connection must still be usable + assert engine_adapter.fetchone("select 1") == (1,) + + +def test_df_copy_cleans_up_staging_table_outside_transaction( + engine_adapter: PostgresEngineAdapter, ctx: TestContext +): + table = ctx.table("outside_txn") + types = _types(a="int") + engine_adapter.create_table(table, types) + source_queries = engine_adapter._df_to_source_queries( + pd.DataFrame({"a": [1, 2]}), types, 10000, table + ) + with source_queries[0] as query: + assert ( + len( + _leftover_staging_tables(engine_adapter, exp.to_table(ctx.schema(TEST_SCHEMA)).name) + ) + == 1 + ) + assert engine_adapter.fetchall(query) == [(1,), (2,)] + assert ( + _leftover_staging_tables(engine_adapter, exp.to_table(ctx.schema(TEST_SCHEMA)).name) == [] + ) + + +def test_df_copy_long_physical_table_name(engine_adapter: PostgresEngineAdapter, ctx: TestContext): + name = ( + "sqlmesh_physical_" + "x" * 30 + "__" + "9" * 12 + ) # 61 characters, a typical physical name + table = ctx.table(name) + assert len(table.name) > 47 + types = _types(a="int") + engine_adapter.create_table(table, types) + engine_adapter.insert_append( + table, pd.DataFrame({"a": [1, 2, 3]}), target_columns_to_types=types + ) + assert engine_adapter.fetchall( + f"SELECT a FROM {table.sql('postgres', identify=True)} ORDER BY a" + ) == [(1,), (2,), (3,)] + assert ( + _leftover_staging_tables(engine_adapter, exp.to_table(ctx.schema(TEST_SCHEMA)).name) == [] + ) + + +def test_df_copy_concurrent_threads( + engine_adapter: PostgresEngineAdapter, ctx: TestContext, config: Config +): + """Every thread has its own connection (and transaction); staging names must not collide.""" + connection_config = config.gateways[ctx.gateway].connection + assert connection_config + # `concurrent_tasks > 1` gives every thread its own connection, like during a plan + adapter = t.cast( + PostgresEngineAdapter, connection_config.create_engine_adapter(concurrent_tasks=6) + ) + types = _types(a="int", b="text") + tables = [ctx.table(f"concurrent_{i}") for i in range(6)] + for table in tables: + engine_adapter.create_table(table, types) + + adapter.DEFAULT_BATCH_SIZE = 300 + + def _load(i: int) -> None: + df = pd.DataFrame({"a": range(i * 1000, i * 1000 + 1000), "b": [f"t{i}"] * 1000}) + adapter.insert_append(tables[i], df, target_columns_to_types=types) + + try: + with ThreadPoolExecutor(max_workers=6) as pool: + list(pool.map(_load, range(6))) + finally: + adapter.close() + + for i, table in enumerate(tables): + rows = engine_adapter.fetchall( + f"SELECT MIN(a), MAX(a), COUNT(*), COUNT(DISTINCT b), MIN(b) FROM {table.sql('postgres', identify=True)}" + ) + assert rows == [(i * 1000, i * 1000 + 999, 1000, 1, f"t{i}")] + assert ( + _leftover_staging_tables(engine_adapter, exp.to_table(ctx.schema(TEST_SCHEMA)).name) == [] + ) + + +def test_df_copy_without_create_privilege_fails_like_other_ddl( + engine_adapter: PostgresEngineAdapter, ctx: TestContext, config: Config +): + schema = exp.to_table(ctx.schema("no_create")).name + engine_adapter.create_schema(schema) + table = exp.to_table(f"{schema}.target") + types = _types(a="int") + engine_adapter.create_table(table, types) + with create_users(engine_adapter, "reader") as roles: + user = roles["reader"]["username"] + engine_adapter.execute(f'GRANT USAGE ON SCHEMA "{schema}" TO "{user}"') + engine_adapter.execute(f'GRANT INSERT, SELECT ON "{schema}"."target" TO "{user}"') + reader = create_engine_adapter_for_role(roles["reader"], ctx, config) + try: + with pytest.raises(Exception, match="permission denied for schema"): + reader.insert_append(table, pd.DataFrame({"a": [1]}), target_columns_to_types=types) + assert reader.fetchone("select 1") == (1,) + finally: + reader.close() + + +def test_df_copy_generated_and_identity_columns( + engine_adapter: PostgresEngineAdapter, ctx: TestContext +): + """Only the listed columns are copied, so identity / generated columns keep working.""" + table = ctx.table("generated") + sql_table = table.sql("postgres", identify=True) + engine_adapter.execute( + f"CREATE TABLE {sql_table} (" + "id INT GENERATED ALWAYS AS IDENTITY, a INT, doubled INT GENERATED ALWAYS AS (a * 2) STORED, b TEXT)" + ) + engine_adapter.insert_append( + table, + pd.DataFrame({"a": [1, 2], "b": ["x", "y"]}), + target_columns_to_types=_types(a="int", b="text"), + ) + assert engine_adapter.fetchall(f"SELECT id, a, doubled, b FROM {sql_table} ORDER BY a") == [ + (1, 1, 2, "x"), + (2, 2, 4, "y"), + ] + + +def test_df_copy_engine_adapter_operations( + engine_adapter: PostgresEngineAdapter, ctx: TestContext, mocker: MockerFixture +): + spy = mocker.spy(PostgresEngineAdapter, "_copy_df") + types = _types(id="int", name="text", ds="text") + + def _read(table: exp.Table) -> t.List[t.Tuple]: + return engine_adapter.fetchall( + f"SELECT * FROM {table.sql('postgres', identify=True)} ORDER BY id" + ) + + # replace_query: creates the table, then replaces its content + replaced = ctx.table("replaced") + engine_adapter.replace_query( + replaced, + pd.DataFrame({"id": [1, 2], "name": ["a", "b"], "ds": ["2020-01-01"] * 2}), + target_columns_to_types=types, + ) + assert _read(replaced) == [(1, "a", "2020-01-01"), (2, "b", "2020-01-01")] + engine_adapter.replace_query( + replaced, + pd.DataFrame({"id": [3], "name": ["c"], "ds": ["2020-01-02"]}), + target_columns_to_types=types, + ) + assert _read(replaced) == [(3, "c", "2020-01-02")] + + # insert_append + engine_adapter.insert_append( + replaced, + pd.DataFrame({"id": [4], "name": [None], "ds": ["2020-01-03"]}), + target_columns_to_types=types, + ) + assert _read(replaced) == [(3, "c", "2020-01-02"), (4, None, "2020-01-03")] + + # ctas from a DataFrame + ctas = ctx.table("ctas") + engine_adapter.ctas( + ctas, + pd.DataFrame({"id": [7], "name": ["z"], "ds": ["2020-01-07"]}), + target_columns_to_types=types, + ) + assert _read(ctas) == [(7, "z", "2020-01-07")] + + # insert_overwrite_by_time_partition + partitioned = ctx.table("partitioned") + engine_adapter.create_table(partitioned, types) + engine_adapter.insert_append( + partitioned, + pd.DataFrame( + {"id": [1, 2, 3], "name": list("abc"), "ds": ["2020-01-01", "2020-01-02", "2020-01-03"]} + ), + target_columns_to_types=types, + ) + engine_adapter.insert_overwrite_by_time_partition( + partitioned, + pd.DataFrame({"id": [20, 30], "name": ["B", "C"], "ds": ["2020-01-02", "2020-01-03"]}), + start="2020-01-02", + end="2020-01-03", + time_formatter=lambda x, _: exp.Literal.string(to_ds(x)), + time_column="ds", + target_columns_to_types=types, + ) + assert _read(partitioned) == [ + (1, "a", "2020-01-01"), + (20, "B", "2020-01-02"), + (30, "C", "2020-01-03"), + ] + + # merge: native (PG >= 15) and the logical implementation used for PG < 15 + for version in (engine_adapter.server_version, (14, 0)): + merged = ctx.table(f"merged_{version[0]}") + engine_adapter.create_table(merged, types) + engine_adapter.insert_append( + merged, + pd.DataFrame({"id": [1, 2], "name": ["a", "b"], "ds": ["x", "x"]}), + target_columns_to_types=types, + ) + engine_adapter.server_version = version # type: ignore + engine_adapter.merge( + merged, + pd.DataFrame({"id": [2, 3], "name": ["B", "C"], "ds": ["y", "y"]}), + target_columns_to_types=types, + unique_key=[exp.to_column("id", quoted=True)], + ) + assert _read(merged) == [(1, "a", "x"), (2, "B", "y"), (3, "C", "y")] + + assert spy.call_count >= 8 + assert ( + _leftover_staging_tables(engine_adapter, exp.to_table(ctx.schema(TEST_SCHEMA)).name) == [] + ) + + +def test_df_copy_python_and_seed_models(ctx: TestContext, tmp_path: Path, mocker: MockerFixture): + engine_adapter = t.cast(PostgresEngineAdapter, ctx.engine_adapter) + spy = mocker.spy(PostgresEngineAdapter, "_copy_df") + schema = exp.to_table(ctx.schema(TEST_SCHEMA)).this + models_dir = tmp_path / "models" + seeds_dir = tmp_path / "seeds" + models_dir.mkdir() + seeds_dir.mkdir() + + (seeds_dir / "seed.csv").write_text('id,name,score\n1,alice,1.5\n2,"b,ob",\n3,\\N,3.25\n') + (models_dir / "seed_model.sql").write_text(f""" + MODEL ( + name {schema}.seed_model, + kind SEED (path '../seeds/seed.csv'), + columns (id int, name text, score double) + ); + """) + (models_dir / "py_model.py").write_text(f""" +import typing as t +from datetime import datetime + +import pandas as pd +from sqlmesh import ExecutionContext, model +from sqlmesh.core.model.kind import ModelKindName + + +@model( + "{schema}.py_model", + kind=dict(name=ModelKindName.INCREMENTAL_BY_TIME_RANGE, time_column="ds", batch_size=1), + columns={{"id": "int", "name": "text", "value": "double", "ds": "text"}}, + start="2020-01-01", + cron="@daily", +) +def execute(context: ExecutionContext, start: datetime, end: datetime, execution_time: datetime, **kwargs: t.Any) -> pd.DataFrame: + ds = start.strftime("%Y-%m-%d") + return pd.DataFrame( + {{ + "id": [1, 2, 3], + "name": [f"n{{ds}}", None, 'q"uote,'], + "value": [0.1, float("nan"), 2.5], + "ds": [ds] * 3, + }} + ) +""") + + sqlmesh = ctx.create_context(path=tmp_path) + with time_machine.travel("2020-01-04 00:00:00"): + sqlmesh.plan(auto_apply=True, no_prompts=True) + + py_rows = engine_adapter.fetchall( + f'SELECT ds, COUNT(*), COUNT(name), COUNT(value), MIN(id) FROM "{schema}"."py_model" GROUP BY ds ORDER BY ds' + ) + assert py_rows == [(f"2020-01-0{d}", 3, 2, 2, 1) for d in (1, 2, 3)] + assert engine_adapter.fetchall( + f'SELECT id, name, score FROM "{schema}"."seed_model" ORDER BY id' + ) == [ + (1, "alice", 1.5), + (2, "b,ob", None), + (3, "\\N", 3.25), + ] + assert spy.call_count >= 4 + assert ( + _leftover_staging_tables(engine_adapter, exp.to_table(ctx.schema(TEST_SCHEMA)).name) == [] + ) + + +@pytest.mark.parametrize("operation", ["replace_query", "insert_append"]) +def test_df_copy_failure_leaves_no_staging_table( + engine_adapter: PostgresEngineAdapter, ctx: TestContext, operation: str +): + """A failing COPY (here: a NUL character) must not leak the staging table, outside a transaction too.""" + table = ctx.table("copy_failure") + types = _types(x="text", y="int") + engine_adapter.create_table(table, types) + df = pd.DataFrame({"x": ["a\x00b"], "y": [1]}) + + with pytest.raises(Exception, match="0x00|null character|NUL"): + getattr(engine_adapter, operation)(table, df, target_columns_to_types=types) + + assert ( + _leftover_staging_tables(engine_adapter, exp.to_table(ctx.schema(TEST_SCHEMA)).name) == [] + ) + assert engine_adapter.fetchone("select 1") == (1,) diff --git a/tests/core/engine_adapter/test_base.py b/tests/core/engine_adapter/test_base.py index 1971ba3bbc..4cff662313 100644 --- a/tests/core/engine_adapter/test_base.py +++ b/tests/core/engine_adapter/test_base.py @@ -3489,6 +3489,40 @@ def test_get_temp_table(mocker: MockerFixture, make_mocked_engine_adapter: t.Cal assert value.sql() == '"catalog"."db"."__temp_test_table_abcdefgh"' +@pytest.mark.parametrize( + "max_length, table_name, expected", + [ + # The limit is only enforced when the engine defines one + ( + None, + "a_very_long_table_name_" + "x" * 40, + "__temp_a_very_long_table_name_" + "x" * 40 + "_abcdefgh", + ), + # Exactly at the limit: unchanged + (30, "t" * 14, "__temp_" + "t" * 14 + "_abcdefgh"), + # Over the limit: prefix and random suffix are kept, the middle of the name is cut + (30, "aaaaaaaa_bbbbbbbb", "__temp_aaaaaaa" + "bbbbbbb" + "_abcdefgh"), + # The limit is so small that nothing of the name fits + (10, "name", "__temp__abcdefgh"), + ], +) +def test_get_temp_table_respects_max_identifier_length( + mocker: MockerFixture, + make_mocked_engine_adapter: t.Callable, + max_length: t.Optional[int], + table_name: str, + expected: str, +): + adapter = make_mocked_engine_adapter(EngineAdapter) + adapter.MAX_IDENTIFIER_LENGTH = max_length + mocker.patch("sqlmesh.core.engine_adapter.base.random_id", return_value="abcdefgh") + + value = adapter._get_temp_table(f"catalog.db.{table_name}") + + assert value.name == expected + assert (value.catalog, value.db) == ("catalog", "db") + + def test_get_data_objects_batching(mocker: MockerFixture, make_mocked_engine_adapter: t.Callable): adapter = make_mocked_engine_adapter(EngineAdapter) adapter._get_data_objects = mocker.Mock(return_value=[]) diff --git a/tests/core/engine_adapter/test_postgres.py b/tests/core/engine_adapter/test_postgres.py index ebcdd03f55..568d7e133c 100644 --- a/tests/core/engine_adapter/test_postgres.py +++ b/tests/core/engine_adapter/test_postgres.py @@ -1,5 +1,10 @@ +import re import typing as t +from datetime import date +from decimal import Decimal +import numpy as np # noqa: TID253 +import pandas as pd # noqa: TID253 import pytest from pytest_mock import MockFixture from pytest_mock.plugin import MockerFixture @@ -282,3 +287,290 @@ def test_sync_grants_config_with_default_schema( "WHERE table_schema = 'public' AND table_name = 'test_table' " "AND grantor = current_role AND grantee <> current_role" ) + + +# -------------------------------------------------------------------------------------------- +# DataFrame bulk loading via COPY +# -------------------------------------------------------------------------------------------- + +_NULL_MARKER_RE = re.compile(r"NULL '([^']+)'") + + +def _capture_copy_calls(adapter: PostgresEngineAdapter) -> t.List[t.Tuple[str, str]]: + """Records (sql, payload) for every `copy_expert` call; the payload is read at call time.""" + calls: t.List[t.Tuple[str, str]] = [] + + def _copy_expert(sql: str, file: t.Any) -> None: + calls.append((sql, file.read())) + + adapter.cursor.copy_expert.side_effect = _copy_expert + return calls + + +@pytest.fixture +def copy_adapter( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture, make_temp_table_name: t.Callable +) -> PostgresEngineAdapter: + adapter = make_mocked_engine_adapter(PostgresEngineAdapter) + mocker.patch.object(adapter, "table_exists", return_value=False) + mocker.patch( + "sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table", + return_value=make_temp_table_name("tbl", "abcdefgh"), + ) + return adapter + + +def test_df_copy_chunks_sql_and_payload(copy_adapter: PostgresEngineAdapter): + calls = _capture_copy_calls(copy_adapter) + df = pd.DataFrame( + { + "a": [1, 2, 3], + "b": ["x", None, 'q"uote, comma'], + "c": [1.5, np.nan, 0.1], + "d": [True, False, True], + "e": pd.to_datetime( + ["2020-01-01 00:00:00.123456", None, "1960-01-01 12:00:00"], # type: ignore + format="ISO8601", + ), + } + ) + types = { + "a": exp.DataType.build("int"), + "b": exp.DataType.build("text"), + "c": exp.DataType.build("double"), + "d": exp.DataType.build("boolean"), + "e": exp.DataType.build("timestamp"), + } + + copy_adapter.DEFAULT_BATCH_SIZE = 2 + copy_adapter.insert_append("tbl", df, target_columns_to_types=types) + + # One COPY per chunk of `batch_size` rows + assert len(calls) == 2 + marker = _NULL_MARKER_RE.search(calls[0][0]).group(1) # type: ignore + assert len(marker) <= 32 # pandas truncates longer `na_rep` values in float columns + expected_sql = ( + 'COPY "__temp_tbl_abcdefgh" ("a", "b", "c", "d", "e") ' + f"FROM STDIN WITH (FORMAT csv, NULL '{marker}')" + ) + assert [sql for sql, _ in calls] == [expected_sql, expected_sql] + assert ( + calls[0][1] + == f"1,x,1.5,True,2020-01-01 00:00:00.123456\n2,{marker},{marker},False,{marker}\n" + ) + assert calls[1][1] == '3,"q""uote, comma",0.1,True,1960-01-01 12:00:00\n' + + assert to_sql_calls(copy_adapter) == [ + 'CREATE TABLE IF NOT EXISTS "__temp_tbl_abcdefgh" ' + '("a" BIGINT, "b" TEXT, "c" DECIMAL, "d" BOOLEAN, "e" TIMESTAMP)', + 'INSERT INTO "tbl" ("a", "b", "c", "d", "e") ' + 'SELECT CAST("a" AS INT) AS "a", CAST("b" AS TEXT) AS "b", CAST("c" AS DOUBLE PRECISION) AS "c", ' + 'CAST("d" AS BOOLEAN) AS "d", CAST("e" AS TIMESTAMP) AS "e" FROM "__temp_tbl_abcdefgh"', + 'DROP TABLE IF EXISTS "__temp_tbl_abcdefgh"', + ] + + +def test_df_copy_batch_size_zero_is_single_chunk(copy_adapter: PostgresEngineAdapter): + calls = _capture_copy_calls(copy_adapter) + df = pd.DataFrame({"a": [1, 2, 3]}) + copy_adapter.DEFAULT_BATCH_SIZE = 0 + copy_adapter.insert_append("tbl", df, target_columns_to_types={"a": exp.DataType.build("int")}) + assert len(calls) == 1 + assert calls[0][1] == "1\n2\n3\n" + + +def test_df_copy_reorders_columns_and_skips_missing_source_columns( + copy_adapter: PostgresEngineAdapter, +): + calls = _capture_copy_calls(copy_adapter) + df = pd.DataFrame({"b": ["x", "y"], "a": [1, 2], "unused": [9, 9]}) + copy_adapter.insert_append( + "tbl", + df, + target_columns_to_types={ + "a": exp.DataType.build("int"), + "b": exp.DataType.build("text"), + "c": exp.DataType.build("int"), + }, + source_columns=["b", "a"], + ) + + # Only the source columns are copied, in the source order + assert len(calls) == 1 + assert calls[0][0].startswith('COPY "__temp_tbl_abcdefgh" ("b", "a") FROM STDIN') + assert calls[0][1] == "x,1\ny,2\n" + assert to_sql_calls(copy_adapter)[1] == ( + 'INSERT INTO "tbl" ("a", "b", "c") SELECT CAST("a" AS INT) AS "a", CAST("b" AS TEXT) AS "b", ' + 'CAST(NULL AS INT) AS "c" FROM "__temp_tbl_abcdefgh"' + ) + + +def test_df_copy_staging_table_is_created_once( + copy_adapter: PostgresEngineAdapter, mocker: MockerFixture +): + calls = _capture_copy_calls(copy_adapter) + df = pd.DataFrame({"a": [1, 2]}) + source_queries = copy_adapter._df_to_source_queries( + df, {"a": exp.DataType.build("int")}, 10000, "tbl" + ) + assert len(source_queries) == 1 + + # Once the staging table exists, calling the factory again must not reload the data + with source_queries[0]: + pass + copy_adapter.table_exists.return_value = True # type: ignore + with source_queries[0]: + pass + assert len(calls) == 1 + + +@pytest.mark.parametrize( + "df, remove_copy_expert", + [ + (pd.DataFrame({"a": [1, 2]}), True), # driver without copy_expert (e.g. pg8000) + (pd.DataFrame({"a": [{"k": 1}, {"k": 2}]}), False), # dict values + (pd.DataFrame({"a": [[1], [2]]}), False), # list values + (pd.DataFrame({"a": ["x", 1]}), False), # mixed str / int + (pd.DataFrame({"a": [b"x", b"y"]}), False), # bytes + (pd.DataFrame({"a": [Decimal("1.5"), Decimal("2")]}), False), # Decimal + (pd.DataFrame({"a": [date(2020, 1, 1), date(2020, 1, 2)]}), False), # date objects + (pd.DataFrame({"a": pd.Categorical(["x", "y"])}), False), # unsupported dtype + (pd.DataFrame({"a": pd.to_timedelta([1, 2], unit="s")}), False), # unsupported dtype + (pd.DataFrame({"a": np.array([1, 2], dtype="uint64")}), False), # may exceed BIGINT + ( + pd.DataFrame({"a": np.array([1, 2], dtype="float32")}), + False, + ), # repr differs from float64 + (pd.DataFrame({"a": ["x", "\\."]}), False), # `\.` can end COPY data in CSV format + ], +) +def test_df_copy_falls_back_to_values( + copy_adapter: PostgresEngineAdapter, df: pd.DataFrame, remove_copy_expert: bool +): + if remove_copy_expert: + del copy_adapter.cursor.copy_expert + + copy_adapter.insert_append("tbl", df, target_columns_to_types={"a": exp.DataType.build("text")}) + + if not remove_copy_expert: + copy_adapter.cursor.copy_expert.assert_not_called() + sql_calls = to_sql_calls(copy_adapter) + assert len(sql_calls) == 1 + assert "VALUES" in sql_calls[0] + assert "__temp_" not in sql_calls[0] + + +def test_df_copy_falls_back_on_null_marker_collision( + copy_adapter: PostgresEngineAdapter, mocker: MockerFixture +): + mocker.patch("sqlmesh.core.engine_adapter.postgres.random_id", return_value="collide") + marker = "__sqlmesh_null_collide__" + df = pd.DataFrame({"a": ["x", marker]}) + + copy_adapter.insert_append("tbl", df, target_columns_to_types={"a": exp.DataType.build("text")}) + + copy_adapter.cursor.copy_expert.assert_not_called() + sql_calls = to_sql_calls(copy_adapter) + assert len(sql_calls) == 1 + assert "VALUES" in sql_calls[0] + + +def test_df_copy_not_used_for_empty_dataframe(copy_adapter: PostgresEngineAdapter): + # The adapter itself must not create a staging table for an empty DataFrame + df = pd.DataFrame({"a": pd.Series([], dtype="int64")}) + source_queries = copy_adapter._df_to_source_queries( + df, {"a": exp.DataType.build("int")}, 10000, "tbl" + ) + assert source_queries == [] + copy_adapter.cursor.copy_expert.assert_not_called() + assert to_sql_calls(copy_adapter) == [] + + +def test_df_copy_driver_call_is_isolated_in_copy_df( + copy_adapter: PostgresEngineAdapter, mocker: MockerFixture +): + """Only `_copy_df` may touch the driver's COPY API (keeps a psycopg3 migration small).""" + copy_df = mocker.patch.object(copy_adapter, "_copy_df") + df = pd.DataFrame({"a": [1, 2]}) + copy_adapter.insert_append("tbl", df, target_columns_to_types={"a": exp.DataType.build("int")}) + copy_df.assert_called_once() + copy_adapter.cursor.copy_expert.assert_not_called() + + +def test_temp_table_name_respects_max_identifier_length( + make_mocked_engine_adapter: t.Callable, +): + adapter = make_mocked_engine_adapter(PostgresEngineAdapter) + table = f"sqlmesh__schema.{'x' * 30}__{'y' * 28}" # 60 character table name + assert len(exp.to_table(table).name) == 60 + + first = adapter._get_temp_table(table) + second = adapter._get_temp_table(table) + + assert first.db == "sqlmesh__schema" + for temp in (first, second): + assert len(temp.name) <= 63 + assert temp.name.startswith("__temp_xxxxxx") # head of the name is kept + assert temp.name.rsplit("_", 1)[0].endswith("yyyyyy") # tail of the name is kept + assert first.name != second.name # random suffix keeps names unique + # Identifier-length validation on DDL must accept the staging name + adapter._check_identifier_length(exp.Create(this=first, kind="TABLE")) + + +def test_temp_table_name_unchanged_when_short(make_mocked_engine_adapter: t.Callable): + adapter = make_mocked_engine_adapter(PostgresEngineAdapter) + name = "z" * 47 # 7 prefix + 47 + 1 + 8 suffix == 63 + temp = adapter._get_temp_table(f"s.{name}") + assert temp.name.startswith(f"__temp_{name}_") + assert len(temp.name) == 63 + + +class _FakeDbError(Exception): + def __init__(self, pgcode: t.Optional[str]) -> None: + super().__init__("boom") + self.pgcode = pgcode + + +@pytest.mark.parametrize( + "pgcode, in_transaction, raises", + [ + ("25P02", True, False), # aborted transaction: its rollback removes the staging table + ("25P02", False, True), + ("42501", True, True), + (None, True, True), + ], +) +def test_drop_copy_staging_table_error_handling( + copy_adapter: PostgresEngineAdapter, + mocker: MockerFixture, + pgcode: t.Optional[str], + in_transaction: bool, + raises: bool, +): + mocker.patch.object(copy_adapter, "drop_table", side_effect=_FakeDbError(pgcode)) + mocker.patch.object( + type(copy_adapter._connection_pool), + "is_transaction_active", + new_callable=mocker.PropertyMock, + return_value=in_transaction, + ) + table = exp.to_table("tbl") + if raises: + with pytest.raises(_FakeDbError): + copy_adapter._drop_copy_staging_table(table) + else: + copy_adapter._drop_copy_staging_table(table) + + +def test_df_copy_failure_drops_staging_table_and_reraises( + copy_adapter: PostgresEngineAdapter, mocker: MockerFixture +): + mocker.patch.object(copy_adapter, "_copy_df", side_effect=ValueError("copy failed")) + df = pd.DataFrame({"a": [1, 2]}) + source_queries = copy_adapter._df_to_source_queries( + df, {"a": exp.DataType.build("int")}, 10000, "tbl" + ) + # `replace_query` calls the factory outside of `with source_query`, so the factory must clean up + with pytest.raises(ValueError, match="copy failed"): + source_queries[0].query_factory() + assert to_sql_calls(copy_adapter)[-1] == 'DROP TABLE IF EXISTS "__temp_tbl_abcdefgh"' diff --git a/tests/core/engine_adapter/test_risingwave.py b/tests/core/engine_adapter/test_risingwave.py index ed3cd77a3f..2711b5af16 100644 --- a/tests/core/engine_adapter/test_risingwave.py +++ b/tests/core/engine_adapter/test_risingwave.py @@ -84,3 +84,18 @@ def test_drop_view(adapter: t.Callable): call('DROP VIEW IF EXISTS "db"."view"'), ] ) + + +def test_df_is_not_loaded_with_copy(adapter: t.Callable): + """RisingWave inherits from Postgres but must keep using the VALUES path.""" + import pandas as pd # noqa: TID253 + + risingwave = adapter + risingwave.insert_append( + "tbl", pd.DataFrame({"a": [1, 2]}), target_columns_to_types={"a": exp.DataType.build("int")} + ) + + risingwave.cursor.copy_expert.assert_not_called() + sql_calls = [call_[0][0] for call_ in risingwave.cursor.execute.call_args_list] + assert len(sql_calls) == 1 + assert "VALUES" in str(sql_calls[0]) and "__temp_" not in str(sql_calls[0])