Skip to content

Commit 5131364

Browse files
blarghmateyclaude
blarghmatey
and
claude
authored
fix(edxorg-s3): fix dlt resource naming, incremental loading, and schema metadata (#2315)
* fix(edxorg-s3): fix dlt resource naming, incremental loading, and schema metadata A @dlt.resource wrapper that *returns* another DltResource causes dlt to discard the outer resource's name, write_disposition, primary_key, and table_format entirely, replacing the resource with the returned one. All tables were therefore written to a destination table named 'read_csv_duckdb' (the transformer's name), making the schema lookup in extract_resource_metadata fail with empty columns and an empty jobs list. Fix loads.py: - Replace the @dlt.resource wrapper + return pattern with .with_name() and .apply_hints() applied directly on the filesystem | read_csv_duckdb pipe. This is the correct way to attach hints to a transformer-based resource. - Pass incremental=incremental('modification_date') natively to filesystem() rather than via apply_hints(). Each per-table resource gets its own modification_date cursor stored under the resource name, so subsequent runs only process files newer than the last successful run. - Fix table_format: was always 'iceberg' even for local/dev where there is no Glue catalog entry; now conditional on DAGSTER_ENVIRONMENT. Fix dagster_assets.py: - Add loader_file_format='parquet' to every dlt.run() call to match the standalone run_pipeline() and produce the correct staging format for filesystem+Iceberg destinations. - Set storage_kind and kinds to 'iceberg'/'filesystem' based on env. data.destination.destination_name returns the named destination ('production'), not the storage type, so catalog badges were misleading. - Explicitly compute dagster/table_name from pipeline.dataset_name because dlt does not populate schema_name in load_info for filesystem+Iceberg destinations, leaving extract_resource_metadata with table_name=None. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(edxorg): pin table_name in apply_hints; dict-wrap TableMetadataSet Explicitly pass table_name=resource_name to apply_hints() so the destination table is named by a concrete hint rather than the with_name() fallback — guarding against any future dlt version that sets an explicit table_name on the read_csv_duckdb transformer. Convert TableMetadataSet(storage_kind=...) to dict() before spreading into the metadata mapping; TableMetadataSet is not guaranteed to implement the mapping protocol, so **TableMetadataSet(...) can raise TypeError at runtime. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent df980d6 commit 5131364

2 files changed

Lines changed: 106 additions & 85 deletions

File tree

‎dg_projects/data_loading/data_loading/defs/edxorg_s3_ingest/dagster_assets.py‎

Lines changed: 30 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -12,44 +12,59 @@
1212
AssetKey,
1313
AssetSpec,
1414
)
15+
from dagster._core.definitions.metadata.metadata_set import TableMetadataSet
1516
from dagster_dlt import DagsterDltResource, DagsterDltTranslator, dlt_assets
1617
from dagster_dlt.translator import DltResourceTranslatorData
1718
from ol_orchestrate.lib.constants import EDXORG_DB_TABLES
1819

1920
from .loads import (
21+
dagster_env,
2022
edxorg_s3_pipeline,
2123
edxorg_s3_source,
2224
edxorg_s3_source_instance,
2325
table_format,
2426
)
2527

28+
_STORAGE_KIND = "iceberg" if dagster_env in ("qa", "production") else "filesystem"
29+
2630

2731
class EdxorgDltTranslator(DagsterDltTranslator):
2832
"""Custom translator for edxorg dlt assets with upstream dependencies."""
2933

3034
def get_asset_spec(self, data: DltResourceTranslatorData) -> AssetSpec:
31-
"""
32-
Map dlt resource to Dagster asset spec with one-to-one upstream deps.
33-
"""
34-
# Get the default spec from parent
3535
default_spec = super().get_asset_spec(data)
3636

3737
# The resource name is raw__edxorg__s3__tables__{table}
3838
resource_name = data.resource.name
39-
40-
# Create asset key with ol_warehouse_raw_data prefix
4139
asset_key = AssetKey(["ol_warehouse_raw_data", resource_name])
4240

43-
# Extract table name and create upstream dependency
4441
deps = []
42+
table_name_meta = {}
4543
if resource_name.startswith("raw__edxorg__s3__tables__"):
46-
table_name = resource_name.replace("raw__edxorg__s3__tables__", "")
47-
# Add non-blocking dependency on upstream partitioned asset
48-
deps = [AssetDep(AssetKey(["edxorg", "raw_data", "db_table", table_name]))]
44+
short_name = resource_name.replace("raw__edxorg__s3__tables__", "")
45+
deps = [AssetDep(AssetKey(["edxorg", "raw_data", "db_table", short_name]))]
46+
# Assemble the Glue catalog fully-qualified name. dlt does not
47+
# populate schema_name in load_info for filesystem+Iceberg
48+
# destinations, so extract_resource_metadata leaves table_name
49+
# as None. Set it explicitly here from the pipeline dataset_name.
50+
if data.pipeline and data.pipeline.dataset_name:
51+
table_name_meta = dict(
52+
TableMetadataSet(
53+
table_name=f"{data.pipeline.dataset_name}.{resource_name}"
54+
)
55+
)
4956

5057
return default_spec.replace_attributes(
5158
key=asset_key,
5259
deps=deps,
60+
# storage_kind defaults to the named destination ("production"),
61+
# not the actual storage type. Override with the real kind.
62+
kinds={"dlt", _STORAGE_KIND},
63+
metadata={
64+
**default_spec.metadata,
65+
**table_name_meta,
66+
**dict(TableMetadataSet(storage_kind=_STORAGE_KIND)),
67+
},
5368
)
5469

5570

@@ -106,7 +121,11 @@ def edxorg_s3_consolidated_tables(
106121
# between tables via AssumeRoleWithWebIdentity (IRSA).
107122
for table_name in selected_tables:
108123
table_source = edxorg_s3_source(tables=[table_name], table_format=table_format)
109-
yield from dlt.run(context=context, dlt_source=table_source)
124+
yield from dlt.run(
125+
context=context,
126+
dlt_source=table_source,
127+
loader_file_format="parquet",
128+
)
110129

111130

112131
__all__ = ["edxorg_s3_consolidated_tables"]

‎dg_projects/data_loading/data_loading/defs/edxorg_s3_ingest/loads.py‎

Lines changed: 76 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,6 @@
2828

2929
import dlt
3030
import s3fs
31-
from dlt.extract import DltResource
3231
from dlt.sources import incremental
3332
from dlt.sources.filesystem import filesystem, read_csv_duckdb
3433
from ol_orchestrate.lib.constants import EDXORG_DB_TABLES
@@ -68,94 +67,97 @@ def edxorg_s3_source(
6867
tables_to_load = tables if tables is not None else EDXORG_DB_TABLES
6968

7069
for table_name in tables_to_load:
71-
# Create a resource for each table with standardized naming
72-
@dlt.resource(
73-
name=f"raw__edxorg__s3__tables__{table_name}",
74-
write_disposition="merge",
70+
resource_name = f"raw__edxorg__s3__tables__{table_name}"
71+
# Pattern to match ONLY prod TSV files for this table
72+
# Matches: db_table/{table_name}/prod/*/*.tsv
73+
# Excludes: db_table/{table_name}/edge/*/*.tsv (edge staging files)
74+
file_glob = f"db_table/{table_name}/prod/**/*.tsv"
75+
76+
# Create an s3fs filesystem WITHOUT explicit credentials.
77+
#
78+
# dlt's AwsCredentials.to_session_credentials() calls
79+
# get_frozen_credentials(), which snapshots the live IRSA
80+
# RefreshableCredentials object as static strings and passes
81+
# them to s3fs.S3FileSystem(key=..., secret=..., token=...).
82+
# Once s3fs holds explicit static strings there is no refresh
83+
# path: when the STS session token expires (~1 hour by default)
84+
# every subsequent GetObject call fails with ExpiredToken.
85+
#
86+
# By constructing S3FileSystem with no explicit key/secret/token
87+
# we bypass dlt's credential snapshotting entirely. aiobotocore
88+
# then uses its own credential chain which, for IRSA pods, invokes
89+
# AioAssumeRoleWithWebIdentityFetcher and wraps the result in
90+
# AioRefreshableCredentials. Those credentials automatically call
91+
# AssumeRoleWithWebIdentity again whenever the token expires, so
92+
# the pipeline can run indefinitely without hitting ExpiredToken.
93+
#
94+
# dlt's filesystem() source accepts an AbstractFileSystem instance
95+
# as its `credentials` parameter and uses it directly, skipping
96+
# its own credential-dispatch logic.
97+
#
98+
# Do NOT pass extract_content=True. That flag reads the entire
99+
# file into Python bytes in a single blocking request before
100+
# passing it to read_csv_duckdb, which can itself exceed a token
101+
# lifetime for very large tables. Leaving it False lets
102+
# read_csv_duckdb stream the content chunk-by-chunk via PyArrow.
103+
fs = s3fs.S3FileSystem()
104+
105+
# Pass incremental directly to filesystem() so dlt tracks the
106+
# modification_date cursor per-resource. Only files newer than
107+
# the cursor from the last successful run are processed.
108+
files = filesystem(
109+
bucket_url=bucket_url,
110+
file_glob=file_glob,
111+
credentials=fs,
112+
incremental=incremental("modification_date"),
113+
)
114+
115+
# Pipe filesystem items through read_csv_duckdb, then rename and
116+
# configure the resulting resource. Using with_name() + apply_hints()
117+
# is required here: a @dlt.resource wrapper that *returns* another
118+
# DltResource causes dlt to replace the outer resource entirely,
119+
# discarding its name, write_disposition, primary_key, and
120+
# table_format hints. The resulting writes land under the
121+
# transformer's own name ("read_csv_duckdb"), making schema and job
122+
# lookups fail. Applying hints directly on the pipe avoids this.
123+
yield (
124+
(
125+
files
126+
| read_csv_duckdb(
127+
use_pyarrow=True,
128+
delimiter="\t", # TSV files use tabs
129+
ignore_errors=True,
130+
# Force all columns to VARCHAR to prevent pyarrow schema mismatches.
131+
# This occurs when DuckDB infers different types for the same column
132+
# across files (e.g., TIMESTAMP vs. VARCHAR for columns that are
133+
# empty in some files). Downstream dbt models handle type casting.
134+
all_varchar=True,
135+
)
136+
)
137+
.with_name(resource_name)
75138
# All edxorg TSV exports include a row_hash column. Use a composite key
76139
# with extracted_course_key because row_hash is computed from only the
77140
# original CSV columns (excluding the extracted_* fields added during
78141
# archive processing), so identical rows from different courses would
79142
# otherwise collide on row_hash alone.
80-
primary_key=["row_hash", "extracted_course_key"],
81-
table_format=table_format,
82-
)
83-
def load_table(table=table_name) -> DltResource:
84-
"""
85-
Load TSV files for a specific table from S3.
86-
87-
Uses incremental loading based on file modification_date to avoid
88-
reprocessing unchanged files on subsequent runs.
89-
90-
Only processes 'prod' source files, excluding 'edge' staging environment.
91-
"""
92-
93-
# Pattern to match ONLY prod TSV files for this table
94-
# Matches: db_table/{table_name}/prod/*/*.tsv (prod/course_id/*.tsv)
95-
# Excludes: db_table/{table_name}/edge/*/*.tsv (edge staging files)
96-
file_glob = f"db_table/{table}/prod/**/*.tsv"
97-
98-
# Create an s3fs filesystem WITHOUT explicit credentials.
99-
#
100-
# dlt's AwsCredentials.to_session_credentials() calls
101-
# get_frozen_credentials(), which snapshots the live IRSA
102-
# RefreshableCredentials object as static strings and passes
103-
# them to s3fs.S3FileSystem(key=..., secret=..., token=...).
104-
# Once s3fs holds explicit static strings there is no refresh
105-
# path: when the STS session token expires (~1 hour by default)
106-
# every subsequent GetObject call fails with ExpiredToken.
107-
#
108-
# By constructing S3FileSystem with no explicit key/secret/token
109-
# we bypass dlt's credential snapshotting entirely. aiobotocore
110-
# then uses its own credential chain which, for IRSA pods, invokes
111-
# AioAssumeRoleWithWebIdentityFetcher and wraps the result in
112-
# AioRefreshableCredentials. Those credentials automatically call
113-
# AssumeRoleWithWebIdentity again whenever the token expires, so
114-
# the pipeline can run indefinitely without hitting ExpiredToken.
115-
#
116-
# dlt's filesystem() source accepts an AbstractFileSystem instance
117-
# as its `credentials` parameter and uses it directly, skipping
118-
# its own credential-dispatch logic.
119-
#
120-
# Do NOT pass extract_content=True. That flag reads the entire
121-
# file into Python bytes in a single blocking request before
122-
# passing it to read_csv_duckdb, which can itself exceed a token
123-
# lifetime for very large tables. Leaving it False lets
124-
# read_csv_duckdb stream the content chunk-by-chunk via PyArrow.
125-
fs = s3fs.S3FileSystem()
126-
files = filesystem(
127-
bucket_url=bucket_url,
128-
file_glob=file_glob,
129-
credentials=fs,
143+
.apply_hints(
144+
table_name=resource_name,
145+
write_disposition="merge",
146+
primary_key=["row_hash", "extracted_course_key"],
147+
table_format=table_format,
130148
)
131-
132-
# Enable incremental loading based on file modification date
133-
# Only processes files modified since the last successful run
134-
files.apply_hints(incremental=incremental("modification_date"))
135-
136-
return files | read_csv_duckdb(
137-
use_pyarrow=True,
138-
delimiter="\t", # TSV files use tabs
139-
ignore_errors=True,
140-
# Force all columns to VARCHAR to prevent pyarrow schema mismatches.
141-
# This occurs when DuckDB infers different types for the same column
142-
# across files (e.g., TIMESTAMP vs. VARCHAR for columns that are
143-
# empty in some files). Downstream dbt models handle type casting.
144-
all_varchar=True,
145-
)
146-
147-
yield load_table
149+
)
148150

149151

150152
# Determine environment and destination
151153
# Use DAGSTER_ENVIRONMENT to determine destination and Glue database
152-
dagster_env = os.getenv("DAGSTER_ENVIRONMENT", "dev")
154+
dagster_env: str = os.getenv("DAGSTER_ENVIRONMENT", "dev")
153155

154156
# Map DAGSTER_ENVIRONMENT to destination config
155157
# dev/ci -> local filesystem (Parquet, no Iceberg, no Glue)
156158
# qa -> filesystem + Iceberg (S3 + ol_warehouse_qa_raw Glue database)
157159
# production -> filesystem + Iceberg (S3 + ol_warehouse_production_raw Glue database)
158-
table_format = "iceberg"
160+
table_format = "iceberg" if dagster_env in ("qa", "production") else "parquet"
159161
if dagster_env in ("qa", "production"):
160162
destination_env = dagster_env
161163
dataset_name = f"ol_warehouse_{dagster_env}_raw"

0 commit comments

Comments
 (0)