import hashlib
import io
import logging
import os
import random
import time
from contextlib import contextmanager
from pathlib import Path
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 ack, 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:
entry = Path(path)
if entry.is_file():
yield entry
else:
for root, _, files in os.walk(entry):
for name in files:
if not name.startswith("."):
yield Path(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 ack the message. Nack without
requeuing if the message requires review (e.g. unexpected messages), for any future dead-letter exchange.
"""
def errback(exc):
# 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(exc, OperationalError) and isinstance(exc.__cause__, errors.DeadlockDetected):
logger.error("Deadlock when consuming %r, requeuing message", body, exc_info=exc)
time.sleep(random.randint(1, 5)) # noqa: S311 # non-cryptographic
nack(state, channel, method.delivery_tag, requeue=True)
elif isinstance(exc, IntegrityError) and isinstance(exc.__cause__, errors.ForeignKeyViolation):
# A foreign-key violation on a collection reference occurs when a collection is deleted (by the wiper)
# while another worker is writing rows that reference it: the concurrent transaction commits those rows,
# so the foreign key fails at COMMIT. The message is obsolete, so skip it, like Collection.DoesNotExist
# below. (Note: Workers that call lock_collection() ack such messages earlier.)
if "collection_id" in (exc.__cause__.diag.constraint_name or ""):
logger.error("Collection deleted while consuming %r, discarding message", body, exc_info=exc)
ack(state, channel, method.delivery_tag)
# 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.error("Unhandled exception when consuming %r, shutting down gracefully", body, exc_info=exc)
add_callback_threadsafe(state.connection, state.interrupt)
# The wiper worker can delete a collection between a worker reading it and reading its parent collection,
# in compilable() and completable().
elif isinstance(exc, Collection.DoesNotExist):
logger.error("Collection deleted while consuming %r, discarding message", body, exc_info=exc)
ack(state, channel, method.delivery_tag)
# 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.
elif isinstance(exc, InvalidFormError | IntegrityError):
logger.error("%s maybe caused by duplicate message %r, discarding", type(exc).__name__, body, exc_info=exc)
ack(state, channel, method.delivery_tag)
# These errors should never occur under normal operations. However, such messages interrupt processing, so they
# are discarded.
elif isinstance(exc, CollectionFile.DoesNotExist | Record.DoesNotExist):
logger.error("Unprocessable message %r, discarding message", body, exc_info=exc)
nack(state, channel, method.delivery_tag, requeue=False)
else:
logger.error("Unhandled exception when consuming %r, shutting down gracefully", body, exc_info=exc)
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 lock_collection(collection_id):
"""
Lock the collection row with KEY SHARE, and return the collection, or ``None`` if it is deleted.
Call this within a transaction that writes rows referencing the collection. KEY SHARE conflicts only with
the wiper worker's FOR UPDATE; as such, other concurrent workers don't block each other.
A worker can call this only if a single transaction encloses all its writes, and if it can abandon that
transaction on a ``None`` return value. Workers that can't call this:
- ``compiler`` creates each step in autocommit mode, interleaved with publishing messages.
- ``file_worker`` can't abandon its transaction: deleting_step() would create a CHECK step for the collection.
"""
collections = Collection.objects.raw("SELECT * FROM collection WHERE id = %s FOR KEY SHARE", [collection_id])
return next(iter(collections), None)
[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: # expected if the wiper worker deleted the collection and its steps
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)