@@ -160,7 +160,10 @@ def get_sql_alchemy_url(
160160 uri_opts = uri_opts ,
161161 )
162162 if self .use_odbc :
163- uri = f"{ uri } ?{ urllib .parse .urlencode (self .uri_args )} "
163+ uri_args = ({k :v for k ,v in self .uri_args .items () if k .lower () != 'database' }
164+ if current_db else
165+ self .uri_args )
166+ uri = f"{ uri } ?{ urllib .parse .urlencode (uri_args )} "
164167 return uri
165168
166169 @property
@@ -939,6 +942,24 @@ def construct_flow_workunits(
939942 aspect = data_flow .as_container_aspect ,
940943 ).as_workunit ()
941944
945+ def _database_names_from_engine (self , engine ):
946+ """
947+ Helper method to get database names from the engine.
948+ This is used to fetch the list of databases in the SQL Server instance.
949+ """
950+ with engine .begin () as conn :
951+ databases = conn .execute (
952+ "SELECT name FROM master.sys.databases WHERE name NOT IN \
953+ ('master', 'model', 'msdb', 'tempdb', 'Resource', \
954+ 'distribution' , 'reportserver', 'reportservertempdb'); "
955+ ).fetchall ()
956+ return [db ["name" ] for db in databases ]
957+
958+ def _inspector_for_database (self , db_name : str ) -> Inspector :
959+ url = self .config .get_sql_alchemy_url (current_db = db_name )
960+ engine = create_engine (url , ** self .config .options )
961+ return inspect (engine )
962+
942963 def get_inspectors (self ) -> Iterable [Inspector ]:
943964 # This method can be overridden in the case that you want to dynamically
944965 # run on multiple databases.
@@ -950,19 +971,10 @@ def get_inspectors(self) -> Iterable[Inspector]:
950971 inspector = inspect (engine )
951972 yield inspector
952973 else :
953- with engine .begin () as conn :
954- databases = conn .execute (
955- "SELECT name FROM master.sys.databases WHERE name NOT IN \
956- ('master', 'model', 'msdb', 'tempdb', 'Resource', \
957- 'distribution' , 'reportserver', 'reportservertempdb'); "
958- ).fetchall ()
959-
960- for db in databases :
961- if self .config .database_pattern .allowed (db ["name" ]):
962- url = self .config .get_sql_alchemy_url (current_db = db ["name" ])
963- engine = create_engine (url , ** self .config .options )
964- inspector = inspect (engine )
965- self .current_database = db ["name" ]
974+ for db_name in self ._database_names_from_engine (engine ):
975+ if self .config .database_pattern .allowed (db_name ):
976+ inspector = self ._inspector_for_database (db_name )
977+ self .current_database = db_name
966978 yield inspector
967979
968980 def get_identifier (
0 commit comments