Source code for datanommer.models

# This file is a part of datanommer, a message sink for fedmsg.
# Copyright (C) 2014, Red Hat, Inc.
#
# This program is free software: you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free Software
# Foundation, either version 3 of the License, or (at your option) any later
# version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along
# with this program.  If not, see <http://www.gnu.org/licenses/>.
import datetime
import importlib.metadata
import json
import logging
import math
import traceback
import uuid
from warnings import warn

from sqlalchemy import (
    and_,
    between,
    Column,
    create_engine,
    DateTime,
    DDL,
    event,
    ForeignKey,
    func,
    Index,
    Integer,
    not_,
    or_,
    select,
    String,
    Table,
    text,
    TypeDecorator,
    Unicode,
    UnicodeText,
    UniqueConstraint,
)
from sqlalchemy.dialects import postgresql
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import (
    declarative_base,
    relationship,
    scoped_session,
    sessionmaker,
    validates,
)
from sqlalchemy.sql import operators

from .view import create_view

try:
    from psycopg2.errors import UniqueViolation
except ImportError:  # pragma: no cover
    from psycopg2.errorcodes import lookup as lookup_error

    UniqueViolation = lookup_error("23505")


__version__ = importlib.metadata.version("datanommer-models")


log = logging.getLogger("datanommer")

maker = sessionmaker()
session = scoped_session(maker)

DeclarativeBase = declarative_base()
DeclarativeBase.query = session.query_property()


[docs] def init(uri=None, alembic_ini=None, engine=None, create=False): """Initialize a connection. Create tables if requested.""" if uri and engine: raise ValueError("uri and engine cannot both be specified") if uri is None and not engine: raise ValueError("One of uri or engine must be specified") if uri and not engine: engine = create_engine(uri, future=True) # We need to hang our own attribute on the sqlalchemy session to stop # ourselves from initializing twice. That is only a problem if the code # calling us isn't consistent. if getattr(session, "_datanommer_initialized", None): log.warning("Session already initialized. Bailing") return session._datanommer_initialized = True maker.configure(bind=engine) DeclarativeBase.query = session.query_property() if create: with engine.begin() as connection: connection.execute(text("CREATE EXTENSION IF NOT EXISTS timescaledb")) DeclarativeBase.metadata.create_all(engine) with engine.begin() as connection: create_view(connection) # Loads the alembic configuration and generates the version table, with # the most recent revision stamped as head if alembic_ini is not None: # pragma: no cover from alembic import command from alembic.config import Config alembic_cfg = Config(alembic_ini) command.stamp(alembic_cfg, "head")
[docs] def add(message): """Take a the fedora-messaging Message and store in the message table. """ headers = message._properties.headers sent_at = headers.get("sent-at", None) if sent_at: # fromisoformat doesn't parse Z suffix (yet) see: # https://discuss.python.org/t/parse-z-timezone-suffix-in-datetime/2220 try: sent_at = datetime.datetime.fromisoformat(sent_at.replace("Z", "+00:00")) except ValueError: log.exception("Failed to parse sent-at timestamp value") return else: sent_at = datetime.datetime.now(tz=datetime.UTC) # Workaround schemas misbehaving try: usernames = message.usernames except Exception: log.exception( "Could not get the list of users from a message on %s with id %s", message.topic, message.id, ) usernames = [] try: packages = message.packages except Exception: log.exception( "Could not get the list of packages from a message on %s with id %s", message.topic, message.id, ) packages = [] Message.create( i=0, msg_id=message.id, topic=message.topic, timestamp=sent_at, msg=message.body, headers=headers, agent_name=getattr(message, "agent_name", None), users=usernames, packages=packages, ) session.commit()
# https://docs.sqlalchemy.org/en/14/core/custom_types.html#marshal-json-strings class _JSONEncodedDict(TypeDecorator): """Represents an immutable structure as a json-encoded string.""" impl = UnicodeText cache_ok = True def process_bind_param(self, value, dialect): if value is not None: value = json.dumps(value) return value def process_result_value(self, value, dialect): if value is not None: value = json.loads(value) return value def coerce_compared_value(self, op, value): # https://docs.sqlalchemy.org/en/14/core/custom_types.html#dealing-with-comparison-operations if op in (operators.like_op, operators.not_like_op): return String() else: return self users_assoc_table = Table( "users_messages", DeclarativeBase.metadata, Column("user_id", ForeignKey("users.id"), primary_key=True), Column("msg_id", Integer, primary_key=True, index=True), Column("msg_timestamp", DateTime, primary_key=True, index=True), ) packages_assoc_table = Table( "packages_messages", DeclarativeBase.metadata, Column("package_id", ForeignKey("packages.id"), primary_key=True), Column("msg_id", Integer, primary_key=True, index=True), Column("msg_timestamp", DateTime, primary_key=True, index=True), )
[docs] class Message(DeclarativeBase): __tablename__ = "messages" __table_args__ = ( UniqueConstraint("msg_id", "timestamp"), Index( "ix_messages_headers", "headers", postgresql_using="gin", postgresql_ops={"headers": "jsonb_path_ops"}, ), ) id = Column(Integer, primary_key=True, autoincrement=True) msg_id = Column(Unicode, nullable=True, default=None, index=True) i = Column(Integer, nullable=False) topic = Column(Unicode, nullable=False, index=True) timestamp = Column(DateTime, nullable=False, index=True, primary_key=True) certificate = Column(UnicodeText) signature = Column(UnicodeText) category = Column(Unicode, nullable=False, index=True) agent_name = Column(Unicode, index=True) crypto = Column(UnicodeText) source_name = Column(Unicode, default="datanommer") source_version = Column(Unicode, default=lambda context: __version__) msg = Column(_JSONEncodedDict, nullable=False) headers = Column(postgresql.JSONB(none_as_null=True)) users = relationship( "User", secondary=users_assoc_table, backref="messages", primaryjoin=lambda: and_( Message.id == users_assoc_table.c.msg_id, Message.timestamp == users_assoc_table.c.msg_timestamp, ), ) packages = relationship( "Package", secondary=packages_assoc_table, backref="messages", primaryjoin=lambda: and_( Message.id == packages_assoc_table.c.msg_id, Message.timestamp == packages_assoc_table.c.msg_timestamp, ), )
[docs] @validates("topic") def get_category(self, key, topic): """Update the category when the topic is set. The method seems... unnatural. But even zzzeek says it's OK to do it: https://stackoverflow.com/a/6442201 """ index = 2 if "VirtualTopic" in topic else 3 try: self.category = topic.split(".")[index] except Exception: traceback.print_exc() self.category = "Unclassified" return topic
[docs] @classmethod def create(cls, **kwargs): users = kwargs.pop("users") packages = kwargs.pop("packages") if not kwargs.get("msg_id"): log.info("Message on %s was received without a msg_id", kwargs["topic"]) kwargs["msg_id"] = str(uuid.uuid4()) obj = cls(**kwargs) try: session.add(obj) session.flush() except IntegrityError as e: if isinstance(e.orig, UniqueViolation): log.warning( "Skipping message from %s with duplicate id: %s", kwargs["topic"], kwargs["msg_id"], ) else: log.exception( "Unknown Integrity Error: message %s with id %s", kwargs["topic"], kwargs["msg_id"], ) session.rollback() return obj._insert_list(User, users_assoc_table, users) obj._insert_list(Package, packages_assoc_table, packages)
def _insert_list(self, rel_class, assoc_table, values): if not values: return assoc_col_name = assoc_table.c[0].name insert_values = [] for name in set(values): attr_obj = rel_class.get_or_create(name) # This would normally be a simple "obj.[users|packages].append(name)" kind # of statement, but here we drop down out of sqlalchemy's ORM and into the # sql abstraction in order to gain a little performance boost. insert_values.append( { assoc_col_name: attr_obj.id, "msg_id": self.id, "msg_timestamp": self.timestamp, } ) session.execute(assoc_table.insert(), insert_values) session.flush()
[docs] @classmethod def from_msg_id(cls, msg_id): return session.execute(select(cls).where(cls.msg_id == msg_id)).scalar_one_or_none()
[docs] def as_dict(self, request=None): return dict( i=self.i, msg_id=self.msg_id, topic=self.topic, timestamp=self.timestamp, certificate=self.certificate, signature=self.signature, agent_name=self.agent_name, username=self.agent_name, # DEPRECATED crypto=self.crypto, msg=self.msg, headers=self.headers, source_name=self.source_name, source_version=self.source_version, users=list(sorted(u.name for u in self.users)), packages=list(sorted(p.name for p in self.packages)), )
[docs] def as_fedora_message_dict(self): headers = self.headers or {} if "sent-at" not in headers: headers["sent-at"] = self.timestamp.astimezone(datetime.UTC).isoformat() return dict( body=self.msg, headers=headers, id=self.msg_id, priority=headers.get("priority", 0), queue=None, topic=self.topic, )
def __json__(self, request=None): warn( "The __json__() method has been renamed to as_dict(), and will be removed " "in the next major version", DeprecationWarning, stacklevel=2, ) return self.as_dict(request) @property def username(self): warn( "The username attribute has been renamed to agent_name, and will be removed " "in the next major version", DeprecationWarning, stacklevel=2, ) return self.agent_name
[docs] @classmethod def make_query( cls, start=None, end=None, msg_id=None, users=None, not_users=None, packages=None, not_packages=None, categories=None, not_categories=None, topics=None, not_topics=None, agents=None, not_agents=None, contains=None, ): """Flexible query interface for messages. Arguments are filters. start and end should be :mod:`datetime` objs. Other filters should be lists of strings. They are applied in a conjunctive-normal-form (CNF) kind of way for example, the following:: users = ['ralph', 'lmacken'] categories = ['bodhi', 'wiki'] should return messages where (user=='ralph' OR user=='lmacken') AND (category=='bodhi' OR category=='wiki') Furthermore, you can use a negative version of each argument. users = ['ralph'] not_categories = ['bodhi', 'wiki'] should return messages where (user == 'ralph') AND NOT (category == 'bodhi' OR category == 'wiki') """ users = users or [] not_users = not_users or [] packages = packages or [] not_packs = not_packages or [] categories = categories or [] not_cats = not_categories or [] topics = topics or [] not_topics = not_topics or [] agents = agents or [] not_agents = not_agents or [] contains = contains or [] Message = cls query = select(Message) # A little argument validation. We could provide some defaults in # these mixed cases.. but instead we'll just leave it up to our caller. if (start is not None and end is None) or (end is not None and start is None): raise ValueError( "Either both start and end must be specified or neither must be specified" ) if start and end: query = query.where(between(Message.timestamp, start, end)) if msg_id: query = query.where(Message.msg_id == msg_id) # Add the four positive filters as necessary if users: query = query.where(or_(*(Message.users.any(User.name == u) for u in users))) if packages: query = query.where(or_(*(Message.packages.any(Package.name == p) for p in packages))) if categories: query = query.where(or_(*(Message.category == category for category in categories))) if topics: query = query.where(or_(*(Message.topic == topic for topic in topics))) if agents: query = query.where(or_(*(Message.agent_name == agent for agent in agents))) if contains: query = query.where(or_(*(Message.msg.like(f"%{contain}%") for contain in contains))) # And then the four negative filters as necessary if not_users: query = query.where(not_(or_(*(Message.users.any(User.name == u) for u in not_users)))) if not_packs: query = query.where( not_(or_(*(Message.packages.any(Package.name == p) for p in not_packs))) ) if not_cats: query = query.where(not_(or_(*(Message.category == category for category in not_cats)))) if not_topics: query = query.where(not_(or_(*(Message.topic == topic for topic in not_topics)))) if not_agents: query = query.where(not_(or_(*(Message.agent_name == agent for agent in not_agents)))) return query
[docs] @classmethod def grep( cls, *, page=1, rows_per_page=100, order="asc", defer=False, **kwargs, ): """Flexible query interface for messages. Arguments are filters. start and end should be :mod:`datetime` objs. Other filters should be lists of strings. They are applied in a conjunctive-normal-form (CNF) kind of way for example, the following:: users = ['ralph', 'lmacken'] categories = ['bodhi', 'wiki'] should return messages where (user=='ralph' OR user=='lmacken') AND (category=='bodhi' OR category=='wiki') Furthermore, you can use a negative version of each argument. users = ['ralph'] not_categories = ['bodhi', 'wiki'] should return messages where (user == 'ralph') AND NOT (category == 'bodhi' OR category == 'wiki') ---- If the `defer` argument evaluates to True, the query won't actually be executed, but a SQLAlchemy query object returned instead. """ query = cls.make_query(**kwargs) # Finally, tag on our pagination arguments Message = cls query_total = query.with_only_columns(func.count(Message.id)) total = None query = query.order_by(getattr(Message.timestamp, order)()) if not rows_per_page: pages = 1 else: total = session.scalar(query_total) pages = math.ceil(total / float(rows_per_page)) query = query.offset(rows_per_page * (page - 1)).limit(rows_per_page) if defer: if total is None: total = session.scalar(query_total) return total, pages, query else: # Execute! messages = session.scalars(query).all() if pages == 1: total = len(messages) return total, pages, messages
[docs] @classmethod def get_first(cls, *, order="asc", **kwargs): """Get the first message matching the regular grep filters.""" query = cls.make_query(**kwargs) query = query.order_by(getattr(Message.timestamp, order)()).limit(1) return session.scalars(query).first()
[docs] class NamedSingleton: id = Column(Integer, primary_key=True, autoincrement=True) name = Column(UnicodeText, index=True, unique=True)
[docs] @classmethod def get_or_create(cls, name): """ Return the instance of the class with the specified name. If it doesn't already exist, create it. """ # Use an in-memory cache to speed things up. if name in cls._cache: # If we cache the instance, SQLAlchemy will run this query anyway because the instance # will be from a different transaction. So just cache the id. return session.get(cls, cls._cache[name]) obj = session.execute(select(cls).where(cls.name == name)).scalar_one_or_none() if obj is None: obj = cls(name=name) session.add(obj) session.flush() cls._cache[name] = obj.id return obj
[docs] @classmethod def clear_cache(cls): cls._cache.clear()
[docs] class User(DeclarativeBase, NamedSingleton): __tablename__ = "users" _cache = {}
[docs] class Package(DeclarativeBase, NamedSingleton): __tablename__ = "packages" _cache = {}
def _setup_hypertable(table_class): event.listen( table_class.__table__, "after_create", DDL(f"SELECT create_hypertable('{table_class.__tablename__}', 'timestamp');"), ) _setup_hypertable(Message)