diff --git a/pkg-py/src/querychat/_datasource.py b/pkg-py/src/querychat/_datasource.py index 0cdb7627..9562b2b8 100644 --- a/pkg-py/src/querychat/_datasource.py +++ b/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() diff --git a/pkg-py/src/querychat/_querychat.py b/pkg-py/src/querychat/_querychat.py index 833bc5de..142791ec 100644 --- a/pkg-py/src/querychat/_querychat.py +++ b/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): """