Fix bug when editing cube node with materialization (#1362)

* Fix bug where editing a cube node with a materialization configured can error out based on there being no lookback_window

* Fix and add new tests

* Lint

Author: shangyian

datajunction-server/datajunction_server/internal/materializations.py

@@ -173,9 +173,14 @@ async def build_cube_materialization_config(
                     if col.semantic_type == SemanticType.DIMENSION
                 ],
                 measures=metrics_expressions,
-                spark=upsert_input.config.spark,
+                spark=upsert_input.config.spark.__root__
+                if hasattr(upsert_input, "config") and upsert_input.config.spark
+                else {},
                 upstream_tables=measures_query.upstream_tables,
                 columns=measures_query.columns,
+                lookback_window=upsert_input.lookback_window
+                if hasattr(upsert_input, "lookback_window")
+                else "",
             )
         return generic_config
     except (KeyError, ValidationError, AttributeError) as exc:  # pragma: no cover

datajunction-server/datajunction_server/internal/nodes.py

@@ -874,11 +874,17 @@ async def update_cube_node(
     ]
     if major_changes and active_materializations:
         for old in active_materializations:
+            # Once we've migrated all materializations to the new format, we should only
+            # be using UpsertCubeMaterialization for cube nodes
+            job_type = MaterializationJobTypeEnum.find_match(old.job)
+            materialization_upsert_class = UpsertMaterialization
+            if job_type == MaterializationJobTypeEnum.DRUID_CUBE:
+                materialization_upsert_class = UpsertCubeMaterialization
             new_cube_revision.materializations.append(
                 await create_new_materialization(
                     session,
                     new_cube_revision,
-                    UpsertMaterialization(
+                    materialization_upsert_class(
                         **MaterializationConfigOutput.from_orm(old).dict(
                             exclude={"job"},
                         ),

datajunction-server/tests/api/cubes_test.py

@@ -1609,6 +1609,68 @@ async def test_updating_cube(
     ]
 
 
+@pytest.mark.asyncio
+async def test_updating_cube_with_existing_cube_materialization(
+    client_with_repairs_cube: AsyncClient,
+    module__query_service_client: QueryServiceClient,
+):
+    """
+    Verify updating a cube with an existing new-style cube materialization
+    """
+    cube_name = "default.repairs_cube__default_incremental_11"
+    await make_a_test_cube(
+        client_with_repairs_cube,
+        cube_name,
+    )
+    response = await client_with_repairs_cube.post(
+        f"/nodes/{cube_name}/columns/default.hard_hat.hire_date/partition",
+        json={
+            "type_": "temporal",
+            "granularity": "day",
+            "format": "yyyyMMdd",
+        },
+    )
+    assert response.status_code in (200, 201)
+    response = await client_with_repairs_cube.post(
+        f"/nodes/{cube_name}/materialization/",
+        json={
+            "job": "druid_cube",
+            "strategy": "incremental_time",
+            "schedule": "@daily",
+            "lookback_window": "1 DAY",
+        },
+    )
+    # Update the cube, but keep the temporal partition column. This should succeed
+    response = await client_with_repairs_cube.patch(
+        f"/nodes/{cube_name}",
+        json={
+            "metrics": ["default.discounted_orders_rate"],
+            "dimensions": ["default.hard_hat.city", "default.hard_hat.hire_date"],
+        },
+    )
+    result = response.json()
+    assert result["version"] == "v2.0"
+
+    # Check that the configured materialization was updated
+    response = await client_with_repairs_cube.get(f"/cubes/{cube_name}/")
+    data = response.json()
+    assert [
+        col["semantic_entity"]
+        for col in data["materializations"][0]["config"]["columns"]
+    ] == [
+        "default.hard_hat.city",
+        "default.hard_hat.hire_date",
+        "default.discounted_orders_rate.default_DOT_discounted_orders_rate",
+    ]
+    assert data["materializations"][0]["job"] == "DruidMetricsCubeMaterializationJob"
+    assert (
+        data["materializations"][0]["name"]
+        == "druid_metrics_cube__incremental_time__default.hard_hat.hire_date"
+    )
+    assert data["materializations"][0]["strategy"] == "incremental_time"
+    assert data["materializations"][0]["schedule"] == "@daily"
+
+
 @pytest.mark.asyncio
 async def test_updating_cube_with_existing_materialization(
     client_with_repairs_cube: AsyncClient,