from urllib.parse import urlparse
from datetime import datetime

import dramatiq
from dramatiq.brokers.redis import RedisBroker
from sqlmodel import Session

from app.db import engine
from app.models import Document, DocumentInput
from app.predict import CategoryPrediction
from app.settings import settings

redis_parameters = urlparse(settings.redis_url)
redis_broker = RedisBroker(
    host=redis_parameters.hostname,
    port=redis_parameters.port,
    username=redis_parameters.username,
    password=redis_parameters.password,
    # Heroku Redis with TLS use self-signed certs, so we need to tinker a bit
    ssl=redis_parameters.scheme == "rediss",
    ssl_cert_reqs=None,
)
dramatiq.set_broker(redis_broker)


category_prediction = CategoryPrediction()


@dramatiq.actor
def ingest_document(document_json: str):
    document = DocumentInput.parse_raw(document_json)
    with Session(engine) as session:
        document_db = session.get(Document, document.id)
        if document_db is None:
            document_db = Document(**document.dict())
        else:
            document_dict = document.dict(exclude_unset=True)
            for key, value in document_dict.items():
                setattr(document_db, key, value)

        document_db.category = category_prediction.predict(document_db.content)
        document_db.updated_at = datetime.utcnow()

        session.add(document_db)
        session.commit()

app.predict.CategoryPrediction dramatiq.set_broker datetime.datetime.utcnow urllib.parse.urlparse sqlmodel.Session app.models.DocumentInput.parse_raw dramatiq.brokers.redis.RedisBroker
[((328, 356), 'urllib.parse.urlparse', 'urlparse', (['settings.redis_url'], {}), '(settings.redis_url)\n', (336, 356), False, 'from urllib.parse import urlparse\n'), ((372, 584), 'dramatiq.brokers.redis.RedisBroker', 'RedisBroker', ([], {'host': 'redis_parameters.hostname', 'port': 'redis_parameters.port', 'username': 'redis_parameters.username', 'password': 'redis_parameters.password', 'ssl': "(redis_parameters.scheme == 'rediss')", 'ssl_cert_reqs': 'None'}), "(host=redis_parameters.hostname, port=redis_parameters.port,\n    username=redis_parameters.username, password=redis_parameters.password,\n    ssl=redis_parameters.scheme == 'rediss', ssl_cert_reqs=None)\n", (383, 584), False, 'from dramatiq.brokers.redis import RedisBroker\n'), ((682, 715), 'dramatiq.set_broker', 'dramatiq.set_broker', (['redis_broker'], {}), '(redis_broker)\n', (701, 715), False, 'import dramatiq\n'), ((740, 760), 'app.predict.CategoryPrediction', 'CategoryPrediction', ([], {}), '()\n', (758, 760), False, 'from app.predict import CategoryPrediction\n'), ((835, 873), 'app.models.DocumentInput.parse_raw', 'DocumentInput.parse_raw', (['document_json'], {}), '(document_json)\n', (858, 873), False, 'from app.models import Document, DocumentInput\n'), ((883, 898), 'sqlmodel.Session', 'Session', (['engine'], {}), '(engine)\n', (890, 898), False, 'from sqlmodel import Session\n'), ((1346, 1363), 'datetime.datetime.utcnow', 'datetime.utcnow', ([], {}), '()\n', (1361, 1363), False, 'from datetime import datetime\n')]