You can not select more than 25 topics
			Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
		
		
		
		
		
			
		
			
				
					117 lines
				
				4.2 KiB
			
		
		
			
		
	
	
					117 lines
				
				4.2 KiB
			| 
								 
											3 years ago
										 
									 | 
							
								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"]
							 | 
						||
| 
								 | 
							
								                            )
							 | 
						||
| 
								 | 
							
								                        )
							 | 
						||
| 
								 | 
							
								                    )
							 |