# Copyright 2017 The Forseti Security Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Database session handling for Forseti Server."""
from builtins import object
from sqlalchemy.orm import sessionmaker
from google.cloud.forseti.common.util import logger
LOGGER = logger.get_logger(__name__)
[docs]class ScopedSession(object):
"""A scoped session is automatically released."""
def __init__(self, session, auto_commit=False, readonly=False):
"""Constructor.
Args:
session (object): Database session to use scope.
auto_commit (bool): Set to true, of commit should automatically
happen upon close.
readonly (bool): whether or not the session is read only.
"""
self.session = session
self.auto_commit = auto_commit
self.readonly = readonly
[docs] def __enter__(self):
"""To support with statement.
Returns:
object: Returns its session.
"""
if self.readonly:
stub_out_flush_operation(self.session)
return self.session
[docs] def __exit__(self, exc_type, value, traceback):
"""To support with statement.
Args:
exc_type (object): Exception type
value (object): Exception value
traceback (object): Traceback if any
"""
try:
if traceback is None and self.auto_commit:
self.session.commit()
finally:
# session.close() is not a must have operation based on the
# sqlalchemy documentation, it's only useful when you are in
# some intermediate state, if the session is readonly, there
# is no need to close the session.
# http://docs.sqlalchemy.org/en/latest/orm/session_transaction.html
if not self.readonly:
self.session.close()
[docs]class ScopedSessionMaker(object):
"""Wraps session maker to create scoped sessions."""
def __init__(self, session_maker, auto_commit=False):
"""Constructor.
Args:
session_maker (object): Session creator.
auto_commit (bool): If set to true, generated sessions will
automatically commit the transaction upon close
"""
self.sessionmaker = session_maker
self.auto_commit = auto_commit
[docs] def __call__(self, *args):
"""Creates a new session.
Args:
*args (list): Will be forwarded to session creator.
Returns:
object: Scoped session.
"""
return ScopedSession(self.sessionmaker(*args), self.auto_commit)
[docs]def create_scoped_sessionmaker(engine):
"""Creates a scoped session maker
Args:
engine (object): Engine to bind session to.
Returns:
object: Scoped session maker.
"""
return ScopedSessionMaker(
sessionmaker(
bind=engine),
auto_commit=True)
[docs]def _abort_readonly():
"""Intercept the flush operation and log a warning message."""
LOGGER.warning('This session is read-only, no flush is allowed.')
[docs]def stub_out_flush_operation(session):
"""Stub out flush operation.
Args:
session (Session): Session to stub out.
Returns:
Session: The session after stubbed out the flush operation.
"""
session.flush = _abort_readonly
return session
[docs]def create_scoped_readonly_session(engine):
"""Creates a readonly scoped session.
Args:
engine (object): Engine to bind session to.
Returns:
object: Scoped session maker.
"""
# Set session autocommit to true to ensure each query gets a new connection
# from the pool.
session = sessionmaker(bind=engine, autocommit=True, autoflush=False)()
return ScopedSession(session, auto_commit=False, readonly=True)
[docs]def create_readonly_session(engine):
"""Creates a readonly session.
Args:
engine (object): Engine to bind session to.
Returns:
object: Scoped session maker.
"""
# Set session autocommit to true to ensure each query gets a new connection
# from the pool.
session = sessionmaker(bind=engine, autocommit=True, autoflush=False)()
stub_out_flush_operation(session)
return session