src/o2ims/service/unit_of_work.py

   1 # Copyright (C) 2021 Wind River Systems, Inc.
   2 #
   3 #  Licensed under the Apache License, Version 2.0 (the "License");
   4 #  you may not use this file except in compliance with the License.
   5 #  You may obtain a copy of the License at
   6 #
   7 #      http://www.apache.org/licenses/LICENSE-2.0
   8 #
   9 #  Unless required by applicable law or agreed to in writing, software
  10 #  distributed under the License is distributed on an "AS IS" BASIS,
  11 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 #  See the License for the specific language governing permissions and
  13 #  limitations under the License.
  14
  15 # pylint: disable=attribute-defined-outside-init
  16 from __future__ import annotations
  17 import abc
  18 from sqlalchemy import create_engine
  19 from sqlalchemy.orm import sessionmaker
  20 from sqlalchemy.orm.session import Session
  21
  22
  23 from o2ims import config
  24 from o2ims.adapter import ocloud_repository
  25
  26
  27 class AbstractUnitOfWork(abc.ABC):
  28     oclouds: ocloud_repository.OcloudRepository
  29
  30     def __enter__(self) -> AbstractUnitOfWork:
  31         return self
  32
  33     def __exit__(self, *args):
  34         self.rollback()
  35
  36     def commit(self):
  37         self._commit()
  38
  39     def collect_new_events(self):
  40         for ocloud in self.oclouds.seen:
  41             while ocloud.events:
  42                 yield ocloud.events.pop(0)
  43
  44     @abc.abstractmethod
  45     def _commit(self):
  46         raise NotImplementedError
  47
  48     @abc.abstractmethod
  49     def rollback(self):
  50         raise NotImplementedError
  51
  52
  53 DEFAULT_SESSION_FACTORY = sessionmaker(
  54     bind=create_engine(
  55         config.get_postgres_uri(),
  56         isolation_level="REPEATABLE READ",
  57     )
  58 )
  59
  60
  61 class SqlAlchemyUnitOfWork(AbstractUnitOfWork):
  62     def __init__(self, session_factory=DEFAULT_SESSION_FACTORY):
  63         self.session_factory = session_factory
  64
  65     def __enter__(self):
  66         self.session = self.session_factory()  # type: Session
  67         self.oclouds = ocloud_repository.OcloudSqlAlchemyRepository(self.session)
  68         return super().__enter__()
  69
  70     def __exit__(self, *args):
  71         super().__exit__(*args)
  72         self.session.close()
  73
  74     def _commit(self):
  75         self.session.commit()
  76
  77     def rollback(self):
  78         self.session.rollback()