Fix!: quote all identifiers in execute except those coming from fetchdf (#1128)

* Fix!: normalize exprs in execute except those coming from fetchdf

* Fix tests

* Set default to false

* Simplify snowflake _insert_append_pandas_df

* Redshift fixup

* Cleanup

* Fix test

* Only quote identifiers in _to_sql

* Snowflake fixups

* Formatting

* Quote identifiers in _to_sql

* PR feedback

* Comment fixup

* Add empty dialect to seed model as well

* Remove default dialects in model level from wursthall project

* Add comment

* Fixup

Author: georgesittas

examples/wursthall/models/db/customer_d.sql

@@ -4,7 +4,6 @@ MODEL (
     time_column (register_ds, '%Y-%m-%d'),
     batch_size 200,
   ),
-  dialect "",
   cron '@daily',
   owner jen,
   start '2022-06-01 00:00:00+00:00',

examples/wursthall/models/db/item_d.sql

@@ -1,7 +1,6 @@
 MODEL (
   name db.item_d,
   kind VIEW,
-  dialect "",
   cron '@daily',
   owner jen,
   start '2022-06-01 00:00:00+00:00',

examples/wursthall/models/db/order_f.py

@@ -5,6 +5,7 @@
 import numpy as np
 import pandas as pd
 from models.src.shared import DATA_START_DATE_STR, set_seed  # type: ignore
+from sqlglot import parse_one
 
 from sqlmesh import ExecutionContext, model
 from sqlmesh.core.model import IncrementalByTimeRangeKind, TimeColumn
@@ -43,27 +44,29 @@ def execute(
     item_d_table_name = context.table("db.item_d")
     order_item_f_table_name = context.table("db.order_item_f")
 
+    # We use parse_one here instead of a raw string because this is a multi-dialect
+    # project and we want to ensure that the resulting query is properly quoted in
+    # the target dialect before executing it
     df_item_d = context.fetchdf(
-        f"""
-        SELECT
-            item_id,
-            item_price
-        FROM {item_d_table_name}
-        """
+        parse_one(f"SELECT item_id, item_price FROM {item_d_table_name}"),
+        quote_identifiers=True,
     )
 
     df_order_item_f = context.fetchdf(
-        f"""
-        SELECT
-            order_id,
-            customer_id,
-            item_id,
-            quantity,
-            order_ds
-        FROM {order_item_f_table_name}
-        WHERE
-            order_ds BETWEEN '{to_ds(start)}' AND '{to_ds(end)}'
-        """
+        parse_one(
+            f"""
+            SELECT
+                order_id,
+                customer_id,
+                item_id,
+                quantity,
+                order_ds
+            FROM {order_item_f_table_name}
+            WHERE
+                order_ds BETWEEN '{to_ds(start)}' AND '{to_ds(end)}'
+            """
+        ),
+        quote_identifiers=True,
     )
 
     df_order_item_f = df_order_item_f.merge(df_item_d, how="inner", on="item_id")

examples/wursthall/models/db/order_item_f.sql

@@ -4,7 +4,6 @@ MODEL (
     time_column (order_ds, '%Y-%m-%d'),
     batch_size 200,
   ),
-  dialect "",
   cron '@daily',
   owner jen,
   start '2022-06-01 00:00:00+00:00',

examples/wursthall/models/src/menu_item_details.sql

@@ -4,4 +4,4 @@ MODEL (
     path '../../seeds/src/menu_item_details.csv',
   ),
   owner jen
-)
+)

examples/wursthall/models/src/order_item_details.py

@@ -7,6 +7,7 @@
 import pandas as pd
 from faker import Faker
 from models.src.shared import DATA_START_DATE_STR, iter_dates, set_seed  # type: ignore
+from sqlglot import parse_one
 
 from sqlmesh import ExecutionContext, model
 from sqlmesh.core.model import IncrementalByTimeRangeKind, TimeColumn
@@ -49,23 +50,26 @@ def execute(
     customer_details_table_name = context.table("src.customer_details")
     menu_item_details_table_name = context.table("src.menu_item_details")
 
+    # We use parse_one here instead of a raw string because this is a multi-dialect
+    # project and we want to ensure that the resulting query is properly quoted in
+    # the target dialect before executing it
     df_customers = context.fetchdf(
-        f"""
-        SELECT
-            id AS customer_id,
-            register_ds
-        FROM {customer_details_table_name}
-        WHERE
-            register_ds <= '{to_ds(end)}'
-        """
+        parse_one(
+            f"""
+            SELECT
+                id AS customer_id,
+                register_ds
+            FROM {customer_details_table_name}
+            WHERE
+                register_ds <= '{to_ds(end)}'
+            """
+        ),
+        quote_identifiers=True,
     )
 
     df_menu_items = context.fetchdf(
-        f"""
-        SELECT
-            id AS item_id
-        FROM {menu_item_details_table_name}
-        """
+        parse_one(f"SELECT id AS item_id FROM {menu_item_details_table_name}"),
+        quote_identifiers=True,
     )
 
     num_menu_items = len(df_menu_items.index)

setup.py

@@ -115,6 +115,7 @@
         ],
         "snowflake": [
             "snowflake-connector-python[pandas,secure-local-storage]",
+            "pyarrow>=10.0.1,<10.1.0",
         ],
         "web": [
             "fastapi==0.100.0",

sqlmesh/core/context.py

@@ -130,27 +130,33 @@ def table(self, model_name: str) -> str:
         """
         return self._model_tables[model_name]
 
-    def fetchdf(self, query: t.Union[exp.Expression, str]) -> pd.DataFrame:
+    def fetchdf(
+        self, query: t.Union[exp.Expression, str], quote_identifiers: bool = False
+    ) -> pd.DataFrame:
         """Fetches a dataframe given a sql string or sqlglot expression.
 
         Args:
             query: SQL string or sqlglot expression.
+            quote_identifiers: Whether to quote all identifiers in the query.
 
         Returns:
             The default dataframe is Pandas, but for Spark a PySpark dataframe is returned.
         """
-        return self.engine_adapter.fetchdf(query)
+        return self.engine_adapter.fetchdf(query, quote_identifiers=quote_identifiers)
 
-    def fetch_pyspark_df(self, query: t.Union[exp.Expression, str]) -> PySparkDataFrame:
+    def fetch_pyspark_df(
+        self, query: t.Union[exp.Expression, str], quote_identifiers: bool = False
+    ) -> PySparkDataFrame:
         """Fetches a PySpark dataframe given a sql string or sqlglot expression.
 
         Args:
             query: SQL string or sqlglot expression.
+            quote_identifiers: Whether to quote all identifiers in the query.
 
         Returns:
             A PySpark dataframe.
         """
-        return self.engine_adapter.fetch_pyspark_df(query)
+        return self.engine_adapter.fetch_pyspark_df(query, quote_identifiers=quote_identifiers)
 
 
 class ExecutionContext(BaseContext):

sqlmesh/core/engine_adapter/base.py

@@ -19,6 +19,7 @@
 from sqlglot import Dialect, exp
 from sqlglot.errors import ErrorLevel
 from sqlglot.helper import ensure_list
+from sqlglot.optimizer.qualify_columns import quote_identifiers
 
 from sqlmesh.core.dialect import pandas_to_sql
 from sqlmesh.core.engine_adapter.shared import DataObject, TransactionType
@@ -698,7 +699,9 @@ def _merge(
         match_expressions: t.List[exp.When],
     ) -> None:
         this = exp.alias_(exp.to_table(target_table), alias=MERGE_TARGET_ALIAS, table=True)
-        using = exp.Subquery(this=source_table, alias=MERGE_SOURCE_ALIAS)
+        using = exp.alias_(
+            exp.Subquery(this=source_table), alias=MERGE_SOURCE_ALIAS, copy=False, table=True
+        )
         self.execute(
             exp.Merge(
                 this=this,
@@ -778,33 +781,49 @@ def fetchone(
         self,
         query: t.Union[exp.Expression, str],
         ignore_unsupported_errors: bool = False,
+        quote_identifiers: bool = False,
     ) -> t.Tuple:
-        self.execute(query, ignore_unsupported_errors=ignore_unsupported_errors)
+        self.execute(
+            query,
+            ignore_unsupported_errors=ignore_unsupported_errors,
+            quote_identifiers=quote_identifiers,
+        )
         return self.cursor.fetchone()
 
     def fetchall(
         self,
         query: t.Union[exp.Expression, str],
         ignore_unsupported_errors: bool = False,
+        quote_identifiers: bool = False,
     ) -> t.List[t.Tuple]:
-        self.execute(query, ignore_unsupported_errors=ignore_unsupported_errors)
+        self.execute(
+            query,
+            ignore_unsupported_errors=ignore_unsupported_errors,
+            quote_identifiers=quote_identifiers,
+        )
         return self.cursor.fetchall()
 
-    def _fetch_native_df(self, query: t.Union[exp.Expression, str]) -> DF:
+    def _fetch_native_df(
+        self, query: t.Union[exp.Expression, str], quote_identifiers: bool = False
+    ) -> DF:
         """Fetches a DataFrame that can be either Pandas or PySpark from the cursor"""
-        self.execute(query)
+        self.execute(query, quote_identifiers=quote_identifiers)
         return self.cursor.fetchdf()
 
-    def fetchdf(self, query: t.Union[exp.Expression, str]) -> pd.DataFrame:
+    def fetchdf(
+        self, query: t.Union[exp.Expression, str], quote_identifiers: bool = False
+    ) -> pd.DataFrame:
         """Fetches a Pandas DataFrame from the cursor"""
-        df = self._fetch_native_df(query)
+        df = self._fetch_native_df(query, quote_identifiers=quote_identifiers)
         if not isinstance(df, pd.DataFrame):
             raise NotImplementedError(
                 "The cursor's `fetch_native_df` method is not returning a pandas DataFrame. Need to update `fetchdf` so a Pandas DataFrame is returned"
             )
         return df
 
-    def fetch_pyspark_df(self, query: t.Union[exp.Expression, str]) -> PySparkDataFrame:
+    def fetch_pyspark_df(
+        self, query: t.Union[exp.Expression, str], quote_identifiers: bool = False
+    ) -> PySparkDataFrame:
         """Fetches a PySpark DataFrame from the cursor"""
         raise NotImplementedError(f"Engine does not support PySpark DataFrames: {type(self)}")
 
@@ -835,6 +854,7 @@ def execute(
         self,
         expressions: t.Union[str, exp.Expression, t.Sequence[exp.Expression]],
         ignore_unsupported_errors: bool = False,
+        quote_identifiers: bool = True,
         **kwargs: t.Any,
     ) -> None:
         """Execute a sql query."""
@@ -843,7 +863,11 @@ def execute(
         )
 
         for e in ensure_list(expressions):
-            sql = self._to_sql(e, **to_sql_kwargs) if isinstance(e, exp.Expression) else e
+            sql = (
+                self._to_sql(e, quote=quote_identifiers, **to_sql_kwargs)
+                if isinstance(e, exp.Expression)
+                else e
+            )
             logger.debug(f"Executing SQL:\n{sql}")
             self.cursor.execute(sql, **kwargs)
 
@@ -882,7 +906,7 @@ def _create_table_properties(
         """Creates a SQLGlot table properties expression for ddl."""
         return None
 
-    def _to_sql(self, e: exp.Expression, **kwargs: t.Any) -> str:
+    def _to_sql(self, expression: exp.Expression, quote: bool = True, **kwargs: t.Any) -> str:
         """
         Converts an expression to a SQL string. Has a set of default kwargs to apply, and then default
         kwargs defined for the given dialect, and then kwargs provided by the user when defining the engine
@@ -896,7 +920,11 @@ def _to_sql(self, e: exp.Expression, **kwargs: t.Any) -> str:
             **self.sql_gen_kwargs,
             **kwargs,
         }
-        return e.sql(**sql_gen_kwargs)  # type: ignore
+
+        if quote:
+            quote_identifiers(expression)
+
+        return expression.sql(**sql_gen_kwargs)  # type: ignore
 
     def _get_data_objects(
         self, schema_name: str, catalog_name: t.Optional[str] = None
@@ -917,9 +945,11 @@ def _get_temp_table(
         """
         table = t.cast(exp.Table, exp.to_table(table).copy())
         table.set("this", exp.to_identifier(f"__temp_{table.name}_{uuid.uuid4().hex}"))
+
         if table_only:
             table.set("db", None)
             table.set("catalog", None)
+
         return table
 
     def _add_where_to_query(self, query: Query, where: t.Optional[exp.Expression]) -> Query:
@@ -933,8 +963,10 @@ def _add_where_to_query(self, query: Query, where: t.Optional[exp.Expression]) -
             .from_(query.subquery("_subquery", copy=False), copy=False)
             .where(where, copy=False)
         )
+
         if with_:
             query.set("with", with_)
+
         return query
 
 

sqlmesh/core/engine_adapter/base_postgres.py

@@ -100,6 +100,7 @@ def execute(
         self,
         expressions: t.Union[str, exp.Expression, t.Sequence[exp.Expression]],
         ignore_unsupported_errors: bool = False,
+        quote_identifiers: bool = True,
         **kwargs: t.Any,
     ) -> None:
         """
@@ -108,7 +109,12 @@ def execute(
 
         Reference: https://www.psycopg.org/psycopg3/docs/basic/transactions.html
         """
-        super().execute(expressions, ignore_unsupported_errors=ignore_unsupported_errors, **kwargs)
+        super().execute(
+            expressions,
+            ignore_unsupported_errors=ignore_unsupported_errors,
+            quote_identifiers=quote_identifiers,
+            **kwargs,
+        )
         if not self._connection_pool.is_transaction_active:
             self._connection_pool.commit()