Source code for thoth.storages.graph.models_base

#!/usr/bin/env python3
# thoth-storages
# Copyright(C) 2019, 2020 Fridolin Pokorny
#
# 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/>.

"""A base and utilities for implementing SQLAlchemy based models."""

import logging
import datetime
from itertools import combinations
from typing import List

from sqlalchemy import Index
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import class_mapper
from sqlalchemy.orm import ColumnProperty
from sqlalchemy.exc import IntegrityError


Base = declarative_base()


_LOGGER = logging.getLogger(__name__)


[docs]class BaseExtension: """Extend base class with additional functionality."""
[docs] @classmethod def get_or_create(cls, session, **kwargs): """Query for the given entity, create if it does not exist yet.""" instance = session.query(cls).filter_by(**kwargs).first() if instance: return instance, True else: try: session.begin_nested() instance = cls(**kwargs) session.add(instance) session.commit() return instance, False except IntegrityError as exc: session.rollback() _LOGGER.warning( "Integrity error on creating a new record; this can be due to " "concurrent writes to database, recovering (attributes: %r): %s", kwargs, str(exc), ) return session.query(cls).filter_by(**kwargs).one(), True
[docs] @classmethod def attribute_names(cls): """Get names of attributes for the given model declaration.""" return [prop.key for prop in class_mapper(cls).iterate_properties if isinstance(prop, ColumnProperty)]
[docs] def to_dict(self, without_id: bool = True) -> dict: """Convert model to a dictionary representation keeping just rows as attributes.""" result = {} for column in self.__table__.columns: if without_id and column.name == "id": continue if getattr(self, column.name) is None: result[column.name] = None else: value = getattr(self, column.name) if isinstance(value, (datetime.datetime, datetime.date)): result[column.name] = str(value) else: result[column.name] = value return result
[docs]def get_python_package_version_index_combinations() -> List[Index]: """Create index for all possible combinations which we can query.""" result = [] _columns_to_variate = ( "os_name", "os_version", "python_version", ) for i in range(0, len(_columns_to_variate) + 1): for j, variation in enumerate(combinations(_columns_to_variate, i)): result.append( Index( f"python_package_version_index_idx_{i}{j}", "package_name", "package_version", *variation, ) ) return result