import hashlib
import io
import logging
import os
import random
import time
from contextlib import contextmanager
from textwrap import fill
import ijson
import simplejson as json
from django.conf import settings
from django.db import IntegrityError, OperationalError, connections, transaction
from psycopg import errors
from yapw.clients import AsyncConsumer, Blocking
from yapw.decorators import decorate
from yapw.methods import add_callback_threadsafe, nack
from process.exceptions import InvalidFormError
from process.models import Collection, CollectionFile, CollectionNote, ProcessingStep, Record
logger = logging.getLogger(__name__)
YAPW_KWARGS = {"url": settings.RABBIT_URL, "exchange": settings.RABBIT_EXCHANGE_NAME, "prefetch_count": 20}
EXTENSION_URL = "https://raw.githubusercontent.com/open-contracting-extensions/ocds_{}_extension/master/extension.json"
[docs]
def wrap(string):
"""Format a long string as a help message, and return it."""
return "\n".join(fill(paragraph, width=78, replace_whitespace=False) for paragraph in string.split("\n"))
[docs]
def walk(paths):
for path in paths:
if os.path.isfile(path):
yield path
else:
for root, _, files in os.walk(path):
for name in files:
if not name.startswith("."):
yield os.path.join(root, name)
[docs]
@contextmanager
def get_publisher():
client = Blocking(**YAPW_KWARGS)
try:
yield client
finally:
client.close()
[docs]
def consume(*args, **kwargs):
client = AsyncConsumer(*args, **kwargs, **YAPW_KWARGS)
client.start()
[docs]
def decorator(decode, callback, state, channel, method, properties, body):
"""
Close the database connections opened by the callback, before returning.
If the callback raises an exception, shut down the client in the main thread, without acknowledgment. For some
exceptions, assume that the same message was delivered twice, log an error, and nack the message.
"""
def errback(exception):
# A deadlock can occur when concurrent transactions INSERT or DELETE the same rows in a different order.
# Requeue the message to retry it. The callback must leave no partial state on a rolled-back transaction
# (see e.g. file_worker and wiper). Sleep first, so that the transaction that won the deadlock can commit,
# and so that concurrent threads retry at different times, to avoid repeating the deadlock.
if isinstance(exception, OperationalError) and isinstance(exception.__cause__, errors.DeadlockDetected):
logger.exception("Deadlock when consuming %r, requeuing", body)
time.sleep(random.randint(1, 5)) # noqa: S311 # non-cryptographic
nack(state, channel, method.delivery_tag, requeue=True)
elif isinstance(exception, IntegrityError) and isinstance(exception.__cause__, errors.ForeignKeyViolation):
# A foreign-key violation on a collection reference occurs when a collection is deleted (by the wiper)
# while another worker is compiling releases into it or storing files for it: the concurrent transaction
# commits rows referencing the now-deleted collection, so the foreign key fails at COMMIT. The message is
# obsolete, so skip it, like Collection.DoesNotExist below.
if "collection_id" in (exception.__cause__.diag.constraint_name or ""):
logger.exception("Collection deleted while consuming %r, skipping", body)
nack(state, channel, method.delivery_tag, requeue=False)
# Any other foreign-key violation means a package_data or data row is still referenced (e.g. by another
# collection's release). It indicates an error in logic or configuration (like toggling DEDUPLICATE_DATA).
else:
logger.exception("Unhandled exception when consuming %r, shutting down gracefully", body)
add_callback_threadsafe(state.connection, state.interrupt)
# These errors should only occur if the RabbitMQ and/or PostgreSQL connection is lost. It's not possible to
# have a transaction that spans both systems, so it's possible to insert a row then fail to ack a message.
#
# That said, we monitor the frequency of these errors via Sentry, to ensure that they are caused by the above
# and not by an error in logic. Their number should not exceed the prefetch count.
#
# InvalidFormError is included, as it may be for a "unique_together" error, which is an integrity error.
#
# Collection.DoesNotExist should only occur in the wiper worker due to a duplicate message. It can also occur
# in the finisher worker if the worker was stopped, and the wiper ran before the finisher.
elif isinstance(exception, InvalidFormError | IntegrityError | Collection.DoesNotExist):
logger.exception("%s maybe caused by duplicate message %r, skipping", type(exception).__name__, body)
nack(state, channel, method.delivery_tag, requeue=False)
# These errors should never occur under normal operations. However, such messages interrupt processing, so they
# are discarded.
elif isinstance(exception, CollectionFile.DoesNotExist | Record.DoesNotExist):
logger.exception("Unprocessable message %r, skipping", body)
nack(state, channel, method.delivery_tag, requeue=False)
else:
logger.exception("Unhandled exception when consuming %r, shutting down gracefully", body)
add_callback_threadsafe(state.connection, state.interrupt)
def finalback():
for conn in connections.all():
conn.close()
decorate(decode, callback, state, channel, method, properties, body, errback, finalback)
[docs]
def get_or_create(model, data):
"""Get or create a PackageData or Data instance."""
if not settings.DEDUPLICATE_DATA:
return model.objects.create(hash_md5="", data=data)
hash_md5 = hashlib.md5( # noqa: S324 # non-cryptographic
json.dumps(data, separators=(",", ":"), sort_keys=True, use_decimal=True).encode("utf-8")
).hexdigest()
try:
# Another transaction is needed here, otherwise a parent transaction catches the integrity error.
with transaction.atomic():
obj, _created = model.objects.get_or_create(hash_md5=hash_md5, defaults={"data": data})
# If another transaction in another thread COMMITs the same data after the SELECT, but before the INSERT.
except IntegrityError:
obj = model.objects.get(hash_md5=hash_md5)
return obj
[docs]
def create_note(collection, code, note, **kwargs):
if isinstance(note, list):
note = "\n".join(note)
CollectionNote(collection=collection, code=code, note=note, **kwargs).save()
[docs]
def create_step(name, collection_id, **kwargs):
ProcessingStep(name=name, collection_id=collection_id, **kwargs).save()
[docs]
@contextmanager
def deleting_step(*args, **kwargs):
"""Delete the named step and run any finish callback only if successful or if the error is expected."""
try:
yield
# Delete the step so that the collection is completable, only if the error was expected.
except (
# See the errback() function in the decorator() function.
InvalidFormError,
IntegrityError,
# See the try/except block in the callback() function of the file_worker worker.
FileNotFoundError,
ijson.common.IncompleteJSONError,
) as exception:
delete_step(*args, **kwargs, exception=exception)
raise
else:
delete_step(*args, **kwargs)
[docs]
def delete_step(name, finish=None, finish_args=(), exception=None, **kwargs):
# kwargs must include collection_id, collection_file_id and/or ocid.
processing_steps = ProcessingStep.objects.filter(name=name, **kwargs)
deleted, _ = processing_steps.delete()
if not deleted:
logger.warning("No such processing step found: %s: %s", name, kwargs)
if finish:
finish(*finish_args, exception=exception)
[docs]
@contextmanager
def create_logger_note(collection, name):
stream = io.StringIO()
handler = logging.StreamHandler(stream)
handler.setLevel(logging.WARNING)
logger = logging.getLogger(name)
logger.addHandler(handler)
yield
logger.removeHandler(handler)
if note := stream.getvalue():
create_note(collection, CollectionNote.Level.WARNING, note)
[docs]
def get_extensions(package):
extensions = set()
package_extensions = package.get("extensions")
if isinstance(package_extensions, list):
extensions = {extension for extension in package_extensions if isinstance(extension, str)}
# The master version of the lots extension depends on OCDS 1.2 or the submission terms extension.
if EXTENSION_URL.format("lots") in extensions:
extensions.add(EXTENSION_URL.format("submissionTerms"))
return frozenset(extensions)