117 lines
4.2 KiB
Python
117 lines
4.2 KiB
Python
|
from sqlalchemy import inspect
|
||
|
|
from sqlalchemy import Integer
|
||
|
|
from ... import create_engine
|
||
|
|
from ... import exc
|
||
|
|
from ...schema import Column
|
||
|
|
from ...schema import DropConstraint
|
||
|
|
from ...schema import ForeignKeyConstraint
|
||
|
|
from ...schema import MetaData
|
||
|
|
from ...schema import Table
|
||
|
|
from ...testing.provision import create_db
|
||
|
|
from ...testing.provision import drop_all_schema_objects_pre_tables
|
||
|
|
from ...testing.provision import drop_db
|
||
|
|
from ...testing.provision import get_temp_table_name
|
||
|
|
from ...testing.provision import log
|
||
|
|
from ...testing.provision import run_reap_dbs
|
||
|
|
from ...testing.provision import temp_table_keyword_args
|
||
|
|
|
||
|
|
|
||
|
|
@create_db.for_db("mssql")
|
||
|
|
def _mssql_create_db(cfg, eng, ident):
|
||
|
|
with eng.connect().execution_options(isolation_level="AUTOCOMMIT") as conn:
|
||
|
|
conn.exec_driver_sql("create database %s" % ident)
|
||
|
|
conn.exec_driver_sql(
|
||
|
|
"ALTER DATABASE %s SET ALLOW_SNAPSHOT_ISOLATION ON" % ident
|
||
|
|
)
|
||
|
|
conn.exec_driver_sql(
|
||
|
|
"ALTER DATABASE %s SET READ_COMMITTED_SNAPSHOT ON" % ident
|
||
|
|
)
|
||
|
|
conn.exec_driver_sql("use %s" % ident)
|
||
|
|
conn.exec_driver_sql("create schema test_schema")
|
||
|
|
conn.exec_driver_sql("create schema test_schema_2")
|
||
|
|
|
||
|
|
|
||
|
|
@drop_db.for_db("mssql")
|
||
|
|
def _mssql_drop_db(cfg, eng, ident):
|
||
|
|
with eng.connect().execution_options(isolation_level="AUTOCOMMIT") as conn:
|
||
|
|
_mssql_drop_ignore(conn, ident)
|
||
|
|
|
||
|
|
|
||
|
|
def _mssql_drop_ignore(conn, ident):
|
||
|
|
try:
|
||
|
|
# typically when this happens, we can't KILL the session anyway,
|
||
|
|
# so let the cleanup process drop the DBs
|
||
|
|
# for row in conn.exec_driver_sql(
|
||
|
|
# "select session_id from sys.dm_exec_sessions "
|
||
|
|
# "where database_id=db_id('%s')" % ident):
|
||
|
|
# log.info("killing SQL server session %s", row['session_id'])
|
||
|
|
# conn.exec_driver_sql("kill %s" % row['session_id'])
|
||
|
|
conn.exec_driver_sql("drop database %s" % ident)
|
||
|
|
log.info("Reaped db: %s", ident)
|
||
|
|
return True
|
||
|
|
except exc.DatabaseError as err:
|
||
|
|
log.warning("couldn't drop db: %s", err)
|
||
|
|
return False
|
||
|
|
|
||
|
|
|
||
|
|
@run_reap_dbs.for_db("mssql")
|
||
|
|
def _reap_mssql_dbs(url, idents):
|
||
|
|
log.info("db reaper connecting to %r", url)
|
||
|
|
eng = create_engine(url)
|
||
|
|
with eng.connect().execution_options(isolation_level="AUTOCOMMIT") as conn:
|
||
|
|
|
||
|
|
log.info("identifiers in file: %s", ", ".join(idents))
|
||
|
|
|
||
|
|
to_reap = conn.exec_driver_sql(
|
||
|
|
"select d.name from sys.databases as d where name "
|
||
|
|
"like 'TEST_%' and not exists (select session_id "
|
||
|
|
"from sys.dm_exec_sessions "
|
||
|
|
"where database_id=d.database_id)"
|
||
|
|
)
|
||
|
|
all_names = {dbname.lower() for (dbname,) in to_reap}
|
||
|
|
to_drop = set()
|
||
|
|
for name in all_names:
|
||
|
|
if name in idents:
|
||
|
|
to_drop.add(name)
|
||
|
|
|
||
|
|
dropped = total = 0
|
||
|
|
for total, dbname in enumerate(to_drop, 1):
|
||
|
|
if _mssql_drop_ignore(conn, dbname):
|
||
|
|
dropped += 1
|
||
|
|
log.info(
|
||
|
|
"Dropped %d out of %d stale databases detected", dropped, total
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@temp_table_keyword_args.for_db("mssql")
|
||
|
|
def _mssql_temp_table_keyword_args(cfg, eng):
|
||
|
|
return {}
|
||
|
|
|
||
|
|
|
||
|
|
@get_temp_table_name.for_db("mssql")
|
||
|
|
def _mssql_get_temp_table_name(cfg, eng, base_name):
|
||
|
|
return "##" + base_name
|
||
|
|
|
||
|
|
|
||
|
|
@drop_all_schema_objects_pre_tables.for_db("mssql")
|
||
|
|
def drop_all_schema_objects_pre_tables(cfg, eng):
|
||
|
|
with eng.connect().execution_options(isolation_level="AUTOCOMMIT") as conn:
|
||
|
|
inspector = inspect(conn)
|
||
|
|
for schema in (None, "dbo", cfg.test_schema, cfg.test_schema_2):
|
||
|
|
for tname in inspector.get_table_names(schema=schema):
|
||
|
|
tb = Table(
|
||
|
|
tname,
|
||
|
|
MetaData(),
|
||
|
|
Column("x", Integer),
|
||
|
|
Column("y", Integer),
|
||
|
|
schema=schema,
|
||
|
|
)
|
||
|
|
for fk in inspect(conn).get_foreign_keys(tname, schema=schema):
|
||
|
|
conn.execute(
|
||
|
|
DropConstraint(
|
||
|
|
ForeignKeyConstraint(
|
||
|
|
[tb.c.x], [tb.c.y], name=fk["name"]
|
||
|
|
)
|
||
|
|
)
|
||
|
|
)
|