diff --git a/usaspending_api/download/filestreaming/download_generation.py b/usaspending_api/download/filestreaming/download_generation.py old mode 100644 new mode 100755 index 244a42ef49..41eaa39857 --- a/usaspending_api/download/filestreaming/download_generation.py +++ b/usaspending_api/download/filestreaming/download_generation.py @@ -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: @@ -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) @@ -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) @@ -158,9 +157,7 @@ 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( @@ -168,8 +165,9 @@ def _process_download_sources( ) -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) @@ -192,16 +190,18 @@ 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) @@ -209,16 +209,18 @@ 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) @@ -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 @@ -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 @@ -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 = {} @@ -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 " 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 @@ -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} ) @@ -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: diff --git a/usaspending_api/download/tests/integration/test_download_generation.py b/usaspending_api/download/tests/integration/test_download_generation.py new file mode 100755 index 0000000000..6854ae9f51 --- /dev/null +++ b/usaspending_api/download/tests/integration/test_download_generation.py @@ -0,0 +1,94 @@ +import re + +import pytest +from django.db.models import F, Value + +from usaspending_api.common.helpers.orm_helpers import ConcatAll, generate_raw_quoted_query +from usaspending_api.download.filestreaming.download_generation import apply_annotations_to_sql +from usaspending_api.search.models import AwardSearch, TransactionSearch + + +@pytest.mark.django_db +def test_apply_annotations_to_sql_single_table(): + selected_columns = ["transaction_id", "action_type_code", "action_type", "date_range"] + queryset = ( + TransactionSearch.objects.values("transaction_id") + .annotate( + action_type_code=F("action_type"), + action_type=F("action_type_description"), + date_range=ConcatAll( + F("period_of_performance_start_date"), Value(" - "), F("period_of_performance_current_end_date") + ), + ) + .values(*selected_columns) + .order_by("transaction_id") + ) + queryset = queryset[:1] + raw_query = generate_raw_quoted_query(queryset) + result_query = apply_annotations_to_sql(raw_query, selected_columns) + expected_query = ( + "SELECT " + '"transaction_search"."transaction_id" AS "transaction_id", ' + '"transaction_search"."action_type" AS "action_type_code", ' + '"transaction_search"."action_type_description" AS "action_type", ' + "CONCAT(" + '"transaction_search"."period_of_performance_start_date", ' + "' - ', " + '"transaction_search"."period_of_performance_current_end_date"' + ') AS "date_range" ' + 'FROM "transaction_search" ' + "ORDER BY 1 ASC LIMIT 1" + ) + assert result_query == expected_query + + +@pytest.mark.django_db +def test_apply_annotations_to_sql_joined_table(): + select_columns = [ + "award_id", + "earliest_transaction_search__solicitation_date", + "latest_transaction_search__action_date", + ] + annotate_columns = ["assistance_award_unique_key", "award_id_fain", "date_range"] + queryset = ( + AwardSearch.objects.values(*select_columns) + .annotate( + assistance_award_unique_key=F("generated_unique_award_id"), + award_id_fain=F("fain"), + date_range=ConcatAll( + F("latest_transaction_search__period_of_performance_start_date"), + Value(" - "), + F("latest_transaction_search__period_of_performance_current_end_date"), + ), + ) + .values(*select_columns, *annotate_columns) + .order_by("award_id") + ) + queryset = queryset[:1] + raw_query = generate_raw_quoted_query(queryset) + result_query = apply_annotations_to_sql(raw_query, [*select_columns, *annotate_columns]) + + # The second reference to the "transaction_search" table should be aliases. Instead of hard coding this value + # it is retrieved via a regex in case it happens to change. + table_alias = re.search(r'"transaction_search" ([^ ]*) ON', raw_query).group(1) + + expected_query = ( + "SELECT " + '"award_search"."award_id" AS "award_id", ' + '"transaction_search"."solicitation_date" AS "earliest_transaction_search__solicitation_date", ' + f'{table_alias}."action_date" AS "latest_transaction_search__action_date", ' + '"award_search"."generated_unique_award_id" AS "assistance_award_unique_key", ' + '"award_search"."fain" AS "award_id_fain", ' + "CONCAT(" + f'{table_alias}."period_of_performance_start_date", ' + "' - ', " + f'{table_alias}."period_of_performance_current_end_date"' + ') AS "date_range" ' + 'FROM "award_search" ' + 'LEFT OUTER JOIN "transaction_search" ON ' + '("award_search"."earliest_transaction_search_id" = "transaction_search"."transaction_id") ' + f'LEFT OUTER JOIN "transaction_search" {table_alias} ' + f'ON ("award_search"."latest_transaction_search_id" = {table_alias}."transaction_id") ' + "ORDER BY 1 ASC LIMIT 1" + ) + assert result_query == expected_query diff --git a/usaspending_api/download/tests/unit/test_sql_annotations.py b/usaspending_api/download/tests/unit/test_sql_annotations.py index 4539c925d0..07ff3bb3d5 100644 --- a/usaspending_api/download/tests/unit/test_sql_annotations.py +++ b/usaspending_api/download/tests/unit/test_sql_annotations.py @@ -2,9 +2,9 @@ # Imports from your apps from usaspending_api.download.filestreaming.download_generation import ( - apply_annotations_to_sql, - _top_level_split, _select_columns, + _top_level_split, + apply_annotations_to_sql, ) @@ -30,7 +30,7 @@ def test_top_level_split_quote_logic(): def test_top_level_split_unsplittable(): unsplittable = "sdoiaghioweg oijfewiwje efoiwnovisper" - with pytest.raises(Exception): + with pytest.raises(ValueError): _top_level_split(unsplittable, "not in this") @@ -48,20 +48,6 @@ def test_select_columns_nesting_logic(): assert _select_columns(cte_sql) == ("WITH cte AS (SELECT inner_things FROM a inner_place)", ["things"]) -def test_apply_annotations_to_sql(): - sql_string = str( - 'SELECT "table"."col1", "table"."col2", (SELECT table2."three" FROM table_two table2 WHERE ' - 'table2."code" = table."othercode") AS "alias_one" FROM table WHERE six = \'something\'' - ) - aliases = ["alias_one", "col1", "col2"] - - annotated_sql = apply_annotations_to_sql(sql_string, aliases) - - annotated_string = 'SELECT (SELECT table2."three" FROM table_two table2 WHERE table2."code" = table."othercode") AS "alias_one", "table"."col1" AS "col1", "table"."col2" AS "col2" FROM table WHERE six = \'something\'' - - assert annotated_sql == annotated_string - - def test_group_by_annotations(): sql = ( "SELECT"