Add tox
[pti/o2.git] / 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 from o2ims import config
23 from o2ims.adapter import ocloud_repository
24
25
26 class AbstractUnitOfWork(abc.ABC):
27     oclouds: ocloud_repository.OcloudRepository
28
29     def __enter__(self):
30         return self
31
32     def __exit__(self, *args):
33         self.rollback()
34
35     def commit(self):
36         self._commit()
37
38     def collect_new_events(self):
39         for ocloud in self.oclouds.seen:
40             while ocloud.events:
41                 yield ocloud.events.pop(0)
42
43     @abc.abstractmethod
44     def _commit(self):
45         raise NotImplementedError
46
47     @abc.abstractmethod
48     def rollback(self):
49         raise NotImplementedError
50
51
52 DEFAULT_SESSION_FACTORY = sessionmaker(
53     bind=create_engine(
54         config.get_postgres_uri(),
55         isolation_level="REPEATABLE READ",
56     )
57 )
58
59
60 class SqlAlchemyUnitOfWork(AbstractUnitOfWork):
61     def __init__(self, session_factory=DEFAULT_SESSION_FACTORY):
62         self.session_factory = session_factory
63
64     def __enter__(self):
65         self.session = self.session_factory()  # type: Session
66         self.oclouds = ocloud_repository\
67             .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()