Skip to content

Commit e65bce1

Browse files
C00ldudeNoonanC00ldudeNoonan
authored
fix: migrate legacy BigQuery parameter calls (#146)
* fix: migrate legacy BigQuery parameter calls * fix: address BigQuery parameter review feedback
1 parent 22149a0 commit e65bce1

14 files changed

Lines changed: 681 additions & 460 deletions

‎macro_agents/src/macro_agents/defs/domains/fomc_transcripts.py‎

Lines changed: 55 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import dagster as dg
99

1010
from macro_agents.defs.analysis.fed_sentiment.resource import FedSentimentResource
11+
from macro_agents.defs.resources.bigquery_query import QueryParameter
1112
from macro_agents.defs.resources.federal_reserve import FederalReserveResource
1213
from macro_agents.defs.resources.gcs import GCSResource
1314
from macro_agents.defs.resources.bigquery_warehouse import BigQueryWarehouseResource
@@ -398,36 +399,70 @@ def process_fomc_transcripts(
398399
page_count = meta["page_count"]
399400

400401
# Update fomc_transcripts with extracted text
401-
conn.query(
402+
bq.execute_query(
402403
"""
403404
UPDATE fomc_transcripts
404-
SET full_text = $1, word_count = $2, page_count = $3,
405+
SET full_text = @full_text,
406+
word_count = @word_count,
407+
page_count = @page_count,
405408
processed_date = CURRENT_TIMESTAMP
406-
WHERE transcript_id = $4
409+
WHERE transcript_id = @transcript_id
407410
""",
408-
[full_text, word_count, page_count, transcript_id], # ty: ignore[invalid-argument-type]
409-
).result()
411+
read_only=False,
412+
params={
413+
"full_text": full_text,
414+
"word_count": word_count,
415+
"page_count": page_count,
416+
"transcript_id": transcript_id,
417+
},
418+
)
410419

411420
# Parse into speaker sections
412421
sections = _parse_transcript_sections(full_text, transcript_id)
413422
for section in sections:
414-
conn.query(
423+
bq.execute_query(
415424
"""
416-
INSERT OR REPLACE INTO transcript_sections
417-
(section_id, transcript_id, section_order, section_type,
418-
speaker, speaker_role, content)
419-
VALUES ($1, $2, $3, $4, $5, $6, $7)
425+
MERGE transcript_sections AS target
426+
USING (
427+
SELECT
428+
@section_id AS section_id,
429+
@transcript_id AS transcript_id,
430+
@section_order AS section_order,
431+
@section_type AS section_type,
432+
@speaker AS speaker,
433+
@speaker_role AS speaker_role,
434+
@content AS content
435+
) AS source
436+
ON target.section_id = source.section_id
437+
WHEN MATCHED THEN UPDATE SET
438+
transcript_id = source.transcript_id,
439+
section_order = source.section_order,
440+
section_type = source.section_type,
441+
speaker = source.speaker,
442+
speaker_role = source.speaker_role,
443+
content = source.content
444+
WHEN NOT MATCHED THEN INSERT (
445+
section_id, transcript_id, section_order,
446+
section_type, speaker, speaker_role, content
447+
) VALUES (
448+
source.section_id, source.transcript_id,
449+
source.section_order, source.section_type,
450+
source.speaker, source.speaker_role, source.content
451+
)
420452
""",
421-
[ # ty: ignore[invalid-argument-type]
422-
section["section_id"],
423-
transcript_id,
424-
section["section_order"],
425-
section["section_type"],
426-
section["speaker"],
427-
section.get("speaker_role"),
428-
section["content"],
429-
],
430-
).result()
453+
read_only=False,
454+
params={
455+
"section_id": section["section_id"],
456+
"transcript_id": transcript_id,
457+
"section_order": section["section_order"],
458+
"section_type": section["section_type"],
459+
"speaker": QueryParameter(section["speaker"], "STRING"),
460+
"speaker_role": QueryParameter(
461+
section.get("speaker_role"), "STRING"
462+
),
463+
"content": section["content"],
464+
},
465+
)
431466

432467
total_sections += len(sections)
433468
context.log.info(

‎macro_agents/src/macro_agents/defs/domains/sec/bi.py‎

Lines changed: 39 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
from macro_agents.defs.domains.sec.tables import ensure_sec_filing_search_terms_table
88
from macro_agents.defs.domains.sec.text import sec_filing_text_extracted
99
from macro_agents.defs.resources.gcs import GCSResource
10+
from macro_agents.defs.resources.bigquery_query import numeric_query_parameter
1011
from macro_agents.defs.resources.bigquery_warehouse import BigQueryWarehouseResource
1112
from macro_agents.defs.resources.sec_edgar import SECEdgarResource
1213
from macro_agents.defs.domains.sec.config import (
@@ -173,26 +174,46 @@ def sec_filing_business_intelligence(
173174
filing_id, signal.category, signal.position
174175
)
175176

176-
conn.query(
177+
bq.execute_query(
177178
"""
178-
INSERT INTO sec_filing_search_terms
179-
(term_id, filing_id, term_category, term_text,
180-
context_text, section_name, confidence_score)
181-
VALUES (?, ?, ?, ?, ?, ?, ?)
182-
ON CONFLICT (term_id) DO UPDATE SET
183-
context_text = EXCLUDED.context_text,
184-
confidence_score = EXCLUDED.confidence_score
179+
MERGE sec_filing_search_terms AS target
180+
USING (
181+
SELECT
182+
@term_id AS term_id,
183+
@filing_id AS filing_id,
184+
@term_category AS term_category,
185+
@term_text AS term_text,
186+
@context_text AS context_text,
187+
@section_name AS section_name,
188+
@confidence_score AS confidence_score
189+
) AS source
190+
ON target.term_id = source.term_id
191+
WHEN MATCHED THEN UPDATE SET
192+
context_text = source.context_text,
193+
confidence_score = source.confidence_score
194+
WHEN NOT MATCHED THEN INSERT (
195+
term_id, filing_id, term_category, term_text,
196+
context_text, section_name, confidence_score
197+
) VALUES (
198+
source.term_id, source.filing_id,
199+
source.term_category, source.term_text,
200+
source.context_text, source.section_name,
201+
source.confidence_score
202+
)
185203
""",
186-
[ # ty: ignore[invalid-argument-type]
187-
term_id,
188-
filing_id,
189-
signal.category,
190-
signal.term[:SIGNAL_TERM_MAX_LENGTH],
191-
signal.context[:SIGNAL_CONTEXT_MAX_LENGTH],
192-
signal.section_name,
193-
signal.confidence_score,
194-
],
195-
).result()
204+
read_only=False,
205+
params={
206+
"term_id": term_id,
207+
"filing_id": filing_id,
208+
"term_category": signal.category,
209+
"term_text": signal.term[:SIGNAL_TERM_MAX_LENGTH],
210+
"context_text": signal.context[:SIGNAL_CONTEXT_MAX_LENGTH],
211+
"section_name": signal.section_name,
212+
"confidence_score": numeric_query_parameter(
213+
signal.confidence_score
214+
),
215+
},
216+
)
196217

197218
# Track category counts
198219
category_counts[signal.category] = (

‎macro_agents/src/macro_agents/defs/domains/sec/documents.py‎

Lines changed: 64 additions & 82 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,8 @@
1010
from macro_agents.defs.domains.sec.config import MAX_ERROR_DETAILS
1111
from macro_agents.defs.domains.sec.filing_downloader import FilingDownloader
1212
from macro_agents.defs.domains.sec.metadata import sec_filing_metadata
13-
from macro_agents.defs.resources.gcs import GCSResource
1413
from macro_agents.defs.resources.bigquery_warehouse import BigQueryWarehouseResource
14+
from macro_agents.defs.resources.gcs import GCSResource
1515
from macro_agents.defs.resources.sec_edgar import SECEdgarResource
1616

1717
sec_filing_documents_partitions = dg.DynamicPartitionsDefinition(
@@ -38,65 +38,57 @@ def sec_filing_documents(
3838
Partition key is the filing_id. Uses FilingDownloader for download logic.
3939
"""
4040
filing_id = context.partition_key
41-
downloader = FilingDownloader(sec_edgar, gcs, context.log)
42-
43-
conn = None
44-
try:
45-
conn = bq.get_connection()
46-
47-
filing_df = downloader.query_filing_by_id(conn, filing_id)
48-
49-
if filing_df.is_empty():
50-
context.log.warning(f"No filing found for filing_id={filing_id}")
51-
return dg.MaterializeResult(
52-
metadata={
53-
"status": "not_found",
54-
"filing_id": filing_id,
55-
"documents_downloaded": 0,
56-
}
57-
)
41+
downloader = FilingDownloader(sec_edgar, gcs, bq, context.log)
42+
43+
filing_df = downloader.query_filing_by_id(filing_id)
44+
45+
if filing_df.is_empty():
46+
context.log.warning(f"No filing found for filing_id={filing_id}")
47+
return dg.MaterializeResult(
48+
metadata={
49+
"status": "not_found",
50+
"filing_id": filing_id,
51+
"documents_downloaded": 0,
52+
}
53+
)
5854

59-
# Process all rows for this filing_id (may span multiple symbols
60-
# if a CIK/accession-derived ID maps to more than one symbol)
61-
rows = filing_df.to_dicts()
62-
total_downloaded = 0
63-
last_result = None
64-
65-
for row in rows:
66-
if downloader.is_filing_processed(conn, filing_id, row["symbol"]):
67-
context.log.debug(
68-
f"Filing {filing_id} already processed for {row['symbol']}"
69-
)
70-
continue
71-
72-
last_result = downloader.download_filing(conn, row)
73-
if last_result.status == "downloaded":
74-
total_downloaded += 1
75-
76-
if last_result is None:
77-
return dg.MaterializeResult(
78-
metadata={
79-
"status": "already_processed",
80-
"filing_id": filing_id,
81-
"symbol": rows[0]["symbol"],
82-
"documents_downloaded": 0,
83-
}
55+
# Process all rows for this filing_id (may span multiple symbols
56+
# if a CIK/accession-derived ID maps to more than one symbol)
57+
rows = filing_df.to_dicts()
58+
total_downloaded = 0
59+
last_result = None
60+
61+
for row in rows:
62+
if downloader.is_filing_processed(filing_id, row["symbol"]):
63+
context.log.debug(
64+
f"Filing {filing_id} already processed for {row['symbol']}"
8465
)
66+
continue
67+
68+
last_result = downloader.download_filing(row)
69+
if last_result.status == "downloaded":
70+
total_downloaded += 1
8571

72+
if last_result is None:
8673
return dg.MaterializeResult(
8774
metadata={
88-
"status": last_result.status,
89-
"filing_id": last_result.filing_id,
90-
"symbol": last_result.symbol,
91-
"gcs_path": last_result.gcs_path,
92-
"documents_downloaded": total_downloaded,
93-
"error": last_result.error,
75+
"status": "already_processed",
76+
"filing_id": filing_id,
77+
"symbol": rows[0]["symbol"],
78+
"documents_downloaded": 0,
9479
}
9580
)
9681

97-
finally:
98-
if conn:
99-
conn.close()
82+
return dg.MaterializeResult(
83+
metadata={
84+
"status": last_result.status,
85+
"filing_id": last_result.filing_id,
86+
"symbol": last_result.symbol,
87+
"gcs_path": last_result.gcs_path,
88+
"documents_downloaded": total_downloaded,
89+
"error": last_result.error,
90+
}
91+
)
10092

10193

10294
@dg.asset(
@@ -117,38 +109,28 @@ def sec_filing_documents_batch(
117109
Queries sec_filings for up to 25 unprocessed filings across all companies,
118110
downloads each primary document to GCS, and marks them processed.
119111
"""
120-
downloader = FilingDownloader(sec_edgar, gcs, context.log)
121-
122-
conn = None
123-
try:
124-
conn = bq.get_connection()
112+
downloader = FilingDownloader(sec_edgar, gcs, bq, context.log)
125113

126-
filings_to_process = downloader.query_unprocessed_filings(conn)
127-
128-
if filings_to_process.is_empty():
129-
context.log.debug("No unprocessed filings to download")
130-
return dg.MaterializeResult(
131-
metadata={"status": "no_filings", "documents_downloaded": 0}
132-
)
133-
134-
context.log.info(f"Downloading documents for {len(filings_to_process)} filings")
135-
136-
result = downloader.download_batch(conn, filings_to_process)
114+
filings_to_process = downloader.query_unprocessed_filings()
137115

116+
if filings_to_process.is_empty():
117+
context.log.debug("No unprocessed filings to download")
138118
return dg.MaterializeResult(
139-
metadata={
140-
"status": "completed",
141-
"documents_downloaded": result.total_downloaded,
142-
"errors": result.total_errors,
143-
"remaining_to_process": result.remaining,
144-
"error_details": (
145-
result.error_details[:MAX_ERROR_DETAILS]
146-
if result.error_details
147-
else []
148-
),
149-
}
119+
metadata={"status": "no_filings", "documents_downloaded": 0}
150120
)
151121

152-
finally:
153-
if conn:
154-
conn.close()
122+
context.log.info(f"Downloading documents for {len(filings_to_process)} filings")
123+
124+
result = downloader.download_batch(filings_to_process)
125+
126+
return dg.MaterializeResult(
127+
metadata={
128+
"status": "completed",
129+
"documents_downloaded": result.total_downloaded,
130+
"errors": result.total_errors,
131+
"remaining_to_process": result.remaining,
132+
"error_details": (
133+
result.error_details[:MAX_ERROR_DETAILS] if result.error_details else []
134+
),
135+
}
136+
)

0 commit comments

Comments
 (0)