import logging
import threading
from functools import cache
from typing import Any
import orjson
import psycopg
from psycopg import sql
from psycopg.rows import dict_row
from psycopg.types.json import set_json_dumps, set_json_loads
from yapw.clients import AsyncConsumer, Blocking
from pelican.util import settings
db_connected = False
db_connection = None
db_cursor_idx = 0
logger = logging.getLogger(__name__)
# RabbitMQ
[docs]
def encode(message: Any, content_type: str | None) -> bytes:
"""
Encode the body of a message for RabbitMQ.
:param message: a decoded message
:param content_type: the message's content type
"""
return orjson.dumps(message)
[docs]
def decode(body: bytes, content_type: str | None) -> Any:
"""
Decode the body of a message from RabbitMQ.
:param message: an encoded message
:param content_type: the message's content type
"""
return orjson.loads(body)
YAPW_KWARGS = {
"url": settings.RABBIT_URL,
"exchange": settings.RABBIT_EXCHANGE_NAME,
"encode": encode,
"decode": decode,
}
[docs]
def consume(*args: Any, prefetch_count=1, **kwargs: Any) -> None:
"""Consume messages from RabbitMQ."""
client = AsyncConsumer(*args, prefetch_count=prefetch_count, **kwargs, **YAPW_KWARGS)
client.start()
[docs]
def publish(*args: Any, **kwargs: Any) -> None:
"""Publish a message to RabbitMQ."""
client = Blocking(**YAPW_KWARGS)
try:
client.publish(*args, **kwargs)
finally:
client.close()
# PostgreSQL
set_json_dumps(orjson.dumps)
set_json_loads(orjson.loads)
[docs]
def get_cursor(name="") -> psycopg.Cursor[dict[str, Any]]:
"""Connect to the database, if needed, and return a database cursor."""
global db_connected, db_connection, db_cursor_idx # noqa: PLW0603
if not db_connected:
db_connection = psycopg.connect(settings.DATABASE_URL, row_factory=dict_row)
db_connected = True
if name:
# https://github.com/django/django/blob/stable/4.2.x/django/db/backends/postgresql/base.py#L469
db_cursor_idx += 1
cursor_name = f"{name}-{threading.current_thread().ident}-{db_cursor_idx}"
# Avoid "named cursor isn't valid anymore". Another option is to use a separate connection.
# https://www.psycopg.org/psycopg3/docs/advanced/cursors.html#server-side-cursors
return db_connection.cursor(name=cursor_name, withhold=True)
return db_connection.cursor()
[docs]
def commit() -> None:
"""Commit the transaction."""
db_connection.commit()
[docs]
def rollback() -> None:
"""Rollback the transaction."""
db_connection.rollback()
[docs]
class State:
IN_PROGRESS = "IN_PROGRESS"
OK = "OK"
[docs]
class Phase:
CONTRACTING_PROCESS = "CONTRACTING_PROCESS"
DATASET = "DATASET"
TIME_VARIANCE = "TIME_VARIANCE"
CHECKED = "CHECKED"
DELETED = "DELETED"
[docs]
def initialize_dataset_state(dataset_id: int) -> None:
"""
Initialize a dataset's progress.
:param dataset_id: the dataset's ID
"""
sql = """\
INSERT INTO progress_monitor_dataset (dataset_id, phase, state, size)
VALUES (%(dataset_id)s, %(phase)s, %(state)s, 0)
"""
with get_cursor() as cursor:
cursor.execute(sql, {"dataset_id": dataset_id, "phase": Phase.CONTRACTING_PROCESS, "state": State.IN_PROGRESS})
[docs]
def update_dataset_state(dataset_id: int, phase: str, state: str, size: int | None = None) -> None:
"""
Update a dataset's progress to the given phase and state.
:param dataset_id: the dataset's ID
:param phase: the phase to be set
:param state: the state to set
:param size: number of data items to process
"""
variables = {"phase": phase, "state": state, "dataset_id": dataset_id}
sql = """\
UPDATE progress_monitor_dataset
SET phase = %(phase)s, state = %(state)s, modified = now()
WHERE dataset_id = %(dataset_id)s
"""
if size:
variables["size"] = size
sql = """\
UPDATE progress_monitor_dataset
SET phase = %(phase)s, state = %(state)s, modified = now(), size = %(size)s
WHERE dataset_id = %(dataset_id)s
"""
with get_cursor() as cursor:
cursor.execute(sql, variables)
[docs]
def initialize_items_state(dataset_id: int, item_ids: list[int]) -> None:
"""
Initialize data items' progress.
:param dataset_id: the dataset's ID
:param item_ids: the data items' IDs
"""
with get_cursor() as cursor:
cursor.executemany(
"""\
INSERT INTO progress_monitor_item (dataset_id, item_id, state)
VALUES (%(dataset_id)s, %(item_id)s, %(state)s)
""",
[{"dataset_id": dataset_id, "item_id": item_id, "state": State.IN_PROGRESS} for item_id in item_ids],
)
[docs]
def update_items_state(dataset_id: int, item_ids: list[int], state: str) -> None:
"""
Update data items' progress to the given state.
:param dataset_id: the dataset's ID
:param item_ids: the data items' IDs
:param state: the state to set
"""
item_ids = list(item_ids)
if not item_ids:
return
with get_cursor() as cursor:
records = sql.SQL(", ").join(sql.SQL("(%s, %s, %s)") for _ in item_ids)
statement = sql.SQL(
"""\
UPDATE progress_monitor_item
SET state = data.state, modified = now()
FROM (VALUES {}) AS data (dataset_id, item_id, state)
WHERE progress_monitor_item.dataset_id = data.dataset_id AND progress_monitor_item.item_id = data.item_id
"""
).format(records)
cursor.execute(statement, [parameter for item_id in item_ids for parameter in (dataset_id, item_id, state)])
[docs]
def get_processed_items_count(dataset_id: int) -> int:
"""
Return the number of items processed.
:param dataset_id: the dataset's ID
"""
with get_cursor() as cursor:
cursor.execute(
"SELECT COUNT(*) cnt FROM progress_monitor_item WHERE dataset_id = %(dataset_id)s AND state = %(state)s",
{"dataset_id": dataset_id, "state": State.OK},
)
return cursor.fetchone()["cnt"]
# The check.dataset worker calls this function when phase=CONTRACTING_PROCESS and state=OK, at which point size is set.
[docs]
@cache
def get_total_items_count(dataset_id: int) -> int:
"""
Return the number of items to process.
:param dataset_id: the dataset's ID
"""
with get_cursor() as cursor:
cursor.execute(
"SELECT size FROM progress_monitor_dataset WHERE dataset_id = %(dataset_id)s", {"dataset_id": dataset_id}
)
return cursor.fetchone()["size"]
[docs]
def get_dataset_progress(dataset_id: int) -> tuple[Any, ...] | None:
"""
Return the dataset's progress.
:param dataset_id: the dataset's ID
"""
with get_cursor() as cursor:
cursor.execute(
"SELECT * FROM progress_monitor_dataset WHERE dataset_id = %(dataset_id)s", {"dataset_id": dataset_id}
)
return cursor.fetchone()