def test_get_multi_by_instance_index(
db, redis_client, redis_flushall, faker, user, application):
cluster_name = 'bar'
key = 'test'
action_type = action_types.UPDATE_CONFIG
_, data_type = action_types[action_type].split('_', 1)
data_type = data_type.lower()
audit_num = 3
for _ in range(audit_num):
extra = {
'application_name': application.application_name,
'cluster_name': 'bar',
'key': 'test',
'old_data': 'old',
'new_data': 'new'
}
action = action_creator.make_action(
action_types.UPDATE_CONFIG, **extra)
AuditLog.create(user.id, faker.ipv4(), action)
audit_logs = AuditLog.get_multi_by_instance_index(
AuditLog.TYPE_CONFIG, application.id, cluster_name, key)
assert len(audit_logs[:]) == audit_num
with DBSession().close_on_exit(False):
DBSession.delete(user)
User._db_session.close()
redis_flushall(redis_client)
audit_logs = AuditLog.get_multi_by_instance_index(
AuditLog.TYPE_CONFIG, application.id, cluster_name, key)
assert not any(getattr(x, 'user') for x in audit_logs)
def test_cache_mixin(mocker):
assert 'team|0' in repr(Team(id=0, team_name="hello"))
assert Team(id=0, team_name='hello').pk_name() == 'id'
assert Team(id=0, team_name='hello').pk_attribute().key == 'id'
t = Team(team_name='test_cache_mixin')
DBSession.add(t)
DBSession.commit()
assert len(Team.mget([t.id], force=True)) == 1
with pytest.raises(NotImplementedError):
CacheMixinBase._cache_client.error
with pytest.raises(NotImplementedError):
CacheMixinBase._db_session.error
with mocker.patch.object(CacheMixinBase, 'RAWDATA_VERSION',
mocker.MagicMock(__str__=lambda x: '233')):
assert Team.gen_raw_key(1) == 'team|1|233'
with mocker.patch.object(Team.__mapper__, 'primary_key',
mocker.MagicMock(return_value=[],
__nonzero__=mocker.MagicMock(
return_value=False))):
assert Team(id=2, team_name='hello').pk_name() is None
assert Team(id=3, team_name='hello').pk_attribute() is None
def get_audit_ids_by_date(cls, target_type, target_id, date):
db = DBSession()
rs = db.query(cls.audit_id).filter(
cls.target_id == target_id,
cls.target_type == target_type,
cast(cls.created_at, DATE) == date
)
return sorted((r[0] for r in rs), reverse=True)
def close_session(response):
try:
DBSession.remove()
except Exception:
logger.exception('Failed to close database during request teardown')
sentry.captureException()
if settings.TESTING:
raise
return response
def get_audit_ids(cls, instance_type, application_id, cluster_name,
instance_key):
"""Get ids of :class:`.audit.AuditLog` by instance information.
Regardless of cluster if the ``cluster_name`` is ``None``
"""
db = DBSession()
rs = db.query(cls.audit_id).filter_by(
instance_type=instance_type, application_id=application_id,
cluster_name=cluster_name, instance_key=instance_key)
return sorted((r[0] for r in rs), reverse=True)
def broken_db(db):
broken_engine = create_engine(
'mysql+pymysql://[email protected]:1/dotdotdot?charset=utf8')
broken_engines = {'master': broken_engine, 'slave': broken_engine}
healthy_engines = dict(db.engines)
try:
DBSession.registry.registry.clear()
DBSession.configure(engines=broken_engines)
yield DBSession()
finally:
DBSession.registry.registry.clear()
DBSession.configure(engines=healthy_engines)
def db(redis_client):
def truncate_tables(session):
for table in reversed(DeclarativeBase.metadata.sorted_tables):
if table.name.startswith('alembic'):
continue
session.execute(table.delete())
session.commit()
try:
with contextlib.closing(DBSession()) as session:
truncate_tables(session)
yield session
finally:
with contextlib.closing(DBSession()) as session:
truncate_tables(session)
def test_comfirm_close_when_exception(mocker):
with mocker.patch.object(
Session, 'commit', mocker.MagicMock(
side_effect=[gevent.Timeout(), ])):
session = DBSession()
session.execute("select 1")
session.flush()
with pytest.raises(gevent.Timeout):
session.commit()
for engine in session.engines.itervalues():
for parent in session.transaction._iterate_parents():
conn = parent._connections.get(engine)
if conn:
assert conn[0].invalidated
session.close()
def test_with_statement_closing(team_point):
def get_team_point():
return DBSession().query(team_point.__class__).get(team_point.id)
assert DBSession().identity_map
with DBSession().close_on_exit(False):
pass
assert DBSession().identity_map
assert get_team_point() is team_point
with DBSession():
pass
assert not DBSession().identity_map
assert get_team_point() is not team_point
def test_admin_user(admin_user):
'''
check the existence of admin user
'''
assert DBSession().query(User) \
.filter_by(username=admin_user.username) \
.first()
def _set_status(self, status):
with DBSession().close_on_exit(False):
self.status = status
self.flush([self.id])
self.get_all_ids.flush()
self.get_ids_by_team.flush(self.team.id)
self.get_id_by_name.flush(self.application_name)
def create_normal(cls, username, password, email=None, is_active=False):
"""Creates a normal user.
:param username: The unique name of user. It can not be the same to
existed users and applications.
:param password: The plain text password of user.
:param email: The optional email address of user.
:param is_active: Only ``True`` is available for now.
:returns User: The new created user.
"""
try:
with DBSession().close_on_exit(False) as db:
instance = cls(username=username,
password=_hash_password(password),
email=email,
is_active=int(bool(is_active)),
huskar_admin=False,
is_app=False)
db.add(instance)
except IntegrityError:
raise NameOccupiedError
cls.flush([instance.id])
cls.get_id_by_name.flush(username)
cls.get_ids_of_normal.flush()
cls.get_id_by_email.flush(email)
return instance
def create(cls, url, hook_type=TYPE_NORMAL):
with DBSession().close_on_exit(False) as db:
instance = cls(url=url, hook_type=hook_type)
db.add(instance)
cls.flush([instance.id])
cls.get_all_ids.flush()
cls.get_ids_by_type.flush(hook_type)
return instance
def create_application(cls, application_name):
stmt = cls.upsert().values(username=application_name,
password=_hash_password(application_name),
email=None,
is_active=True,
is_app=True)
with DBSession().close_on_exit(False) as db:
rs = db.execute(stmt)
instance = cls.get(rs.lastrowid)
DBSession().refresh(instance)
cls.get_id_by_name.flush(application_name)
cls.get_ids_of_normal.flush()
cls.flush([instance.id])
return instance
def get_ids(cls, application_id, webhook_id, action_type):
conds = {name: value for name, value in [
('application_id', application_id),
('webhook_id', webhook_id),
('action_type', action_type),
] if value is not None}
rs = DBSession().query(cls.id).filter_by(**conds).all()
return [r[0] for r in rs]
def _set_active_status(self, is_active):
"""Set value of `is_active` column to
disable or active this user.
"""
with DBSession().close_on_exit(False):
self.is_active = is_active
self.__class__.get_id_by_name.flush(self.username)
self.__class__.get_ids_of_normal.flush()
self.__class__.flush([self.id])
def delete(cls, team_id):
team = cls.get(team_id)
with DBSession().close_on_exit(False) as db:
team_will_be_deleted.send(cls, db=db, team_id=team_id)
for user_id in TeamAdmin.get_user_ids(team_id):
TeamAdmin.discard(team_id, user_id)
db.query(cls).filter_by(id=team_id).delete()
if team is not None: # can be skip safety (instance will be erased)
cls.get_id_by_name.flush(team.team_name)
cls.get_all_ids.flush()
cls.flush([team_id])
def create(cls, application_id, webhook_id, action_type):
with DBSession().close_on_exit(False) as db:
instance = cls(webhook_id=webhook_id,
application_id=application_id,
action_type=action_type)
db.add(instance)
cls.flush([cls.id])
cls.get_id.flush(application_id, webhook_id, action_type)
cls.get_ids.flush(application_id, webhook_id, action_type)
cls.get_ids.flush(application_id, webhook_id, None)
cls.get_ids.flush(application_id, None, None)
return instance
def delete(self):
with DBSession().close_on_exit(False) as db:
db.delete(self)
self.flush([self.id])
self.get_id.flush(
self.application_id, self.webhook_id, self.action_type)
self.get_ids.flush(
self.application_id, self.webhook_id, self.action_type)
self.get_ids.flush(
self.application_id, self.webhook_id, None)
self.get_ids.flush(
self.application_id, None, None)
def test_close_sessions_extra():
session = DBManager()
session.close_sessions(should_close_connection=False)
session.close_sessions(should_close_connection=True)
db_manager = DBManager()
DBSession = db_manager.get_session('default')
class ThisExc(Exception):
pass
with pytest.raises(ThisExc):
with DBSession() as session:
session.execute('select 1')
_invalidate_connections(session)
raise ThisExc('boom!')
session = DBSession()
assert session.transaction.is_active
session.transaction = None
db_manager.close_sessions(should_close_connection=True)
def main():
db = DBSession()
linters = [
check_users(db),
]
for code, info in itertools.chain.from_iterable(linters):
desc = CODES[code]
print(u'%s\t%s\t%s' % (code, desc, info), file=sys.stderr)
else:
sys.exit(0)
sys.exit(1)
def delete(cls, application_id):
application_id = int(application_id)
with DBSession().close_on_exit(False) as db:
application = cls.get(application_id)
auth_set = ApplicationAuth.search_by(application_id=application_id)
for auth in auth_set:
db.delete(auth)
db.flush()
db.delete(application)
cls.get_id_by_name.flush(application.application_name)
cls.get_ids_by_team.flush(application.team_id)
cls.get_all_ids.flush()
cls.flush([application_id])
def discard(cls, authority, user_id, application_id):
"""Deletes an application auth if it does exist.
:param authority: ``AUTHORITY_READ`` or ``AUTHORITY_WRITE``
:param user_id: The id of owner.
:param application_id: The id of application.
"""
assert authority in Authority
stmt = cls.__table__.delete().where(
(cls.authority == authority.value) & (cls.user_id == user_id)
& (cls.application_id == application_id))
with DBSession().close_on_exit(False) as db:
db.execute(stmt)
cls.flush_by(authority, user_id, application_id)
def test_with_statement_on_exc():
class ThisExc(Exception):
pass
trans_orig = DBSession().transaction
with pytest.raises(ThisExc):
with DBSession() as session:
session.execute('select 1')
_invalidate_connections(session)
raise ThisExc('boom!')
session = DBSession()
assert session.transaction.is_active
assert session.transaction is not trans_orig
session.commit()
session.close()
def ensure(cls, authority, user_id, application_id):
"""Creates an application auth if it does not exist.
:param authority: ``AUTHORITY_READ`` or ``AUTHORITY_WRITE``
:param user_id: The id of owner.
:param application_id: The id of application.
"""
assert authority in Authority
stmt = cls.upsert().values(authority=authority.value,
user_id=user_id,
application_id=application_id)
with DBSession().close_on_exit(False) as db:
db.execute(stmt)
cls.flush_by(authority, user_id, application_id)
def create(cls, application_name, team_id):
try:
with DBSession().close_on_exit(False) as db:
cls.check_default_user(application_name)
instance = cls(application_name=application_name,
team_id=team_id)
db.add(instance)
except IntegrityError:
raise NameOccupiedError
cls.get_id_by_name.flush(application_name)
cls.get_ids_by_team.flush(team_id)
cls.get_all_ids.flush()
instance.setup_default_auth()
instance.setup_default_zpath()
return instance
def team_point():
team_point = Team(team_name='233')
try:
DBSession().query(Team).delete()
DBSession().commit()
DBSession().add(team_point)
DBSession().commit()
yield team_point
finally:
DBSession().query(Team).delete()
DBSession().commit()
def search_ids_by(cls, authority, user_id, application_id):
if authority is not None:
assert authority in Authority
authority = authority.value
conditions = {
name: value
for name, value in [
('authority', authority),
('user_id', user_id),
('application_id', application_id),
] if value is not None
}
rs = DBSession().query(cls.id).filter_by(**conditions).order_by(
cls.id.desc())
return [r[0] for r in rs]
def create(cls, name, desc=None):
"""Creates a new team.
:param name: The unique name of team.
:param desc: The readable name of team.
:returns: The instance of :class:`Team`.
"""
try:
with DBSession().close_on_exit(False) as db:
instance = cls(team_name=name, team_desc=desc or name)
db.add(instance)
except IntegrityError:
raise NameOccupiedError
cls.flush([instance.id])
cls.get_id_by_name.flush(instance.team_name)
cls.get_all_ids.flush()
return instance
def create(cls, user_id, remote_addr, action, rollback_to=None):
"""Creates a new audit log.
:param int user_id: The id of operator.
:param str remote_addr: The IP address of operator.
:param tuple action: Use :meth:`action_creator.make_action` to create
this tuple.
:param int rollback_to: Optional. Default is :obj:`None`. If this is
a rollback operation, pass id of the audit log
want to rollback.
:returns: The created instance.
"""
if rollback_to is not None and cls.get(rollback_to) is None:
raise ValueError('rollback_to is not a valid id')
action_type, action_data, action_indices = action
trace_all_application_events(action_type, action_data)
action_data = json.dumps(action_data,
sort_keys=True,
ensure_ascii=False)
if len(action_data) >= cls.MAX_AUDIT_LENGTH:
_publish_new_action(user_id, remote_addr, action)
raise AuditLogTooLongError('The audit log is too long.')
instance = cls(user_id=user_id,
remote_addr=remote_addr,
action_type=action_type,
action_data=action_data,
rollback_id=rollback_to)
try:
with DBSession().close_on_exit(False) as db:
db.add(instance)
db.flush()
for args in action_indices:
create_index(db, instance.id, instance.created_at, args)
except SQLAlchemyError:
_publish_new_action(user_id, remote_addr, action)
raise AuditLogLostError()
date = instance.created_at.date()
for index_args in action_indices:
flush_index_cache(date, index_args)
_publish_new_action(user_id, remote_addr, action)
return instance