[dagster-airlift] complete revamp #23369

dpeng817 · 2024-08-01T22:03:31Z

Complete project revamp.

dpeng817 · 2024-08-01T22:03:49Z

[dagster-airlift] complete revamp #23369 : 2 dependent PRs (#23370 , #23410 ) 👈
[dagster-airlift] mwaa auth backend #23344
[dagster-airlift] AirflowInstance #23343
[dagster-airlift] Fix materialization reporting order for tasks #23326
[dagster-airlift] improve unit test speed #23324
[dagster-airlift] actually put examples under test #23300
master

This stack of pull requests is managed by Graphite. Learn more about stacking.

Join @dpeng817 and the rest of your teammates on Graphite

schrockn · 2024-08-02T16:22:46Z

examples/experimental/dagster-airlift/dagster_airlift/airflow_utils.py

+# Each task_id.yaml will look something like this:
+# - asset: my/asset/key
+#   tags: {}
+#   metadata: {}
+#   deps: [upstream/asset/key]
+#   compute_pointer: my.compute.pointer
+# - asset: my/other/asset/key


this is totally configurable on a per-blueprint basis, correct?

schrockn · 2024-08-02T16:24:29Z

...s/experimental/dagster-airlift/dagster_airlift_tests/airflow_migration_project/migrated.yaml

@@ -0,0 +1,5 @@
+dags:


can easily change this later but I do think we will want one file per dag

schrockn · 2024-08-02T16:25:36Z

examples/experimental/dagster-airlift/dagster_airlift/airflow_utils.py

+class AssetComputePointer(BaseModel):
+    module_name: str
+    fn_name: str
+
+
+class AssetInfo(BaseModel):
+    asset_key: str
+    tags: Dict[str, str] = {}
+    metadata: Dict[str, Any] = {}
+    deps: List[str] = []
+    asset_compute_pointer: Optional[AssetComputePointer] = None
+
+
+class MultiAssetSpec(BaseModel):
+    asset_infos: List[AssetInfo]
+
+
+class MigratingAssetsBlueprint(BaseModel):
+    type: Literal["external_asset"] = "external_asset"
+    asset_infos: List[AssetInfo]
+    migrated: bool
+    compute_kind: str
+    name: str
+
+    def build_defs(self) -> Definitions:
+        asset_specs = [
+            AssetSpec(
+                key=AssetKey.from_user_string(info.asset_key),
+                metadata=info.metadata,
+                tags=info.tags,
+                deps=[AssetDep(AssetKey.from_user_string(dep)) for dep in info.deps],
+            )
+            for info in self.asset_infos
+        ]
+
+        @multi_asset(
+            specs=asset_specs,
+            name=self.name,
+            compute_kind=self.compute_kind,
+        )
+        def _external_asset() -> None:
+            if self.migrated:
+                for info in self.asset_infos:
+                    ptr = check.not_none(info.asset_compute_pointer)
+                    module = import_module(ptr.module_name)
+                    compute_fn = getattr(module, ptr.fn_name)
+                    compute_fn()
+
+        return Definitions(assets=[_external_asset])
+
+
+class TaskMigrationState(BaseModel):
+    task_id: str
+    migrated: bool


For our own sake, I think we should split these into different files, and make it super clear what we expect to be in the airflow process versus the dagster process. We will probably want to have different modules eventually.

Heavily agreed about the different modules. I figure for a lot of people adding a heavyweight dependency to dagster within their airflow code will be a non-starter. I've been thinking that airlift will consist of the dagster component, and then a lightweight-pipes esque package with the only dependency being requests. Maybe we could even get away without that.

I'll put them into their own file for now.

schrockn · 2024-08-02T16:26:28Z

examples/experimental/dagster-airlift/dagster_airlift/airflow_utils.py

+class MigratingAssetsBlueprint(BaseModel):
+    type: Literal["external_asset"] = "external_asset"
+    asset_infos: List[AssetInfo]
+    migrated: bool
+    compute_kind: str
+    name: str
+
+    def build_defs(self) -> Definitions:
+        asset_specs = [
+            AssetSpec(
+                key=AssetKey.from_user_string(info.asset_key),
+                metadata=info.metadata,
+                tags=info.tags,
+                deps=[AssetDep(AssetKey.from_user_string(dep)) for dep in info.deps],
+            )
+            for info in self.asset_infos
+        ]
+
+        @multi_asset(
+            specs=asset_specs,
+            name=self.name,
+            compute_kind=self.compute_kind,
+        )
+        def _external_asset() -> None:
+            if self.migrated:
+                for info in self.asset_infos:
+                    ptr = check.not_none(info.asset_compute_pointer)
+                    module = import_module(ptr.module_name)
+                    compute_fn = getattr(module, ptr.fn_name)
+                    compute_fn()
+
+        return Definitions(assets=[_external_asset])


I'd also like to see the pattern where have an existing blueprint (e.g. dbt) and have it only return specs when it is not migrated.

schrockn

This one probably worth talking through live.

schrockn · 2024-08-02T17:15:05Z

examples/experimental/dagster-airlift/dagster_airlift/airflow_utils.py

+#   metadata: {}
+#   deps: [upstream/asset/key]
+#   compute_pointer: my.compute.pointer
+# - asset: my/other/asset/key


type: python python_fn: my_module.python_fn assets: - key: tags: metadata: deps:

schrockn · 2024-08-02T17:17:42Z

examples/experimental/dagster-airlift/dagster_airlift/airflow_utils.py

+            if self.migrated:
+                for info in self.asset_infos:
+                    ptr = check.not_none(info.asset_compute_pointer)
+                    module = import_module(ptr.module_name)
+                    compute_fn = getattr(module, ptr.fn_name)
+                    compute_fn()


accessing the attr should happen in the factory function, not the body of the multi-asset

schrockn

Flushing through comments

schrockn · 2024-08-05T20:19:38Z

examples/experimental/dagster-airlift/dagster_airlift/dbt/multi_asset.py

+class FromEnvVar(BaseModel):
+    env_var: str
+
+
+class DbtProjectDefs(Blueprint):
+    type: Literal["dbt_project"] = "dbt_project"
+    dbt_project_path: FromEnvVar
+    group: Optional[str] = None
+
+    def build_defs(self) -> Definitions:
+        source_file_name_without_ext = self.source_file_name.split(".")[0]
+        dbt_project_path = Path(os.environ[self.dbt_project_path.env_var])
+        dbt_manifest_path = dbt_project_path / "target" / "manifest.json"
+


Would prefer to see FromEnvVar go away so we don't rely on it

schrockn · 2024-08-05T20:20:23Z

examples/experimental/dagster-airlift/examples/peering-with-dbt/tox.ini

+  -e ../../../../../python_modules/libraries/dagster-databricks
+  -e ../../../../../python_modules/libraries/dagster-pandas
+  -e ../../../../../python_modules/libraries/dagster-pyspark
+  -e ../../../../../python_modules/libraries/dagster-spark


why do we need all of these?

schrockn

Ok given that this is experimental we can fix forward stuff, but

Eliminating unnecessary dependencies very important (e.g. early design partners who do not use databricks should not install dagster-databricks)
We should cut our dependency on blueprints.

dpeng817 mentioned this pull request Aug 1, 2024

[dagster-airlift] mark dags as migrating #23370

Merged

dpeng817 marked this pull request as ready for review August 1, 2024 22:17

dpeng817 requested a review from schrockn August 2, 2024 16:21

schrockn reviewed Aug 2, 2024

View reviewed changes

schrockn requested changes Aug 2, 2024

View reviewed changes

schrockn reviewed Aug 2, 2024

View reviewed changes

This was referenced Aug 2, 2024

[dagster-airlift] dagster operator #23386

Merged

[dagster-airlift] airflow op swizzler #23402

Merged

dpeng817 force-pushed the dpeng817/mwaa_auth_backend branch from 1eb3e6b to 587e720 Compare August 3, 2024 23:59

dpeng817 force-pushed the dpeng817/airflow_blueprints branch from 4ff80f5 to de7ca4a Compare August 3, 2024 23:59

dpeng817 force-pushed the dpeng817/mwaa_auth_backend branch from 587e720 to 2c7d6eb Compare August 4, 2024 00:03

dpeng817 force-pushed the dpeng817/airflow_blueprints branch from de7ca4a to a57a8c1 Compare August 4, 2024 00:04

dpeng817 force-pushed the dpeng817/mwaa_auth_backend branch from 2c7d6eb to 3247e12 Compare August 5, 2024 02:08

dpeng817 force-pushed the dpeng817/airflow_blueprints branch from a57a8c1 to 7c7c8b4 Compare August 5, 2024 02:08

dpeng817 mentioned this pull request Aug 5, 2024

[dagster-airlift] full revamp #23410

Merged

dpeng817 changed the title ~~[dagster-airlift] airflow migration blueprints~~ [dagster-airlift] airflow migration tooling + defs from yaml Aug 5, 2024

dpeng817 changed the title ~~[dagster-airlift] airflow migration tooling + defs from yaml~~ [dagster-airlift] complete revamp Aug 5, 2024

dpeng817 force-pushed the dpeng817/airflow_blueprints branch 2 times, most recently from e745d9a to 15efcd4 Compare August 5, 2024 16:22

dpeng817 force-pushed the dpeng817/airflow_blueprints branch 2 times, most recently from 1ad114f to a70502a Compare August 5, 2024 17:54

Base automatically changed from dpeng817/mwaa_auth_backend to master August 5, 2024 19:37

schrockn reviewed Aug 5, 2024

View reviewed changes

schrockn approved these changes Aug 5, 2024

View reviewed changes

dpeng817 force-pushed the dpeng817/airflow_blueprints branch from a70502a to ba0ad3f Compare August 5, 2024 20:53

[dagster-airlift] airflow migration blueprints

6ea9e6b

dpeng817 force-pushed the dpeng817/airflow_blueprints branch from ba0ad3f to 6ea9e6b Compare August 5, 2024 21:05

dpeng817 merged commit bb9aec7 into master Aug 5, 2024
1 check passed

dpeng817 deleted the dpeng817/airflow_blueprints branch August 5, 2024 21:21

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[dagster-airlift] complete revamp #23369

[dagster-airlift] complete revamp #23369

dpeng817 commented Aug 1, 2024 •

edited

Loading

dpeng817 commented Aug 1, 2024 •

edited

Loading

schrockn Aug 2, 2024

schrockn Aug 2, 2024

schrockn Aug 2, 2024

dpeng817 Aug 4, 2024

schrockn Aug 2, 2024

schrockn left a comment

schrockn Aug 2, 2024

schrockn Aug 2, 2024

schrockn left a comment

schrockn Aug 5, 2024

schrockn Aug 5, 2024

schrockn left a comment

[dagster-airlift] complete revamp #23369

[dagster-airlift] complete revamp #23369

Conversation

dpeng817 commented Aug 1, 2024 • edited Loading

dpeng817 commented Aug 1, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

schrockn left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

schrockn left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

schrockn left a comment

Choose a reason for hiding this comment

dpeng817 commented Aug 1, 2024 •

edited

Loading

dpeng817 commented Aug 1, 2024 •

edited

Loading