Optimize the resource release of SQL query execution

Author: junjie.miao

db_query/tools/sql_query.py

@@ -47,21 +47,14 @@ def _invoke(
         output_format = tool_parameters.get("output_format", "markdown").lower()
 
         try:
-            db = DbUtil(db_type=db_type,
+            with DbUtil(db_type=db_type,
                         username=db_username, password=db_password,
                         host=db_host, port=db_port,
-                        database=db_name, properties=db_properties)
+                        database=db_name, properties=db_properties) as db:
+                records = db.run_query(query_sql)
         except Exception as e:
-            message = "Database connection creation exception."
-            logging.exception(message)
-            raise Exception(message + " {}".format(e))
-
-        try:
-            records = db.run_query(query_sql)
-        except Exception as e:
-            message = "SQL query execution exception."
-            logging.exception(message)
-            raise Exception(message + " {}".format(e))
+            logging.exception("SQL query execution failed: %s", str(e))
+            raise RuntimeError(f"Error executing SQL: {e}") from e
 
         if output_format == "json":
             yield self.create_json_message({"records": records})

db_query_pre_auth/provider/db_query.py

@@ -24,13 +24,13 @@ def _validate_credentials(self, credentials: dict) -> None:
             raise ValueError("Database port can be empty or fill with integer value")
         db_name = credentials.get("db_name", "")
         db_properties = credentials.get("db_properties", "")
+
         try:
-            db = DbUtil(db_type=db_type,
+            with DbUtil(db_type=db_type,
                         username=db_username, password=db_password,
                         host=db_host, port=db_port,
-                        database=db_name, properties=db_properties)
-            db.run_query(db.test_sql())
+                        database=db_name, properties=db_properties) as db:
+                db.run_query(db.test_sql())
         except Exception as e:
-            message = "Database connection failure."
-            logging.exception(message)
-            raise Exception(message + " {}".format(e))
+            logging.exception("Database connection failed: %s", str(e))
+            raise RuntimeError(f"Error connection database: {e}") from e

db_query_pre_auth/tools/sql_query.py

@@ -47,21 +47,14 @@ def _invoke(
         output_format = tool_parameters.get("output_format", "markdown").lower()
 
         try:
-            db = DbUtil(db_type=db_type,
+            with DbUtil(db_type=db_type,
                         username=db_username, password=db_password,
                         host=db_host, port=db_port,
-                        database=db_name, properties=db_properties)
+                        database=db_name, properties=db_properties) as db:
+                records = db.run_query(query_sql)
         except Exception as e:
-            message = "Database connection creation exception."
-            logging.exception(message)
-            raise Exception(message + " {}".format(e))
-
-        try:
-            records = db.run_query(query_sql)
-        except Exception as e:
-            message = "SQL query execution exception."
-            logging.exception(message)
-            raise Exception(message + " {}".format(e))
+            logging.exception("SQL query execution failed: %s", str(e))
+            raise RuntimeError(f"Error executing SQL: {e}") from e
 
         if output_format == "json":
             yield self.create_json_message({"records": records})