diff --git a/pyproject.toml b/pyproject.toml index 9df860f2b..e762d7bd8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -174,6 +174,11 @@ archiver-mongodb = [ "pymongo>=4.0", ] +# QuestDB archiver connector backend +archiver-questdb = [ + "asyncpg>=0.29", +] + # Facility knowledge / ontology support (OKF markdown bundles, RDF reasoning) knowledge = [ "rdflib>=7.0", @@ -246,6 +251,7 @@ module = [ "pyepics.*", "podman.*", "pymongo.*", + "asyncpg.*", ] ignore_missing_imports = true diff --git a/src/osprey/connectors/archiver/questdb_archiver_connector.py b/src/osprey/connectors/archiver/questdb_archiver_connector.py new file mode 100644 index 000000000..d1fbaa09c --- /dev/null +++ b/src/osprey/connectors/archiver/questdb_archiver_connector.py @@ -0,0 +1,430 @@ +""" +QuestDB archiver connector for historical PV data retrieval. + +Provides interface to QuestDB instances containing archived PV data. +Reads data over the PostgreSQL wire protocol (port 8812) using asyncpg. + +Assumed schema (table and column names are configurable): + + CREATE TABLE pv_archive ( + ts TIMESTAMP, -- designated timestamp + pv_name SYMBOL, -- PV identifier + value DOUBLE + ) TIMESTAMP(ts) PARTITION BY DAY; + +Note on identifiers vs values: asyncpg's parameterized queries ($1, $2, ...) +safely bind *values* (PV names, timestamps). Table/column *names* come only +from connector config (never user input at request time) and are validated +against a strict allow-list pattern before being interpolated into SQL. +""" + +import asyncio +import re +from datetime import UTC, datetime +from typing import Any + +import pandas as pd + +from osprey.connectors.archiver.base import ArchiverConnector, ArchiverMetadata +from osprey.utils.logger import get_logger + +logger = get_logger("questdb_archiver_connector") + +# Identifiers (table/column names) are only ever validated against this +# pattern -- they never come from per-request user input, only from the +# connect() config, but we validate anyway since they're interpolated +# directly into SQL (asyncpg has no way to parameterize identifiers). +_IDENTIFIER_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") + + +def _validate_identifier(name: str, label: str) -> str: + """Ensure a table/column name is a safe SQL identifier before use.""" + if not _IDENTIFIER_RE.match(name): + raise ValueError(f"Invalid {label} '{name}': must be a plain SQL identifier") + return name + + +def _to_utc(dt: datetime) -> datetime: + """Normalize a datetime to UTC, assuming naive datetimes are already UTC.""" + if dt.tzinfo is None: + return dt.replace(tzinfo=UTC) + return dt.astimezone(UTC) + + +def _ms_to_sample_unit(precision_ms: int) -> str: + """Convert milliseconds to a QuestDB SAMPLE BY unit string.""" + if precision_ms < 1000: + return f"{precision_ms}T" + seconds = precision_ms // 1000 + if seconds < 60: + return f"{seconds}s" + minutes = seconds // 60 + if minutes < 60: + return f"{minutes}m" + hours = minutes // 60 + if hours < 24: + return f"{hours}h" + return f"{hours // 24}d" + + +class QuestDBArchiverConnector(ArchiverConnector): + """ + QuestDB archiver connector for historical PV data. + + Connects to QuestDB over its PostgreSQL wire protocol using asyncpg. + Queries a single narrow (long-format) table and pivots results into + the wide DataFrame shape required by the ArchiverConnector interface. + + All PV names and timestamps are passed as bound query parameters + ($1, $2, ...) rather than interpolated into the SQL string, so values + containing quotes or other special characters cannot break or inject + into the query. Table/column names come only from connector config and + are validated against a strict identifier pattern. + + Example: + >>> config = { + >>> 'host': 'questdb.als.lbl.gov', + >>> 'port': 8812, + >>> 'database': 'qdb', + >>> 'username': 'admin', + >>> 'password_env': 'QUESTDB_PASSWORD', + >>> 'table': 'pv_archive', + >>> } + >>> connector = QuestDBArchiverConnector() + >>> await connector.connect(config) + >>> df = await connector.get_data( + >>> pv_list=['BEAM:CURRENT'], + >>> start_date=datetime(2024, 1, 1), + >>> end_date=datetime(2024, 1, 2) + >>> ) + """ + + def __init__(self): + self._connected = False + self._pool = None + self._timeout = 60 + + # Schema config (overridable via connect config) + self._table = "pv_archive" + self._pv_col = "pv_name" + self._val_col = "value" + self._ts_col = "ts" + + async def connect(self, config: dict[str, Any]) -> None: + """ + Open an asyncpg connection pool to QuestDB. + + Args: + config: Configuration with keys: + - host: QuestDB host (required) + - port: PostgreSQL wire protocol port (default: 8812) + - database: Database name (default: 'qdb') + - username: QuestDB username (required) + - password_env: Environment variable name for password (required) + - timeout: Default timeout in seconds (default: 60) + - table: Table name (default: 'pv_archive') + - pv_column: PV name column (default: 'pv_name') + - value_column: Value column (default: 'value') + - ts_column: Timestamp column (default: 'ts') + + Raises: + ImportError: If asyncpg is not installed + ValueError: If required config values are missing or schema + names are not valid SQL identifiers + ConnectionError: If connection cannot be established + """ + try: + import asyncpg + except ImportError as e: + raise ImportError( + "asyncpg is required for QuestDB archiver. Install with: pip install asyncpg" + ) from e + + import os + + host = config.get("host") + if not host: + raise ValueError("host is required for QuestDB archiver") + + username = config.get("username") + if not username: + raise ValueError("username is required for QuestDB archiver") + + password_env = config.get("password_env") + if not password_env: + raise ValueError("password_env is required for QuestDB archiver") + + password = os.getenv(password_env) + if not password: + raise ValueError( + f"Environment variable '{password_env}' not set. " + "Password is required for QuestDB authentication." + ) + + port = config.get("port", 8812) + database = config.get("database", "qdb") + self._timeout = config.get("timeout", 60) + + # Schema overrides -- validated since they get interpolated into SQL + # as identifiers (asyncpg can only parameterize values, not names). + self._table = _validate_identifier(config.get("table", self._table), "table name") + self._pv_col = _validate_identifier(config.get("pv_column", self._pv_col), "pv_column") + self._val_col = _validate_identifier( + config.get("value_column", self._val_col), "value_column" + ) + self._ts_col = _validate_identifier(config.get("ts_column", self._ts_col), "ts_column") + + try: + self._pool = await asyncpg.create_pool( + host=host, + port=port, + database=database, + user=username, + password=password, + command_timeout=self._timeout, + ) + + # Smoke test + async with self._pool.acquire() as conn: + await conn.fetchval("SELECT 1") + + self._connected = True + logger.debug(f"QuestDB archiver connector initialized: {host}:{port}/{database}") + + except asyncpg.PostgresError as e: + await self.disconnect() + raise ConnectionError( + f"Cannot connect to QuestDB at {host}:{port}. " + "Please check connectivity and authentication." + ) from e + except OSError as e: + await self.disconnect() + raise ConnectionError(f"QuestDB connection failed: {e}") from e + + async def disconnect(self) -> None: + """Close the asyncpg connection pool. Safe to call when not connected.""" + if self._pool: + try: + await self._pool.close() + except Exception as e: + logger.warning(f"Error closing QuestDB connection pool: {e}") + + self._pool = None + self._connected = False + logger.debug("QuestDB archiver connector disconnected") + + async def get_data( + self, + pv_list: list[str], + start_date: datetime, + end_date: datetime, + precision_ms: int = 1000, + timeout: int | None = None, + ) -> pd.DataFrame: + """ + Retrieve historical data from QuestDB. + + Queries the configured table for the given PVs and time range. + When precision_ms > 0, uses QuestDB's SAMPLE BY for server-side + downsampling. Results are pivoted from long to wide format. + + PV names and timestamps are passed as bound parameters, not + interpolated into the SQL string. start_date/end_date are + normalized to UTC before querying; naive datetimes are assumed + to already be in UTC. + + Args: + pv_list: List of PV names to retrieve + start_date: Start of time range + end_date: End of time range + precision_ms: Time precision in milliseconds (0 = raw rows) + timeout: Optional timeout in seconds + + Returns: + DataFrame with a UTC datetime index and one column per PV. + PVs with no data in range are present as NaN columns. + + Raises: + RuntimeError: If archiver not connected + TypeError: If start_date or end_date are not datetime objects + ValueError: If pv_list is empty or time range is invalid + TimeoutError: If operation times out + ConnectionError: If QuestDB cannot be reached + """ + timeout = timeout or self._timeout + + if not self._connected or self._pool is None: + raise RuntimeError("QuestDB archiver not connected") + + if not isinstance(start_date, datetime): + raise TypeError(f"start_date must be a datetime object, got {type(start_date)}") + if not isinstance(end_date, datetime): + raise TypeError(f"end_date must be a datetime object, got {type(end_date)}") + + if not pv_list: + raise ValueError("pv_list cannot be empty") + + start_utc = _to_utc(start_date) + end_utc = _to_utc(end_date) + if start_utc >= end_utc: + raise ValueError("start_date must be before end_date") + + # PV names and timestamps are bound as $1, $2, ... -- never + # interpolated into the SQL string. Table/column names are + # identifiers validated in connect() and are safe to interpolate. + pv_placeholders = ", ".join(f"${i + 3}" for i in range(len(pv_list))) + + if precision_ms > 0: + sample_unit = _ms_to_sample_unit(precision_ms) + sql = ( + f"SELECT {self._ts_col}, {self._pv_col}, " + f"avg({self._val_col}) AS {self._val_col} " + f"FROM {self._table} " + f"WHERE {self._ts_col} BETWEEN $1 AND $2 " + f" AND {self._pv_col} IN ({pv_placeholders}) " + f"SAMPLE BY {sample_unit} ALIGN TO CALENDAR;" + ) + else: + sql = ( + f"SELECT {self._ts_col}, {self._pv_col}, {self._val_col} " + f"FROM {self._table} " + f"WHERE {self._ts_col} BETWEEN $1 AND $2 " + f" AND {self._pv_col} IN ({pv_placeholders}) " + f"ORDER BY {self._ts_col};" + ) + + params = [start_utc, end_utc, *pv_list] + + async def fetch(): + async with self._pool.acquire() as conn: + return await conn.fetch(sql, *params) + + try: + records = await asyncio.wait_for(fetch(), timeout=float(timeout)) + except TimeoutError as e: + raise TimeoutError(f"QuestDB query timed out after {timeout}s") from e + except Exception as e: + # Network/connection issues surface here distinctly from a + # plain timeout so callers can tell the two apart. + raise ConnectionError(f"QuestDB query failed: {e}") from e + + if not records: + logger.debug(f"No data found in range {start_utc} to {end_utc}") + return pd.DataFrame(index=pd.DatetimeIndex([], tz="UTC"), columns=pv_list) + + raw_df = pd.DataFrame(records, columns=[self._ts_col, self._pv_col, self._val_col]) + raw_df[self._ts_col] = pd.to_datetime(raw_df[self._ts_col], utc=True) + + df = raw_df.pivot_table( + index=self._ts_col, + columns=self._pv_col, + values=self._val_col, + aggfunc="mean", + ) + df.index.name = "datetime" + df.columns.name = None + + # Ensure all requested PVs are present + for pv in pv_list: + if pv not in df.columns: + df[pv] = float("nan") + + logger.debug(f"Retrieved QuestDB data: {len(df)} points for {len(pv_list)} PVs") + return df[pv_list] + + async def get_metadata(self, pv_name: str) -> ArchiverMetadata: + """ + Get archiving metadata for a PV. + + Args: + pv_name: Name of the process variable + + Returns: + ArchiverMetadata with archiving information + + Raises: + RuntimeError: If archiver not connected + ValueError: If pv_name is empty + """ + if not self._connected or self._pool is None: + raise RuntimeError("QuestDB archiver not connected") + + if not pv_name: + raise ValueError("pv_name cannot be empty") + + sql = ( + f"SELECT " + f" min({self._ts_col}) AS archival_start, " + f" max({self._ts_col}) AS archival_end, " + f" count() AS sample_count, " + f" datediff('ms', min({self._ts_col}), max({self._ts_col})) " + f" / nullif(count() - 1, 0) AS avg_period_ms " + f"FROM {self._table} " + f"WHERE {self._pv_col} = $1;" + ) + + try: + async with self._pool.acquire() as conn: + row = await conn.fetchrow(sql, pv_name) + except Exception as e: + # A query failure here means we genuinely don't know the + # archival status -- report it rather than silently claiming + # the PV isn't archived, which would be misleading. + raise ConnectionError(f"QuestDB metadata query failed for {pv_name}: {e}") from e + + if not row or row["sample_count"] == 0: + return ArchiverMetadata(pv_name=pv_name, is_archived=False) + + start = pd.to_datetime(row["archival_start"], utc=True).to_pydatetime() + end = pd.to_datetime(row["archival_end"], utc=True).to_pydatetime() + avg_ms = row["avg_period_ms"] + sampling_s = float(avg_ms) / 1000.0 if avg_ms is not None else None + + return ArchiverMetadata( + pv_name=pv_name, + is_archived=True, + archival_start=start, + archival_end=end, + sampling_period=sampling_s, + ) + + async def check_availability(self, pv_names: list[str]) -> dict[str, bool]: + """ + Check which PVs have archived data. + + Uses a single batched query for efficiency. + + Args: + pv_names: List of PV names to check + + Returns: + Dictionary mapping PV name to availability status + + Raises: + RuntimeError: If archiver not connected + ConnectionError: If the availability query fails + """ + if not self._connected or self._pool is None: + raise RuntimeError("QuestDB archiver not connected") + + if not pv_names: + return {} + + placeholders = ", ".join(f"${i + 1}" for i in range(len(pv_names))) + sql = ( + f"SELECT DISTINCT {self._pv_col} " + f"FROM {self._table} " + f"WHERE {self._pv_col} IN ({placeholders});" + ) + + try: + async with self._pool.acquire() as conn: + records = await conn.fetch(sql, *pv_names) + except Exception as e: + # Surface the failure rather than reporting every PV as + # unavailable, which would look identical to "checked, none + # archived" and could mislead a caller. + raise ConnectionError(f"QuestDB availability query failed: {e}") from e + + found = {row[0] for row in records} + return {pv: (pv in found) for pv in pv_names} diff --git a/src/osprey/connectors/factory.py b/src/osprey/connectors/factory.py index 1e5316bbe..de84fa1d2 100644 --- a/src/osprey/connectors/factory.py +++ b/src/osprey/connectors/factory.py @@ -257,3 +257,14 @@ def register_builtin_connectors() -> None: ConnectorFactory.register_archiver(types.EPICS_ARCHIVER, EPICSArchiverConnector) if _mongo_available: ConnectorFactory.register_archiver(types.MONGODB_ARCHIVER, MongoDBArchiverConnector) + + +try: + from osprey.connectors.archiver.questdb_archiver_connector import QuestDBArchiverConnector + + _questdb_available = True +except ImportError: + _questdb_available = False + +if _questdb_available: + ConnectorFactory.register_archiver(types.QUESTDB_ARCHIVER, QuestDBArchiverConnector) diff --git a/src/osprey/connectors/types.py b/src/osprey/connectors/types.py index 50bc818b5..7e8fb43da 100644 --- a/src/osprey/connectors/types.py +++ b/src/osprey/connectors/types.py @@ -13,7 +13,8 @@ MOCK_ARCHIVER = "mock_archiver" EPICS_ARCHIVER = "epics_archiver" MONGODB_ARCHIVER = "mongodb_archiver" +QUESTDB_ARCHIVER = "questdb_archiver" # -- CLI choice lists (only types with implementations) -- CLI_CONTROL_SYSTEM_TYPES = [MOCK, EPICS] -CLI_ARCHIVER_TYPES = [MOCK_ARCHIVER, EPICS_ARCHIVER, MONGODB_ARCHIVER] +CLI_ARCHIVER_TYPES = [MOCK_ARCHIVER, EPICS_ARCHIVER, MONGODB_ARCHIVER, QUESTDB_ARCHIVER] diff --git a/tests/connectors/test_questdb_archiver_connector.py b/tests/connectors/test_questdb_archiver_connector.py new file mode 100644 index 000000000..cd54a7b74 --- /dev/null +++ b/tests/connectors/test_questdb_archiver_connector.py @@ -0,0 +1,454 @@ +"""Tests for QuestDB archiver connector.""" + +import asyncio +from datetime import datetime +from unittest.mock import AsyncMock, MagicMock, patch + +import pandas as pd +import pytest + +from osprey.connectors.archiver.base import ArchiverMetadata +from osprey.connectors.archiver.questdb_archiver_connector import QuestDBArchiverConnector + +asyncpg = pytest.importorskip("asyncpg", reason="asyncpg not installed") + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_record(ts: str, pv: str, value: float) -> MagicMock: + """Simulate an asyncpg Record row.""" + record = MagicMock() + record.__iter__ = MagicMock(return_value=iter([ts, pv, value])) + record.__getitem__ = MagicMock( + side_effect=lambda k: {"ts": ts, "pv_name": pv, "value": value}[k] + ) + return record + + +def _make_pool(fetch_return=None, fetchrow_return=None, fetchval_return=1): + """Build a mock asyncpg pool.""" + mock_conn = AsyncMock() + mock_conn.fetch = AsyncMock(return_value=fetch_return or []) + mock_conn.fetchrow = AsyncMock(return_value=fetchrow_return) + mock_conn.fetchval = AsyncMock(return_value=fetchval_return) + + mock_ctx = AsyncMock() + mock_ctx.__aenter__ = AsyncMock(return_value=mock_conn) + mock_ctx.__aexit__ = AsyncMock(return_value=False) + + mock_pool = AsyncMock() + mock_pool.acquire = MagicMock(return_value=mock_ctx) + mock_pool.close = AsyncMock() + return mock_pool, mock_conn + + +def _make_connector_with_pool(pool): + """Return a QuestDBArchiverConnector with pool injected directly.""" + connector = QuestDBArchiverConnector() + connector._pool = pool + connector._connected = True + return connector + + +# --------------------------------------------------------------------------- +# Connect / Disconnect +# --------------------------------------------------------------------------- + + +class TestConnectDisconnectLifecycle: + @pytest.mark.asyncio + async def test_connect_missing_host_raises_value_error(self): + connector = QuestDBArchiverConnector() + with pytest.raises(ValueError, match="host"): + await connector.connect({"username": "admin", "password_env": "PW"}) + + @pytest.mark.asyncio + async def test_connect_missing_username_raises_value_error(self): + connector = QuestDBArchiverConnector() + with pytest.raises(ValueError, match="username"): + await connector.connect({"host": "localhost", "password_env": "PW"}) + + @pytest.mark.asyncio + async def test_connect_missing_password_env_raises_value_error(self): + connector = QuestDBArchiverConnector() + with pytest.raises(ValueError, match="password_env"): + await connector.connect({"host": "localhost", "username": "admin"}) + + @pytest.mark.asyncio + async def test_connect_unset_password_env_raises_value_error(self, monkeypatch): + monkeypatch.delenv("QUESTDB_PW", raising=False) + connector = QuestDBArchiverConnector() + with pytest.raises(ValueError, match="not set"): + await connector.connect( + { + "host": "localhost", + "username": "admin", + "password_env": "QUESTDB_PW", + } + ) + + @pytest.mark.asyncio + async def test_connect_success_sets_connected(self, monkeypatch): + monkeypatch.setenv("QUESTDB_PW", "secret") + pool, _ = _make_pool() + + with patch("asyncpg.create_pool", AsyncMock(return_value=pool)): + connector = QuestDBArchiverConnector() + await connector.connect( + { + "host": "localhost", + "username": "admin", + "password_env": "QUESTDB_PW", + } + ) + + assert connector._connected is True + await connector.disconnect() + + @pytest.mark.asyncio + async def test_connect_schema_overrides_applied(self, monkeypatch): + monkeypatch.setenv("QUESTDB_PW", "secret") + pool, _ = _make_pool() + + with patch("asyncpg.create_pool", AsyncMock(return_value=pool)): + connector = QuestDBArchiverConnector() + await connector.connect( + { + "host": "localhost", + "username": "admin", + "password_env": "QUESTDB_PW", + "table": "beam_data", + "pv_column": "channel", + "value_column": "reading", + "ts_column": "timestamp", + } + ) + + assert connector._table == "beam_data" + assert connector._pv_col == "channel" + assert connector._val_col == "reading" + assert connector._ts_col == "timestamp" + await connector.disconnect() + + @pytest.mark.asyncio + async def test_disconnect_closes_pool(self): + pool, _ = _make_pool() + connector = _make_connector_with_pool(pool) + await connector.disconnect() + pool.close.assert_awaited_once() + assert connector._connected is False + assert connector._pool is None + + @pytest.mark.asyncio + async def test_disconnect_when_not_connected_is_safe(self): + connector = QuestDBArchiverConnector() + await connector.disconnect() + assert connector._connected is False + + +# --------------------------------------------------------------------------- +# get_data +# --------------------------------------------------------------------------- + + +class TestGetDataMethod: + @pytest.mark.asyncio + async def test_not_connected_raises_runtime_error(self): + connector = QuestDBArchiverConnector() + with pytest.raises(RuntimeError, match="not connected"): + await connector.get_data( + pv_list=["BEAM:CURRENT"], + start_date=datetime(2024, 1, 1), + end_date=datetime(2024, 1, 2), + ) + + @pytest.mark.asyncio + async def test_empty_pv_list_raises_value_error(self): + pool, _ = _make_pool() + connector = _make_connector_with_pool(pool) + with pytest.raises(ValueError, match="pv_list"): + await connector.get_data( + pv_list=[], + start_date=datetime(2024, 1, 1), + end_date=datetime(2024, 1, 2), + ) + + @pytest.mark.asyncio + async def test_invalid_time_range_raises_value_error(self): + pool, _ = _make_pool() + connector = _make_connector_with_pool(pool) + with pytest.raises(ValueError, match="start_date"): + await connector.get_data( + pv_list=["PV:X"], + start_date=datetime(2024, 1, 2), + end_date=datetime(2024, 1, 1), + ) + + @pytest.mark.asyncio + async def test_returns_dataframe_with_datetime_index(self): + rows = [ + ("2024-01-01T00:00:00Z", "BEAM:CURRENT", 499.8), + ("2024-01-01T00:00:01Z", "BEAM:CURRENT", 499.7), + ] + pool, conn = _make_pool(fetch_return=rows) + connector = _make_connector_with_pool(pool) + + df = await connector.get_data( + pv_list=["BEAM:CURRENT"], + start_date=datetime(2024, 1, 1), + end_date=datetime(2024, 1, 1, 1), + ) + + assert isinstance(df, pd.DataFrame) + assert isinstance(df.index, pd.DatetimeIndex) + assert "BEAM:CURRENT" in df.columns + + @pytest.mark.asyncio + async def test_correct_values_returned(self): + rows = [ + ("2024-01-01T00:00:00Z", "PV:X", 1.0), + ("2024-01-01T00:00:01Z", "PV:X", 2.0), + ("2024-01-01T00:00:02Z", "PV:X", 3.0), + ] + pool, _ = _make_pool(fetch_return=rows) + connector = _make_connector_with_pool(pool) + + df = await connector.get_data( + pv_list=["PV:X"], + start_date=datetime(2024, 1, 1), + end_date=datetime(2024, 1, 1, 1), + ) + + assert list(df["PV:X"]) == [1.0, 2.0, 3.0] + + @pytest.mark.asyncio + async def test_multi_pv_one_column_each(self): + rows = [ + ("2024-01-01T00:00:00Z", "PV:1", 1.0), + ("2024-01-01T00:00:00Z", "PV:2", 2.0), + ] + pool, _ = _make_pool(fetch_return=rows) + connector = _make_connector_with_pool(pool) + + df = await connector.get_data( + pv_list=["PV:1", "PV:2"], + start_date=datetime(2024, 1, 1), + end_date=datetime(2024, 1, 1, 1), + ) + + assert "PV:1" in df.columns + assert "PV:2" in df.columns + + @pytest.mark.asyncio + async def test_absent_pv_filled_with_nan(self): + rows = [("2024-01-01T00:00:00Z", "PV:A", 5.0)] + pool, _ = _make_pool(fetch_return=rows) + connector = _make_connector_with_pool(pool) + + df = await connector.get_data( + pv_list=["PV:A", "PV:MISSING"], + start_date=datetime(2024, 1, 1), + end_date=datetime(2024, 1, 1, 1), + ) + + assert "PV:MISSING" in df.columns + assert df["PV:MISSING"].isna().all() + + @pytest.mark.asyncio + async def test_empty_response_returns_empty_dataframe(self): + pool, _ = _make_pool(fetch_return=[]) + connector = _make_connector_with_pool(pool) + + df = await connector.get_data( + pv_list=["BEAM:CURRENT"], + start_date=datetime(2024, 1, 1), + end_date=datetime(2024, 1, 1, 1), + ) + + assert isinstance(df, pd.DataFrame) + assert len(df) == 0 + + @pytest.mark.asyncio + async def test_timeout_raises_timeout_error(self): + async def slow_fetch(*args, **kwargs): + await asyncio.sleep(10) + + pool, conn = _make_pool() + conn.fetch = slow_fetch + connector = _make_connector_with_pool(pool) + + with pytest.raises(TimeoutError): + await connector.get_data( + pv_list=["BEAM:CURRENT"], + start_date=datetime(2024, 1, 1), + end_date=datetime(2024, 1, 2), + timeout=1, + ) + + +# --------------------------------------------------------------------------- +# get_metadata +# --------------------------------------------------------------------------- + + +class TestGetMetadataMethod: + @pytest.mark.asyncio + async def test_not_connected_raises_runtime_error(self): + connector = QuestDBArchiverConnector() + with pytest.raises(RuntimeError, match="not connected"): + await connector.get_metadata("BEAM:CURRENT") + + @pytest.mark.asyncio + async def test_empty_pv_name_raises_value_error(self): + pool, _ = _make_pool() + connector = _make_connector_with_pool(pool) + with pytest.raises(ValueError): + await connector.get_metadata("") + + @pytest.mark.asyncio + async def test_returns_archiver_metadata(self): + row = MagicMock() + row.__getitem__ = MagicMock( + side_effect=lambda k: { + "archival_start": "2024-01-01T00:00:00Z", + "archival_end": "2024-01-02T00:00:00Z", + "sample_count": 86400, + "avg_period_ms": 1000, + }[k] + ) + + pool, _ = _make_pool(fetchrow_return=row) + connector = _make_connector_with_pool(pool) + metadata = await connector.get_metadata("BEAM:CURRENT") + + assert isinstance(metadata, ArchiverMetadata) + assert metadata.pv_name == "BEAM:CURRENT" + assert metadata.is_archived is True + assert metadata.sampling_period == pytest.approx(1.0) + + @pytest.mark.asyncio + async def test_not_archived_when_no_rows(self): + row = MagicMock() + row.__getitem__ = MagicMock( + side_effect=lambda k: { + "archival_start": None, + "archival_end": None, + "sample_count": 0, + "avg_period_ms": None, + }[k] + ) + + pool, _ = _make_pool(fetchrow_return=row) + connector = _make_connector_with_pool(pool) + metadata = await connector.get_metadata("NONEXISTENT:PV") + + assert metadata.is_archived is False + + +# --------------------------------------------------------------------------- +# check_availability +# --------------------------------------------------------------------------- + + +class TestCheckAvailability: + @pytest.mark.asyncio + async def test_not_connected_raises_runtime_error(self): + connector = QuestDBArchiverConnector() + with pytest.raises(RuntimeError, match="not connected"): + await connector.check_availability(["PV:X"]) + + @pytest.mark.asyncio + async def test_empty_list_returns_empty_dict(self): + pool, _ = _make_pool() + connector = _make_connector_with_pool(pool) + result = await connector.check_availability([]) + assert result == {} + + @pytest.mark.asyncio + async def test_present_pvs_true_absent_false(self): + rows = [("PV:1",), ("PV:2",)] + pool, conn = _make_pool(fetch_return=rows) + connector = _make_connector_with_pool(pool) + result = await connector.check_availability(["PV:1", "PV:2", "PV:MISSING"]) + + assert result["PV:1"] is True + assert result["PV:2"] is True + assert result["PV:MISSING"] is False + + @pytest.mark.asyncio + async def test_single_batched_query(self): + pool, conn = _make_pool(fetch_return=[]) + connector = _make_connector_with_pool(pool) + await connector.check_availability(["PV:1", "PV:2", "PV:3"]) + assert conn.fetch.call_count == 1 + + +# --------------------------------------------------------------------------- +# Factory integration +# --------------------------------------------------------------------------- + + +class TestFactoryIntegration: + @pytest.fixture(autouse=True) + def setup_factory(self): + from osprey.connectors.factory import ConnectorFactory + + ConnectorFactory.register_archiver("questdb_archiver", QuestDBArchiverConnector) + yield + ConnectorFactory._archiver_connectors.pop("questdb_archiver", None) + + @pytest.mark.asyncio + async def test_factory_creates_questdb_archiver_connector(self, monkeypatch): + monkeypatch.setenv("QUESTDB_PW", "secret") + pool, _ = _make_pool() + + with patch("asyncpg.create_pool", AsyncMock(return_value=pool)): + from osprey.connectors.factory import ConnectorFactory + + config = { + "type": "questdb_archiver", + "questdb_archiver": { + "host": "localhost", + "username": "admin", + "password_env": "QUESTDB_PW", + }, + } + connector = await ConnectorFactory.create_archiver_connector(config) + + assert isinstance(connector, QuestDBArchiverConnector) + assert connector._connected is True + await connector.disconnect() + + +# --------------------------------------------------------------------------- +# Integration test (requires live QuestDB) +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +class TestQuestDBIntegration: + @pytest.mark.asyncio + async def test_connect_and_query_live_instance(self, monkeypatch): + """ + Integration test against a real QuestDB instance. + + Start one locally with: + docker run -p 8812:8812 -p 9000:9000 questdb/questdb + + Set environment variable QUESTDB_PW=quest before running. + """ + monkeypatch.setenv("QUESTDB_PW", "quest") + connector = QuestDBArchiverConnector() + await connector.connect( + { + "host": "localhost", + "port": 8812, + "username": "admin", + "password_env": "QUESTDB_PW", + } + ) + assert connector._connected is True + await connector.disconnect()