From 6787602820b000bb92803761dad24879ac171fc2 Mon Sep 17 00:00:00 2001 From: mikeqfu Date: Mon, 9 Jan 2023 11:55:16 +0000 Subject: [PATCH] Modify dbms.py for SQLAlchemy 2.0 --- pyhelpers/dbms.py | 75 +++++++++++++++++++++++------------------------ 1 file changed, 37 insertions(+), 38 deletions(-) diff --git a/pyhelpers/dbms.py b/pyhelpers/dbms.py index 4ae5e18..1a6fa8e 100644 --- a/pyhelpers/dbms.py +++ b/pyhelpers/dbms.py @@ -524,10 +524,10 @@ def get_database_names(self, names_only=True): 'postgres'] """ - conn = self.engine.connect(close_with_result=True) - rslt = conn.execute('SELECT datname FROM pg_database;') + with self.engine.connect() as connection: + result = connection.execute('SELECT datname FROM pg_database;') - db_names = rslt.fetchall() + db_names = result.fetchall() if names_only: database_names = list(itertools.chain(*db_names)) else: @@ -574,10 +574,11 @@ def database_exists(self, database_name=None): db_name = self.database_name if database_name is None else str(database_name) - conn = self.engine.connect(close_with_result=True) - rslt = conn.execute( - f"SELECT EXISTS(SELECT datname FROM pg_catalog.pg_database WHERE datname='{db_name}');") - db_exists = bool(rslt.fetchone()[0]) + with self.engine.connect() as connection: + result = connection.execute( + f"SELECT EXISTS(SELECT datname FROM pg_catalog.pg_database WHERE datname='{db_name}');") + + db_exists = bool(result.fetchone()[0]) return db_exists @@ -1066,15 +1067,14 @@ def get_schema_names(self, include_all=False, names_only=True, verbose=False): "AND nspname NOT LIKE 'pg_toast%%' " \ "AND nspname NOT LIKE 'pg_temp%%' " - query = f"SELECT s.nspname, s.oid, u.usename FROM pg_catalog.pg_namespace s " \ - f"JOIN pg_catalog.pg_user u ON u.usesysid = s.nspowner" \ - f"{condition}" \ - f"ORDER BY s.nspname;" - - conn = self.engine.connect(close_with_result=True) - rslt = conn.execute(query) + with self.engine.connect() as connection: + query = f"SELECT s.nspname, s.oid, u.usename FROM pg_catalog.pg_namespace s " \ + f"JOIN pg_catalog.pg_user u ON u.usesysid = s.nspowner" \ + f"{condition}" \ + f"ORDER BY s.nspname;" + result = connection.execute(query) - schema_info_ = rslt.fetchall() + schema_info_ = result.fetchall() if not schema_info_: schema_info = None if verbose: @@ -1226,12 +1226,12 @@ def table_exists(self, table_name, schema_name=None): schema_name_ = self._schema_name(schema_name=schema_name) - rslt = self.engine.execute( + result_ = self.engine.execute( f"SELECT EXISTS(" f"SELECT * FROM information_schema.tables " f"WHERE table_schema='{schema_name_}' AND table_name='{table_name}');") - result = rslt.fetchone()[0] + result = result_.fetchone()[0] return result @@ -1492,13 +1492,13 @@ def get_table_names(self, schema_name=None, verbose=False): print("The schema \"{}\" does not exist.".format(schema_name_)) else: - rslt = self.engine.execute( + result_ = self.engine.execute( "SELECT table_name FROM information_schema.tables " "WHERE table_schema='{}' AND table_type='BASE TABLE';".format(schema_name_)) - temp_list = rslt.fetchall() + result = result_.fetchall() - table_list = [x[0] for x in temp_list] + table_list = [x[0] for x in result] return table_list @@ -2336,13 +2336,12 @@ def __init__(self, host=None, port=None, username=None, password=None, database_ try: if reconnect_db: self.engine = sqlalchemy.create_engine(url=self.engine.url, isolation_level='AUTOCOMMIT') - test_conn = self.engine.connect() - test_conn.close() - del test_conn + with self.engine.connect() as test_conn: + test_conn.close() if verbose: print("Successfully.") except Exception as e: - print("Failed. {}".format(e)) + print(f"Failed. {e}.") def _database_name(self, database_name=None): """ @@ -2573,7 +2572,7 @@ def create_connection(self, database_name=None, close_with_result=False, mode=No >>> db_conn.closed True - >>> db_conn = mssql.create_connection(close_with_result=True) + >>> db_conn = mssql.create_connection() >>> db_conn.should_close_with_result True >>> db_conn.closed @@ -2661,10 +2660,10 @@ def get_database_names(self, names_only=True): 'msdb'] """ - conn = self.engine.connect(close_with_result=True) - rslt = conn.execute('SELECT name, database_id, create_date FROM sys.databases;') + with self.engine.connect() as connection: + result = connection.execute('SELECT name, database_id, create_date FROM sys.databases;') - db_names = rslt.fetchall() + db_names = result.fetchall() if names_only: database_names = [x[0] for x in db_names] else: @@ -2712,15 +2711,15 @@ def database_exists(self, database_name=None): # noinspection PyBroadException try: - rslt = self.engine.execute( + result_ = self.engine.execute( f"IF (EXISTS (SELECT name FROM master.sys.databases WHERE name='{db_name}')) " f"SELECT 1 ELSE SELECT 0") except Exception: - rslt = self.engine.execute( + result_ = self.engine.execute( f"SELECT COUNT(*) FROM master.sys.databases " f"WHERE '[' + name + ']' = '{db_name}' OR name = '{db_name}';") - result = bool(rslt.fetchone()[0]) + result = bool(result_.fetchone()[0]) return result @@ -2968,11 +2967,11 @@ def schema_exists(self, schema_name): schema_name_ = self._schema_name(schema_name=schema_name) - rslt = self.engine.execute( + result_ = self.engine.execute( f"IF EXISTS (SELECT name FROM sys.schemas WHERE name='{schema_name_}') " f"SELECT 1 ELSE SELECT 0") - result = bool(rslt.fetchone()[0]) + result = bool(result_.fetchone()[0]) return result @@ -3165,13 +3164,13 @@ def table_exists(self, table_name, schema_name=None): schema_name_ = self._schema_name(schema_name=schema_name) - rslt = self.engine.execute( + result_ = self.engine.execute( f"IF (EXISTS (" f"SELECT * FROM INFORMATION_SCHEMA.TABLES " f"WHERE TABLE_SCHEMA = '{schema_name_}' AND TABLE_NAME = '{table_name}'))" f"SELECT 1 ELSE SELECT 0") - result = bool(rslt.fetchone()[0]) + result = bool(result_.fetchone()[0]) return result @@ -3236,10 +3235,10 @@ def get_row_count(self, table_name, schema_name=None): table_name_ = self._table_name(table_name=table_name, schema_name=schema_name) - conn = self.engine.connect(close_with_result=True) - rslt = conn.execute(f'SELECT COUNT(*) FROM {table_name_};') + with self.engine.connect() as connection: + result = connection.execute(f'SELECT COUNT(*) FROM {table_name_};') - row_count = rslt.fetchone()[0] + row_count = result.fetchone()[0] return row_count