Update sling assets

.gitignore

@@ -172,5 +172,5 @@ cython_debug/
 tmp*/
 
 # dbt
-dbt_project/example.duckdb
+dbt_project/example.duckdb*
 dbt_project/logs

hooli-demo-assets/hooli_demo_assets/assets/sling.py

@@ -1,12 +1,25 @@
-from dagster import AssetSpec
-from dagster_embedded_elt.sling import build_sling_asset, SlingMode
-
-
-raw_location = build_sling_asset(
-    asset_spec=AssetSpec(key=["locations"]),
-    source_stream="s3://hooli-demo/embedded-elt/",
-    target_object= "RAW_DATA.LOCATIONS",
-    mode=SlingMode.FULL_REFRESH,
-    source_options={"format": "csv"},
-    sling_resource_key="s3_to_snowflake_resource"
-)
+from dagster_embedded_elt.sling import (
+   sling_assets,
+   SlingResource,
+)
+from dagster_embedded_elt.sling.dagster_sling_translator import DagsterSlingTranslator
+
+from hooli_demo_assets.resources import replication_config
+
+
+class CustomSlingTranslator(DagsterSlingTranslator):
+   def __init__(self, target_prefix="RAW_DATA"):
+        super().__init__(target_prefix=target_prefix)
+
+   def get_group_name(self, stream_definition):
+       return "RAW_DATA"
+
+
+@sling_assets(
+   replication_config=replication_config,
+   dagster_sling_translator=CustomSlingTranslator(),
+)
+def my_sling_assets(context, sling: SlingResource):
+   yield from sling.replicate(context=context)
+   for row in sling.stream_raw_logs():
+       context.log.info(row)

hooli-demo-assets/hooli_demo_assets/definitions.py

@@ -1,20 +1,18 @@
 from dagster import (
-    Definitions,
-    load_assets_from_modules,
-
+   Definitions,  
 )
 
-from hooli_demo_assets.assets import sling
+from hooli_demo_assets.assets.sling import my_sling_assets
 from hooli_demo_assets.jobs import daily_sling_job
-from hooli_demo_assets.resources import get_env, resource_def
-from hooli_demo_assets.schedules import daily_sling_assets 
+from hooli_demo_assets.resources import sling_resource
+from hooli_demo_assets.schedules import daily_sling_assets
 
 
-sling_assets = load_assets_from_modules([sling], key_prefix="RAW_DATA", group_name="RAW_DATA")
-
 defs = Definitions(
-    assets=[*sling_assets],
-    schedules=[daily_sling_assets],
-    jobs=[daily_sling_job],
-    resources=resource_def[get_env()],
-)
+   assets=[my_sling_assets],
+   schedules=[daily_sling_assets],
+   jobs=[daily_sling_job],
+   resources={
+       "sling": sling_resource
+   },
+)

hooli-demo-assets/hooli_demo_assets/jobs/__init__.py

@@ -1,8 +1,10 @@
 from dagster import AssetSelection, define_asset_job
 
+
 raw_location_by_day = AssetSelection.keys(["RAW_DATA", "locations"])
 
+
 daily_sling_job = define_asset_job(
-    name="daily_sling_job",
-    selection=raw_location_by_day,
-)
+   name="daily_sling_job",
+   selection=raw_location_by_day,
+)

hooli-demo-assets/hooli_demo_assets/resources/__init__.py

@@ -1,66 +1,99 @@
 import os
+from pathlib import Path
 
 from dagster import EnvVar
-from pathlib import Path
 from dagster_embedded_elt.sling import (
-    SlingResource,
-    SlingSourceConnection,
-    SlingTargetConnection,
+   SlingResource,
+   SlingConnectionResource,
 )
 
+
 def get_env():
-    if os.getenv("DAGSTER_CLOUD_IS_BRANCH_DEPLOYMENT", "") == "1":
-        return "BRANCH"
-    if os.getenv("DAGSTER_CLOUD_DEPLOYMENT_NAME", "") == "data-eng-prod":
-        return "PROD"
-    return "LOCAL"
+   if os.getenv("DAGSTER_CLOUD_IS_BRANCH_DEPLOYMENT", "") == "1":
+       return "BRANCH"
+   if os.getenv("DAGSTER_CLOUD_DEPLOYMENT_NAME", "") == "data-eng-prod":
+       return "PROD"
+   return "LOCAL"
+
 
-# embedded elt source and target
+# Path for duckdb connection - needed for local dev
 current_file_path = Path(__file__)
 parent_dir_path = current_file_path.parent.parent.parent.parent
 DUCKDB_PATH = parent_dir_path / "dbt_project" / "example.duckdb"
 
-source = SlingSourceConnection(
-        type="s3",
-        bucket=EnvVar("AWS_S3_BUCKET"),
-        region=EnvVar("AWS_REGION"),
-        access_key_id=EnvVar("AWS_ACCESS_KEY_ID"),
-        secret_access_key=EnvVar("AWS_SECRET_ACCESS_KEY"),
-)
-local_target = SlingTargetConnection(
-        type="duckdb",
-        instance=f"{DUCKDB_PATH}",
-)
-branch_target = SlingTargetConnection(
-        type="snowflake",
-        host=EnvVar("SNOWFLAKE_HOST"),
-        database="DEMO_DB2_BRANCH",
-        user=EnvVar("SNOWFLAKE_USER"),
-        password=EnvVar("SNOWFLAKE_PASSWORD"),
-        role=EnvVar("SNOWFLAKE_ROLE"),
-)
-prod_target = SlingTargetConnection(
-        type="snowflake",
-        host=EnvVar("SNOWFLAKE_HOST"),
-        database="DEMO_DB2",
-        user=EnvVar("SNOWFLAKE_USER"),
-        password=EnvVar("SNOWFLAKE_PASSWORD"),
-        role=EnvVar("SNOWFLAKE_ROLE"),
-)
 
-resource_def = {
-    "LOCAL": {
-         "s3_to_snowflake_resource":
-         SlingResource(source_connection=source, target_connection=local_target),
-    },
-    "BRANCH": {
-         "s3_to_snowflake_resource": SlingResource(
-            source_connection=source, target_connection=branch_target
-        ),
-    },
-    "PROD": {
-        "s3_to_snowflake_resource": SlingResource(
-            source_connection=source, target_connection=prod_target
-        ),
-    },
-}
+# Alternatively a replication.yaml file can be used
+def create_replication_config(env: str):
+   # Determine the target database dynamically based on the environment
+   target_database = {
+       'LOCAL': 'DUCKDB',
+       'BRANCH': 'SNOWFLAKE_BRANCH',
+       'PROD': 'SNOWFLAKE_PROD'
+   }.get(env, 'DUCKDB')  # Default to DUCKDB if no environment matches
+
+   replication_config = {
+       "source": "S3",
+       "target": target_database,
+       "defaults": {
+           "mode": "full-refresh",
+           "object": "{stream_file_folder}.{stream_file_name}",
+           "source_options": {
+               "format": "csv"
+           }
+       },
+       "streams": {
+           "s3://hooli-demo/embedded-elt/locations.csv": {
+               "object": "locations"
+           }
+       }
+   }
+   return replication_config
+
+
+def create_sling_resource(env: str):
+   # Dynamically generate connection based on enviornment
+   connections = [
+       SlingConnectionResource(
+           name="S3",
+           type="s3",
+           bucket=EnvVar("AWS_S3_BUCKET"),
+           region=EnvVar("AWS_REGION"),
+           access_key_id=EnvVar("AWS_ACCESS_KEY_ID"),
+           secret_access_key=EnvVar("AWS_SECRET_ACCESS_KEY"),
+       )
+   ]
+   if env == 'LOCAL':
+       connections.append(SlingConnectionResource(
+           name="DUCKDB",
+           type="duckdb",
+           instance=f"{DUCKDB_PATH}",
+           schema="RAW_DATA",
+       ))
+   elif env == 'BRANCH':
+       connections.append(SlingConnectionResource(
+           name="SNOWFLAKE_BRANCH",
+           type="snowflake",
+           host=EnvVar("SNOWFLAKE_HOST"),
+           user=EnvVar("SNOWFLAKE_USER"),
+           password=EnvVar("SNOWFLAKE_PASSWORD"),
+           role=EnvVar("SNOWFLAKE_ROLE"),
+           database="DEMO_DB2_BRANCH",
+           schema="RAW_DATA",
+       ))
+   elif env == 'PROD':
+       connections.append(SlingConnectionResource(
+           name="SNOWFLAKE_PROD",
+           type="snowflake",
+           host=EnvVar("SNOWFLAKE_HOST"),
+           user=EnvVar("SNOWFLAKE_USER"),
+           password=EnvVar("SNOWFLAKE_PASSWORD"),
+           role=EnvVar("SNOWFLAKE_ROLE"),
+           database="DEMO_DB2",
+           schema="RAW_DATA",
+       ))
+   return SlingResource(connections=connections)
+
+
+env = get_env()
+replication_config = create_replication_config(env)
+sling_resource = create_sling_resource(env)

hooli-demo-assets/hooli_demo_assets/schedules/__init__.py

@@ -1,7 +1,8 @@
 from dagster import ScheduleDefinition
 from hooli_demo_assets.jobs import daily_sling_job
 
+
 daily_sling_assets = ScheduleDefinition(
-    job=daily_sling_job,
-    cron_schedule="0 0 * * *", # every day at midnight
-)
+   job=daily_sling_job,
+   cron_schedule="0 0 * * *", # every day at midnight
+)