#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#-------------------------------------------------------------------------------
# This file is part of Mentat system (https://mentat.cesnet.cz/).
#
# Copyright (C) since 2011 CESNET, z.s.p.o (http://www.ces.net/)
# Use of this source is governed by the MIT license, see LICENSE file.
#-------------------------------------------------------------------------------
"""
Database storage abstraction layer. The current implementation is based on the
awesome `SQLAlchemy <http://www.sqlalchemy.org/>`__ library.
.. warning::
Current implementation is for optimalization purposes using some advanced
schema features provided by the `PostgreSQL <https://www.postgresql.org/>`__
database and thus no other engines are currently supported.
"""
__author__ = "Jan Mach <jan.mach@cesnet.cz>"
__credits__ = "Pavel Kácha <pavel.kacha@cesnet.cz>, Andrea Kropáčová <andrea.kropacova@cesnet.cz>"
import copy
import sqlalchemy
#
# Custom libraries
#
from mentat.const import CKEY_CORE_DATABASE, CKEY_CORE_DATABASE_SQLSTORAGE
from mentat.datatype.sqldb import MODEL
_MANAGER = None
[docs]class RetryingQuery(sqlalchemy.orm.Query):
"""
An override of SQLAlchemy's Query class, allowing for recovery from a lost DB
connection.
"""
def _execute_and_instances(self, querycontext):
for _ in range(2):
try:
return super()._execute_and_instances(querycontext)
except sqlalchemy.exc.OperationalError:
self.session.close()
continue
[docs]class StorageService:
"""
Proxy object for working with persistent SQL storages. Maintains and provides
access to database engine, session maker and session proxies.
"""
def __init__(self, **enginecfg):
"""
Open and cache connection to SQL storage. The connection arguments for
database engine are passed directly to :py:func:`sqlalchemy.engine_from_config`
method. The appropriate prefix is the empty string.
:param enginecfg: Connection arguments.
"""
self.dbengine = sqlalchemy.engine_from_config(enginecfg, prefix = '')
self.sessionmaker = sqlalchemy.orm.scoped_session(sqlalchemy.orm.sessionmaker(bind = self.dbengine))
self._session = self.sessionmaker(query_cls = RetryingQuery)
def __del__(self):
self.close()
[docs] def close(self):
"""
Close current database connection.
"""
self.session_close()
self.dbengine = None
@property
def session(self):
"""
"""
if not self._session:
self._session = self.sessionmaker(query_cls = RetryingQuery)
return self._session
[docs] def session_close(self):
"""
"""
if self._session:
self.sessionmaker.remove()
self._session = None
[docs] def session_new(self):
"""
Close existing session and create new fresh database session.
"""
self.session_close()
return self.session
[docs] def database_create(self):
"""
Create database SQL schema.
"""
MODEL.metadata.create_all(self.dbengine)
[docs] def database_drop(self):
"""
Drop database SQL schema.
"""
MODEL.metadata.drop_all(self.dbengine)
[docs]class StorageServiceManager:
"""
Class representing a custom _StorageServiceManager_ capable of understanding
and parsing Mentat system core configurations and enabling easy access to
preconfigured database collections via indirrect handles (identifiers).
"""
def __init__(self, core_config, updates = None):
"""
Initialize a _StorageServiceManager_ proxy object with full core configuration
tree structure.
:param dict core_config: Mentat core configuration structure.
:param dict updates: Optional configuration updates (same structure as ``core_config``).
"""
self._dbconfig = {}
self._storage = None
self._configure_dbconfig(core_config, updates)
def _configure_dbconfig(self, core_config, updates):
"""
Internal sub-initialization helper: Configure database structure parameters
and optionally merge them with additional updates.
:param dict core_config: Mentat core configuration structure.
:param dict updates: Optional configuration updates (same structure as ``core_config``).
"""
self._dbconfig = copy.deepcopy(core_config[CKEY_CORE_DATABASE][CKEY_CORE_DATABASE_SQLSTORAGE])
if updates and CKEY_CORE_DATABASE in updates and CKEY_CORE_DATABASE_SQLSTORAGE in updates[CKEY_CORE_DATABASE]:
self._dbconfig.update(
updates[CKEY_CORE_DATABASE][CKEY_CORE_DATABASE_SQLSTORAGE]
)
#---------------------------------------------------------------------------
[docs] def close(self):
"""
Close internal storage connection.
"""
if self._storage:
self._storage.close()
self._storage = None
[docs] def service(self):
"""
Return handle to storage connection service according to internal configurations.
:return: Reference to storage service.
:rtype: mentat.services.sqlstorage.StorageService
"""
if not self._storage:
self._storage = StorageService(**self._dbconfig)
return self._storage
#-------------------------------------------------------------------------------
[docs]def init(core_config, updates = None):
"""
(Re-)Initialize :py:class:`mentat.services.sqlstorage.StorageServiceManager`
instance at module level and store the refence within module.
:param dict core_config: Mentat core configuration structure.
:param dict updates: Optional configuration updates (same structure as ``core_config``).
"""
global _MANAGER # pylint: disable=locally-disabled,global-statement
_MANAGER = StorageServiceManager(core_config, updates)
[docs]def set_manager(new_manager):
"""
Set manager from outside the module. This should be used only when you know
exactly what you are doing.
"""
global _MANAGER # pylint: disable=locally-disabled,global-statement
_MANAGER = new_manager
[docs]def manager():
"""
Obtain reference to :py:class:`mentat.services.sqlstorage.StorageServiceManager`
instance stored at module level.
:return: Storage service manager reference.
:rtype: mentat.services.sqlstorage.StorageServiceManager
"""
return _MANAGER
[docs]def service():
"""
Obtain reference to :py:class:`mentat.services.sqlstorage.StorageService`
instance from module level manager.
:return: Storage service reference.
:rtype: mentat.services.sqlstorage.StorageServiceManager
"""
return manager().service()
[docs]def close():
"""
Close database connection on :py:class:`mentat.services.sqlstorage.StorageService`
instance from module level manager.
"""
return manager().close()