Repository: TobikoData/sqlmesh
Pull request #2808 (2024-06-21 20:01:10)

Fix: Support Airflow URLs without trailing slashes


Patch:
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
index f5c3e444..bbd656fb 100644
--- a/sqlmesh/core/model/definition.py
+++ b/sqlmesh/core/model/definition.py
@@ -423,8 +423,12 @@ class _Model(ModelMeta, frozen=True):
         query = self.render_query_or_raise(**render_kwarg).copy()
 
         for select_or_set_op in query.find_all(exp.Select, exp.SetOperation):
-            cte = select_or_set_op.find_ancestor(exp.With, exp.Select, exp.Subquery)
-            skip_limit = isinstance(cte, exp.With) and cte.recursive
+            skip_limit = False
+            ancestor = select_or_set_op.parent
+            while ancestor and not skip_limit:
+                if isinstance(ancestor, exp.With) and ancestor.recursive:
+                    skip_limit = True
+                ancestor = ancestor.parent
 
             if isinstance(select_or_set_op, exp.Select) and select_or_set_op.args.get("from"):
                 select_or_set_op.where(exp.false(), copy=False)
diff --git a/sqlmesh/schedulers/airflow/client.py b/sqlmesh/schedulers/airflow/client.py
index f2d1a657..e0d93be0 100644
--- a/sqlmesh/schedulers/airflow/client.py
+++ b/sqlmesh/schedulers/airflow/client.py
@@ -40,6 +40,8 @@ VERSIONS_PATH = f"{common.SQLMESH_API_BASE_PATH}/versions"
 class BaseAirflowClient(abc.ABC):
     def __init__(self, airflow_url: str, console: t.Optional[Console]):
         self._airflow_url = airflow_url
+        if not self._airflow_url.endswith("/"):
+            self._airflow_url += "/"
         self._console = console
 
     @property


Test patch:
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index 10d5c00b..687fbc22 100644
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -2163,6 +2163,45 @@ def test_model_ctas_query():
         == 'WITH RECURSIVE "a" AS (SELECT * FROM "x" AS "x" WHERE FALSE), "b" AS (SELECT * FROM "a" AS "a" WHERE FALSE UNION ALL SELECT * FROM "a" AS "a" WHERE FALSE) SELECT * FROM "b" AS "b" WHERE FALSE LIMIT 0'
     )
 
+    expressions = d.parse(
+        """
+        MODEL (name `a-b-c.table`, kind FULL, dialect bigquery);
+        WITH RECURSIVE a AS (
+            SELECT * FROM (SELECT * FROM (SELECT * FROM x))
+        ), b AS (
+            SELECT * FROM a UNION ALL SELECT * FROM a
+        )
+        SELECT * FROM b
+
+    """
+    )
+
+    assert (
+        load_sql_based_model(expressions, dialect="bigquery").ctas_query().sql()
+        == 'WITH RECURSIVE "a" AS (SELECT * FROM (SELECT * FROM (SELECT * FROM "x" AS "x" WHERE FALSE) AS "_q_0" WHERE FALSE) AS "_q_1" WHERE FALSE), "b" AS (SELECT * FROM "a" AS "a" WHERE FALSE UNION ALL SELECT * FROM "a" AS "a" WHERE FALSE) SELECT * FROM "b" AS "b" WHERE FALSE LIMIT 0'
+    )
+
+    expressions = d.parse(
+        """
+        MODEL (name `a-b-c.table`, kind FULL, dialect bigquery);
+        WITH RECURSIVE a AS (
+            WITH nested_a AS (
+                SELECT * FROM (SELECT * FROM (SELECT * FROM x))
+            )
+            SELECT * FROM nested_a
+        ), b AS (
+            SELECT * FROM a UNION ALL SELECT * FROM a
+        )
+        SELECT * FROM b
+
+    """
+    )
+
+    assert (
+        load_sql_based_model(expressions, dialect="bigquery").ctas_query().sql()
+        == 'WITH RECURSIVE "a" AS (WITH "nested_a" AS (SELECT * FROM (SELECT * FROM (SELECT * FROM "x" AS "x" WHERE FALSE) AS "_q_0" WHERE FALSE) AS "_q_1" WHERE FALSE) SELECT * FROM "nested_a" AS "nested_a" WHERE FALSE), "b" AS (SELECT * FROM "a" AS "a" WHERE FALSE UNION ALL SELECT * FROM "a" AS "a" WHERE FALSE) SELECT * FROM "b" AS "b" WHERE FALSE LIMIT 0'
+    )
+
 
 def test_is_breaking_change():
     model = create_external_model("a", columns={"a": "int", "limit": "int"})
diff --git a/tests/schedulers/airflow/test_client.py b/tests/schedulers/airflow/test_client.py
index 9dd7cd88..5384dc74 100644
--- a/tests/schedulers/airflow/test_client.py
+++ b/tests/schedulers/airflow/test_client.py
@@ -466,3 +466,18 @@ def test_get_variable(mocker: MockerFixture):
     assert client.get_variable("test_key") == "test_value"
 
     get_variable_mock.assert_called_once_with("http://localhost:8080/api/v1/variables/test_key")
+
+
+def test_url_no_trailing_slash(mocker: MockerFixture, snapshot: Snapshot):
+    get_variable_response_mock = mocker.Mock()
+    get_variable_response_mock.status_code = 200
+    get_variable_response_mock.json.return_value = {"value": "test_value", "key": "test_key"}
+    get_variable_mock = mocker.patch("requests.Session.get")
+    get_variable_mock.return_value = get_variable_response_mock
+
+    client = AirflowClient(airflow_url="http://localhost:8080/prefix", session=requests.Session())
+    client.get_variable("test_key")
+
+    get_variable_mock.assert_called_once_with(
+        "http://localhost:8080/prefix/api/v1/variables/test_key"
+    )