fix: bigquery snowflake source columns support

Author: eakmanrq

.circleci/continue_config.yml

@@ -306,17 +306,17 @@ workflows:
             parameters:
               engine:
                 - snowflake
-                - databricks
-                - redshift
+#                - databricks
+#                - redshift
                 - bigquery
-                - clickhouse-cloud
-                - athena
-                - fabric
-                - gcp-postgres
-          filters:
-            branches:
-              only:
-                - main
+#                - clickhouse-cloud
+#                - athena
+#                - fabric
+#                - gcp-postgres
+#          filters:
+#            branches:
+#              only:
+#                - main
       - ui_style
       - ui_test
       - vscode_test

sqlmesh/core/engine_adapter/bigquery.py

@@ -169,17 +169,18 @@ def _df_to_source_queries(
         )
 
         def query_factory() -> Query:
-            if bigframes_pd and isinstance(df, bigframes_pd.DataFrame):
-                df.to_gbq(
+            ordered_df = df[list(source_columns_to_types)]
+            if bigframes_pd and isinstance(ordered_df, bigframes_pd.DataFrame):
+                ordered_df.to_gbq(
                     f"{temp_bq_table.project}.{temp_bq_table.dataset_id}.{temp_bq_table.table_id}",
                     if_exists="replace",
                 )
             elif not self.table_exists(temp_table):
                 # Make mypy happy
-                assert isinstance(df, pd.DataFrame)
+                assert isinstance(ordered_df, pd.DataFrame)
                 self._db_call(self.client.create_table, table=temp_bq_table, exists_ok=False)
                 result = self.__load_pandas_to_table(
-                    temp_bq_table, df, source_columns_to_types, replace=False
+                    temp_bq_table, ordered_df, source_columns_to_types, replace=False
                 )
                 if result.errors:
                     raise SQLMeshError(result.errors)

sqlmesh/core/engine_adapter/snowflake.py

@@ -354,6 +354,8 @@ def query_factory() -> Query:
                 else None
             )
 
+            ordered_df = df[list(source_columns_to_types)]
+
             if is_snowpark_dataframe:
                 temp_table.set("catalog", database)
 
@@ -362,11 +364,11 @@ def query_factory() -> Query:
                 # 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
+                    col.startswith('"') and col.endswith('"') for col in ordered_df.columns
                 )
-                local_df = df
+                local_df = ordered_df
                 if not columns_already_quoted:
-                    local_df = df.rename(
+                    local_df = ordered_df.rename(
                         {
                             col: exp.to_identifier(col).sql(dialect=self.dialect, identify=True)
                             for col in source_columns_to_types
@@ -375,7 +377,7 @@ def query_factory() -> Query:
                 local_df.createOrReplaceTempView(
                     temp_table.sql(dialect=self.dialect, identify=True)
                 )  # type: ignore
-            elif isinstance(df, pd.DataFrame):
+            elif isinstance(ordered_df, pd.DataFrame):
                 from snowflake.connector.pandas_tools import write_pandas
 
                 # Workaround for https://github.com/snowflakedb/snowflake-connector-python/issues/1034
@@ -388,16 +390,16 @@ def query_factory() -> Query:
 
                 # See: https://stackoverflow.com/a/75627721
                 for column, kind in source_columns_to_types.items():
-                    if is_datetime64_any_dtype(df.dtypes[column]):
+                    if is_datetime64_any_dtype(ordered_df.dtypes[column]):
                         if kind.is_type("date"):  # type: ignore
-                            df[column] = pd.to_datetime(df[column]).dt.date  # type: ignore
-                        elif getattr(df.dtypes[column], "tz", None) is not None:  # type: ignore
-                            df[column] = pd.to_datetime(df[column]).dt.strftime(
+                            ordered_df[column] = pd.to_datetime(ordered_df[column]).dt.date  # type: ignore
+                        elif getattr(ordered_df.dtypes[column], "tz", None) is not None:  # type: ignore
+                            ordered_df[column] = pd.to_datetime(ordered_df[column]).dt.strftime(
                                 "%Y-%m-%d %H:%M:%S.%f%z"
                             )  # type: ignore
                         # https://github.com/snowflakedb/snowflake-connector-python/issues/1677
                         else:  # type: ignore
-                            df[column] = pd.to_datetime(df[column]).dt.strftime(
+                            ordered_df[column] = pd.to_datetime(ordered_df[column]).dt.strftime(
                                 "%Y-%m-%d %H:%M:%S.%f"
                             )  # type: ignore
 
@@ -407,7 +409,7 @@ def query_factory() -> Query:
 
                 write_pandas(
                     self._connection_pool.get(),
-                    df,
+                    ordered_df,
                     temp_table.name,
                     schema=temp_table.db or None,
                     database=database.sql(dialect=self.dialect) if database else None,
@@ -417,7 +419,7 @@ def query_factory() -> Query:
                 )
             else:
                 raise SQLMeshError(
-                    f"Unknown dataframe type: {type(df)} for {target_table}. Expecting pandas or snowpark."
+                    f"Unknown dataframe type: {type(ordered_df)} for {target_table}. Expecting pandas or snowpark."
                 )
 
             return exp.select(

tests/core/engine_adapter/integration/config.yaml

@@ -155,7 +155,7 @@ gateways:
     connection:
       type: bigquery
       method: service-account
-      keyfile: {{ env_var('BIGQUERY_KEYFILE') }}
+      keyfile: { { env_var('BIGQUERY_KEYFILE') } }
       check_import: false
     state_connection:
       type: duckdb