Repository: TobikoData/sqlmesh
Pull request #2807 (2024-06-21 19:51:01)

Docs: Add custom materializations guide


Patch:
diff --git a/docs/guides/custom_materializations.md b/docs/guides/custom_materializations.md
new file mode 100644
index 00000000..4391db88
--- /dev/null
+++ b/docs/guides/custom_materializations.md
@@ -0,0 +1,153 @@
+# Custom materializations guide
+
+SQLMesh supports a variety of [model kinds](../concepts/models/model_kinds.md) to capture the most common semantics of how transformations can be evaluated and materialized.
+
+There are times, however, when a specific use case doesn't align with any of the supported materialization strategies. For scenarios like this, SQLMesh allows users to create their own materialization implementation using Python.
+
+Please note that this is an advanced feature and should only be considered if all other approaches to addressing a use case have been exhausted.
+
+## Creating a materialization
+
+The fastest way to add a new custom materialization is to add a new `.py` file with the implementation to the `materializations/` folder of the project. SQLMesh will automatically import all Python modules in this folder at project load time and register the custom materializations accordingly.
+
+To create a custom materialization strategy, you need to inherit the `CustomMaterialization` base class and, at a very minimum, provide an implementation for the `insert` method.
+
+For example, a simple custom full-refresh materialization strategy might look like the following:
+
+```python linenums="1"
+from __future__ import annotations
+
+import typing as t
+
+from sqlmesh import CustomMaterialization, Model
+
+if t.TYPE_CHECKING:
+    from sqlmesh import QueryOrDF
+
+
+class CustomFullMaterialization(CustomMaterialization):
+    NAME = "my_custom_full"
+
+    def insert(
+        self,
+        table_name: str,
+        query_or_df: QueryOrDF,
+        model: Model,
+        is_first_insert: bool,
+        **kwargs: t.Any,
+    ) -> None:
+        self.adapter.replace_query(table_name, query_or_df)
+
+```
+
+Let's unpack the above implementation:
+
+* `NAME` - determines the name of the custom materialization. This name will be used in model definitions to reference a specific strategy. If not specified, the name of the class will be used instead.
+* The `insert` method comes with the following arguments:
+    * `table_name` - the name of a target table (or a view) into which the data should be inserted.
+    * `query_or_df` - a query (a SQLGlot expression) or a DataFrame (pandas, PySpark, or Snowpark) instance which has to be inserted.
+    * `model` - the associated model definition object which can be used to get any model parameters as well as custom materialization settings.
+    * `is_first_insert` - whether this is the first insert for the current version of the model.
+    * `kwargs` - contains additional and future arguments.
+* The `self.adapter` instance is used to interact with the target engine. It comes with a set of useful high-level APIs like `replace_query`, `create_table`, and `table_exists`, but also supports execution of arbitrary SQL expressions with its `execute` method.
+
+You can also control how the associated data objects (tables, views, etc.) are created and deleted by overriding the `create` and `delete` methods accordingly:
+
+```python linenums="1"
+from __future__ import annotations
+
+import typing as t
+
+from sqlmesh import CustomMaterialization, Model
+
+
+class CustomFullMaterialization(CustomMaterialization):
+    ...
+
+    def create(
+        self,
+        table_name: str,
+        model: Model,
+        is_table_deployable: bool,
+        **render_kwargs: t.Any,
+    ) -> None:
+        # Custom creation logic.
+
+    def delete(self, name: str, **kwargs: t.Any) -> None:
+        # Custom deletion logic.
+```
+
+## Using custom materializations in models
+
+In order to use the newly created materialization, use the special model kind `CUSTOM`:
+
+```sql linenums="1"
+MODEL (
+  name my_db.my_model,
+  kind CUSTOM (materialization 'my_custom_full')
+);
+```
+
+The name of the materialization strategy is provided in the `materialization` attribute of the `CUSTOM` kind.
+
+Additionally, you can provide an optional list of arbitrary key-value pairs in the `materialization_properties` attribute:
+
+```sql linenums="1"
+MODEL (
+  name my_db.my_model,
+  kind CUSTOM (
+    materialization 'my_custom_full',
+    materialization_properties (
+      'config_key' = 'config_value'
+    )
+  )
+);
+```
+
+These properties can be accessed with the model reference within the materialization implementation:
+
+```python linenums="1" hl_lines="12"
+class CustomFullMaterialization(CustomMaterialization):
+    NAME = "my_custom_full"
+
+    def insert(
+        self,
+        table_name: str,
+        query_or_df: QueryOrDF,
+        model: Model,
+        is_first_insert: bool,
+        **kwargs: t.Any,
+    ) -> None:
+        config_value = model.custom_materialization_properties["config_key"]
+        # Proceed with implementing the insertion logic.
+```
+
+## Packaging custom materializations
+
+To share custom materializations across multiple SQLMesh projects, you need to create and publish a Python package containing your implementation.
+
+When using SQLMesh with Airflow or other external schedulers, note that the `materializations/` folder might not be available on the Airflow cluster side. Therefore, you'll need a package that can be installed there.
+
+Custom materializations can be packaged into a Python package and exposed via [setuptools entrypoints](https://packaging.python.org/en/latest/guides/creating-and-discovering-plugins/#using-package-metadata) mechanism. Once the package is installed, SQLMesh will automatically load custom materializations from the entrypoint list.
+
+If your custom materialization class is defined in the `my_package/my_materialization.py` module, you can expose it as an entry point in the `pyproject.toml` file as follows:
+
+```toml
+[project.entry-points."sqlmesh.materializations"]
+my_materialization = "my_package.my_materialization:CustomFullMaterialization"
+```
+
+Or in `setup.py`:
+
+```python
+setup(
+    ...,
+    entry_points={
+        "sqlmesh.materializations": [
+            "my_materialization = my_package.my_materialization:CustomFullMaterialization",
+        ],
+    },
+)
+```
+
+Refer to the [custom_materializations](https://github.com/TobikoData/sqlmesh/tree/main/examples/custom_materializations) package example for more details.
diff --git a/mkdocs.yml b/mkdocs.yml
index 04167f13..8bdd8f51 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -22,6 +22,7 @@ nav:
       - guides/notifications.md
       - guides/migrations.md
       - guides/table_migration.md
+      - guides/custom_materializations.md
     - Project content:
       - guides/models.md
       - guides/incremental_time.md
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/core/snapshot/evaluator.py b/sqlmesh/core/snapshot/evaluator.py
index 70e1d29e..81214ac3 100644
--- a/sqlmesh/core/snapshot/evaluator.py
+++ b/sqlmesh/core/snapshot/evaluator.py
@@ -965,7 +965,7 @@ class EvaluationStrategy(abc.ABC):
         Args:
             table_name: The name of a table or a view.
             model: The target model.
-            is_table_deployable: Whether the table that is being created is deployable.
+            is_table_deployable: Whether the table that is being created is deployable (can be deployed to the production environment).
             render_kwargs: Additional kwargs for node rendering.
         """
 
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"
+    )