Merge pull request #406 from TEAMSchools/348-use-configurableiomanager

replace iomanager in defs
TEAMSchools · Dec 5, 2023 · 13a7e45 · 13a7e45
2 parents 4eb925b + 76e6c4f
commit 13a7e45
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 63 deletions.
diff --git a/src/teamster/kippcamden/definitions.py b/src/teamster/kippcamden/definitions.py
@@ -1,11 +1,11 @@
 from dagster import Definitions, EnvVar, load_assets_from_modules
 from dagster_dbt import DbtCliResource
-from dagster_gcp import BigQueryResource, GCSPickleIOManager, GCSResource
+from dagster_gcp import BigQueryResource, GCSResource
 from dagster_k8s import k8s_job_executor
 
 from teamster import GCS_PROJECT_NAME
 from teamster.core.deanslist.resources import DeansListResource
-from teamster.core.google.io.resources import gcs_io_manager
+from teamster.core.google.storage.io_manager import GCSIOManager
 from teamster.core.sqlalchemy.resources import OracleResource, SqlAlchemyEngineResource
 from teamster.core.ssh.resources import SSHResource
 from teamster.kippcamden import (
@@ -36,20 +36,16 @@
     sensors=[*powerschool.sensors, *edplan.sensors, *titan.sensors],
     resources={
         "gcs": GCS_RESOURCE,
-        "io_manager": GCSPickleIOManager(
-            gcs=GCS_RESOURCE, gcs_bucket=f"teamster-{CODE_LOCATION}"
+        "io_manager": GCSIOManager(
+            gcs=GCS_RESOURCE,
+            gcs_bucket=f"teamster-{CODE_LOCATION}",
+            object_type="pickle",
         ),
-        "io_manager_gcs_avro": gcs_io_manager.configured(
-            config_or_config_fn={
-                "gcs_bucket": f"teamster-{CODE_LOCATION}",
-                "io_format": "avro",
-            }
+        "io_manager_gcs_avro": GCSIOManager(
+            gcs=GCS_RESOURCE, gcs_bucket=f"teamster-{CODE_LOCATION}", object_type="avro"
         ),
-        "io_manager_gcs_file": gcs_io_manager.configured(
-            config_or_config_fn={
-                "gcs_bucket": f"teamster-{CODE_LOCATION}",
-                "io_format": "filepath",
-            }
+        "io_manager_gcs_file": GCSIOManager(
+            gcs=GCS_RESOURCE, gcs_bucket=f"teamster-{CODE_LOCATION}", object_type="file"
         ),
         "dbt_cli": DbtCliResource(project_dir=f"src/dbt/{CODE_LOCATION}"),
         "db_bigquery": BigQueryResource(project=GCS_PROJECT_NAME),

diff --git a/src/teamster/kippmiami/definitions.py b/src/teamster/kippmiami/definitions.py
@@ -1,11 +1,11 @@
 from dagster import Definitions, EnvVar, load_assets_from_modules
 from dagster_dbt import DbtCliResource
-from dagster_gcp import BigQueryResource, GCSPickleIOManager, GCSResource
+from dagster_gcp import BigQueryResource, GCSResource
 from dagster_k8s import k8s_job_executor
 
 from teamster import GCS_PROJECT_NAME
 from teamster.core.deanslist.resources import DeansListResource
-from teamster.core.google.io.resources import gcs_io_manager
+from teamster.core.google.storage.io_manager import GCSIOManager
 from teamster.core.sqlalchemy.resources import OracleResource, SqlAlchemyEngineResource
 from teamster.core.ssh.resources import SSHResource
 from teamster.kippmiami import (
@@ -36,20 +36,16 @@
     sensors=[*powerschool.sensors, *renlearn.sensors, *iready.sensors],
     resources={
         "gcs": GCS_RESOURCE,
-        "io_manager": GCSPickleIOManager(
-            gcs=GCS_RESOURCE, gcs_bucket=f"teamster-{CODE_LOCATION}"
+        "io_manager": GCSIOManager(
+            gcs=GCS_RESOURCE,
+            gcs_bucket=f"teamster-{CODE_LOCATION}",
+            object_type="pickle",
         ),
-        "io_manager_gcs_avro": gcs_io_manager.configured(
-            config_or_config_fn={
-                "gcs_bucket": f"teamster-{CODE_LOCATION}",
-                "io_format": "avro",
-            }
+        "io_manager_gcs_avro": GCSIOManager(
+            gcs=GCS_RESOURCE, gcs_bucket=f"teamster-{CODE_LOCATION}", object_type="avro"
         ),
-        "io_manager_gcs_file": gcs_io_manager.configured(
-            config_or_config_fn={
-                "gcs_bucket": f"teamster-{CODE_LOCATION}",
-                "io_format": "filepath",
-            }
+        "io_manager_gcs_file": GCSIOManager(
+            gcs=GCS_RESOURCE, gcs_bucket=f"teamster-{CODE_LOCATION}", object_type="file"
         ),
         "dbt_cli": DbtCliResource(project_dir=f"src/dbt/{CODE_LOCATION}"),
         "db_bigquery": BigQueryResource(project=GCS_PROJECT_NAME),

diff --git a/src/teamster/kippnewark/definitions.py b/src/teamster/kippnewark/definitions.py
@@ -1,11 +1,11 @@
 from dagster import Definitions, EnvVar, load_assets_from_modules
 from dagster_dbt import DbtCliResource
-from dagster_gcp import BigQueryResource, GCSPickleIOManager, GCSResource
+from dagster_gcp import BigQueryResource, GCSResource
 from dagster_k8s import k8s_job_executor
 
 from teamster import GCS_PROJECT_NAME
 from teamster.core.deanslist.resources import DeansListResource
-from teamster.core.google.io.resources import gcs_io_manager
+from teamster.core.google.storage.io_manager import GCSIOManager
 from teamster.core.sqlalchemy.resources import OracleResource, SqlAlchemyEngineResource
 from teamster.core.ssh.resources import SSHResource
 from teamster.kippnewark import (
@@ -54,20 +54,16 @@
     ],
     resources={
         "gcs": GCS_RESOURCE,
-        "io_manager": GCSPickleIOManager(
-            gcs=GCS_RESOURCE, gcs_bucket=f"teamster-{CODE_LOCATION}"
+        "io_manager": GCSIOManager(
+            gcs=GCS_RESOURCE,
+            gcs_bucket=f"teamster-{CODE_LOCATION}",
+            object_type="pickle",
         ),
-        "io_manager_gcs_avro": gcs_io_manager.configured(
-            config_or_config_fn={
-                "gcs_bucket": f"teamster-{CODE_LOCATION}",
-                "io_format": "avro",
-            }
+        "io_manager_gcs_avro": GCSIOManager(
+            gcs=GCS_RESOURCE, gcs_bucket=f"teamster-{CODE_LOCATION}", object_type="avro"
         ),
-        "io_manager_gcs_file": gcs_io_manager.configured(
-            config_or_config_fn={
-                "gcs_bucket": f"teamster-{CODE_LOCATION}",
-                "io_format": "filepath",
-            }
+        "io_manager_gcs_file": GCSIOManager(
+            gcs=GCS_RESOURCE, gcs_bucket=f"teamster-{CODE_LOCATION}", object_type="file"
         ),
         "dbt_cli": DbtCliResource(project_dir=f"src/dbt/{CODE_LOCATION}"),
         "db_bigquery": BigQueryResource(project=GCS_PROJECT_NAME),

diff --git a/src/teamster/kipptaf/definitions.py b/src/teamster/kipptaf/definitions.py
@@ -2,7 +2,7 @@
 from dagster_airbyte import AirbyteCloudResource
 from dagster_dbt import DbtCliResource
 from dagster_fivetran import FivetranResource
-from dagster_gcp import BigQueryResource, GCSPickleIOManager, GCSResource
+from dagster_gcp import BigQueryResource, GCSResource
 from dagster_k8s import k8s_job_executor
 
 from teamster import GCS_PROJECT_NAME
@@ -14,8 +14,8 @@
 from teamster.core.amplify.resources import MClassResource
 from teamster.core.google.directory.resources import GoogleDirectoryResource
 from teamster.core.google.forms.resources import GoogleFormsResource
-from teamster.core.google.io.resources import gcs_io_manager
 from teamster.core.google.sheets.resources import GoogleSheetsResource
+from teamster.core.google.storage.io_manager import GCSIOManager
 from teamster.core.ldap.resources import LdapResource
 from teamster.core.schoolmint.grow.resources import SchoolMintGrowResource
 from teamster.core.smartrecruiters.resources import SmartRecruitersResource
@@ -95,14 +95,13 @@
     ],
     resources={
         "gcs": GCS_RESOURCE,
-        "io_manager": GCSPickleIOManager(
-            gcs=GCS_RESOURCE, gcs_bucket=f"teamster-{CODE_LOCATION}"
-        ),
-        "io_manager_gcs_avro": gcs_io_manager.configured(
-            config_or_config_fn={
-                "gcs_bucket": f"teamster-{CODE_LOCATION}",
-                "io_format": "avro",
-            }
+        "io_manager": GCSIOManager(
+            gcs=GCS_RESOURCE,
+            gcs_bucket=f"teamster-{CODE_LOCATION}",
+            object_type="pickle",
+        ),
+        "io_manager_gcs_avro": GCSIOManager(
+            gcs=GCS_RESOURCE, gcs_bucket=f"teamster-{CODE_LOCATION}", object_type="avro"
         ),
         "dbt_cli": DbtCliResource(project_dir=f"src/dbt/{CODE_LOCATION}"),
         "db_bigquery": BigQueryResource(project=GCS_PROJECT_NAME),

diff --git a/src/teamster/staging/definitions.py b/src/teamster/staging/definitions.py
@@ -2,7 +2,7 @@
 from dagster_gcp import GCSResource
 
 from teamster import GCS_PROJECT_NAME
-from teamster.core.google.io.resources import gcs_io_manager
+from teamster.core.google.storage.io_manager import GCSIOManager
 from teamster.core.sqlalchemy.resources import OracleResource, SqlAlchemyEngineResource
 from teamster.core.ssh.resources import SSHResource
 from teamster.staging import CODE_LOCATION, powerschool
@@ -13,17 +13,16 @@
     assets=load_assets_from_modules(modules=[powerschool]),
     resources={
         "gcs": GCS_RESOURCE,
-        "io_manager_gcs_avro": gcs_io_manager.configured(
-            config_or_config_fn={
-                "gcs_bucket": f"teamster-{CODE_LOCATION}",
-                "io_format": "avro",
-            }
+        "io_manager": GCSIOManager(
+            gcs=GCS_RESOURCE,
+            gcs_bucket=f"teamster-{CODE_LOCATION}",
+            object_type="pickle",
         ),
-        "io_manager_gcs_file": gcs_io_manager.configured(
-            config_or_config_fn={
-                "gcs_bucket": f"teamster-{CODE_LOCATION}",
-                "io_format": "filepath",
-            }
+        "io_manager_gcs_avro": GCSIOManager(
+            gcs=GCS_RESOURCE, gcs_bucket=f"teamster-{CODE_LOCATION}", object_type="avro"
+        ),
+        "io_manager_gcs_file": GCSIOManager(
+            gcs=GCS_RESOURCE, gcs_bucket=f"teamster-{CODE_LOCATION}", object_type="file"
         ),
         "db_powerschool": OracleResource(
             engine=SqlAlchemyEngineResource(