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
