Experimental flag to attach row count metadata as part of `dagster-dbt` execution #21542

benpankow · 2024-04-30T21:50:21Z

Summary

Adds a new fetch_table_metadata experimental flag to DbtCliResource.cli which allows us to fetch dagster/total_row_count (introduced in #21524) to dbt-built tables:

@dbt_assets(manifest=dbt_manifest)
def jaffle_shop_dbt_assets(
    context: AssetExecutionContext,
    dbt: DbtCliResource,
):
    yield from dbt.cli(
        ["build"],
        context=context,
        fetch_table_metadata=True,
    ).stream()

Under the hood, this PR uses dbt's dbt.adapters.base.impl.BaseAdapter abstraction to let Dagster connect to the user's warehouse using the dbt-provided credentials. Right now, we just run a simple select count(*) on the tables specified in each AssetMaterialization and Output, but this lays some groundwork we could use for fetching other data as well.

There are a few caveats:

When using duckdb, we wait for the dbt run to conclude, since duckdb does not allow simultaneous connections when a write connection is open (e.g. when dbt is running)
We don't query row counts on views, since they may include non-trivial sql which could be expensive to query

Test Plan

Tested locally w/ duckdb, bigquery, and snowflake. Introduced basic pytest test to test against duckdb.

benpankow · 2024-04-30T21:50:34Z

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

Join @benpankow and the rest of your teammates on Graphite

rexledesma

Main thing here is that I'm not sold on this PostprocessingTask abstraction. Let's just define a single stateful transformation function.

python_modules/libraries/dagster-dbt/dagster_dbt/core/resources_v2.py

rexledesma · 2024-05-01T22:07:02Z

python_modules/libraries/dagster-dbt/dagster_dbt/core/resources_v2.py

+ from dbt.adapters.duckdb.credentials import DuckDBCredentials
+
+ if isinstance(config.credentials, DuckDBCredentials):
+ if not config.credentials.config_options:
+ config.credentials.config_options = {}
+ config.credentials.config_options["access_mode"] = "READ_ONLY"


Could you explain what happening here?

attempts to get duckdb adapter to play nice - concurrently, duckdb only accepts either:
a) one read-write connection
b) many read-only connections

this lets us parallelize our accesses for postprocessing in the duckdb case - sadly don't think we can run any of these tasks simultaneously to the dbt build since that's read-write

rexledesma · 2024-05-01T22:11:47Z

python_modules/libraries/dagster-dbt/dagster_dbt/core/utils.py

@@ -282,3 +283,20 @@ def parse_manifest(path: str, target_path: str = DEFAULT_DBT_TARGET_PATH) -> Map
 return json.load(file)
 except FileNotFoundError:
 raise DagsterDbtCliOutputsNotFoundError(path=manifest_path)
+
+
+def get_future_completion_state_or_err(futures: List[Union[Future, Any]]) -> bool:


Could we instead use https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.wait with some timeout for similar functionality?

Probably, but I think we'd still have to do some custom logic around bookkeeping on event log/Future order. Right now the waiting/timeout is handled in the main loop in stream since we just want to wait on the next incomplete event log entry's Future (so we can yield it to user code). This way we don't have to worry about yielding things in the wrong order.

rexledesma · 2024-05-01T22:13:15Z

python_modules/libraries/dagster-dbt/dagster_dbt/core/resources_v2.py

+ # they are emitted by dbt.
+ output_events_and_futures: List[Union[Future, DbtDagsterEventType]] = []
+
+ with pushd(str(self.project_dir)):


Could you explain the pushd here?

needed to get the duckdb adapter to work correctly - it tries to access the db file assuming the working directory is the project dir - seems like the DBT cli does this in its process, but we need to do the same if we want to use the adapter ourselves. Added comment.

python_modules/libraries/dagster-dbt/dagster_dbt/core/resources_v2.py

rexledesma · 2024-05-01T22:27:09Z

python_modules/libraries/dagster-dbt/dagster_dbt/core/resources_v2.py

@@ -543,6 +613,7 @@ def run(
 raise_on_error: bool,
 context: Optional[OpExecutionContext],
 adapter: Optional[BaseAdapter],
+ postprocessing_tasks: Optional[List[DbtPostprocessingAsyncTask]],


Just so we have maximum flexibility, we shouldn't expose this API to the user yet. Once we have confidence in this scheme working for out-of-the-box metadata (e.g. row count), then we can consider exposing it to the user so they also have more flexibility.

Updated. We still expose this functionality to the user, but now it's an @experimental_param flag

queue

python_modules/libraries/dagster-dbt/dagster_dbt/core/resources_v2.py

rexledesma

Main comments:

Divest from the boolean pattern, to more of a composable builder pattern
The with_metadata change should be a separate PR. There should be systematic refactoring of our event inheritance to ensure metadata can be overridden in an ergonomic way.

python_modules/libraries/dagster-dbt/dagster_dbt/core/resources_v2.py

rexledesma · 2024-05-06T15:01:28Z

python_modules/dagster/dagster/_core/definitions/events.py

@@ -149,6 +149,17 @@ def __eq__(self, other: object) -> bool:
 and self.tags == other.tags
 )

+ def with_metadata(self, metadata: Optional[Mapping[str, RawMetadataValue]]) -> "Output":


We should separate this out into its own PR, since this is a core lib change.

python_modules/libraries/dagster-dbt/dagster_dbt/core/resources_v2.py

rexledesma · 2024-05-06T15:08:17Z

python_modules/libraries/dagster-dbt/dagster_dbt/core/resources_v2.py

+ if isinstance(event, Output):
+ return event.with_metadata(metadata={**event.metadata, **additional_metadata})
+ else:
+ return event._replace(metadata={**event.metadata, **additional_metadata})


This signals that the core Dagster XXXEvent classes should have a more unified inheritance pattern.

We should file an issue and mark this as one code path to consolidate.

rexledesma · 2024-05-06T15:10:30Z

python_modules/libraries/dagster-dbt/dagster_dbt/core/resources_v2.py

+ if all_work_complete and event_to_emit_idx >= len(output_events_and_futures):
+ break
+
+ if event_to_emit_idx < len(output_events_and_futures):


nit: invert conditional logic, there's a lot of nesting here

benpankow · 2024-05-06T19:36:34Z

Updated to builder pattern

queue

python_modules/libraries/dagster-dbt/dagster_dbt/core/resources_v2.py

rexledesma

I assume docs + API docstring examples are coming in a follow up PR

rexledesma · 2024-05-07T17:46:46Z

...s/libraries/dagster-dbt/dagster_dbt_tests/core/dbt_packages/test_row_count_postprocessing.py

I don't think this should be in the dbt_packages/ directory since this functionality doesn't rely on the dagster dbt package.

## Summary Adds a small util method to recreate an `Output` object with a new set of metadata keys and values. Used in #21542. ## Test Plan New unit test.

## Summary Adds a small util method to recreate an `Output` object with a new set of metadata keys and values. Used in dagster-io#21542. ## Test Plan New unit test.

…t` execution (dagster-io#21542) ## Summary Adds a new `fetch_table_metadata` experimental flag to `DbtCliResource.cli` which allows us to fetch `dagster/total_row_count` (introduced in dagster-io#21524) to dbt-built tables: ```python @dbt_assets(manifest=dbt_manifest) def jaffle_shop_dbt_assets( context: AssetExecutionContext, dbt: DbtCliResource, ): yield from dbt.cli( ["build"], context=context, fetch_table_metadata=True, ).stream() ``` <img width="534" alt="Screenshot 2024-05-03 at 11 03 19 AM" src="https://github.com/dagster-io/dagster/assets/10215173/c3e64633-5fc3-44e4-99e3-601f0c7a0856"> Under the hood, this PR uses dbt's `dbt.adapters.base.impl.BaseAdapter` abstraction to let Dagster connect to the user's warehouse using the dbt-provided credentials. Right now, we just run a simple `select count(*)` on the tables specified in each `AssetMaterialization` and `Output`, but this lays some groundwork we could use for fetching other data as well. There are a few caveats: - When using duckdb, we wait for the dbt run to conclude, since duckdb does not allow simultaneous connections when a write connection is open (e.g. when dbt is running) - We don't query row counts on views, since they may include non-trivial sql which could be expensive to query ## Test Plan Tested locally w/ duckdb, bigquery, and snowflake. Introduced basic pytest test to test against duckdb.

…n't require an import of a module that might not be present Summary: Attempt to resolve #23952: Instead of attempting to import duckdb and logging a warning on importerror, do a bit of introspection on the relevant class and only try to import it if the class name matches. Test Plan: could use some help with this from reviewers, I don't see any duckdb-specific tests in #21542

…n't require an import of a module that might not be present (#24346) Summary: Attempt to resolve #23952: Instead of attempting to import duckdb and logging a warning on importerror, do a bit of introspection on the relevant class and only try to import it if the class name matches. Test Plan: could use some help with this from reviewers, I don't see any duckdb-specific tests in #21542 ## Summary & Motivation ## How I Tested These Changes ## Changelog Insert changelog entry or "NOCHANGELOG" here. - [ ] `NEW` _(added new feature or capability)_ - [ ] `BUGFIX` _(fixed a bug)_ - [ ] `DOCS` _(added or updated documentation)_

benpankow changed the base branch from master to benpankow/row-count-meta May 1, 2024 20:17

benpankow force-pushed the benpankow/threadpool-thing branch from 8e18c51 to 4329f78 Compare May 1, 2024 20:17

benpankow mentioned this pull request May 1, 2024

Standardize existing row count metadata to use dagster/row_count key #21524

Merged

benpankow changed the title ~~test~~ [wip] Pull row count as part of dagster-dbt execution May 1, 2024

benpankow changed the base branch from benpankow/row-count-meta to persist-run-for-production-2-schedule May 1, 2024 20:21

benpankow changed the base branch from persist-run-for-production-2-schedule to benpankow/row-count-meta May 1, 2024 20:21

benpankow requested review from rexledesma and sryza May 1, 2024 21:35

rexledesma previously requested changes May 1, 2024

View reviewed changes

rexledesma reviewed May 1, 2024

View reviewed changes

benpankow force-pushed the benpankow/row-count-meta branch from 563fccc to 13fd45e Compare May 2, 2024 20:46

benpankow force-pushed the benpankow/threadpool-thing branch 2 times, most recently from 4befca2 to 46ffcb3 Compare May 2, 2024 23:16

benpankow force-pushed the benpankow/row-count-meta branch from 1b5c0c7 to 375cf8b Compare May 3, 2024 18:13

benpankow force-pushed the benpankow/threadpool-thing branch from 2a58b25 to 7b14290 Compare May 3, 2024 18:13

benpankow marked this pull request as ready for review May 3, 2024 18:17

benpankow changed the title ~~[wip] Pull row count as part of dagster-dbt execution~~ Experimental flag to attach row count metadata as part of dagster-dbt execution May 3, 2024

benpankow force-pushed the benpankow/row-count-meta branch from 375cf8b to c3ae946 Compare May 3, 2024 18:37

benpankow force-pushed the benpankow/threadpool-thing branch 2 times, most recently from a35aa1b to 0a44fe7 Compare May 3, 2024 21:29

benpankow requested a review from rexledesma May 3, 2024 21:40

sryza reviewed May 5, 2024

View reviewed changes

python_modules/libraries/dagster-dbt/dagster_dbt/core/resources_v2.py Outdated Show resolved Hide resolved

rexledesma previously requested changes May 6, 2024

View reviewed changes

benpankow mentioned this pull request May 6, 2024

add small util to modify Output metadata #21666

Merged

benpankow force-pushed the benpankow/row-count-meta branch from d70aa25 to 9d957ab Compare May 6, 2024 18:54

benpankow force-pushed the benpankow/threadpool-thing branch from ff9b5e9 to 9a27b13 Compare May 6, 2024 18:54

benpankow changed the base branch from benpankow/row-count-meta to benpankow/output-with-metadata May 6, 2024 18:57

benpankow force-pushed the benpankow/threadpool-thing branch from 9a27b13 to 596b0bf Compare May 6, 2024 18:57

benpankow mentioned this pull request May 6, 2024

unify metadata replacement api across event types #21667

Merged

benpankow requested review from rexledesma and sryza May 6, 2024 19:36

rexledesma reviewed May 6, 2024

View reviewed changes

python_modules/libraries/dagster-dbt/dagster_dbt/core/resources_v2.py Outdated Show resolved Hide resolved

python_modules/libraries/dagster-dbt/dagster_dbt/core/resources_v2.py Outdated Show resolved Hide resolved

benpankow requested a review from rexledesma May 7, 2024 17:21

rexledesma approved these changes May 7, 2024

View reviewed changes

benpankow force-pushed the benpankow/output-with-metadata branch from ac37a1b to 58bcadb Compare May 7, 2024 20:29

benpankow force-pushed the benpankow/threadpool-thing branch from 4d12f63 to 437bdc3 Compare May 7, 2024 20:29

benpankow force-pushed the benpankow/output-with-metadata branch from 58bcadb to b9d8a75 Compare May 7, 2024 22:11

benpankow force-pushed the benpankow/threadpool-thing branch from 437bdc3 to ee88d8c Compare May 7, 2024 22:11

benpankow force-pushed the benpankow/output-with-metadata branch from b9d8a75 to 5e94d59 Compare May 8, 2024 23:17

benpankow force-pushed the benpankow/threadpool-thing branch from d616ef9 to 46998df Compare May 8, 2024 23:18

Base automatically changed from benpankow/output-with-metadata to master May 8, 2024 23:44

benpankow added a commit that referenced this pull request May 8, 2024

add small util to modify Output metadata (#21666)

3f5e6c4

## Summary Adds a small util method to recreate an `Output` object with a new set of metadata keys and values. Used in #21542. ## Test Plan New unit test.

benpankow added 5 commits May 9, 2024 10:36

total_row_count to row_count

5c33c09

test

4369597

update some comments, label count

dc3e1ae

respond to some comments

527ae4b

move test

9b3a0e2

benpankow force-pushed the benpankow/threadpool-thing branch from 46998df to 9b3a0e2 Compare May 9, 2024 17:37

stream in main thread if no postprocessing

f7fd476

benpankow merged commit 536685d into master May 9, 2024
1 check was pending

benpankow deleted the benpankow/threadpool-thing branch May 9, 2024 20:15

gibsondan mentioned this pull request Sep 10, 2024

Gate duckdb-specific instance check in dagster-dbt in a way that doesn't require an import of a module that might not be present #24346

Merged

3 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Experimental flag to attach row count metadata as part of `dagster-dbt` execution #21542

Experimental flag to attach row count metadata as part of `dagster-dbt` execution #21542

benpankow commented Apr 30, 2024 •

edited

Loading

benpankow commented Apr 30, 2024 •

edited

Loading

rexledesma left a comment

rexledesma May 1, 2024

benpankow May 1, 2024

rexledesma May 1, 2024

benpankow May 3, 2024

rexledesma May 1, 2024

benpankow May 1, 2024 •

edited

Loading

rexledesma May 1, 2024

benpankow May 3, 2024

rexledesma left a comment

rexledesma May 6, 2024

benpankow May 6, 2024

rexledesma May 6, 2024

benpankow May 6, 2024

rexledesma May 6, 2024

benpankow commented May 6, 2024

rexledesma left a comment

rexledesma May 7, 2024

Experimental flag to attach row count metadata as part of dagster-dbt execution #21542

Experimental flag to attach row count metadata as part of dagster-dbt execution #21542

Conversation

benpankow commented Apr 30, 2024 • edited Loading

Summary

Test Plan

benpankow commented Apr 30, 2024 • edited Loading

rexledesma left a 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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

benpankow May 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

rexledesma left a 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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

benpankow commented May 6, 2024

rexledesma left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Experimental flag to attach row count metadata as part of `dagster-dbt` execution #21542

Experimental flag to attach row count metadata as part of `dagster-dbt` execution #21542

benpankow commented Apr 30, 2024 •

edited

Loading

benpankow commented Apr 30, 2024 •

edited

Loading

benpankow May 1, 2024 •

edited

Loading