diff --git a/src/pudl/transform/vcerare.py b/src/pudl/transform/vcerare.py index 6ee9e6ce8..840651f95 100644 --- a/src/pudl/transform/vcerare.py +++ b/src/pudl/transform/vcerare.py @@ -9,6 +9,7 @@ import pyarrow as pa import pyarrow.parquet as pq from dagster import ( + AssetCheckExecutionContext, AssetCheckResult, asset, asset_check, @@ -359,16 +360,23 @@ def _load_duckdb_table(): blocking=True, description="Check that row count matches expected.", ) -def check_rows() -> AssetCheckResult: +def check_rows(context: AssetCheckExecutionContext) -> AssetCheckResult: """Check rows.""" logger.info("Check VCE RARE hourly table is the expected length") + + # Define row counts for fast/full etl + row_counts = { + "etl_full": 136437000, + "etl_fast": 27287400, + } + vce = _load_duckdb_table() # noqa: F841 (length,) = duckdb.query("SELECT COUNT(*) FROM vce").fetchone() - if length != 136437000: + if (expected_length := row_counts[context.op_execution_context.job_name]) != length: return AssetCheckResult( passed=False, description="Table unexpected length", - metadata={"table_length": length, "expected_length": 136437000}, + metadata={"table_length": length, "expected_length": expected_length}, ) return AssetCheckResult(passed=True)