def test_enqueued_fake_message_type(): q = get_sqs_queue() q.send_message("1235", {"attr1": "val1", "attr2": 111}) msg = get_sqs_queue()._messages().pop() assert isinstance(msg, FakeSQSMessage) assert msg.body == "1235" assert msg.message_attributes.get("attr2") == 111
def download( self, file_name, prime_award_types=None, agency=None, sub_agency=None, date_type=None, start_date=None, end_date=None, columns=[], file_format="csv", monthly_download=False, cleanup=False, use_sqs=False, ): date_range = {} if start_date: date_range["start_date"] = start_date if end_date: date_range["end_date"] = end_date json_request = { "constraint_type": "year", "filters": { "prime_award_types": prime_award_types, "agency": str(agency), "date_type": date_type, "date_range": date_range, }, "columns": columns, "file_format": file_format, } download_viewset = YearLimitedDownloadViewSet() download_viewset.process_filters(json_request) validated_request = validate_award_request(json_request) download_job = DownloadJob.objects.create( job_status_id=JOB_STATUS_DICT["ready"], file_name=file_name, json_request=json.dumps(order_nested_object(validated_request)), monthly_download=True, ) if not use_sqs: # Note: Because of the line below, it's advised to only run this script on a separate instance as this will # modify your bulk download settings. settings.BULK_DOWNLOAD_S3_BUCKET_NAME = settings.MONTHLY_DOWNLOAD_S3_BUCKET_NAME download_generation.generate_download(download_job=download_job) if cleanup: # Get all the files that have the same prefix except for the update date file_name_prefix = file_name[: -12] # subtracting the 'YYYYMMDD.zip' for key in self.bucket.objects.filter(Prefix=file_name_prefix): if key.key == file_name: # ignore the one we just uploaded continue key.delete() logger.info("Deleting {} from bucket".format(key.key)) else: queue = get_sqs_queue( queue_name=settings.BULK_DOWNLOAD_SQS_QUEUE_NAME) queue.send_message(MessageBody=str(download_job.download_job_id))
def handle(self, *args, **options): # Configure Tracer to drop traces of polls of the queue that have been flagged as uninteresting DatadogEagerlyDropTraceFilter.activate() queue = get_sqs_queue() log_job_message(logger=logger, message="Starting SQS polling", job_type=JOB_TYPE) message_found = None keep_polling = True while keep_polling: # Start a Datadog Trace for this poll iter to capture activity in APM with tracer.trace(name=f"job.{JOB_TYPE}", service="bulk-download", resource=queue.url, span_type=SpanTypes.WORKER) as span: # Set True to add trace to App Analytics: # - https://docs.datadoghq.com/tracing/app_analytics/?tab=python#custom-instrumentation span.set_tag(ANALYTICS_SAMPLE_RATE_KEY, 1.0) # Setup dispatcher that coordinates job activity on SQS dispatcher = SQSWorkDispatcher( queue, worker_process_name=JOB_TYPE, worker_can_start_child_processes=True) try: # Check the queue for work and hand it to the given processing function message_found = dispatcher.dispatch(download_service_app) # Mark the job as failed if: there was an error processing the download; retries after interrupt # are not allowed; or all retries have been exhausted # If the job is interrupted by an OS signal, the dispatcher's signal handling logic will log and # handle this case # Retries are allowed or denied by the SQS queue's RedrivePolicy config # That is, if maxReceiveCount > 1 in the policy, then retries are allowed # - if queue retries are allowed, the queue message will retry to the max allowed by the queue # - As coded, no cleanup should be needed to retry a download # - the psql -o will overwrite the output file # - the zip will use 'w' write mode to create from scratch each time # The worker function controls the maximum allowed runtime of the job except (QueueWorkerProcessError, QueueWorkDispatcherError) as exc: _handle_queue_error(exc) if not message_found: # Flag the the Datadog trace for dropping, since no trace-worthy activity happened on this poll DatadogEagerlyDropTraceFilter.drop(span) # When you receive an empty response from the queue, wait before trying again time.sleep(1) # If this process is exiting, don't poll for more work keep_polling = not dispatcher.is_exiting
def test_fake_receive_message(): q = get_sqs_queue() q.send_message("1235", {"attr1": "val1", "attr2": 111}) q.send_message("2222") q.send_message("3333") q.send_message("4444", {"attr1": "v1", "attr2": "v2"}) msgs = q.receive_messages(10) assert len(msgs) == 1 assert msgs[0].body == "1235" assert msgs[0].message_attributes.get("attr2") == 111
def _patch_get_sqs_queue(fake_sqs_queue, monkeypatch): """ Patch the use of get_sqs_queue HERE, in tests in this module (defined by ``__name__``), where the function is imported as a new name in this module. Chaining to fixture ``fake_sqs_queue`` will also take care of fake-queue purging and cleanup. """ monkeypatch.setattr(f"{__name__}.get_sqs_queue", get_unittest_fake_sqs_queue) should_be_fake = get_sqs_queue() # Check that the patch worked, and the fake unit test queue is returned assert should_be_fake.url.split("/")[-1] == UNITTEST_FAKE_QUEUE_NAME
def process_request(self, download_job: DownloadJob): if settings.IS_LOCAL and settings.RUN_LOCAL_DOWNLOAD_IN_PROCESS: # Eagerly execute the download in this running process download_generation.generate_download(download_job) else: # Send a SQS message that will be processed by another server which will eventually run # download_generation.generate_download(download_source) (see download_sqs_worker.py) write_to_log( message= f"Passing download_job {download_job.download_job_id} to SQS", download_job=download_job) queue = get_sqs_queue( queue_name=settings.BULK_DOWNLOAD_SQS_QUEUE_NAME) queue.send_message(MessageBody=str(download_job.download_job_id))
def test_fake_receive_messages(): q = get_sqs_queue() q.send_message("1235", {"attr1": "val1", "attr2": 111}) q.send_message("2222") q.send_message("3333") q.send_message("4444", {"attr1": "v1", "attr2": "v2"}) msgs = q.receive_messages(10, MaxNumberOfMessages=10) assert len(msgs) == 4 assert msgs[0].body == "1235" assert msgs[0].message_attributes.get("attr2") == 111 assert msgs[3].body == "4444" assert msgs[3].message_attributes.get("attr2") == "v2" q.purge() q.send_message("1235", {"attr1": "val1", "attr2": 111}) q.send_message("2222") q.send_message("3333") q.send_message("4444", {"attr1": "v1", "attr2": "v2"}) msgs = q.receive_messages(10, MaxNumberOfMessages=2) assert len(msgs) == 2 assert msgs[0].body == "1235" assert msgs[0].message_attributes.get("attr2") == 111 assert msgs[1].body == "2222"
def test_delete_fake_received_message(): q = get_sqs_queue() q.send_message("1235", {"attr1": "val1", "attr2": 111}) q.send_message("2222") q.send_message("3333") q.send_message("4444", {"attr1": "v1", "attr2": "v2"}) assert len(get_sqs_queue()._messages()) == 4 bodies = [m.body for m in get_sqs_queue()._messages()] assert "1235" in bodies msgs = q.receive_messages(10) assert len(msgs) == 1 assert len(get_sqs_queue()._messages()) == 4 msg = msgs[0] assert msg.body == "1235" assert msg.message_attributes.get("attr2") == 111 msg.delete() assert len(get_sqs_queue()._messages()) == 3 bodies = [m.body for m in get_sqs_queue()._messages()] assert "1235" not in bodies
def push_job_to_queue(self): # Candidate for separate object or file queue = get_sqs_queue(queue_name=settings.BULK_DOWNLOAD_SQS_QUEUE_NAME) queue.send_message(MessageBody=str(self.download_job.download_job_id))
def test_fake_send_message(): assert len(get_sqs_queue()._messages()) == 0 q = get_sqs_queue() q.send_message("1235", {"attr1": "val1", "attr2": 111}) assert len(get_sqs_queue()._messages()) == 1
def test_getting_fake_dead_letter_queue(): q = get_sqs_queue(queue_name=UNITTEST_FAKE_DEAD_LETTER_QUEUE_NAME) assert q.url == _FakeStatelessLoggingSQSDeadLetterQueue.url
def test_unittest_gets_fake_queue(): q = get_sqs_queue() assert q.url.split("/")[-1] == UNITTEST_FAKE_QUEUE_NAME