From 28dba5a6dbdf8885a993a78a02268554f2829019 Mon Sep 17 00:00:00 2001 From: caterryan Date: Tue, 14 Jul 2026 09:50:18 -0500 Subject: [PATCH] feat(plugin-utils): query across databases --- influxdata-plugin-utils/CHANGELOG.md | 4 +- influxdata-plugin-utils/README.md | 25 +++- .../influxdata_plugin_utils/introspection.py | 74 ++++++++-- .../tests/test_introspection.py | 131 ++++++++++++++++++ 4 files changed, 220 insertions(+), 14 deletions(-) create mode 100644 influxdata-plugin-utils/tests/test_introspection.py diff --git a/influxdata-plugin-utils/CHANGELOG.md b/influxdata-plugin-utils/CHANGELOG.md index 6af573f..a16b09c 100644 --- a/influxdata-plugin-utils/CHANGELOG.md +++ b/influxdata-plugin-utils/CHANGELOG.md @@ -13,6 +13,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `write.write_data` — optional `database` parameter for writing to another database. +- `introspection` — optional `database` parameter for schema helpers and + `query_window`. - `parsing.parse_timedelta` — `ms` (milliseconds) and `us` (microseconds) duration units. @@ -39,4 +41,4 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 [Unreleased]: https://github.com/influxdata/influxdb3_plugins/compare/utils-v0.2.0...HEAD [0.2.0]: https://github.com/influxdata/influxdb3_plugins/compare/utils-v0.1.0...utils-v0.2.0 -[0.1.0]: https://github.com/influxdata/influxdb3_plugins/releases/tag/utils-v0.1.0 \ No newline at end of file +[0.1.0]: https://github.com/influxdata/influxdb3_plugins/releases/tag/utils-v0.1.0 diff --git a/influxdata-plugin-utils/README.md b/influxdata-plugin-utils/README.md index 7b88099..c60d321 100644 --- a/influxdata-plugin-utils/README.md +++ b/influxdata-plugin-utils/README.md @@ -19,7 +19,7 @@ pip install -e influxdata-plugin-utils | Module | What it provides | |-----------------|-----------------------------------------------------------------------------------------------------------------------------| | `config` | `load_plugin_config(args, validators)` (dynaconf-backed), `resolve_plugin_dir()`, `resolve_path()`, re-exported `Validator` | -| `introspection` | `get_table_names()`, `get_tag_names()`, `get_field_names()`, `query_window()` | +| `introspection` | `get_table_names()`, `get_tag_names()`, `get_field_names()`, `query_window()` with optional `database=` | | `parsing` | `parse_timedelta()`, `parse_timestamp_ns()`, `parse_int()`, `parse_bool()`, `parse_delimited_list()`, `parse_key_value()` | | `cache` | `cached(influxdb3_local, key, producer, ttl_seconds=3600)` | | `write` | `build_line()`, `build_line_typed()`, `add_field_with_type()`, `write_data()`, `BatchLines` | @@ -69,7 +69,28 @@ write_data(influxdb3_local, lines) # batched + retried by default # write_data(influxdb3_local, lines, no_sync=True) # write_sync API (3.8+) ``` +## Cross-database queries + +On InfluxDB versions that support processing-engine cross-database queries, +the introspection helpers accept `database=` and pass it through to +`influxdb3_local.query`. +Cached schema results are separated per database. + +```python +from influxdata_plugin_utils.introspection import get_field_names, query_window + +fields = get_field_names(influxdb3_local, "cpu", database="source_db") +rows = query_window( + influxdb3_local, + "cpu", + start=start, + end=end, + columns=fields, + database="source_db", +) +``` + ## License Licensed under either of [Apache License 2.0](LICENSE-APACHE) or -[MIT license](LICENSE-MIT) at your option. \ No newline at end of file +[MIT license](LICENSE-MIT) at your option. diff --git a/influxdata-plugin-utils/src/influxdata_plugin_utils/introspection.py b/influxdata-plugin-utils/src/influxdata_plugin_utils/introspection.py index b348c04..83ff0e2 100644 --- a/influxdata-plugin-utils/src/influxdata_plugin_utils/introspection.py +++ b/influxdata-plugin-utils/src/influxdata_plugin_utils/introspection.py @@ -2,6 +2,8 @@ Each helper takes ``influxdb3_local`` explicitly and may optionally use the TTL cache. Queries mirror the patterns already used across the plugins. +Helpers accept an optional ``database`` argument for engines that support +cross-database plugin queries. """ from .cache import cached @@ -22,13 +24,40 @@ def _quote_identifier(identifier: str) -> str: return '"' + identifier.replace('"', '""') + '"' +def _query( + influxdb3_local, + query: str, + args: dict | None = None, + database: str | None = None, +): + """Run a plugin query, preserving the old call shape when no database is set.""" + if database is None: + if args is None: + return influxdb3_local.query(query) + return influxdb3_local.query(query, args) + if args is None: + return influxdb3_local.query(query, database=database) + return influxdb3_local.query(query, args, database=database) + + +def _cache_key(base: str, database: str | None) -> str: + """Return a cache key, keeping existing default-database keys unchanged.""" + if database is None: + return base + return f"{base}:database:{database}" + + def get_table_names( - influxdb3_local, *, use_cache: bool = True, ttl_seconds: int = 3600 + influxdb3_local, + *, + database: str | None = None, + use_cache: bool = True, + ttl_seconds: int = 3600, ) -> list[str]: """Return base table names via ``SHOW TABLES``.""" def producer() -> list[str]: - rows = influxdb3_local.query("SHOW TABLES") + rows = _query(influxdb3_local, "SHOW TABLES", database=database) return [ row["table_name"] for row in rows @@ -36,12 +65,22 @@ def producer() -> list[str]: ] if use_cache: - return cached(influxdb3_local, "shared:tables", producer, ttl_seconds=ttl_seconds) + return cached( + influxdb3_local, + _cache_key("shared:tables", database), + producer, + ttl_seconds=ttl_seconds, + ) return producer() def get_tag_names( - influxdb3_local, table: str, *, use_cache: bool = True, ttl_seconds: int = 3600 + influxdb3_local, + table: str, + *, + database: str | None = None, + use_cache: bool = True, + ttl_seconds: int = 3600, ) -> list[str]: """Return tag column names of a table (``Dictionary(Int32, Utf8)`` columns).""" @@ -50,14 +89,20 @@ def producer() -> list[str]: "SELECT column_name FROM information_schema.columns " "WHERE table_name = $table AND data_type = $data_type" ) - rows = influxdb3_local.query( - query, {"table": table, "data_type": _TAG_DATA_TYPE} + rows = _query( + influxdb3_local, + query, + {"table": table, "data_type": _TAG_DATA_TYPE}, + database=database, ) return [row["column_name"] for row in rows] if use_cache: return cached( - influxdb3_local, f"shared:tags:{table}", producer, ttl_seconds=ttl_seconds + influxdb3_local, + _cache_key(f"shared:tags:{table}", database), + producer, + ttl_seconds=ttl_seconds, ) return producer() @@ -67,6 +112,7 @@ def get_field_names( table: str, *, numeric_only: bool = False, + database: str | None = None, use_cache: bool = True, ttl_seconds: int = 3600, ) -> list[str]: @@ -80,7 +126,7 @@ def producer() -> list[str]: "SELECT column_name, data_type FROM information_schema.columns " "WHERE table_name = $table" ) - rows = influxdb3_local.query(query, {"table": table}) + rows = _query(influxdb3_local, query, {"table": table}, database=database) names: list[str] = [] for row in rows: name = row["column_name"] @@ -93,7 +139,7 @@ def producer() -> list[str]: return names if use_cache: - key = f"shared:fields:{table}:{int(numeric_only)}" + key = _cache_key(f"shared:fields:{table}:{int(numeric_only)}", database) return cached(influxdb3_local, key, producer, ttl_seconds=ttl_seconds) return producer() @@ -105,6 +151,7 @@ def query_window( start, end, columns: list[str] | None = None, + database: str | None = None, ) -> list[dict]: """Run a basic ``time``-window query and return rows (``[]`` if none). @@ -118,5 +165,10 @@ def query_window( f"SELECT {selected} FROM {_quote_identifier(table)} " "WHERE time >= $start AND time < $end ORDER BY time" ) - rows = influxdb3_local.query(query, {"start": start, "end": end}) - return rows or [] \ No newline at end of file + rows = _query( + influxdb3_local, + query, + {"start": start, "end": end}, + database=database, + ) + return rows or [] diff --git a/influxdata-plugin-utils/tests/test_introspection.py b/influxdata-plugin-utils/tests/test_introspection.py new file mode 100644 index 0000000..1b846a8 --- /dev/null +++ b/influxdata-plugin-utils/tests/test_introspection.py @@ -0,0 +1,131 @@ +"""Tests for influxdata_plugin_utils.introspection.""" + +from influxdata_plugin_utils.introspection import ( + get_field_names, + get_table_names, + get_tag_names, + query_window, +) + + +class FakeCache: + def __init__(self): + self.values = {} + self.ttls = {} + + def get(self, key): + return self.values.get(key) + + def put(self, key, value, ttl_seconds): + self.values[key] = value + self.ttls[key] = ttl_seconds + + +class FakeInfluxDB: + def __init__(self, responder): + self.cache = FakeCache() + self.calls = [] + self._responder = responder + + def query(self, query, args=None, *, database=None): + self.calls.append({"query": query, "args": args, "database": database}) + return self._responder(query, args, database) + + +def test_get_table_names_passes_database_and_separates_cache(): + def responder(query, args, database): + return [ + {"table_name": f"{database}_cpu", "table_type": "BASE TABLE"}, + {"table_name": f"{database}_view", "table_type": "VIEW"}, + ] + + influxdb3_local = FakeInfluxDB(responder) + + assert get_table_names(influxdb3_local, database="db_a") == ["db_a_cpu"] + assert get_table_names(influxdb3_local, database="db_a") == ["db_a_cpu"] + assert get_table_names(influxdb3_local, database="db_b") == ["db_b_cpu"] + + assert [call["database"] for call in influxdb3_local.calls] == ["db_a", "db_b"] + assert "shared:tables:database:db_a" in influxdb3_local.cache.values + assert "shared:tables:database:db_b" in influxdb3_local.cache.values + + +def test_default_database_keeps_existing_table_cache_key(): + influxdb3_local = FakeInfluxDB( + lambda query, args, database: [ + {"table_name": "cpu", "table_type": "BASE TABLE"} + ] + ) + + assert get_table_names(influxdb3_local) == ["cpu"] + + assert influxdb3_local.calls[0]["database"] is None + assert "shared:tables" in influxdb3_local.cache.values + + +def test_get_tag_names_passes_database_and_separates_cache(): + def responder(query, args, database): + assert args == { + "table": "cpu", + "data_type": "Dictionary(Int32, Utf8)", + } + return [{"column_name": f"{database}_host"}] + + influxdb3_local = FakeInfluxDB(responder) + + assert get_tag_names(influxdb3_local, "cpu", database="db_a") == ["db_a_host"] + assert get_tag_names(influxdb3_local, "cpu", database="db_b") == ["db_b_host"] + + assert [call["database"] for call in influxdb3_local.calls] == ["db_a", "db_b"] + assert "shared:tags:cpu:database:db_a" in influxdb3_local.cache.values + assert "shared:tags:cpu:database:db_b" in influxdb3_local.cache.values + + +def test_get_field_names_passes_database_and_separates_cache(): + def responder(query, args, database): + assert args == {"table": "cpu"} + return [ + {"column_name": "time", "data_type": "Timestamp"}, + {"column_name": "host", "data_type": "Dictionary(Int32, Utf8)"}, + {"column_name": f"{database}_usage", "data_type": "Float64"}, + {"column_name": f"{database}_state", "data_type": "Utf8"}, + ] + + influxdb3_local = FakeInfluxDB(responder) + + assert get_field_names(influxdb3_local, "cpu", database="db_a") == [ + "db_a_usage", + "db_a_state", + ] + assert get_field_names( + influxdb3_local, "cpu", numeric_only=True, database="db_b" + ) == ["db_b_usage"] + + assert [call["database"] for call in influxdb3_local.calls] == ["db_a", "db_b"] + assert "shared:fields:cpu:0:database:db_a" in influxdb3_local.cache.values + assert "shared:fields:cpu:1:database:db_b" in influxdb3_local.cache.values + + +def test_query_window_passes_database(): + def responder(query, args, database): + assert '"cpu"' in query + assert '"usage"' in query + assert args == { + "start": "2026-01-01T00:00:00Z", + "end": "2026-01-02T00:00:00Z", + } + assert database == "db_a" + return [{"usage": 1.2}] + + influxdb3_local = FakeInfluxDB(responder) + + assert query_window( + influxdb3_local, + "cpu", + start="2026-01-01T00:00:00Z", + end="2026-01-02T00:00:00Z", + columns=["usage"], + database="db_a", + ) == [{"usage": 1.2}] + + assert influxdb3_local.calls[0]["database"] == "db_a"