Ftr/dev 12639 update custom download

usaspending_api/common/experimental_api_flags.py

@@ -21,6 +21,8 @@
 EXPERIMENTAL_API_HEADER = "HTTP_X_EXPERIMENTAL_API"
 ELASTICSEARCH_HEADER_VALUE = "elasticsearch"
 
+DOWNLOAD_API_HEADER = "HTTP-X-DOWNLOAD-API"
+DOWNLOAD_HEADER_VALUE = "download"
 
 def is_experimental_elasticsearch_api(request: Request) -> bool:
     """
@@ -29,6 +31,12 @@ def is_experimental_elasticsearch_api(request: Request) -> bool:
     return request.META.get(EXPERIMENTAL_API_HEADER) == ELASTICSEARCH_HEADER_VALUE
 
 
+def is_experimental_download_api(request: Request) -> bool:
+    """
+    Returns True or False depending on if the expected_header_value matches what is sent with the request
+    """
+    return request.headers.get(DOWNLOAD_API_HEADER) == DOWNLOAD_HEADER_VALUE
+
 def mirror_request_to_elasticsearch(request: Union[HttpRequest, Request]):
     """Duplicate request and send-again against this server, with the ES header attached to mirror
     non-elasticsearch load against elasticsearch for load testing

usaspending_api/download/tests/integration/test_download_accounts.py

@@ -497,3 +497,29 @@ def test_empty_array_filter_fail(client, download_test_data):
     assert (
         "Field 'filters|def_codes' value '[]' is below min '1' items" in resp.json()["detail"]
     ), "Incorrect error message"
+
+@pytest.mark.django_db(databases=[settings.DOWNLOAD_DB_ALIAS, settings.DEFAULT_DB_ALIAS])
+def test_file_c_spark_download(client, download_test_data):
+    download_generation.retrieve_db_string = Mock(return_value=get_database_dsn_string())
+
+    resp = client.post(
+        "/api/v2/download/accounts/",
+        content_type="application/json",
+        data=json.dumps(
+            {
+                "account_level": "federal_account",
+                "filters": { "budget_function": "all", "agency": "all",
+                            "submission_types": [
+                                "account_balances",
+                                "object_class_program_activity",
+                                "award_financial"
+                            ],
+                            "fy": "2021",
+                            "period": 12
+                },
+                "file_format": "csv"
+            }
+        ),
+    )
+
+    assert resp.status_code == status.HTTP_200_OK

usaspending_api/download/v2/base_download_viewset.py

@@ -13,7 +13,9 @@
 from usaspending_api.broker.lookups import EXTERNAL_DATA_TYPE_DICT
 from usaspending_api.broker.models import ExternalDataLoadDate
 from usaspending_api.common.api_versioning import API_TRANSFORM_FUNCTIONS, api_transformations
+from usaspending_api.common.experimental_api_flags import is_experimental_download_api
 from usaspending_api.common.helpers.dict_helpers import order_nested_object
+from usaspending_api.common.spark.jobs import DatabricksStrategy, LocalStrategy, SparkJobs
 from usaspending_api.common.sqs.sqs_handler import get_sqs_queue
 from usaspending_api.download.download_utils import create_unique_filename, log_new_download_job
 from usaspending_api.download.filestreaming import download_generation
@@ -24,7 +26,7 @@
 from usaspending_api.download.v2.request_validations import DownloadValidatorBase
 from usaspending_api.routers.replicas import ReadReplicaRouter
 from usaspending_api.submissions.models import DABSSubmissionWindowSchedule
-
+from usaspending_api.settings import IS_LOCAL
 
 @api_transformations(api_version=settings.API_VERSION, function_list=API_TRANSFORM_FUNCTIONS)
 class BaseDownloadViewSet(APIView):
@@ -68,7 +70,34 @@ def post(
         )
 
         log_new_download_job(request, download_job)
-        self.process_request(download_job)
+        if (
+            is_experimental_download_api(request)
+            and json_request["request_type"] == "account"
+            and "award_financial" in json_request["download_types"]
+        ):
+            # run spark download with only award_financial in download type
+            str_to_json_original = json.loads(download_job.json_request).copy()
+            str_to_json_award_financial = json.loads(download_job.json_request)
+            str_to_json_award_financial['download_types'] = ['award_financial']
+            str_to_json_award_financial['filters']['submission_types'] = ['award_financial']
+            download_job.json_request = json.dumps(str_to_json_award_financial)
+            download_job.save()
+            # goes to spark for File C account download
+            self.process_account_download_in_spark(download_job=download_job)
+            # remove award_financial from json request to run download with remaining types
+            if len(json_request["download_types"]) > 1:
+                str_to_json_original['filters']['submission_types'].remove('award_financial')
+                str_to_json_original['download_types'].remove('award_financial')
+                download_job.json_request = json.dumps(str_to_json_original)
+                download_job.save()
+                self.process_request(download_job)
+                # add award_financial back for the final response
+                str_to_json_original['filters']['submission_types'].append('award_financial')
+                str_to_json_original['download_types'].append('award_financial')
+                download_job.json_request = json.dumps(str_to_json_original)
+                download_job.save()
+        else:
+            self.process_request(download_job)
 
         return self.get_download_response(file_name=final_output_zip_name)
 
@@ -85,6 +114,15 @@ def process_request(self, download_job: DownloadJob):
             queue = get_sqs_queue(queue_name=settings.BULK_DOWNLOAD_SQS_QUEUE_NAME)
             queue.send_message(MessageBody=str(download_job.download_job_id))
 
+    def process_account_download_in_spark(self, download_job: DownloadJob):
+        """
+        Process File C downloads through spark instead of sqs for better performance
+        """
+        spark_jobs = SparkJobs(LocalStrategy()) if IS_LOCAL else SparkJobs(DatabricksStrategy())
+        spark_jobs.start(job_name="download_delta_table-award_search",
+                         command_name="generate_spark_download",
+                         command_options=[f"--download-job-id={download_job.download_job_id}"])
+
     def get_download_response(self, file_name: str):
         """
         Generate download response which encompasses various elements to provide accurate status for state of a