use dagster-insights

dagster-io · Sep 20, 2023 · a79041b · a79041b
1 parent 328d825
commit a79041b
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 256 deletions.
diff --git a/purina_usage/assets/dbt_snowflake/__init__.py b/purina_usage/assets/dbt_snowflake/__init__.py
@@ -4,7 +4,21 @@
 from dagster_dbt import DbtCliResource, dbt_assets
 
 from dagster import OpExecutionContext
-from ...utils import dagster_insights
+from gql.transport.requests import RequestsHTTPTransport
+from dagster_insights import DagsterInsightsClient, SnowflakeConnectionDetails
+
+transport = RequestsHTTPTransport(
+    url="http://localhost:3000/test/staging/graphql",
+    use_json=True,
+    timeout=300,
+    headers={"Dagster-Cloud-Api-Token": "user:test:joe"},
+)
+dagster_insights = DagsterInsightsClient(
+    organization_id="test",
+    deployment="test",
+    cloud_user_token="",
+    transport=transport,
+)
 
 dbt_project_dir = Path(__file__).joinpath("..", "..", "..", "..", "dbt_project").resolve()
 dbt_cli_resource = DbtCliResource(
@@ -24,4 +38,14 @@ def dbt_snowflake_assets(context: OpExecutionContext, dbt: DbtCliResource):
 
     run_results = dbt_cli_invocation.get_artifact("run_results.json")
     manifest = dbt_cli_invocation.get_artifact("manifest.json")
-    dagster_insights.store_dbt_adapter_metrics(context, manifest, run_results)
+    dagster_insights.store_dbt_adapter_metrics(
+        context,
+        manifest,
+        run_results,
+        snowflake_connection_details=SnowflakeConnectionDetails(
+            user=os.getenv("SNOWFLAKE_USER", ""),
+            password=os.getenv("SNOWFLAKE_PASSWORD", ""),
+            account="na94824.us-east-1",
+            warehouse="DEVELOPMENT",
+        ),
+    )
diff --git a/purina_usage/assets/raw_data/__init__.py b/purina_usage/assets/raw_data/__init__.py
@@ -4,8 +4,22 @@
 from typing import Generator, Any
 from dagster import OpExecutionContext, Output
 
+from gql.transport.requests import RequestsHTTPTransport
 from purina_usage.utils import random_data
-from ...utils import dagster_insights
+from dagster_insights import DagsterInsightsClient, DagsterInsightsMetric
+
+transport = RequestsHTTPTransport(
+    url="http://localhost:3000/test/staging/graphql",
+    use_json=True,
+    timeout=300,
+    headers={"Dagster-Cloud-Api-Token": "user:test:joe"},
+)
+dagster_insights = DagsterInsightsClient(
+    organization_id="test",
+    deployment="test",
+    cloud_user_token="",
+    transport=transport,
+)
 
 
 @asset(compute_kind="random")
@@ -24,7 +38,7 @@ def users(context: OpExecutionContext) -> Generator[Output[pd.DataFrame], Any, A
     dagster_insights.put_context_metrics(
         context,
         metrics=[
-            dagster_insights.Metric(
+            DagsterInsightsMetric(
                 metric_name="rows_affected",
                 metric_value=len(data),
             )
@@ -43,7 +57,7 @@ def orders(context: OpExecutionContext) -> Generator[Output[pd.DataFrame], Any,
     dagster_insights.put_context_metrics(
         context,
         metrics=[
-            dagster_insights.Metric(
+            DagsterInsightsMetric(
                 metric_name="rows_affected",
                 metric_value=len(data),
             )

diff --git a/purina_usage/utils/dagster_insights.py b/purina_usage/utils/dagster_insights.py
diff --git a/setup.py b/setup.py
@@ -6,6 +6,7 @@
     install_requires=[
         "dagster",
         "dagster-cloud",
+        "dagster-insights",
         "pandas",
         "snowflake-connector-python[pandas]",
         "snowflake.sqlalchemy",