Enable K8s Job Executor (#2148)

opensource-observer · Sep 16, 2024 · 25c1ea4 · 25c1ea4
1 parent b741446
commit 25c1ea4
Show file tree

Hide file tree

Showing 6 changed files with 82 additions and 56 deletions.
diff --git a/ops/helm-charts/oso-dagster/Chart.yaml b/ops/helm-charts/oso-dagster/Chart.yaml
@@ -3,7 +3,7 @@ name: oso-dagster
 description: Extension of the dagster template
 
 type: application
-version: 0.4.3
+version: 0.5.0
 appVersion: "1.0.0"
 dependencies:
 - name: dagster

diff --git a/ops/helm-charts/oso-dagster/templates/config-map.yaml b/ops/helm-charts/oso-dagster/templates/config-map.yaml
@@ -15,6 +15,7 @@ data:
   DAGSTER_ALERTS_BASE_URL: "{{ .Values.alerts.baseUrl }}"
   DAGSTER_HTTP_CACHE: "{{ .Values.cache.uri }}"
   DAGSTER_ENABLE_BIGQUERY: "1"
+  DAGSTER_ENABLE_K8S_EXECUTOR: "1"
   DAGSTER_SQLMESH_GATEWAY: clickhouse
   SQLMESH_POSTGRES_INSTANCE_CONNECTION_STRING: "{{ .Values.configMap.secretPrefix }}-{{ .Values.secretmanagerKeys.sqlmeshPostgresInstanceStringConnectionString }}"
   SQLMESH_POSTGRES_USER: "{{ .Values.configMap.secretPrefix }}-{{ .Values.secretmanagerKeys.sqlmeshPostgresUser }}"

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -63,6 +63,7 @@ google-cloud-bigquery-storage = "^2.25.0"
 dagster-sqlmesh = "0.2.0.dev3"
 google-auth = "^2.34.0"
 pillow = "^10.4.0"
+dagster-k8s = "^0.24.7"
 
 
 [tool.poetry.scripts]

diff --git a/warehouse/oso_dagster/constants.py b/warehouse/oso_dagster/constants.py
@@ -87,3 +87,5 @@ def get_project_id():
 
 sqlmesh_dir = os.getenv("DAGSTER_SQLMESH_DIR", str(repo_dir.joinpath("warehouse/metrics_mesh")))
 sqlmesh_gateway = os.getenv("DAGSTER_SQLMESH_GATEWAY", "local")
+
+enable_k8s_executor = os.getenv("DAGSTER_ENABLE_K8S_EXECUTOR", "false").lower() in ["true", "1"]
diff --git a/warehouse/oso_dagster/definitions.py b/warehouse/oso_dagster/definitions.py
@@ -4,6 +4,7 @@
 from dagster_dbt import DbtCliResource
 from dagster_gcp import BigQueryResource, GCSResource
 from dagster_sqlmesh import SQLMeshContextConfig, SQLMeshResource
+from dagster_k8s import k8s_job_executor
 from dagster_embedded_elt.dlt import DagsterDltResource
 from dotenv import load_dotenv
 from . import constants
@@ -116,13 +117,18 @@ def load_definitions():
             project_dir=os.fspath(constants.main_dbt_project_dir), target=target
         )
 
+    extra_kwargs = {}
+    if constants.enable_k8s_executor:
+        extra_kwargs["executor"] = k8s_job_executor
+
     return Definitions(
         assets=asset_factories.assets,
         schedules=all_schedules,
         jobs=asset_factories.jobs,
         asset_checks=asset_factories.checks,
         sensors=asset_factories.sensors,
         resources=resources,
+        **extra_kwargs,
     )