feat(pkg-py): Add .cleanup() method (#121)

* Close #106. Add cleanup method to Python

* Address feedback

* More docstring polish

Author: cpsievert

pkg-py/src/querychat/_datasource.py

@@ -83,6 +83,20 @@ def get_data(self) -> pd.DataFrame:
         """
         ...
 
+    @abstractmethod
+    def cleanup(self) -> None:
+        """
+        Clean up resources associated with the data source.
+
+        This method should clean up any connections or resources used by the
+        data source.
+
+        Returns
+        -------
+        None
+
+        """
+
 
 class DataFrameSource(DataSource):
     """A DataSource implementation that wraps a pandas DataFrame using DuckDB."""
@@ -212,6 +226,18 @@ def get_data(self) -> pd.DataFrame:
         # TODO(@gadenbuie): This should just return `self._df` and not a pandas DataFrame
         return self._df.lazy().collect().to_pandas()
 
+    def cleanup(self) -> None:
+        """
+        Close the DuckDB connection.
+
+        Returns
+        -------
+        None
+
+        """
+        if self._conn:
+            self._conn.close()
+
 
 class SQLAlchemySource(DataSource):
     """
@@ -440,3 +466,15 @@ def _get_sql_type_name(self, type_: sqltypes.TypeEngine) -> str:  # noqa: PLR091
     def _get_connection(self) -> Connection:
         """Get a connection to use for queries."""
         return self._engine.connect()
+
+    def cleanup(self) -> None:
+        """
+        Dispose of the SQLAlchemy engine.
+
+        Returns
+        -------
+        None
+
+        """
+        if self._engine:
+            self._engine.dispose()

pkg-py/src/querychat/_querychat.py

@@ -273,6 +273,20 @@ def data_source(self):
         """
         return self._data_source
 
+    def cleanup(self) -> None:
+        """
+        Clean up resources associated with the data source.
+
+        Call this method when you are done using the QueryChat object to close
+        database connections and avoid resource leaks.
+
+        Returns
+        -------
+        None
+
+        """
+        self._data_source.cleanup()
+
 
 class QueryChat(QueryChatBase):
     """