Skip to content

[DEV-15635] Update download regex and add test case #4713

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sethstoudenmier merged 8 commits into from
Aug 3, 2026
+202 −114
Merged

[DEV-15635] Update download regex and add test case #4713

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
d49a143
[DEV-15635] Update download regex and add test case
sethstoudenmier Jul 30, 2026
1b473a8
[DEV-15635] Remove old test and unnecessary baker statement
sethstoudenmier Jul 30, 2026
d5d2f31
[DEV-15635] linter changes
sethstoudenmier Jul 30, 2026
b8585b4
[DEV-15635] handle assigned table aliases when table is joined more t…
sethstoudenmier Jul 30, 2026
d436110
Merge branch 'qat' into fix/dev-15635-adjust-download-sql-regex
james-at-kc Aug 3, 2026
52a8325
Merge branch 'qat' of https://github.com/fedspendingtransparency/usas…
sethstoudenmier Aug 3, 2026
d9a205c
[DEV-14608] fix typo
sethstoudenmier Aug 3, 2026
1608630
Merge branch 'fix/dev-15635-adjust-download-sql-regex' of https://git…
sethstoudenmier Aug 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
202 changes: 105 additions & 97 deletions usaspending_api/download/filestreaming/download_generation.py
100644 → 100755
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,26 +86,30 @@ def generate_download(download_job: DownloadJob, origination: Optional[str] = No
def _log_download_start(download_job: DownloadJob, file_name: str, request_type: str) -> None:
"""Log the start of download generation with tracing"""
with SubprocessTrace(
name=f"job.{JOB_TYPE}.generate_download_{request_type}",
kind=SpanKind.INTERNAL,
service="bulk-download",
name=f"job.{JOB_TYPE}.generate_download_{request_type}",
kind=SpanKind.INTERNAL,
service="bulk-download",
) as main_trace:
main_trace.set_attributes({
"service": "bulk-download",
"span_type": "Internal",
"job_type": str(JOB_TYPE),
"message": "Creating data archive files from the download job object",
"download_job_id": str(download_job.download_job_id),
"download_job_status": str(download_job.job_status.name),
"download_file_name": str(download_job.file_name),
"download_file_size": download_job.file_size if download_job.file_size is not None else 0,
"number_of_rows": download_job.number_of_rows if download_job.number_of_rows is not None else 0,
"number_of_columns": download_job.number_of_columns if download_job.number_of_columns is not None else 0,
"error_message": download_job.error_message if download_job.error_message else "",
"monthly_download": str(download_job.monthly_download),
"json_request": str(download_job.json_request) if download_job.json_request else "",
"file_name": str(file_name),
})
main_trace.set_attributes(
{
"service": "bulk-download",
"span_type": "Internal",
"job_type": str(JOB_TYPE),
"message": "Creating data archive files from the download job object",
"download_job_id": str(download_job.download_job_id),
"download_job_status": str(download_job.job_status.name),
"download_file_name": str(download_job.file_name),
"download_file_size": download_job.file_size if download_job.file_size is not None else 0,
"number_of_rows": download_job.number_of_rows if download_job.number_of_rows is not None else 0,
"number_of_columns": download_job.number_of_columns
if download_job.number_of_columns is not None
else 0,
"error_message": download_job.error_message if download_job.error_message else "",
"monthly_download": str(download_job.monthly_download),
"json_request": str(download_job.json_request) if download_job.json_request else "",
"file_name": str(file_name),
}
)


def _validate_and_setup_download(download_job: DownloadJob, json_request: dict, file_name: str) -> str:
Expand All @@ -118,9 +122,9 @@ def _validate_and_setup_download(download_job: DownloadJob, json_request: dict,
f"{MAX_DOWNLOAD_LIMIT} records"
)
with SubprocessTrace(
name=f"job.{JOB_TYPE}.generate_download_{json_request.get('request_type')}",
kind=SpanKind.INTERNAL,
service="bulk-download",
name=f"job.{JOB_TYPE}.generate_download_{json_request.get('request_type')}",
kind=SpanKind.INTERNAL,
service="bulk-download",
) as limit_exceeded:
limit_exceeded.set_attributes({"message": msg, "limit": limit})
raise Exception(msg)
Expand All @@ -139,12 +143,7 @@ def _validate_and_setup_download(download_job: DownloadJob, json_request: dict,
return working_dir


def _process_download_sources(
download_job: DownloadJob,
json_request: dict,
working_dir: str,
file_name: str
) -> None:
def _process_download_sources(download_job: DownloadJob, json_request: dict, working_dir: str, file_name: str) -> None:
"""Process all download sources and generate data files"""
columns = json_request.get("columns", None)
limit = json_request.get("limit", None)
Expand All @@ -158,18 +157,17 @@ def _process_download_sources(
for source in sources:
source_column_count = len(source.columns(columns))
if source_column_count == 0:
create_empty_data_file(
source, download_job, working_dir, piid, assistance_id, zip_file_path, file_format
)
create_empty_data_file(source, download_job, working_dir, piid, assistance_id, zip_file_path, file_format)
else:
download_job.number_of_columns += source_column_count
parse_source(
source, columns, download_job, working_dir, piid, assistance_id, zip_file_path, limit, file_format
)


def _add_supplemental_files(download_job: DownloadJob, json_request: dict, working_dir: str,
zip_file_path: str) -> None:
def _add_supplemental_files(
download_job: DownloadJob, json_request: dict, working_dir: str, zip_file_path: str
) -> None:
"""Add data dictionary and file description if requested"""
if json_request.get("include_data_dictionary"):
add_data_dictionary_to_zip(working_dir, zip_file_path)
Expand All @@ -192,33 +190,37 @@ def _handle_invalid_parameter_exception(download_job: DownloadJob, e: InvalidPar
"""Handle InvalidParameterException with tracing"""
exc_msg = "InvalidParameterException was raised while attempting to process the DownloadJob"
with SubprocessTrace(
name=f"job.{JOB_TYPE}.generate_download",
kind=SpanKind.INTERNAL,
service="bulk-download",
name=f"job.{JOB_TYPE}.generate_download",
kind=SpanKind.INTERNAL,
service="bulk-download",
) as error_span:
error_span.set_attributes({
"service": "bulk-download",
"span_type": "Internal",
"message": exc_msg,
"error": str(e),
})
error_span.set_attributes(
{
"service": "bulk-download",
"span_type": "Internal",
"message": exc_msg,
"error": str(e),
}
)
fail_download(download_job, e, exc_msg)


def _handle_general_exception(download_job: DownloadJob, e: Exception) -> None:
"""Handle general exceptions with tracing"""
exc_msg = "An exception was raised while attempting to process the DownloadJob"
with SubprocessTrace(
name=f"job.{JOB_TYPE}.process_download_job",
kind=SpanKind.INTERNAL,
service="bulk-download",
name=f"job.{JOB_TYPE}.process_download_job",
kind=SpanKind.INTERNAL,
service="bulk-download",
) as error_span:
error_span.set_attributes({
"service": "bulk-download",
"span_type": "Internal",
"message": exc_msg,
"error": str(e),
})
error_span.set_attributes(
{
"service": "bulk-download",
"span_type": "Internal",
"message": exc_msg,
"error": str(e),
}
)
fail_download(download_job, e, exc_msg)


Expand All @@ -241,30 +243,33 @@ def _upload_to_s3(download_job: DownloadJob, file_name: str) -> None:
zip_file_path = settings.CSV_LOCAL_PATH + file_name

with SubprocessTrace(
name=f"job.{JOB_TYPE}.download.s3",
kind=SpanKind.INTERNAL,
service="bulk-download",
name=f"job.{JOB_TYPE}.download.s3",
kind=SpanKind.INTERNAL,
service="bulk-download",
) as span:
span.set_attributes({
"service": "bulk-download",
"span_type": "Internal",
"resource": f"s3://{CONFIG.BULK_DOWNLOAD_S3_BUCKET_NAME}",
"message": "Push file to S3 bucket, if not local",
})
span.set_attributes(
{
"service": "bulk-download",
"span_type": "Internal",
"resource": f"s3://{CONFIG.BULK_DOWNLOAD_S3_BUCKET_NAME}",
"message": "Push file to S3 bucket, if not local",
}
)

with SubprocessTrace(
name=f"job.{JOB_TYPE}.s3.command",
kind=SpanKind.SERVER,
service="bulk-download",
name=f"job.{JOB_TYPE}.s3.command",
kind=SpanKind.SERVER,
service="bulk-download",
) as s3_span:
s3_span.set_attributes({
"service": "aws.s3",
"span_type": "WEB",
"resource": ".".join([
multipart_upload.__module__,
(multipart_upload.__qualname__ or multipart_upload.__name__)
]),
})
s3_span.set_attributes(
{
"service": "aws.s3",
"span_type": "WEB",
"resource": ".".join(
[multipart_upload.__module__, (multipart_upload.__qualname__ or multipart_upload.__name__)]
),
}
)

try:
bucket = CONFIG.BULK_DOWNLOAD_S3_BUCKET_NAME
Expand All @@ -273,22 +278,23 @@ def _upload_to_s3(download_job: DownloadJob, file_name: str) -> None:
start_uploading = time.perf_counter()
multipart_upload(bucket, region, zip_file_path, os.path.basename(zip_file_path))
write_to_log(
message=f"Uploading took {time.perf_counter() - start_uploading:.2f}s",
download_job=download_job
message=f"Uploading took {time.perf_counter() - start_uploading:.2f}s", download_job=download_job
)
except Exception as e:
exc_msg = "An exception was raised while attempting to upload the file"
with SubprocessTrace(
name=f"job.{JOB_TYPE}.upload_file_to_aws",
kind=SpanKind.SERVER,
service="bulk-download",
name=f"job.{JOB_TYPE}.upload_file_to_aws",
kind=SpanKind.SERVER,
service="bulk-download",
) as error_span:
error_span.set_attributes({
"service": "bulk-download",
"span_type": "Internal",
"message": exc_msg,
"error": str(e),
})
error_span.set_attributes(
{
"service": "bulk-download",
"span_type": "Internal",
"message": exc_msg,
"error": str(e),
}
)
fail_download(download_job, e, exc_msg)
if isinstance(e, InvalidParameterException):
raise InvalidParameterException(e) from e
Expand Down Expand Up @@ -745,15 +751,21 @@ def apply_annotations_to_sql(

cte_sql, select_statements = _select_columns(raw_query)

DIRECT_SELECT_QUERY_REGEX = r'^[^ ]*\."[^"]*"$' # Django is pretty consistent with how it prints out queries
# Django is pretty consistent with how it prints out queries
# !!! IMPORTANT: This can change when changing versions of Django !!!
DIRECT_SELECT_QUERY_REGEX = r'^("?[^"|\.]*"?\."[^"]*") AS "[^"]*"$'
# Create a list from the non-derived values between SELECT and FROM
selects_list = [val for val in select_statements if re.search(DIRECT_SELECT_QUERY_REGEX, val)]
selects_list = [
re.fullmatch(DIRECT_SELECT_QUERY_REGEX, val).group(1)
for val in select_statements
if re.fullmatch(DIRECT_SELECT_QUERY_REGEX, val)
]

# Create a list from the derived values between SELECT and FROM
aliased_list = [
(idx + 1, val)
for idx, val in enumerate(select_statements)
if not re.search(DIRECT_SELECT_QUERY_REGEX, val.strip())
if not re.fullmatch(DIRECT_SELECT_QUERY_REGEX, val.strip())
]
deriv_dict = {}

Expand All @@ -768,20 +780,16 @@ def apply_annotations_to_sql(
for idx, val in aliased_list:
split_string = _top_level_split(val, " AS ")
alias = split_string[1].replace('"', "").replace(",", "").strip()
if alias not in aliases:
raise Exception(f'alias "{alias}" not found!')
col_select = split_string[0]
if alias in aliases:
deriv_dict[alias] = col_select
if annotated_group_by_columns and alias in annotated_group_by_columns:
group_by_to_replace.append((idx, col_select))
else:
# Django 5.2 emits an "AS <query_path>" alias on every values()-projected field
# (e.g. "treasury_account__funding_toptier_agency__name"). The query path is not a caller-known alias, so
# strip it and use the expression just like a direct SELECT in older Django
selects_list.append(col_select.strip())
deriv_dict[alias] = col_select
if annotated_group_by_columns and alias in annotated_group_by_columns:
group_by_to_replace.append((idx, col_select))

if group_by_to_replace:
first_half_query, second_half_query = _top_level_split(raw_query, " GROUP BY ")
# Fist we replace the position with a valid parameter we can use for formatting
# First we replace the position with a valid parameter we can use for formatting
for idx, _ in group_by_to_replace:
# It is assumed that all non-positional values in the GROUP BY are column names meaning the first number
# matching the value of "idx" will be the position. However, this is not guaranteed in the rest of the
Expand All @@ -795,7 +803,7 @@ def apply_annotations_to_sql(
elif second_half_query.startswith(f"{idx},"):
second_half_query = second_half_query.replace(f"{idx},", f" {{idx_{idx}}},", 1)
elif second_half_query.endswith(f" {idx}"):
second_half_query = second_half_query[:-1 * len(f" {idx}")] + f" {{idx_{idx}}}"
second_half_query = second_half_query[: -1 * len(f" {idx}")] + f" {{idx_{idx}}}"
second_half_query = second_half_query.format(
**{f"idx_{idx}": col_select for idx, col_select in group_by_to_replace}
)
Expand Down Expand Up @@ -880,7 +888,7 @@ def _top_level_split(sql: str, splitter: str) -> str:
if parens_depth == 0:
if sql[index : index + len(splitter)] == splitter:
return [sql[:index], sql[index + len(splitter) :]]
raise Exception(f"SQL string ${sql} cannot be split on ${splitter}")
raise ValueError(f"SQL string ${sql} cannot be split on ${splitter}")


def execute_psql(temp_sql_file_path: str, source_path: str, download_job: DownloadJob) -> None:
Expand Down
Loading
Loading