diff --git a/src/hflow/catalog_ui.py b/src/hflow/catalog_ui.py index f63661a5e..2cc6cc99f 100644 --- a/src/hflow/catalog_ui.py +++ b/src/hflow/catalog_ui.py @@ -86,11 +86,22 @@ def _start_duckdb_ui_server(catalog_connection: duckdb.DuckDBPyConnection, port: ) from error +def _ui_extension_is_loaded(catalog_connection: duckdb.DuckDBPyConnection) -> bool: + row = catalog_connection.execute( + "SELECT count(*) FROM duckdb_extensions() WHERE extension_name = 'ui' AND loaded" + ).fetchone() + return row is not None and int(row[0]) > 0 + + def _stop_duckdb_ui_server( catalog_connection: duckdb.DuckDBPyConnection, server_started: bool ) -> None: try: - if server_started: + # stop_ui_server only exists once the ui extension is loaded. Calling it + # without the extension makes DuckDB auto-install the extension from the + # network -- seconds on a slow connection, and pointless: there is no + # server to stop if the extension was never loaded. + if server_started and _ui_extension_is_loaded(catalog_connection): catalog_connection.execute("CALL stop_ui_server()") except duckdb.Error: # The process is already leaving. A server that stopped independently diff --git a/tests/test_catalog_ui.py b/tests/test_catalog_ui.py index 2df5a788d..0126ebcde 100644 --- a/tests/test_catalog_ui.py +++ b/tests/test_catalog_ui.py @@ -123,6 +123,100 @@ def test_catalog_ui_refuses_a_port_of_the_wrong_type( replace(settings, port=port) +class _RecordingCatalogConnection: + """DuckDB connection proxy that records every statement it executes.""" + + def __init__(self, inner: duckdb.DuckDBPyConnection) -> None: + self._inner = inner + self.executed_statements: list[str] = [] + self.closed = False + + def execute(self, sql: str, *parameters: object) -> object: + self.executed_statements.append(sql) + return self._inner.execute(sql, *parameters) + + def close(self) -> None: + self.closed = True + self._inner.close() + + def __getattr__(self, name: str) -> object: + return getattr(self._inner, name) + + +def _serve_until_shutdown( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ui_extension_loaded: bool, +) -> _RecordingCatalogConnection: + shutdown_event = Event() + server_started_event = Event() + connections: list[_RecordingCatalogConnection] = [] + background_failures: list[BaseException] = [] + + real_open_catalog_connection = catalog_ui.open_catalog_connection + + def record_server_start(catalog_connection: duckdb.DuckDBPyConnection, port: int) -> None: + assert port == catalog_ui.DEFAULT_CATALOG_UI_PORT + server_started_event.set() + + def open_recorded_connection(catalog_root: Path) -> _RecordingCatalogConnection: + connection = _RecordingCatalogConnection(real_open_catalog_connection(catalog_root)) + connections.append(connection) + return connection + + monkeypatch.setattr(catalog_ui, "open_catalog_connection", open_recorded_connection) + monkeypatch.setattr(catalog_ui, "_start_duckdb_ui_server", record_server_start) + monkeypatch.setattr( + catalog_ui, + "_ui_extension_is_loaded", + lambda _catalog_connection: ui_extension_loaded, + ) + + def run_catalog_ui() -> None: + try: + catalog_ui.serve_catalog_ui( + catalog_ui.CatalogUiSettings( + catalog_root=tmp_path / "catalog", + open_browser=False, + catalog_poll_interval_seconds=0.01, + ), + shutdown_event=shutdown_event, + ) + except BaseException as error: + background_failures.append(error) + + catalog_ui_thread = Thread(target=run_catalog_ui) + catalog_ui_thread.start() + try: + assert server_started_event.wait(timeout=2) + finally: + shutdown_event.set() + catalog_ui_thread.join(timeout=2) + + assert not catalog_ui_thread.is_alive() + assert background_failures == [] + (connection,) = connections + return connection + + +def test_catalog_ui_shutdown_does_not_stop_a_server_the_ui_extension_never_loaded( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + connection = _serve_until_shutdown(tmp_path, monkeypatch, ui_extension_loaded=False) + + assert "CALL stop_ui_server()" not in connection.executed_statements + assert connection.closed + + +def test_catalog_ui_shutdown_stops_the_server_when_the_ui_extension_is_loaded( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + connection = _serve_until_shutdown(tmp_path, monkeypatch, ui_extension_loaded=True) + + assert "CALL stop_ui_server()" in connection.executed_statements + assert connection.closed + + @pytest.mark.parametrize("port", [1, 65535]) def test_catalog_ui_accepts_port_boundaries(tmp_path: Path, port: int) -> None: settings = catalog_ui.CatalogUiSettings(catalog_root=tmp_path / "catalog", port=port)