sqlmesh 0.183.1.dev11__tar.gz → 0.183.1.dev13__tar.gz

{sqlmesh-0.183.1.dev11 → sqlmesh-0.183.1.dev13}/.circleci/continue_config.yml

@@ -281,7 +281,6 @@ workflows:
                 - trino
                 - spark
                 - clickhouse
-                - clickhouse-cluster
                 - risingwave
       - engine_tests_cloud:
           name: cloud_engine_<< matrix.engine >>

{sqlmesh-0.183.1.dev11 → sqlmesh-0.183.1.dev13}/Makefile

@@ -126,54 +126,51 @@ engine-%-down:
 ##################
 
 clickhouse-test: engine-clickhouse-up
-	pytest -n auto -x -m "clickhouse" --retries 3 --junitxml=test-results/junit-clickhouse.xml
-
-clickhouse-cluster-test: engine-clickhouse-up
-	pytest -n auto -x -m "clickhouse_cluster" --retries 3 --junitxml=test-results/junit-clickhouse-cluster.xml
+	pytest -n auto -m "clickhouse" --retries 3 --junitxml=test-results/junit-clickhouse.xml
 
 duckdb-test: engine-duckdb-install
-	pytest -n auto -x -m "duckdb" --retries 3 --junitxml=test-results/junit-duckdb.xml
+	pytest -n auto -m "duckdb" --retries 3 --junitxml=test-results/junit-duckdb.xml
 
 mssql-test: engine-mssql-up
-	pytest -n auto -x -m "mssql" --retries 3 --junitxml=test-results/junit-mssql.xml
+	pytest -n auto -m "mssql" --retries 3 --junitxml=test-results/junit-mssql.xml
 
 mysql-test: engine-mysql-up
-	pytest -n auto -x -m "mysql" --retries 3 --junitxml=test-results/junit-mysql.xml
+	pytest -n auto -m "mysql" --retries 3 --junitxml=test-results/junit-mysql.xml
 
 postgres-test: engine-postgres-up
-	pytest -n auto -x -m "postgres" --retries 3 --junitxml=test-results/junit-postgres.xml
+	pytest -n auto -m "postgres" --retries 3 --junitxml=test-results/junit-postgres.xml
 
 spark-test: engine-spark-up
-	pytest -n auto -x -m "spark or pyspark" --retries 3 --junitxml=test-results/junit-spark.xml
+	pytest -n auto -m "spark" --retries 3 --junitxml=test-results/junit-spark.xml
 
 trino-test: engine-trino-up
-	pytest -n auto -x -m "trino or trino_iceberg or trino_delta or trino_nessie" --retries 3 --junitxml=test-results/junit-trino.xml
+	pytest -n auto -m "trino" --retries 3 --junitxml=test-results/junit-trino.xml
 
 risingwave-test: engine-risingwave-up
-	pytest -n auto -x -m "risingwave" --retries 3 --junitxml=test-results/junit-risingwave.xml
+	pytest -n auto -m "risingwave" --retries 3 --junitxml=test-results/junit-risingwave.xml
 
 #################
 # Cloud Engines #
 #################
 
 snowflake-test: guard-SNOWFLAKE_ACCOUNT guard-SNOWFLAKE_WAREHOUSE guard-SNOWFLAKE_DATABASE guard-SNOWFLAKE_USER guard-SNOWFLAKE_PASSWORD engine-snowflake-install
-	pytest -n auto -x -m "snowflake" --retries 3 --junitxml=test-results/junit-snowflake.xml
+	pytest -n auto -m "snowflake" --retries 3 --junitxml=test-results/junit-snowflake.xml
 
 bigquery-test: guard-BIGQUERY_KEYFILE engine-bigquery-install
-	pytest -n auto -x -m "bigquery" --retries 3 --junitxml=test-results/junit-bigquery.xml
+	pytest -n auto -m "bigquery" --retries 3 --junitxml=test-results/junit-bigquery.xml
 
 databricks-test: guard-DATABRICKS_CATALOG guard-DATABRICKS_SERVER_HOSTNAME guard-DATABRICKS_HTTP_PATH guard-DATABRICKS_ACCESS_TOKEN guard-DATABRICKS_CONNECT_VERSION engine-databricks-install
 	pip install 'databricks-connect==${DATABRICKS_CONNECT_VERSION}'
-	pytest -n auto -x -m "databricks" --retries 3 --junitxml=test-results/junit-databricks.xml
+	pytest -n auto -m "databricks" --retries 3 --junitxml=test-results/junit-databricks.xml
 
 redshift-test: guard-REDSHIFT_HOST guard-REDSHIFT_USER guard-REDSHIFT_PASSWORD guard-REDSHIFT_DATABASE engine-redshift-install
-	pytest -n auto -x -m "redshift" --retries 3 --junitxml=test-results/junit-redshift.xml
+	pytest -n auto -m "redshift" --retries 3 --junitxml=test-results/junit-redshift.xml
 
 clickhouse-cloud-test: guard-CLICKHOUSE_CLOUD_HOST guard-CLICKHOUSE_CLOUD_USERNAME guard-CLICKHOUSE_CLOUD_PASSWORD engine-clickhouse-install
 	pytest -n 1 -m "clickhouse_cloud" --retries 3 --junitxml=test-results/junit-clickhouse-cloud.xml
 
 athena-test: guard-AWS_ACCESS_KEY_ID guard-AWS_SECRET_ACCESS_KEY guard-ATHENA_S3_WAREHOUSE_LOCATION engine-athena-install
-	pytest -n auto -x -m "athena" --retries 3 --retry-delay 10 --junitxml=test-results/junit-athena.xml
+	pytest -n auto -m "athena" --retries 3 --retry-delay 10 --junitxml=test-results/junit-athena.xml
 
 vscode_settings:
 	mkdir -p .vscode

{sqlmesh-0.183.1.dev11 → sqlmesh-0.183.1.dev13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlmesh
-Version: 0.183.1.dev11
+Version: 0.183.1.dev13
 Summary: Next-generation data transformation framework
 Author-email: "TobikoData Inc." <engineering@tobikodata.com>
 License:                                  Apache License

{sqlmesh-0.183.1.dev11 → sqlmesh-0.183.1.dev13}/pytest.ini

@@ -16,13 +16,11 @@ markers =
     github: test for Github CI/CD bot
     jupyter: tests for Jupyter integration
     web: tests for web UI
-    spark_pyspark: test for Spark with PySpark dependency
     # Engine Adapters
     engine: test all engine adapters
     athena: test for Athena
     bigquery: test for BigQuery
-    clickhouse: test for Clickhouse (standalone mode)
-    clickhouse_cluster: test for Clickhouse (cluster mode)
+    clickhouse: test for Clickhouse (standalone mode / cluster mode)
     clickhouse_cloud: test for Clickhouse (cloud mode)
     databricks: test for Databricks
     duckdb: test for DuckDB
@@ -33,9 +31,8 @@ markers =
     redshift: test for Redshift
     snowflake: test for Snowflake
     spark: test for Spark
-    trino: test for Trino (Hive connector)
-    trino_iceberg: test for Trino (Iceberg connector)
-    trino_delta: test for Trino (Delta connector)
+    trino: test for Trino (all connectors)
+    risingwave: test for Risingwave
 addopts = -n 0 --dist=loadgroup
 
 asyncio_default_fixture_loop_scope = session

{sqlmesh-0.183.1.dev11 → sqlmesh-0.183.1.dev13}/sqlmesh/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.183.1.dev11'
-__version_tuple__ = version_tuple = (0, 183, 1, 'dev11')
+__version__ = version = '0.183.1.dev13'
+__version_tuple__ = version_tuple = (0, 183, 1, 'dev13')

{sqlmesh-0.183.1.dev11 → sqlmesh-0.183.1.dev13}/sqlmesh/core/engine_adapter/base.py

@@ -1008,6 +1008,8 @@ class EngineAdapter:
         if materialized_properties and not materialized:
             raise SQLMeshError("Materialized properties are only supported for materialized views")
 
+        query_or_df = self._native_df_to_pandas_df(query_or_df)
+
         if isinstance(query_or_df, pd.DataFrame):
             values: t.List[t.Tuple[t.Any, ...]] = list(
                 query_or_df.itertuples(index=False, name=None)
@@ -2002,6 +2004,19 @@ class EngineAdapter:
             self.execute(query, quote_identifiers=quote_identifiers)
             return self.cursor.fetchdf()
 
+    def _native_df_to_pandas_df(
+        self,
+        query_or_df: QueryOrDF,
+    ) -> t.Union[Query, pd.DataFrame]:
+        """
+        Take a "native" DataFrame (eg Pyspark, Bigframe, Snowpark etc) and convert it to Pandas
+        """
+        if isinstance(query_or_df, (exp.Query, exp.DerivedTable, pd.DataFrame)):
+            return query_or_df
+
+        # EngineAdapter subclasses that have native DataFrame types should override this
+        raise NotImplementedError(f"Unable to convert {type(query_or_df)} to Pandas")
+
     def fetchdf(
         self, query: t.Union[exp.Expression, str], quote_identifiers: bool = False
     ) -> pd.DataFrame:

{sqlmesh-0.183.1.dev11 → sqlmesh-0.183.1.dev13}/sqlmesh/core/engine_adapter/bigquery.py

@@ -26,6 +26,7 @@ from sqlmesh.core.schema_diff import SchemaDiffer
 from sqlmesh.utils import optional_import
 from sqlmesh.utils.date import to_datetime
 from sqlmesh.utils.errors import SQLMeshError
+from sqlmesh.utils.pandas import columns_to_types_from_dtypes
 
 if t.TYPE_CHECKING:
     from google.api_core.retry import Retry
@@ -1107,6 +1108,39 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin, ClusteredByMixin, Row
     def _normalize_nested_value(self, col: exp.Expression) -> exp.Expression:
         return exp.func("TO_JSON_STRING", col, dialect=self.dialect)
 
+    @t.overload
+    def _columns_to_types(
+        self, query_or_df: DF, columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None
+    ) -> t.Dict[str, exp.DataType]: ...
+
+    @t.overload
+    def _columns_to_types(
+        self, query_or_df: Query, columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None
+    ) -> t.Optional[t.Dict[str, exp.DataType]]: ...
+
+    def _columns_to_types(
+        self, query_or_df: QueryOrDF, columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None
+    ) -> t.Optional[t.Dict[str, exp.DataType]]:
+        if (
+            not columns_to_types
+            and bigframes
+            and isinstance(query_or_df, bigframes.dataframe.DataFrame)
+        ):
+            # using dry_run=True attempts to prevent the DataFrame from being materialized just to read the column types from it
+            dtypes = query_or_df.to_pandas(dry_run=True).columnDtypes
+            return columns_to_types_from_dtypes(dtypes.items())
+
+        return super()._columns_to_types(query_or_df, columns_to_types)
+
+    def _native_df_to_pandas_df(
+        self,
+        query_or_df: QueryOrDF,
+    ) -> t.Union[Query, pd.DataFrame]:
+        if bigframes and isinstance(query_or_df, bigframes.dataframe.DataFrame):
+            return query_or_df.to_pandas()
+
+        return super()._native_df_to_pandas_df(query_or_df)
+
     @property
     def _query_data(self) -> t.Any:
         return self._connection_pool.get_attribute("query_data")

{sqlmesh-0.183.1.dev11 → sqlmesh-0.183.1.dev13}/sqlmesh/core/engine_adapter/snowflake.py

@@ -27,6 +27,7 @@ from sqlmesh.core.engine_adapter.shared import (
 from sqlmesh.core.schema_diff import SchemaDiffer
 from sqlmesh.utils import optional_import
 from sqlmesh.utils.errors import SQLMeshError
+from sqlmesh.utils.pandas import columns_to_types_from_dtypes
 
 logger = logging.getLogger(__name__)
 snowpark = optional_import("snowflake.snowpark")
@@ -298,13 +299,23 @@ class SnowflakeEngineAdapter(GetCurrentCatalogFromFunctionMixin, ClusteredByMixi
 
             if is_snowpark_dataframe:
                 temp_table.set("catalog", database)
-                df_renamed = df.rename(
-                    {
-                        col: exp.to_identifier(col).sql(dialect=self.dialect, identify=True)
-                        for col in columns_to_types
-                    }
-                )  # type: ignore
-                df_renamed.createOrReplaceTempView(
+
+                # only quote columns if they arent already quoted
+                # if the Snowpark dataframe was created from a Pandas dataframe via snowpark.create_dataframe(pandas_df),
+                # then they will be quoted already. But if the Snowpark dataframe was created manually by the user, then the
+                # columns may not be quoted
+                columns_already_quoted = all(
+                    col.startswith('"') and col.endswith('"') for col in df.columns
+                )
+                local_df = df
+                if not columns_already_quoted:
+                    local_df = df.rename(
+                        {
+                            col: exp.to_identifier(col).sql(dialect=self.dialect, identify=True)
+                            for col in columns_to_types
+                        }
+                    )  # type: ignore
+                local_df.createOrReplaceTempView(
                     temp_table.sql(dialect=self.dialect, identify=True)
                 )  # type: ignore
             elif isinstance(df, pd.DataFrame):
@@ -356,6 +367,11 @@ class SnowflakeEngineAdapter(GetCurrentCatalogFromFunctionMixin, ClusteredByMixi
 
         def cleanup() -> None:
             if is_snowpark_dataframe:
+                if hasattr(df, "table_name"):
+                    if isinstance(df.table_name, str):
+                        # created by the Snowpark library if the Snowpark DataFrame was created from a Pandas DataFrame
+                        # (if the Snowpark DataFrame was created via native means then there is no 'table_name' property and no temp table)
+                        self.drop_table(df.table_name)
                 self.drop_view(temp_table)
             else:
                 self.drop_table(temp_table)
@@ -381,6 +397,15 @@ class SnowflakeEngineAdapter(GetCurrentCatalogFromFunctionMixin, ClusteredByMixi
             columns = self.cursor._result_set.batches[0].column_names
             return pd.DataFrame([dict(zip(columns, row)) for row in rows])
 
+    def _native_df_to_pandas_df(
+        self,
+        query_or_df: QueryOrDF,
+    ) -> t.Union[Query, pd.DataFrame]:
+        if snowpark and isinstance(query_or_df, snowpark.DataFrame):
+            return query_or_df.to_pandas()
+
+        return super()._native_df_to_pandas_df(query_or_df)
+
     def _get_data_objects(
         self, schema_name: SchemaName, object_names: t.Optional[t.Set[str]] = None
     ) -> t.List[DataObject]:
@@ -426,6 +451,10 @@ class SnowflakeEngineAdapter(GetCurrentCatalogFromFunctionMixin, ClusteredByMixi
         if object_names:
             query = query.where(exp.column("TABLE_NAME").isin(*object_names))
 
+        # exclude SNOWPARK_TEMP_TABLE tables that are managed by the Snowpark library and are an implementation
+        # detail of dealing with DataFrame's
+        query = query.where(exp.column("TABLE_NAME").like("SNOWPARK_TEMP_TABLE%").not_())
+
         df = self.fetchdf(query, quote_identifiers=True)
         if df.empty:
             return []
@@ -537,3 +566,21 @@ class SnowflakeEngineAdapter(GetCurrentCatalogFromFunctionMixin, ClusteredByMixi
             clone_kwargs=clone_kwargs,
             **kwargs,
         )
+
+    @t.overload
+    def _columns_to_types(
+        self, query_or_df: DF, columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None
+    ) -> t.Dict[str, exp.DataType]: ...
+
+    @t.overload
+    def _columns_to_types(
+        self, query_or_df: Query, columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None
+    ) -> t.Optional[t.Dict[str, exp.DataType]]: ...
+
+    def _columns_to_types(
+        self, query_or_df: QueryOrDF, columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None
+    ) -> t.Optional[t.Dict[str, exp.DataType]]:
+        if not columns_to_types and snowpark and isinstance(query_or_df, snowpark.DataFrame):
+            return columns_to_types_from_dtypes(query_or_df.sample(n=1).to_pandas().dtypes.items())
+
+        return super()._columns_to_types(query_or_df, columns_to_types)

{sqlmesh-0.183.1.dev11 → sqlmesh-0.183.1.dev13}/sqlmesh/core/engine_adapter/spark.py

@@ -382,51 +382,14 @@ class SparkEngineAdapter(
             partitioned_by=[exp.column(x) for x in primary_key] if primary_key else None,
         )
 
-    def create_view(
+    def _native_df_to_pandas_df(
         self,
-        view_name: TableName,
         query_or_df: QueryOrDF,
-        columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
-        replace: bool = True,
-        materialized: bool = False,
-        materialized_properties: t.Optional[t.Dict[str, t.Any]] = None,
-        table_description: t.Optional[str] = None,
-        column_descriptions: t.Optional[t.Dict[str, str]] = None,
-        view_properties: t.Optional[t.Dict[str, exp.Expression]] = None,
-        **create_kwargs: t.Any,
-    ) -> None:
-        """Create a view with a query or dataframe.
-
-        If a dataframe is passed in, it will be converted into a literal values statement.
-        This should only be done if the dataframe is very small!
-
-        Args:
-            view_name: The view name.
-            query_or_df: A query or dataframe.
-            columns_to_types: Columns to use in the view statement.
-            replace: Whether or not to replace an existing view - defaults to True.
-            materialized: Whether or not the view should be materialized - defaults to False.
-            materialized_properties: Optional materialized view properties to add to the view.
-            table_description: Optional table description from MODEL DDL.
-            column_descriptions: Optional column descriptions from model query.
-            create_kwargs: Additional kwargs to pass into the Create expression
-            view_properties: Optional view properties to add to the view.
-        """
-        pyspark_df = self.try_get_pyspark_df(query_or_df)
-        if pyspark_df:
-            query_or_df = pyspark_df.toPandas()
-        super().create_view(
-            view_name,
-            query_or_df,
-            columns_to_types,
-            replace,
-            materialized,
-            materialized_properties,
-            table_description,
-            column_descriptions,
-            view_properties=view_properties,
-            **create_kwargs,
-        )
+    ) -> t.Union[Query, pd.DataFrame]:
+        if pyspark_df := self.try_get_pyspark_df(query_or_df):
+            return pyspark_df.toPandas()
+
+        return super()._native_df_to_pandas_df(query_or_df)
 
     def _create_table(
         self,

{sqlmesh-0.183.1.dev11 → sqlmesh-0.183.1.dev13}/sqlmesh/core/macros.py

@@ -1087,17 +1087,17 @@ def haversine_distance(
 @macro()
 def pivot(
     evaluator: MacroEvaluator,
-    column: exp.Column,
-    values: t.Union[exp.Array, exp.Tuple],
-    alias: exp.Boolean = exp.true(),
-    agg: exp.Literal = exp.Literal.string("SUM"),
-    cmp: exp.Literal = exp.Literal.string("="),
-    prefix: exp.Literal = exp.Literal.string(""),
-    suffix: exp.Literal = exp.Literal.string(""),
-    then_value: exp.Literal = exp.Literal.number(1),
-    else_value: exp.Literal = exp.Literal.number(0),
-    quote: exp.Boolean = exp.true(),
-    distinct: exp.Boolean = exp.false(),
+    column: SQL,
+    values: t.List[SQL],
+    alias: bool = True,
+    agg: SQL = SQL("SUM"),
+    cmp: SQL = SQL("="),
+    prefix: SQL = SQL(""),
+    suffix: SQL = SQL(""),
+    then_value: SQL = SQL("1"),
+    else_value: SQL = SQL("0"),
+    quote: bool = True,
+    distinct: bool = False,
 ) -> t.List[exp.Expression]:
     """Returns a list of projections as a result of pivoting the given column on the given values.
 
@@ -1109,15 +1109,19 @@ def pivot(
         'SELECT date_day, SUM(CASE WHEN status = \\'cancelled\\' THEN 1 ELSE 0 END) AS "\\'cancelled\\'", SUM(CASE WHEN status = \\'completed\\' THEN 1 ELSE 0 END) AS "\\'completed\\'" FROM rides GROUP BY 1'
     """
     aggregates: t.List[exp.Expression] = []
-    for value in values.expressions:
-        proj = f"{agg.this}("
-        if distinct.this:
+    for value in values:
+        proj = f"{agg}("
+        if distinct:
             proj += "DISTINCT "
-        proj += f"CASE WHEN {column} {cmp.this} {value} THEN {then_value} ELSE {else_value} END) "
+
+        proj += f"CASE WHEN {column} {cmp} {value} THEN {then_value} ELSE {else_value} END) "
         node = evaluator.parse_one(proj)
-        if alias.this:
-            node = node.as_(f"{prefix.this}{value}{suffix.this}", quoted=quote.this, copy=False)
+
+        if alias:
+            node = node.as_(f"{prefix}{value}{suffix}", quoted=quote, copy=False)
+
         aggregates.append(node)
+
     return aggregates
 
 

{sqlmesh-0.183.1.dev11 → sqlmesh-0.183.1.dev13}/sqlmesh/utils/pandas.py

@@ -27,13 +27,20 @@ PANDAS_TYPE_MAPPINGS = {
     pd.Float32Dtype(): exp.DataType.build("float"),
     pd.Float64Dtype(): exp.DataType.build("double"),
     pd.StringDtype(): exp.DataType.build("text"),  # type: ignore
+    pd.StringDtype("pyarrow"): exp.DataType.build("text"),
     pd.BooleanDtype(): exp.DataType.build("boolean"),
 }
 
 
 def columns_to_types_from_df(df: pd.DataFrame) -> t.Dict[str, exp.DataType]:
+    return columns_to_types_from_dtypes(df.dtypes.items())
+
+
+def columns_to_types_from_dtypes(
+    dtypes: t.Iterable[t.Tuple[t.Hashable, t.Any]],
+) -> t.Dict[str, exp.DataType]:
     result = {}
-    for column_name, column_type in df.dtypes.items():
+    for column_name, column_type in dtypes:
         exp_type: t.Optional[exp.DataType] = None
         if hasattr(pd, "DatetimeTZDtype") and isinstance(column_type, pd.DatetimeTZDtype):
             exp_type = exp.DataType.build("timestamptz")

{sqlmesh-0.183.1.dev11 → sqlmesh-0.183.1.dev13}/sqlmesh.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlmesh
-Version: 0.183.1.dev11
+Version: 0.183.1.dev13
 Summary: Next-generation data transformation framework
 Author-email: "TobikoData Inc." <engineering@tobikodata.com>
 License:                                  Apache License

{sqlmesh-0.183.1.dev11 → sqlmesh-0.183.1.dev13}/tests/core/engine_adapter/integration/__init__.py

@@ -22,6 +22,9 @@ from sqlmesh.utils import random_id
 from sqlmesh.utils.date import to_ds
 from sqlmesh.utils.pydantic import PydanticModel
 from tests.utils.pandas import compare_dataframes
+from dataclasses import dataclass
+from _pytest.mark import MarkDecorator
+from _pytest.mark.structures import ParameterSet
 
 if t.TYPE_CHECKING:
     from sqlmesh.core._typing import TableName
@@ -30,6 +33,120 @@ if t.TYPE_CHECKING:
 TEST_SCHEMA = "test_schema"
 
 
+@dataclass
+class IntegrationTestEngine:
+    engine: str
+    catalog_types: t.Optional[t.List[str]] = None
+    native_dataframe_type: t.Optional[str] = None
+    cloud: bool = False
+
+    @property
+    def dialect(self) -> str:
+        return self.engine.split("_", maxsplit=1)[0]
+
+    @property
+    def pytest_marks(self) -> t.List[MarkDecorator]:
+        marks = [getattr(pytest.mark, self.engine), pytest.mark.engine]
+        if self.cloud:
+            marks.append(pytest.mark.remote)
+        else:
+            marks.append(pytest.mark.docker)
+        if self.engine == "duckdb":
+            marks.extend(
+                [
+                    # run the duckdb tests in `make cicd-test` as well
+                    pytest.mark.slow,
+                    # the duckdb tests cannot run concurrently because many of them point at the same files
+                    # and duckdb does not support multi process read/write on the same files
+                    # ref: https://duckdb.org/docs/connect/concurrency.html#writing-to-duckdb-from-multiple-processes
+                    pytest.mark.xdist_group("engine_integration_duckdb"),
+                ]
+            )
+        return marks
+
+
+ENGINES = [
+    # Docker engines that can be locally tested
+    IntegrationTestEngine("duckdb"),
+    IntegrationTestEngine("postgres"),
+    IntegrationTestEngine("mysql"),
+    IntegrationTestEngine("mssql"),
+    IntegrationTestEngine("trino", catalog_types=["hive", "iceberg", "delta", "nessie"]),
+    IntegrationTestEngine("spark", native_dataframe_type="pyspark"),
+    IntegrationTestEngine("clickhouse", catalog_types=["standalone", "cluster"]),
+    IntegrationTestEngine("risingwave"),
+    # Cloud engines that need paid accounts / special credentials
+    IntegrationTestEngine("clickhouse_cloud", cloud=True),
+    IntegrationTestEngine("redshift", cloud=True),
+    IntegrationTestEngine("athena", catalog_types=["hive", "iceberg"], cloud=True),
+    IntegrationTestEngine("bigquery", native_dataframe_type="bigframe", cloud=True),
+    IntegrationTestEngine("databricks", native_dataframe_type="pyspark", cloud=True),
+    IntegrationTestEngine("snowflake", native_dataframe_type="snowpark", cloud=True),
+]
+
+ENGINES_BY_NAME = {e.engine: e for e in ENGINES}
+
+
+def generate_pytest_params(
+    engines: t.Union[IntegrationTestEngine, t.List[IntegrationTestEngine]],
+    query: bool = True,
+    df: bool = False,
+    show_variant_in_test_id: bool = True,
+) -> t.Iterable[ParameterSet]:
+    """
+    The engine adapter tests have a bunch of variants:
+     - Per engine for engines that dont have pluggable catalogs
+     - Per engine per catalog type for engines that have pluggable catalogs
+
+    In addition, many engine adapter functions take either a SQL Query or a DataFrame so we need to test both combinations.
+    For the methods that take a DataFrame:
+     - Every engine takes a Pandas DataFrame
+     - A small subset of engines also take their own engine-specific DataFrame (eg Bigframe, Snowpark, Pyspark)
+
+    This function controls the parameter generation so that:
+     - Tests that only need to test SQL queries only get called once per engine/catalog
+     - Tests that only need to test DataFrame's get called once for Pandas Dataframe's and once for each engine-specific DataFrame
+     - Tests that need to test both SQL Queries and DataFrame's get called once for every combination of (engine, catalog, *(query, pandas df, native df))
+
+    The goal is to prevent needing to code this kind of logic into tests:
+
+    > if test_type == "df":
+    >    pytest.skip("Test only needs to run for query")
+
+    As well as make it easier to generate the right combinations for new databases / catalogs / DataFrame implementations
+    """
+    if not isinstance(engines, list):
+        engines = [engines]
+
+    for engine in engines:
+        catalogs = engine.catalog_types if engine.catalog_types else [""]
+        for catalog in catalogs:
+            gateway = (
+                f"inttest_{engine.engine}_{catalog}" if catalog else f"inttest_{engine.engine}"
+            )
+            if engine.engine == "athena":
+                # athena only has a single gateway defined, not a gateway per catalog
+                gateway = f"inttest_athena"
+
+            variants = []
+            if query:
+                variants.append("query")
+            if df:
+                variants.append("df-pandas")
+                if engine.native_dataframe_type:
+                    variants.append(f"df-{engine.native_dataframe_type}")
+
+            test_id = f"{engine.engine}_{catalog}" if catalog else f"{engine.engine}"
+            default_table_format = catalog
+
+            for variant in variants:
+                yield pytest.param(
+                    (engine, gateway, variant, default_table_format),
+                    marks=engine.pytest_marks,
+                    id=f"[{variant}]{test_id}" if show_variant_in_test_id else test_id,
+                )
+
+
 class MetadataResults(PydanticModel):
     tables: t.List[str] = []
     views: t.List[str] = []
@@ -77,7 +194,7 @@ class TestContext:
         is_remote: bool = False,
         columns_to_types: t.Optional[t.Dict[str, t.Union[str, exp.DataType]]] = None,
     ):
-        self.test_type = test_type
+        self._test_type = test_type
         self.engine_adapter = engine_adapter
         self.mark = mark
         self.gateway = gateway
@@ -89,6 +206,17 @@ class TestContext:
             str
         ] = []  # keep track of any schemas returned from self.schema() / self.table() so we can drop them at the end
 
+    @property
+    def test_type(self) -> str:
+        return "df" if self._test_type.startswith("df") else "query"
+
+    @property
+    def df_type(self) -> t.Optional[str]:
+        if self.test_type == "df":
+            # the 'pandas' part of 'df-pandas'
+            return self._test_type.split("-", maxsplit=1)[1]
+        return None
+
     @property
     def columns_to_types(self):
         if self._columns_to_types is None:
@@ -193,7 +321,7 @@ class TestContext:
         return data
 
     def init(self):
-        if self.test_type == "pyspark" and not hasattr(self.engine_adapter, "is_pyspark_df"):
+        if self.df_type == "pyspark" and not hasattr(self.engine_adapter, "is_pyspark_df"):
             pytest.skip(f"Engine adapter {self.engine_adapter} doesn't support pyspark")
         self._init_engine_adapter()
 
@@ -210,9 +338,20 @@ class TestContext:
                 batch_end=sys.maxsize,
                 columns_to_types=columns_to_types,
             )
-        if self.test_type == "pyspark":
-            return self.engine_adapter.spark.createDataFrame(data)  # type: ignore
-        return self._format_df(data, to_datetime=self.dialect != "trino")
+        if self.test_type == "df":
+            formatted_df = self._format_df(data, to_datetime=self.dialect != "trino")
+            if self.df_type == "pandas":
+                return formatted_df
+            if self.df_type == "pyspark":
+                return self.engine_adapter.spark.createDataFrame(formatted_df)  # type: ignore
+            if self.df_type == "bigframe":
+                return self.engine_adapter.bigframe.read_pandas(formatted_df)  # type: ignore
+            if self.df_type == "snowpark":
+                return self.engine_adapter.snowpark.create_dataframe(formatted_df)  # type: ignore
+
+            raise ValueError(f"Unknown DF type: {self.df_type}")
+
+        raise ValueError(f"Unknown test type: {self.test_type}")
 
     def output_data(self, data: pd.DataFrame) -> pd.DataFrame:
         return self._format_df(data)
@@ -568,6 +707,10 @@ class TestContext:
                 schema_name=schema_name, ignore_if_not_exists=True, cascade=True
             )
 
+        if snowpark := self.engine_adapter.snowpark:
+            # ensure that the next test gets a fresh Snowpark session
+            snowpark.close()
+
     def upsert_sql_model(self, model_definition: str) -> t.Tuple[Context, SqlModel]:
         if not self._context:
             self._context = self.create_context()

{sqlmesh-0.183.1.dev11 → sqlmesh-0.183.1.dev13}/tests/core/engine_adapter/integration/config.yaml

@@ -8,7 +8,7 @@ gateways:
         testing: 'testing.duckdb'
 
   # Databases with docker images available
-  inttest_trino:
+  inttest_trino_hive:
     connection:
       type: trino
       host: {{ env_var('DOCKER_HOSTNAME', 'localhost') }}
@@ -81,7 +81,7 @@ gateways:
       password: mysql
       port: 3306
       charset: utf8
-  inttest_clickhouse:
+  inttest_clickhouse_standalone:
     connection:
       type: clickhouse
       host: {{ env_var('DOCKER_HOSTNAME', 'localhost') }}