Dev/refactor

CHANGELOG.md

@@ -8,6 +8,7 @@
 * [Doc] Document --persist-replace in API section (#539)
 * [Fix] Fixed CI issue by updating `invalid_connection_string_duckdb` in `test_magic.py` (#631)
 * [Fix] Refactored `ResultSet` to lazy loading (#470)
+* [Fix] Error when executing multiple SQL statements when using DuckDB with `autopandas` on (#674)
 * [Fix] Removed `WITH` when a snippet does not have a dependency (#657)
 
 ## 0.7.9 (2023-06-19)

src/sql/connection.py

@@ -203,6 +203,8 @@ def __init__(self, engine, alias=None):
         )
 
         self.dialect = engine.url.get_dialect()
+
+        # TODO: delete this!
         self.session = self._create_session(engine, self.url)
 
         self.connections[alias or self.url] = self
@@ -651,7 +653,7 @@ def execute(self, query, with_=None):
         Executes SQL query on a given connection
         """
         query = self._prepare_query(query, with_)
-        return self.session.execute(query)
+        return self.engine.execute(query)
 
 
 atexit.register(Connection.close_all, verbose=True)
@@ -672,11 +674,15 @@ def __init__(self, connection, engine):
             }
         )
 
+    # TODO: need to close the cursor
     def execute(self, query):
         cur = self.engine.cursor()
         cur.execute(query)
         return cur
 
+    def commit(self):
+        self.engine.commit()
+
 
 class CustomConnection(Connection):
     """
@@ -698,6 +704,7 @@ def __init__(self, payload, engine=None, alias=None):
         self.name = connection_name_
         self.dialect = connection_name_
         self.session = CustomSession(self, engine)
+        self.engine = self.session
 
         self.connections[alias or connection_name_] = self
 

src/sql/inspect.py

@@ -18,7 +18,7 @@ def _get_inspector(conn):
     if not Connection.current:
         raise exceptions.RuntimeError("No active connection")
     else:
-        return inspect(Connection.current.session)
+        return inspect(Connection.current.engine)
 
 
 class DatabaseInspection:
@@ -239,11 +239,15 @@ def __init__(self, table_name, schema=None) -> None:
         columns_query_result = sql.run.raw_run(
             Connection.current, f"SELECT * FROM {table_name} WHERE 1=0"
         )
+
         if Connection.is_custom_connection():
             columns = [i[0] for i in columns_query_result.description]
         else:
             columns = columns_query_result.keys()
 
+        # TODO: abstract it internally
+        columns_query_result.close()
+
         table_stats = dict({})
         columns_to_include_in_report = set()
         columns_with_styles = []