Mobilizon-Reshare-condividi.../mobilizon_bots/storage/query.py

125 lines
4.0 KiB
Python

from typing import Iterable, Optional, List
import sys
from arrow import Arrow
from tortoise.queryset import QuerySet
from tortoise.transactions import atomic
from mobilizon_bots.event.event import MobilizonEvent, EventPublicationStatus
from mobilizon_bots.models.event import Event
from mobilizon_bots.models.publication import Publication, PublicationStatus
from mobilizon_bots.models.publisher import Publisher
from mobilizon_bots.publishers.coordinator import PublisherCoordinatorReport
# This is due to Tortoise community fixtures to
# set up and tear down a DB instance for Pytest.
# See: https://github.com/tortoise/tortoise-orm/issues/419#issuecomment-696991745
# and: https://docs.pytest.org/en/stable/example/simple.html
CONNECTION_NAME = "models" if "pytest" in sys.modules else None
async def prefetch_event_relations(queryset: QuerySet[Event]) -> List[Event]:
return (
await queryset.prefetch_related("publications__publisher")
.order_by("begin_datetime")
.distinct()
)
def _add_date_window(
query, from_date: Optional[Arrow] = None, to_date: Optional[Arrow] = None,
):
if from_date:
query = query.filter(end_datetime__gt=from_date.datetime)
if to_date:
query = query.filter(end_datetime__lt=to_date.datetime)
return query
async def events_with_status(
status: List[EventPublicationStatus],
from_date: Optional[Arrow] = None,
to_date: Optional[Arrow] = None,
) -> Iterable[MobilizonEvent]:
def _filter_event_with_status(event: Event) -> bool:
# This computes the status client-side instead of running in the DB. It shouldn't pose a performance problem
# in the short term, but should be moved to the query if possible.
event_status = MobilizonEvent.compute_status(list(event.publications))
return event_status in status
query = Event.all()
_add_date_window(query, from_date, to_date)
return map(
MobilizonEvent.from_model,
filter(_filter_event_with_status, await prefetch_event_relations(query),),
)
async def get_all_events(
from_date: Optional[Arrow] = None, to_date: Optional[Arrow] = None,
) -> Iterable[MobilizonEvent]:
return map(
MobilizonEvent.from_model,
await prefetch_event_relations(
_add_date_window(Event.all(), from_date, to_date)
),
)
async def get_published_events() -> Iterable[MobilizonEvent]:
return await events_with_status([EventPublicationStatus.COMPLETED])
async def get_unpublished_events() -> Iterable[MobilizonEvent]:
return await events_with_status([EventPublicationStatus.WAITING])
async def save_event(event):
event_model = event.to_model()
await event_model.save()
return event_model
async def save_publication(publisher_name, event_model, status: PublicationStatus):
publisher = await Publisher.filter(name=publisher_name).first()
await Publication.create(
status=status, event_id=event_model.id, publisher_id=publisher.id,
)
@atomic(CONNECTION_NAME)
async def create_unpublished_events(
unpublished_mobilizon_events: Iterable[MobilizonEvent],
active_publishers: Iterable[str],
) -> None:
# We store only new events, i.e. events whose mobilizon_id wasn't found in the DB.
unpublished_event_models = set(
map(lambda event: event.mobilizon_id, await get_unpublished_events())
)
unpublished_events = list(
filter(
lambda event: event.mobilizon_id not in unpublished_event_models,
unpublished_mobilizon_events,
)
)
for event in unpublished_events:
event_model = await save_event(event)
for publisher in active_publishers:
await save_publication(
publisher, event_model, status=PublicationStatus.WAITING
)
async def create_publisher(name: str, account_ref: Optional[str] = None) -> None:
await Publisher.create(name=name, account_ref=account_ref)
async def save_publication_report(publication_report: PublisherCoordinatorReport):
for publisher_report in publication_report:
pass