Feat: Enable Iceberg support for Snowflake (#4262)

Author: erindru

.circleci/config.yml

@@ -17,7 +17,7 @@ on_tag_filter: &on_tag_filter
       only: /^v.+/
 
 orbs:
-  path-filtering: circleci/path-filtering@1.1.0
+  path-filtering: circleci/path-filtering@1.2.0
 
 jobs:
   vscode-extension-setup:
@@ -116,6 +116,7 @@ workflows:
             pytest.ini|setup.cfg|setup.py python true
             \.circleci/.*|Makefile|\.pre-commit-config\.yaml common true
             vscode/extensions/.* vscode true
+          tag: "3.9"
 
       - vscode-extension-setup:
           <<: *on_main_or_tag_filter

docs/integrations/engines/snowflake.md

@@ -525,7 +525,7 @@ MODEL (
   ),
 );
 ```
- 
+
 ## Custom View and Table types
 
 SQLMesh supports custom view and table types for Snowflake models. You can apply these modifiers to either the physical layer or virtual layer of a model using the `physical_properties` and `virtual_properties` attributes respectively. For example:
@@ -535,8 +535,8 @@ SQLMesh supports custom view and table types for Snowflake models. You can apply
 A table can be exposed through a `SECURE` view in the virtual layer by specifying the `creatable_type` property and setting it to `SECURE`:
 
 ```sql linenums="1"
-Model (
-  name = schema_name.model_name,
+MODEL (
+  name schema_name.model_name,
   virtual_properties (
       creatable_type = SECURE
   )
@@ -550,8 +550,8 @@ SELECT a FROM schema_name.model_b;
 A model can use a `TRANSIENT` table in the physical layer by specifying the `creatable_type` property and setting it to `TRANSIENT`:
 
 ```sql linenums="1"
-Model (
-  name = schema_name.model_name,
+MODEL (
+  name schema_name.model_name,
   physical_properties (
       creatable_type = TRANSIENT
   )
@@ -560,6 +560,52 @@ Model (
 SELECT a FROM schema_name.model_b;
 ```
 
+### Iceberg Tables
+
+In order for Snowflake to be able to create an Iceberg table, there must be an [External Volume](https://docs.snowflake.com/en/user-guide/tables-iceberg-configure-external-volume) configured to store the Iceberg table data on.
+
+Once that is configured, you can create a model backed by an Iceberg table by using `table_format iceberg` like so:
+
+```sql linenums="1" hl_lines="4 6-7"
+MODEL (
+  name schema_name.model_name,
+  kind FULL,
+  table_format iceberg,
+  physical_properties (
+    catalog = 'snowflake',
+    external_volume = '<external volume name>'
+  )
+);
+```
+
+To prevent having to specify `catalog = 'snowflake'` and `external_volume = '<external volume name>'` on every model, see the Snowflake documentation for:
+
+  - [Configuring a default Catalog](https://docs.snowflake.com/en/user-guide/tables-iceberg-configure-catalog-integration#set-a-default-catalog-at-the-account-database-or-schema-level)
+  - [Configuring a default External Volume](https://docs.snowflake.com/en/user-guide/tables-iceberg-configure-external-volume#set-a-default-external-volume-at-the-account-database-or-schema-level)
+
+Alternatively you can also use [model defaults](../../guides/configuration.md#model-defaults) to set defaults at the SQLMesh level instead.
+
+To utilize the wide variety of [optional properties](https://docs.snowflake.com/en/sql-reference/sql/create-iceberg-table-snowflake#optional-parameters) that Snowflake makes available for Iceberg tables, simply specify them as `physical_properties`:
+
+```sql linenums="1" hl_lines="8"
+MODEL (
+  name schema_name.model_name,
+  kind FULL,
+  table_format iceberg,
+  physical_properties (
+    catalog = 'snowflake',
+    external_volume = 'my_external_volume',
+    base_location = 'my/product_reviews/'
+  )
+);
+```
+
+!!! warning "External catalogs"
+
+    Setting `catalog = 'snowflake'` to use Snowflake's internal catalog is a good default because SQLMesh needs to be able to write to the tables it's managing and Snowflake [does not support](https://docs.snowflake.com/en/user-guide/tables-iceberg#catalog-options) writing to Iceberg tables configured under external catalogs.
+
+    You can however still reference a table from an external catalog in your model as a normal [external table](../../concepts/models/external_models.md).
+
 ## Troubleshooting
 
 ### Frequent Authentication Prompts

sqlmesh/core/engine_adapter/snowflake.py

@@ -112,6 +112,38 @@ def snowpark(self) -> t.Optional[SnowparkSession]:
     def catalog_support(self) -> CatalogSupport:
         return CatalogSupport.FULL_SUPPORT
 
+    def _create_table(
+        self,
+        table_name_or_schema: t.Union[exp.Schema, TableName],
+        expression: t.Optional[exp.Expression],
+        exists: bool = True,
+        replace: bool = False,
+        columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
+        table_description: t.Optional[str] = None,
+        column_descriptions: t.Optional[t.Dict[str, str]] = None,
+        table_kind: t.Optional[str] = None,
+        **kwargs: t.Any,
+    ) -> None:
+        table_format = kwargs.get("table_format")
+        if table_format and isinstance(table_format, str):
+            table_format = table_format.upper()
+            if not table_kind:
+                table_kind = f"{table_format} TABLE"
+            elif table_kind == self.MANAGED_TABLE_KIND:
+                table_kind = f"DYNAMIC {table_format} TABLE"
+
+        super()._create_table(
+            table_name_or_schema=table_name_or_schema,
+            expression=expression,
+            exists=exists,
+            replace=replace,
+            columns_to_types=columns_to_types,
+            table_description=table_description,
+            column_descriptions=column_descriptions,
+            table_kind=table_kind,
+            **kwargs,
+        )
+
     def create_managed_table(
         self,
         table_name: TableName,
@@ -230,7 +262,8 @@ def _build_table_properties_exp(
         if table_properties:
             table_properties = {k.upper(): v for k, v in table_properties.items()}
             # if we are creating a non-dynamic table; remove any properties that are only valid for dynamic tables
-            if table_kind != self.MANAGED_TABLE_KIND:
+            # this is necessary because we create "normal" tables from the same managed model definition for dev previews and the "normal" tables dont support these parameters
+            if "DYNAMIC" not in (table_kind or "").upper():
                 for prop in {"WAREHOUSE", "TARGET_LAG", "REFRESH_MODE", "INITIALIZE"}:
                     table_properties.pop(prop, None)
 

sqlmesh/core/snapshot/evaluator.py

@@ -2098,6 +2098,7 @@ def create(
                 table_properties=kwargs.get("physical_properties", model.physical_properties),
                 table_description=model.description,
                 column_descriptions=model.column_descriptions,
+                table_format=model.table_format,
             )
         elif not is_table_deployable:
             # Only create the dev preview table as a normal table.
@@ -2134,6 +2135,7 @@ def insert(
                 table_properties=kwargs.get("physical_properties", model.physical_properties),
                 table_description=model.description,
                 column_descriptions=model.column_descriptions,
+                table_format=model.table_format,
             )
         elif not is_snapshot_deployable:
             # Snapshot isnt deployable; update the preview table instead

tests/core/engine_adapter/integration/test_integration_snowflake.py

@@ -163,3 +163,50 @@ def _get_data_object(table: exp.Table) -> DataObject:
 
     metadata = _get_data_object(target_table_1)
     assert not metadata.is_clustered
+
+
+def test_create_iceberg_table(ctx: TestContext, engine_adapter: SnowflakeEngineAdapter) -> None:
+    # Note: this test relies on a default Catalog and External Volume being configured in Snowflake
+    # ref: https://docs.snowflake.com/en/user-guide/tables-iceberg-configure-catalog-integration#set-a-default-catalog-at-the-account-database-or-schema-level
+    # ref: https://docs.snowflake.com/en/user-guide/tables-iceberg-configure-external-volume#set-a-default-external-volume-at-the-account-database-or-schema-level
+    # This has been done on the Snowflake account used by CI
+
+    model_name = ctx.table("TEST")
+    managed_model_name = ctx.table("TEST_DYNAMIC")
+    sqlmesh = ctx.create_context()
+
+    model = load_sql_based_model(
+        d.parse(f"""
+            MODEL (
+                name {model_name},
+                kind FULL,
+                table_format iceberg,
+                dialect 'snowflake'
+            );
+
+            select 1 as "ID", 'foo' as "NAME";
+            """)
+    )
+
+    managed_model = load_sql_based_model(
+        d.parse(f"""
+            MODEL (
+                name {managed_model_name},
+                kind MANAGED,
+                physical_properties (
+                    target_lag = '20 minutes'
+                ),
+                table_format iceberg,
+                dialect 'snowflake'
+            );
+
+            select "ID", "NAME" from {model_name};
+            """)
+    )
+
+    sqlmesh.upsert_model(model)
+    sqlmesh.upsert_model(managed_model)
+
+    result = sqlmesh.plan(auto_apply=True)
+
+    assert len(result.new_snapshots) == 2

tests/core/engine_adapter/test_snowflake.py

@@ -19,6 +19,13 @@
 pytestmark = [pytest.mark.engine, pytest.mark.snowflake]
 
 
+@pytest.fixture
+def snowflake_mocked_engine_adapter(
+    make_mocked_engine_adapter: t.Callable,
+) -> SnowflakeEngineAdapter:
+    return make_mocked_engine_adapter(SnowflakeEngineAdapter)
+
+
 def test_get_temp_table(mocker: MockerFixture, make_mocked_engine_adapter: t.Callable):
     adapter = make_mocked_engine_adapter(SnowflakeEngineAdapter)
 
@@ -270,11 +277,25 @@ def test_create_managed_table(make_mocked_engine_adapter: t.Callable, mocker: Mo
         },
     )
 
+    # table_format=iceberg
+    adapter.create_managed_table(
+        table_name="test_table",
+        query=query,
+        columns_to_types=columns_to_types,
+        table_properties={
+            "target_lag": exp.Literal.string("20 minutes"),
+            "catalog": exp.Literal.string("snowflake"),
+            "external_volume": exp.Literal.string("test"),
+        },
+        table_format="iceberg",
+    )
+
     assert to_sql_calls(adapter) == [
         """CREATE OR REPLACE DYNAMIC TABLE "test_table" TARGET_LAG='20 minutes' WAREHOUSE="default_warehouse" AS SELECT CAST("a" AS INT) AS "a", CAST("b" AS INT) AS "b" FROM (SELECT "a", "b" FROM "source_table") AS "_subquery\"""",
         """CREATE OR REPLACE DYNAMIC TABLE "test_table" TARGET_LAG='20 minutes' WAREHOUSE="foo" AS SELECT CAST("a" AS INT) AS "a", CAST("b" AS INT) AS "b" FROM (SELECT "a", "b" FROM "source_table") AS "_subquery\"""",
         """CREATE OR REPLACE DYNAMIC TABLE "test_table" CLUSTER BY ("a") TARGET_LAG='20 minutes' WAREHOUSE="default_warehouse" AS SELECT CAST("a" AS INT) AS "a", CAST("b" AS INT) AS "b" FROM (SELECT "a", "b" FROM "source_table") AS "_subquery\"""",
         """CREATE OR REPLACE DYNAMIC TABLE "test_table" TARGET_LAG='20 minutes' REFRESH_MODE='auto' INITIALIZE='on_create' WAREHOUSE="default_warehouse" AS SELECT CAST("a" AS INT) AS "a", CAST("b" AS INT) AS "b" FROM (SELECT "a", "b" FROM "source_table") AS "_subquery\"""",
+        """CREATE OR REPLACE DYNAMIC ICEBERG TABLE "test_table" TARGET_LAG='20 minutes' CATALOG='snowflake' EXTERNAL_VOLUME='test' WAREHOUSE="default_warehouse" AS SELECT CAST("a" AS INT) AS "a", CAST("b" AS INT) AS "b" FROM (SELECT "a", "b" FROM "source_table") AS "_subquery\"""",
     ]
 
 
@@ -666,3 +687,44 @@ def test_clone_table(mocker: MockerFixture, make_mocked_engine_adapter: t.Callab
     adapter.cursor.execute.assert_called_once_with(
         'CREATE TABLE "target_table" CLONE "source_table"'
     )
+
+
+def test_table_format_iceberg(snowflake_mocked_engine_adapter: SnowflakeEngineAdapter) -> None:
+    adapter = snowflake_mocked_engine_adapter
+
+    model = load_sql_based_model(
+        expressions=d.parse("""
+        MODEL (
+            name test.table,
+            kind full,
+            table_format iceberg,
+            physical_properties (
+              catalog = 'snowflake',
+              external_volume = 'test'
+            )
+        );
+        SELECT a::INT;
+        """)
+    )
+    assert isinstance(model, SqlModel)
+    assert model.table_format == "iceberg"
+
+    adapter.create_table(
+        table_name=model.name,
+        columns_to_types=model.columns_to_types_or_raise,
+        table_format=model.table_format,
+        table_properties=model.physical_properties,
+    )
+
+    adapter.ctas(
+        table_name=model.name,
+        query_or_df=model.render_query_or_raise(),
+        columns_to_types=model.columns_to_types_or_raise,
+        table_format=model.table_format,
+        table_properties=model.physical_properties,
+    )
+
+    assert to_sql_calls(adapter) == [
+        'CREATE ICEBERG TABLE IF NOT EXISTS "test"."table" ("a" INT) CATALOG=\'snowflake\' EXTERNAL_VOLUME=\'test\'',
+        'CREATE ICEBERG TABLE IF NOT EXISTS "test"."table" CATALOG=\'snowflake\' EXTERNAL_VOLUME=\'test\' AS SELECT CAST("a" AS INT) AS "a" FROM (SELECT CAST("a" AS INT) AS "a") AS "_subquery"',
+    ]

tests/core/test_snapshot_evaluator.py

@@ -3367,6 +3367,7 @@ def test_create_managed(adapter_mock, make_snapshot, mocker: MockerFixture):
         table_properties=model.physical_properties,
         table_description=model.description,
         column_descriptions=model.column_descriptions,
+        table_format=None,
     )