Repository: ploomber/jupysql
Pull request #318 (2023-03-24 22:44:23)

Drop sqlglot for generating CTEs
## Describe your changes

In the original issue, we need to do:
1. Need some research to see what dialects support backticks, and make a list
2. Remove the sqlglot & add condition on using non-backticks or backticks jinja template

Result:
1. We can check if the implementation of the corresponding dialect in sqlglot has IDENTIFIERS = ["`"] attribute, which is confirmed by sqlglot team
2. We prepare two template in this PR. One is using `"`, another one is using ` (backtick), we will dynamic choose the right template by checking the implementation of the sqlglot.dialect class

## Issue number

Closes https://github.com/ploomber/jupysql/issues/307

## Checklist before requesting a review

- [ ] Performed a self-review of my code
- [ ] Formatted my code with [`pkgmt format`](https://github.com/ploomber/contributing/blob/main/CONTRIBUTING.md#lintingformatting)
- [ ] Added [tests](https://github.com/ploomber/contributing/blob/main/CONTRIBUTING.md#testing) (when necessary).
- [ ] Added documentation in the [docstring](https://github.com/ploomber/contributing/blob/main/CONTRIBUTING.md#documenting-changes-and-new-features) and [changelog](https://github.com/ploomber/contributing/blob/main/CONTRIBUTING.md#changelog) (when needed)



<!-- readthedocs-preview jupysql start -->
----
:books: Documentation preview :books:: https://jupysql--318.org.readthedocs.build/en/318/

<!-- readthedocs-preview jupysql end -->

tonykploomber: > Added minor notes.
> 
> General question - I see you're using this logic
> 
> ```
> identifier = "`" if is_dialect_support_backtick else '"'
> ```
> 
> Is there any chance the default identifier (`"`) won't work?

Good question, I don't have fully confidence to answer you 100% yes or no.

But originally before we use sqlglot, we use simple with_template includes `"`
```
with_template = Template(
    """WITH{% for name in with_ %} "{{name}}" AS ({{saved[name]._query}})\
{{ "," if not loop.last }}{% endfor %}{{query}}"""
)
```

We haven't met any issue then, now in this PR we just detect if the dialect support ` and we use backtick template, I think the impact should be not much
edublancas: what's the status of this? This is blocking the 0.7 release so I'd like to merge this as soon as we can.

Patch:
diff --git a/src/sql/connection.py b/src/sql/connection.py
index 916033a..7faf8eb 100644
--- a/src/sql/connection.py
+++ b/src/sql/connection.py
@@ -44,9 +44,7 @@ MISSING_PACKAGE_LIST_EXCEPT_MATCHERS = {
     "pymssql": "pymssql",
 }
 
-DIALECT_NAME_SQLALCHEMY_TO_SQLGLOT_MAPPING = {
-    "postgresql": "postgres",
-}
+DIALECT_NAME_SQLALCHEMY_TO_SQLGLOT_MAPPING = {"postgresql": "postgres", "mssql": "tsql"}
 
 
 def extract_module_name_from_ModuleNotFoundError(e):
@@ -352,8 +350,16 @@ class Connection:
             conn.session.close()
 
     @classmethod
-    def _get_curr_connection_info(cls):
-        """Returns the dialect, driver, and database server version info"""
+    def _get_curr_sqlalchemy_connection_info(cls):
+        """Get the dialect, driver, and database server version info of current
+        connected dialect
+
+        Returns
+        -------
+        dict
+            The dictionary which contains the SQLAlchemy-based dialect
+            information, or None if there is no current connection.
+        """
         if not cls.current:
             return None
 
@@ -365,14 +371,59 @@ class Connection:
         }
 
     @classmethod
-    def _transpile_query(cls, query):
-        if not cls.current:
-            return query
-        connection_info = cls._get_curr_connection_info()
+    def _get_curr_sqlglot_dialect(cls):
+        """Get the dialect name in sqlglot package scope
+
+        Returns
+        -------
+        str
+            Available dialect in sqlglot package, see more:
+            https://github.com/tobymao/sqlglot/blob/main/sqlglot/dialects/dialect.py
+        """
+        connection_info = cls._get_curr_sqlalchemy_connection_info()
+        if not connection_info:
+            return None
+
+        return DIALECT_NAME_SQLALCHEMY_TO_SQLGLOT_MAPPING.get(
+            connection_info["dialect"], connection_info["dialect"]
+        )
+
+    @classmethod
+    def is_use_backtick_template(cls):
+        """Get if the dialect support backtick (`) syntax as identifier
+
+        Returns
+        -------
+        bool
+            Indicate if the dialect can use backtick identifier in the SQL clause
+        """
+        cur_dialect = cls._get_curr_sqlglot_dialect()
+        if not cur_dialect:
+            return False
         try:
-            write_dialect = DIALECT_NAME_SQLALCHEMY_TO_SQLGLOT_MAPPING.get(
-                connection_info["dialect"], connection_info["dialect"]
+            return (
+                "`" in sqlglot.Dialect.get_or_raise(cur_dialect).Tokenizer.IDENTIFIERS
             )
+        except (ValueError, AttributeError, TypeError):
+            return False
+
+    @classmethod
+    def _transpile_query(cls, query):
+        """Translate the given SQL clause that's compatible to current connected
+        dialect by sqlglot
+
+        Parameters
+        ----------
+        query : str
+            Original SQL clause
+
+        Returns
+        -------
+        str
+            SQL clause that's compatible to current connected dialect
+        """
+        write_dialect = cls._get_curr_sqlglot_dialect()
+        try:
             query = sqlglot.parse_one(query).sql(dialect=write_dialect)
         finally:
             return query
@@ -386,7 +437,7 @@ class Connection:
         """
         identifiers = ["", '"']
         try:
-            connection_info = cls._get_curr_connection_info()
+            connection_info = cls._get_curr_sqlalchemy_connection_info()
             if connection_info:
                 cur_dialect = connection_info["dialect"]
                 identifiers_ = sqlglot.Dialect.get_or_raise(
diff --git a/src/sql/magic.py b/src/sql/magic.py
index ad1b08c..ff86d5b 100644
--- a/src/sql/magic.py
+++ b/src/sql/magic.py
@@ -343,7 +343,7 @@ class SqlMagic(Magics, Configurable):
         )
         payload[
             "connection_info"
-        ] = sql.connection.Connection._get_curr_connection_info()
+        ] = sql.connection.Connection._get_curr_sqlalchemy_connection_info()
         if args.persist:
             return self._persist_dataframe(
                 command.sql, conn, user_ns, append=False, index=not args.no_index
diff --git a/src/sql/plot.py b/src/sql/plot.py
index 48b6347..0a1483b 100644
--- a/src/sql/plot.py
+++ b/src/sql/plot.py
@@ -246,7 +246,9 @@ def boxplot(payload, table, column, *, orient="v", with_=None, conn=None, ax=Non
     if not conn:
         conn = sql.connection.Connection.current.session
 
-    payload["connection_info"] = sql.connection.Connection._get_curr_connection_info()
+    payload[
+        "connection_info"
+    ] = sql.connection.Connection._get_curr_sqlalchemy_connection_info()
 
     ax = plt.gca()
     vert = orient == "v"
@@ -362,7 +364,9 @@ def histogram(
     .. plot:: ../examples/plot_histogram_many.py
     """
     ax = ax or plt.gca()
-    payload["connection_info"] = sql.connection.Connection._get_curr_connection_info()
+    payload[
+        "connection_info"
+    ] = sql.connection.Connection._get_curr_sqlalchemy_connection_info()
     if category:
         if isinstance(column, list):
             if len(column) > 1:
diff --git a/src/sql/run.py b/src/sql/run.py
index bc56506..250739a 100644
--- a/src/sql/run.py
+++ b/src/sql/run.py
@@ -180,7 +180,7 @@ class ResultSet(list, ColumnGuesserMixin):
         frame = pd.DataFrame(self, columns=(self and self.keys) or [])
         payload[
             "connection_info"
-        ] = sql.connection.Connection._get_curr_connection_info()
+        ] = sql.connection.Connection._get_curr_sqlalchemy_connection_info()
         return frame
 
     @telemetry.log_call("polars-data-frame")
diff --git a/src/sql/store.py b/src/sql/store.py
index d98fcf6..8b50e99 100644
--- a/src/sql/store.py
+++ b/src/sql/store.py
@@ -1,7 +1,8 @@
 from typing import Iterator, Iterable
 from collections.abc import MutableMapping
+from jinja2 import Template
 from ploomber_core.exceptions import modify_exceptions
-from sqlglot import parse_one
+import sql.connection
 from IPython.core.error import UsageError
 import warnings
 import difflib
@@ -98,14 +99,26 @@ class SQLQuery:
             )
 
     def __str__(self) -> str:
-        if self._query:
-            with_all = _get_dependencies(self._store, self._with_)
-            parsed_res = parse_one(self._query)
-            for with_key in with_all:
-                with_query = self._store._data[with_key]._query
-                parsed_res = parsed_res.with_(with_key, with_query)
-            return str(parsed_res.sql())
-        return ""
+        """
+        Since some dialects don't support " (double quote) symbol, we will
+        replace to the ' (backtick) symbol if it's supported
+        """
+        with_clause_template = Template(
+            """WITH{% for name in with_ %} "{{name}}" AS ({{saved[name]._query}})\
+{{ "," if not loop.last }}{% endfor %}{{query}}"""
+        )
+        with_clause_template_backtick = Template(
+            """WITH{% for name in with_ %} `{{name}}` AS ({{saved[name]._query}})\
+{{ "," if not loop.last }}{% endfor %}{{query}}"""
+        )
+        is_use_backtick = sql.connection.Connection.is_use_backtick_template()
+        with_all = _get_dependencies(self._store, self._with_)
+        template = (
+            with_clause_template_backtick if is_use_backtick else with_clause_template
+        )
+        return template.render(
+            query=self._query, saved=self._store._data, with_=with_all
+        )
 
 
 def _get_dependencies(store, keys):


Test patch:
diff --git a/src/tests/test_command.py b/src/tests/test_command.py
index e0ee9f5..d5d57b1 100644
--- a/src/tests/test_command.py
+++ b/src/tests/test_command.py
@@ -78,7 +78,8 @@ def test_parsed_sql_when_using_with(ip, sql_magic):
         sql_magic, ip.user_ns, line="--with author_one", cell="SELECT * FROM author_one"
     )
 
-    sql = "WITH author_one AS (SELECT * FROM author LIMIT 1) SELECT * FROM author_one"
+    sql = "WITH `author_one` AS (\n\n        SELECT * FROM author LIMIT 1\n        )\n\
+SELECT * FROM author_one"
 
     sql_original = "\nSELECT * FROM author_one"
 
diff --git a/src/tests/test_compose.py b/src/tests/test_compose.py
index 098f092..3fd16f0 100644
--- a/src/tests/test_compose.py
+++ b/src/tests/test_compose.py
@@ -12,7 +12,7 @@ def test_compose(ip):
     )
 
     result = ip.run_cell("%sqlrender final").result
-    expected = "WITH author_sub AS (SELECT last_name FROM author\
- WHERE year_of_death > 1900) SELECT last_name FROM author_sub"
+    expected = "WITH `author_sub` AS (\nSELECT last_name \
+FROM author WHERE year_of_death > 1900)\nSELECT last_name FROM author_sub;"
 
     assert result == expected
diff --git a/src/tests/test_connection.py b/src/tests/test_connection.py
index 982d890..84dffdc 100644
--- a/src/tests/test_connection.py
+++ b/src/tests/test_connection.py
@@ -2,8 +2,10 @@ import sys
 from unittest.mock import Mock, patch
 import pytest
 from sqlalchemy.engine import Engine
+import sql.connection
 from sql.connection import Connection
 from IPython.core.error import UsageError
+import sqlglot
 
 
 @pytest.fixture
@@ -36,15 +38,146 @@ def test_alias(cleanup):
     assert list(Connection.connections) == ["some-alias"]
 
 
-def test_get_curr_connection_info(mock_postgres):
+def test_get_curr_sqlalchemy_connection_info(mock_postgres):
     Connection.from_connect_str("postgresql://user:topsecret@somedomain.com/db")
-    assert Connection._get_curr_connection_info() == {
+    assert Connection._get_curr_sqlalchemy_connection_info() == {
         "dialect": "postgresql",
         "driver": "psycopg2",
         "server_version_info": None,
     }
 
 
+def test_get_curr_sqlglot_dialect_no_curr_connection(monkeypatch):
+    monkeypatch.setattr(
+        Connection, "_get_curr_sqlalchemy_connection_info", lambda: None
+    )
+    assert Connection._get_curr_sqlglot_dialect() is None
+
+
+@pytest.mark.parametrize(
+    "sqlalchemy_connection_info, expected_sqlglot_dialect",
+    [
+        (
+            {
+                "dialect": "duckdb",
+                "driver": "duckdb_engine",
+                "server_version_info": [8, 0],
+            },
+            "duckdb",
+        ),
+        (
+            {
+                "dialect": "mysql",
+                "driver": "pymysql",
+                "server_version_info": [10, 10, 3, 10, 3],
+            },
+            "mysql",
+        ),
+        # sqlalchemy and sqlglot have different dialect name, test the mapping dict
+        (
+            {
+                "dialect": "sqlalchemy_mock_dialect_name",
+                "driver": "sqlalchemy_mock_driver_name",
+                "server_version_info": [0],
+            },
+            "sqlglot_mock_dialect",
+        ),
+        (
+            {
+                "dialect": "only_support_in_sqlalchemy_dialect",
+                "driver": "sqlalchemy_mock_driver_name",
+                "server_version_info": [0],
+            },
+            "only_support_in_sqlalchemy_dialect",
+        ),
+    ],
+)
+def test_get_curr_sqlglot_dialect(
+    monkeypatch, sqlalchemy_connection_info, expected_sqlglot_dialect
+):
+    """To test if we can get the dialect name in sqlglot package scope
+
+    Args:
+        monkeypatch (fixture): A convenient fixture for monkey-patching
+        sqlalchemy_connection_info (dict): The metadata about the current dialect
+        expected_sqlglot_dialect (str): Expected sqlglot dialect name
+    """
+    monkeypatch.setattr(
+        Connection,
+        "_get_curr_sqlalchemy_connection_info",
+        lambda: sqlalchemy_connection_info,
+    )
+    monkeypatch.setattr(
+        sql.connection,
+        "DIALECT_NAME_SQLALCHEMY_TO_SQLGLOT_MAPPING",
+        {"sqlalchemy_mock_dialect_name": "sqlglot_mock_dialect"},
+    )
+    assert Connection._get_curr_sqlglot_dialect() == expected_sqlglot_dialect
+
+
+@pytest.mark.parametrize(
+    "cur_dialect, expected_support_backtick",
+    [
+        ("mysql", True),
+        ("sqlite", True),
+        ("postgres", False),
+    ],
+)
+def test_is_use_backtick_template(monkeypatch, cur_dialect, expected_support_backtick):
+    """To test if we can get the backtick supportive information from different dialects
+
+    Args:
+        monkeypatch (fixture): A convenient fixture for monkey-patching
+        cur_dialect (bool): Patched dialect name
+        expected_support_backtick (bool): Excepted boolean value to indicate
+        if the dialect supports backtick identifier
+    """
+    monkeypatch.setattr(Connection, "_get_curr_sqlglot_dialect", lambda: cur_dialect)
+    assert Connection.is_use_backtick_template() == expected_support_backtick
+
+
+def test_is_use_backtick_template_sqlglot_missing_dialect_ValueError(monkeypatch):
+    """Since accessing missing dialect will raise ValueError from sqlglot, we assume
+    that's not support case
+    """
+    monkeypatch.setattr(
+        Connection, "_get_curr_sqlglot_dialect", lambda: "something_weird_dialect"
+    )
+    assert Connection.is_use_backtick_template() is False
+
+
+def test_is_use_backtick_template_sqlglot_missing_tokenizer_AttributeError(monkeypatch):
+    """Since accessing the dialect without Tokenizer Class will raise AttributeError
+    from sqlglot, we assume that's not support case
+    """
+    monkeypatch.setattr(Connection, "_get_curr_sqlglot_dialect", lambda: "mysql")
+    monkeypatch.setattr(sqlglot.Dialect.get_or_raise("mysql"), "Tokenizer", None)
+
+    assert Connection.is_use_backtick_template() is False
+
+
+def test_is_use_backtick_template_sqlglot_missing_identifiers_TypeError(monkeypatch):
+    """Since accessing the IDENTIFIERS list of the dialect's Tokenizer Class
+    will raise TypeError from sqlglot, we assume that's not support case
+    """
+    monkeypatch.setattr(Connection, "_get_curr_sqlglot_dialect", lambda: "mysql")
+    monkeypatch.setattr(
+        sqlglot.Dialect.get_or_raise("mysql").Tokenizer, "IDENTIFIERS", None
+    )
+    assert Connection.is_use_backtick_template() is False
+
+
+def test_is_use_backtick_template_sqlglot_empty_identifiers(monkeypatch):
+    """Since looking up the "`" symbol in IDENTIFIERS list of the dialect's
+    Tokenizer Class will raise TypeError from sqlglot, we assume that's not support case
+    """
+    monkeypatch.setattr(Connection, "_get_curr_sqlglot_dialect", lambda: "mysql")
+    monkeypatch.setattr(
+        sqlglot.Dialect.get_or_raise("mysql").Tokenizer, "IDENTIFIERS", []
+    )
+    assert Connection.is_use_backtick_template() is False
+
+
 # Mock the missing package
 # Ref: https://stackoverflow.com/a/28361013
 def test_missing_duckdb_dependencies(cleanup, monkeypatch):
@@ -96,4 +229,4 @@ def test_missing_driver(
 
 def test_no_current_connection_and_get_info(monkeypatch):
     monkeypatch.setattr(Connection, "current", None)
-    assert Connection._get_curr_connection_info() is None
+    assert Connection._get_curr_sqlalchemy_connection_info() is None
diff --git a/src/tests/test_store.py b/src/tests/test_store.py
index 9522d7f..c3371bb 100644
--- a/src/tests/test_store.py
+++ b/src/tests/test_store.py
@@ -74,6 +74,10 @@ def test_key():
         store.store("first", "SELECT * FROM first WHERE x > 20", with_=["first"])
 
 
+@pytest.mark.parametrize(
+    "is_dialect_support_backtick",
+    [(True), (False)],
+)
 @pytest.mark.parametrize(
     "with_",
     [
@@ -89,7 +93,26 @@ def test_key():
         "redundant-end",
     ],
 )
-def test_serial(with_):
+def test_serial(with_, is_dialect_support_backtick, monkeypatch):
+    """To test if SQLStore can store multiple with sql clause
+    and parse into final combined sql clause
+
+    Parameters
+    ----------
+    with_ : string
+        The key to use in with sql clause
+    is_dialect_support_backtick : bool
+        If the current connected dialect support `(backtick) syntax
+    monkeypatch : Monkeypatch
+        A convenient fixture for monkey-patching
+    """
+    monkeypatch.setattr(
+        Connection,
+        "is_use_backtick_template",
+        lambda: is_dialect_support_backtick,
+    )
+    identifier = "`" if is_dialect_support_backtick else '"'
+
     store = SQLStore()
     store.store("first", "SELECT * FROM a WHERE x > 10")
     store.store("second", "SELECT * FROM first WHERE x > 20", with_=["first"])
@@ -97,16 +120,44 @@ def test_serial(with_):
     store.store("third", "SELECT * FROM second WHERE x > 30", with_=["second", "first"])
 
     result = store.render("SELECT * FROM third", with_=with_)
+
     assert (
-        str(result) == "WITH first AS (SELECT * FROM a WHERE x > 10), "
-        "second AS (SELECT * FROM first WHERE x > 20), "
-        "third AS (SELECT * FROM second WHERE x > 30) SELECT * FROM third"
+        str(result)
+        == "WITH {0}first{0} AS (SELECT * FROM a WHERE x > 10), \
+{0}second{0} AS (SELECT * FROM first WHERE x > 20), \
+{0}third{0} AS (SELECT * FROM second WHERE x > 30)SELECT * FROM third".format(
+            identifier
+        )
     )
 
 
-def test_branch_root():
-    store = SQLStore()
+@pytest.mark.parametrize(
+    "is_dialect_support_backtick",
+    [(True), (False)],
+)
+def test_branch_root(is_dialect_support_backtick, monkeypatch):
+    """To test if SQLStore can store multiple with sql clause,
+    but with each with clause has it's own sub-query.
+    To see if SQLStore can parse into final combined sql clause
+
+    Parameters
+    ----------
+    with_ : string
+        The key to use in with sql clause
+    is_dialect_support_backtick : bool
+        If the current connected dialect support `(backtick) syntax
+    monkeypatch : Monkeypatch
+        A convenient fixture for monkey-patching
+    """
+
+    monkeypatch.setattr(
+        Connection,
+        "is_use_backtick_template",
+        lambda: is_dialect_support_backtick,
+    )
+    identifier = "`" if is_dialect_support_backtick else '"'
 
+    store = SQLStore()
     store.store("first_a", "SELECT * FROM a WHERE x > 10")
     store.store("second_a", "SELECT * FROM first_a WHERE x > 20", with_=["first_a"])
     store.store("third_a", "SELECT * FROM second_a WHERE x > 30", with_=["second_a"])
@@ -115,14 +166,41 @@ def test_branch_root():
 
     result = store.render("SELECT * FROM third", with_=["third_a", "first_b"])
     assert (
-        str(result) == "WITH first_a AS (SELECT * FROM a WHERE x > 10), "
-        "second_a AS (SELECT * FROM first_a WHERE x > 20), "
-        "third_a AS (SELECT * FROM second_a WHERE x > 30), "
-        "first_b AS (SELECT * FROM b WHERE y > 10) SELECT * FROM third"
+        str(result)
+        == "WITH {0}first_a{0} AS (SELECT * FROM a WHERE x > 10), \
+{0}second_a{0} AS (SELECT * FROM first_a WHERE x > 20), \
+{0}third_a{0} AS (SELECT * FROM second_a WHERE x > 30), \
+{0}first_b{0} AS (SELECT * FROM b WHERE y > 10)SELECT * FROM third".format(
+            identifier
+        )
     )
 
 
-def test_branch_root_reverse_final_with():
+@pytest.mark.parametrize(
+    "is_dialect_support_backtick",
+    [(True), (False)],
+)
+def test_branch_root_reverse_final_with(is_dialect_support_backtick, monkeypatch):
+    """To test if SQLStore can store multiple with sql clause,
+    but with different reverse order in with_ parameter.
+    To see if SQLStore can parse into final combined sql clause
+
+    Parameters
+    ----------
+    with_ : string
+        The key to use in with sql clause
+    is_dialect_support_backtick : bool
+        If the current connected dialect support `(backtick) syntax
+    monkeypatch : Monkeypatch
+        A convenient fixture for monkey-patching
+    """
+    monkeypatch.setattr(
+        Connection,
+        "is_use_backtick_template",
+        lambda: is_dialect_support_backtick,
+    )
+    identifier = "`" if is_dialect_support_backtick else '"'
+
     store = SQLStore()
 
     store.store("first_a", "SELECT * FROM a WHERE x > 10")
@@ -133,14 +211,39 @@ def test_branch_root_reverse_final_with():
 
     result = store.render("SELECT * FROM third", with_=["first_b", "third_a"])
     assert (
-        str(result) == "WITH first_a AS (SELECT * FROM a WHERE x > 10), "
-        "second_a AS (SELECT * FROM first_a WHERE x > 20), "
-        "first_b AS (SELECT * FROM b WHERE y > 10), "
-        "third_a AS (SELECT * FROM second_a WHERE x > 30) SELECT * FROM third"
+        str(result)
+        == "WITH {0}first_a{0} AS (SELECT * FROM a WHERE x > 10), \
+{0}second_a{0} AS (SELECT * FROM first_a WHERE x > 20), \
+{0}first_b{0} AS (SELECT * FROM b WHERE y > 10), \
+{0}third_a{0} AS (SELECT * FROM second_a WHERE x > 30)SELECT * FROM third".format(
+            identifier
+        )
     )
 
 
-def test_branch():
+@pytest.mark.parametrize(
+    "is_dialect_support_backtick",
+    [(True), (False)],
+)
+def test_branch(is_dialect_support_backtick, monkeypatch):
+    """To test if SQLStore can store multiple with sql clause,
+    but some sub-queries have same with_ dependency.
+    To see if SQLStore can parse into final combined sql clause
+
+    Parameters
+    ----------
+    with_ : string
+        The key to use in with sql clause
+    monkeypatch : Monkeypatch
+        A convenient fixture for monkey-patching
+    """
+    monkeypatch.setattr(
+        Connection,
+        "is_use_backtick_template",
+        lambda: is_dialect_support_backtick,
+    )
+    identifier = "`" if is_dialect_support_backtick else '"'
+
     store = SQLStore()
 
     store.store("first_a", "SELECT * FROM a WHERE x > 10")
@@ -151,8 +254,11 @@ def test_branch():
 
     result = store.render("SELECT * FROM third", with_=["first_b", "third_a"])
     assert (
-        str(result) == "WITH first_a AS (SELECT * FROM a WHERE x > 10), "
-        "second_a AS (SELECT * FROM first_a WHERE x > 20), "
-        "first_b AS (SELECT * FROM second_a WHERE y > 10), "
-        "third_a AS (SELECT * FROM second_a WHERE x > 30) SELECT * FROM third"
+        str(result)
+        == "WITH {0}first_a{0} AS (SELECT * FROM a WHERE x > 10), \
+{0}second_a{0} AS (SELECT * FROM first_a WHERE x > 20), \
+{0}first_b{0} AS (SELECT * FROM second_a WHERE y > 10), \
+{0}third_a{0} AS (SELECT * FROM second_a WHERE x > 30)SELECT * FROM third".format(
+            identifier
+        )
     )