From 9b23842aa0fb3bf582ea376bf5213f16f9b49baa Mon Sep 17 00:00:00 2001 From: Chris Munns Date: Mon, 2 Nov 2026 17:48:58 -0500 Subject: [PATCH 1/2] Collect column expressions and quote largest_tables names Ports dev #210 and #213. Each PostgreSQL column records default_expression and generated_expression. largest_tables quotes each relation name with format('%I.%I'), and partition_analysis no longer lists index partitions. --- AGENTS.md | 6 ++ README.md | 2 + docs/workload-bundle.md | 2 +- docs/workload_capture.md | 5 ++ .../analyzers/performance_analyzer.py | 10 +-- .../database/analyzers/schema_analyzer.py | 20 +++--- planetscale_discovery/workload/bundle.py | 2 +- tests/unit/test_performance_analyzer.py | 16 +++++ tests/unit/test_schema_analyzer.py | 62 +++++++++++++++++-- 9 files changed, 107 insertions(+), 18 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index d6145fe..ed837e7 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -184,6 +184,12 @@ statistics, infrastructure topology, and user and role names. password is used for the connection and is never written to the output. Literal values from queries are not collected either, unless you turn on `capture_log`. +Schema metadata includes the text of definitions written into the schema: view +and function bodies, CHECK constraints, index, policy and partition +expressions, and each column's default and generated expression. A definition +can hold a literal its author wrote, such as `DEFAULT 'pending'`, but never a +value read from a row. + Where the tool reports a statement, it replaces every literal value with a placeholder and removes comments. It records the shape of a query, not the data in it. diff --git a/README.md b/README.md index 6eb4fd0..faab0ec 100644 --- a/README.md +++ b/README.md @@ -387,6 +387,8 @@ workload capture also collects the values described below. **What is NOT collected:** Table contents, customer records, application code, or credentials. Passwords are used only for connection and are never stored in output. Literal values from queries are not collected either, unless you turn on `capture_log`. +**Schema definitions.** Schema metadata includes the text of definitions written into your schema: view and function bodies, CHECK constraints, index, policy and partition expressions, and each column's default and generated expression. A definition can hold a literal that its author wrote, such as `DEFAULT 'pending'`. None of it is a value read from a row. The migration team needs this text to hand your schema back to you with a sharding design applied. + **Query text.** Where the tool reports a statement — a long-running transaction, or the two sides of a lock — it records the statement with literal values replaced by `$N` placeholders, so `SELECT * FROM orders WHERE email = $1` rather than the address itself. Comments are removed. This captures the shape of a query rather than the data in it. Review the report before sharing it outside your organization. **A workload capture records values from your data in two files.** The bundle README names each one: diff --git a/docs/workload-bundle.md b/docs/workload-bundle.md index 52afa69..3d20013 100644 --- a/docs/workload-bundle.md +++ b/docs/workload-bundle.md @@ -36,7 +36,7 @@ a customer. | --- | --- | | `workload.sql` | the query log: one statement per record, `;`-terminated, with `-- neki:` metric headers | | `schema.sql` | sequences, tables, indexes, constraints, views, and foreign keys last | -| `plantest_counts.json` | row count per table, plus NDV and null fraction for indexed and constrained columns | +| `plantest_counts.json` | row count per table, plus NDV and null fraction for the indexed and constrained columns the capture role may `SELECT` ([why](#column-statistics)) | | `plantest_counts-card.json` | identical bytes, under the second name the planning tools open | | `manifest.json` | when the capture ran, and over which intervals | | `column_stats.json` | per-column value distribution, for judging a shard key | diff --git a/docs/workload_capture.md b/docs/workload_capture.md index 607bbdb..320d17c 100644 --- a/docs/workload_capture.md +++ b/docs/workload_capture.md @@ -509,6 +509,11 @@ the counts and frequencies. For the columns you declare in [step 6](#6-declare-the-columns-your-migration-engineer-names), it also reads `most_common_vals`, the most common values themselves. +PostgreSQL shows a column in `pg_stats` only to a role that may `SELECT` it. +The capture role holds `pg_monitor` and nothing else, so the counts cover only +the columns you grant it, which in practice are the declared ones. See +[Grant the capture role each column](#grant-the-capture-role-each-column). + It also calls `pg_total_relation_size`, `pg_relation_size` and `pg_indexes_size` on each table and index in scope. These report how many bytes an object occupies on disk. They read no row, take no lock, and need no diff --git a/planetscale_discovery/database/analyzers/performance_analyzer.py b/planetscale_discovery/database/analyzers/performance_analyzer.py index c28f280..b3cb70e 100644 --- a/planetscale_discovery/database/analyzers/performance_analyzer.py +++ b/planetscale_discovery/database/analyzers/performance_analyzer.py @@ -403,13 +403,13 @@ def _get_resource_utilization(self) -> Dict[str, Any]: SELECT schemaname, tablename, - pg_total_relation_size(schemaname||'.'||tablename) as total_size_bytes, - pg_size_pretty(pg_total_relation_size(schemaname||'.'||tablename)) as total_size_human, - pg_relation_size(schemaname||'.'||tablename) as table_size_bytes, - pg_size_pretty(pg_relation_size(schemaname||'.'||tablename)) as table_size_human + pg_total_relation_size(format('%I.%I', schemaname, tablename)) as total_size_bytes, + pg_size_pretty(pg_total_relation_size(format('%I.%I', schemaname, tablename))) as total_size_human, + pg_relation_size(format('%I.%I', schemaname, tablename)) as table_size_bytes, + pg_size_pretty(pg_relation_size(format('%I.%I', schemaname, tablename))) as table_size_human FROM pg_tables WHERE schemaname NOT IN ('information_schema', 'pg_catalog') - ORDER BY pg_total_relation_size(schemaname||'.'||tablename) DESC + ORDER BY pg_total_relation_size(format('%I.%I', schemaname, tablename)) DESC LIMIT 20 """) largest_tables = [dict(row) for row in cursor.fetchall()] diff --git a/planetscale_discovery/database/analyzers/schema_analyzer.py b/planetscale_discovery/database/analyzers/schema_analyzer.py index 39ea41f..960b6b2 100644 --- a/planetscale_discovery/database/analyzers/schema_analyzer.py +++ b/planetscale_discovery/database/analyzers/schema_analyzer.py @@ -274,9 +274,8 @@ def _get_table_analysis(self) -> List[Dict[str, Any]]: identity_col = ( "a.attidentity" if has_declarative_partitioning else "''" ) + " as identity_kind" - generated_col = ( - "a.attgenerated" if version_num >= 120000 else "''" - ) + " as generated_kind" + generated_attr = "a.attgenerated" if version_num >= 120000 else "''" + generated_col = generated_attr + " as generated_kind" with self.connection.cursor() as cursor: # Build query based on version @@ -323,8 +322,9 @@ def _get_table_analysis(self) -> List[Dict[str, Any]]: # Get column information. # default_sequence is the sequence that the column default # calls when the default is one nextval() call and nothing - # else. The default text stays in the database: the query - # selects only the sequence name, from pg_depend. + # else. A generated column keeps its expression in the + # default slot, so the expression goes to generated_expression + # and never to default_expression. try: with self.connection.cursor() as detail_cursor: detail_cursor.execute( @@ -363,7 +363,13 @@ def _get_table_analysis(self) -> List[Dict[str, Any]]: AND pg_get_expr(d.adbin, d.adrelid) ~ '^nextval\\(''(?:[^'']|'''')+''::regclass\\)$' LIMIT 1 - ) as default_sequence + ) as default_sequence, + CASE WHEN {generated_attr} = '' + THEN pg_get_expr(d.adbin, d.adrelid) + END as default_expression, + CASE WHEN {generated_attr} <> '' + THEN pg_get_expr(d.adbin, d.adrelid) + END as generated_expression FROM pg_attribute a LEFT JOIN pg_attrdef d ON (a.attrelid = d.adrelid AND a.attnum = d.adnum) WHERE a.attrelid = %s @@ -732,7 +738,7 @@ def _get_partition_analysis(self) -> List[Dict[str, Any]]: pg_size_pretty(pg_relation_size(child_class.oid)) as partition_size_human, child_class.reltuples::bigint as estimated_rows FROM pg_inherits i - JOIN pg_class parent_class ON i.inhparent = parent_class.oid + JOIN pg_class parent_class ON i.inhparent = parent_class.oid AND parent_class.relkind = 'p' JOIN pg_namespace parent_ns ON parent_class.relnamespace = parent_ns.oid JOIN pg_class child_class ON i.inhrelid = child_class.oid JOIN pg_namespace child_ns ON child_class.relnamespace = child_ns.oid diff --git a/planetscale_discovery/workload/bundle.py b/planetscale_discovery/workload/bundle.py index eda07bc..74aef4f 100644 --- a/planetscale_discovery/workload/bundle.py +++ b/planetscale_discovery/workload/bundle.py @@ -1119,7 +1119,7 @@ def _render_readme(summary: Dict[str, Any], out_dir: Path) -> str: | --- | --- | | `workload.sql` | the queries, each with how often it ran and how long it took | | `schema.sql` | your tables, indexes and views, as `CREATE` statements | -| `plantest_counts.json` | how many rows each table holds, plus NDV and null fraction for indexed and constrained columns | +| `plantest_counts.json` | how many rows each table holds, plus NDV and null fraction for the indexed and constrained columns the capture role may `SELECT` | | `plantest_counts-card.json` | the same content, under a second name the planning tools look for | | `manifest.json` | when this capture ran, and over which intervals | | `column_stats.json` | how evenly the values in each column are spread | diff --git a/tests/unit/test_performance_analyzer.py b/tests/unit/test_performance_analyzer.py index 385192e..e8e647a 100644 --- a/tests/unit/test_performance_analyzer.py +++ b/tests/unit/test_performance_analyzer.py @@ -466,6 +466,22 @@ def test_resource_utilization_success(self, analyzer, mock_connection): assert len(result["largest_tables"]) == 1 assert result["temp_file_usage"] == [] + def test_resource_utilization_quotes_table_names(self, analyzer, mock_connection): + """A table name that needs quoting, like "User", must not fail the + section: an unquoted name folds to user, which does not exist.""" + _, cursor = mock_connection + cursor.fetchall.side_effect = [[], [], []] + + analyzer._get_resource_utilization() + + largest_sql = next( + c.args[0] + for c in cursor.execute.call_args_list + if "FROM pg_tables" in c.args[0] + ) + assert "format('%I.%I', schemaname, tablename)" in largest_sql + assert "schemaname||'.'||tablename" not in largest_sql + def test_resource_utilization_error(self, analyzer, mock_connection): """Test resource utilization graceful degradation on error.""" _, cursor = mock_connection diff --git a/tests/unit/test_schema_analyzer.py b/tests/unit/test_schema_analyzer.py index 03f4a34..654734f 100644 --- a/tests/unit/test_schema_analyzer.py +++ b/tests/unit/test_schema_analyzer.py @@ -840,7 +840,7 @@ def test_get_table_analysis_column_identity_and_generated( def test_get_table_analysis_column_default_sequence( self, analyzer, mock_connection ): - """Column query names the sequence a nextval() default calls, never the default text""" + """Column query names the sequence a nextval() default calls""" _, cursor = mock_connection cursor.fetchall.return_value = [ {"table_oid": 16385, "table_name": "users", "table_type": "r"} @@ -860,12 +860,45 @@ def test_get_table_analysis_column_default_sequence( assert "dep.objid = d.oid" in column_sql # Only a default that is one nextval() call names a sequence. assert r"~ '^nextval\(''(?:[^'']|'''')+''::regclass\)$'" in column_sql - # The default text appears only in that filter. It is never a - # selected column, so it cannot reach the report. - assert column_sql.count("pg_get_expr") == 1 assert "AND pg_get_expr(d.adbin, d.adrelid)" in column_sql assert "adsrc" not in column_sql + @pytest.mark.parametrize( + "version_num, generated_attr", + [(180000, "a.attgenerated"), (120000, "a.attgenerated"), (110000, "''")], + ) + def test_get_table_analysis_column_default_and_generated_expressions( + self, analyzer, mock_connection, version_num, generated_attr + ): + """A generated column's expression never reads as a plain DEFAULT, and + a server before PostgreSQL 12 is never asked for attgenerated""" + _, cursor = mock_connection + cursor.fetchall.return_value = [ + {"table_oid": 16385, "table_name": "users", "table_type": "r"} + ] + + with patch.object( + analyzer, "_get_server_version_num", return_value=version_num + ): + analyzer._get_table_analysis() + + column_sql = " ".join( + next( + c.args[0] + for c in cursor.execute.call_args_list + if "FROM pg_attribute a" in c.args[0] + ).split() + ) + assert ( + f"CASE WHEN {generated_attr} = '' " + "THEN pg_get_expr(d.adbin, d.adrelid) END as default_expression" + ) in column_sql + # Any generated kind: 's' (stored) and, from PostgreSQL 18, 'v' (virtual). + assert ( + f"CASE WHEN {generated_attr} <> '' " + "THEN pg_get_expr(d.adbin, d.adrelid) END as generated_expression" + ) in column_sql + # --------------------------------------------------------------- # _get_index_analysis() # --------------------------------------------------------------- @@ -1344,6 +1377,27 @@ def cursor_factory(): assert result[0]["parent_table"] == "events" assert result[0]["partition_name"] == "events_2024_01" + def test_get_partition_analysis_excludes_index_partitions( + self, analyzer, mock_connection + ): + """An index on a partitioned table has partitions of its own, with + relispartition true. Only children of a partitioned table count.""" + _, cursor = mock_connection + cursor.fetchall.return_value = [] + + with patch.object(analyzer, "_get_server_version_num", return_value=170000): + analyzer._get_partition_analysis() + + partition_sql = next( + c.args[0] + for c in cursor.execute.call_args_list + if "child_class.relispartition = true" in c.args[0] + ) + assert ( + "JOIN pg_class parent_class ON i.inhparent = parent_class.oid " + "AND parent_class.relkind = 'p'" + ) in partition_sql + def test_get_partition_analysis_pre_pg10(self, mock_connection): """Test partition analysis returns empty on pre-PG10""" connection, _ = mock_connection From a74046d661236f4f63d25694c1c98ebdca178f64 Mon Sep 17 00:00:00 2001 From: Chris Munns Date: Mon, 2 Nov 2026 17:52:27 -0500 Subject: [PATCH 2/2] Add initial Citus support to database discovery Ports dev #211. When the citus extension is installed, discovery reports each table's citus.table_type and fills the existing table fields with logical values read from every shard placement. A new citus section reports nodes, colocation groups, placements and storage. The Citus test fixture FakeConnection gains a closed attribute, which this repo's _collect_schema reads. --- AGENTS.md | 12 + README.md | 19 + docs/citus.md | 195 ++ docs/cleanup.md | 12 + docs/data_size_analysis.md | 9 + docs/output-format.md | 85 +- docs/performance_considerations.md | 29 + docs/providers/azure.md | 2 +- docs/troubleshooting.md | 33 + planetscale_discovery/cli.py | 32 +- .../database/analyzers/citus_analyzer.py | 219 +++ .../database/analyzers/data_size_analyzer.py | 23 + .../database/analyzers/feature_analyzer.py | 25 +- .../analyzers/performance_analyzer.py | 173 +- .../database/analyzers/schema_analyzer.py | 175 +- .../database/analyzers/security_analyzer.py | 18 +- planetscale_discovery/database/citus.py | 1613 +++++++++++++++++ planetscale_discovery/database/discovery.py | 60 +- test-data/citus-coverage.sql | 669 +++++++ tests/fixtures/database_responses.py | 512 ++++++ tests/integration/test_citus_discovery.py | 902 +++++++++ tests/unit/test_citus_analyzer.py | 176 ++ tests/unit/test_citus_detection.py | 648 +++++++ tests/unit/test_citus_shard_statistics.py | 1456 +++++++++++++++ tests/unit/test_cli.py | 142 ++ tests/unit/test_data_size_analyzer.py | 51 + tests/unit/test_database_discovery.py | 122 ++ tests/unit/test_feature_analyzer.py | 34 + tests/unit/test_schema_analyzer.py | 63 + tests/unit/test_security_analyzer.py | 30 + tests/unit/test_workload_citus_insulation.py | 37 + 31 files changed, 7539 insertions(+), 37 deletions(-) create mode 100644 docs/citus.md create mode 100644 planetscale_discovery/database/analyzers/citus_analyzer.py create mode 100644 planetscale_discovery/database/citus.py create mode 100644 test-data/citus-coverage.sql create mode 100644 tests/integration/test_citus_discovery.py create mode 100644 tests/unit/test_citus_analyzer.py create mode 100644 tests/unit/test_citus_detection.py create mode 100644 tests/unit/test_citus_shard_statistics.py create mode 100644 tests/unit/test_workload_citus_insulation.py diff --git a/AGENTS.md b/AGENTS.md index ed837e7..bd0f238 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -107,6 +107,7 @@ send it. | Database engine (`postgres` or `mysql`) | `engine:` in `config.yaml` | [MySQL Setup](docs/mysql.md) for MySQL/Vitess | | Cloud or hosting provider | `providers:` in `config.yaml` | [Provider guides](docs/providers/) | | Large column and LOB analysis | optional module | [Data Size Analysis](docs/data_size_analysis.md) | +| Citus extension | detected automatically, nothing to set. With Citus tables, every active primary node must be reachable during the run | [Citus Extension Support](docs/citus.md) | Supported providers: AWS, GCP, Azure, Supabase, Heroku, Neon, PlanetScale, and Snowflake. @@ -214,6 +215,16 @@ All analysis runs locally. The tool sends nothing to an external service. Permission errors are expected on a managed service, and the tool continues with what it can read. A warning in the report is normal and is not a failed run. +A database with the Citus extension is the exception. A missing Citus +privilege, a connection to a worker instead of the coordinator, and, on a +database with Citus tables, a node or shard that is unavailable each stop +database discovery with the message `Citus discovery stopped`. That is a failed +run, not a warning. Database discovery alone exits with code 1. A combined run +with a cloud provider configured exits with code 0 and writes the report with +`database_results` set to `null`, so check `database_results` rather than the +exit code. The remedy for each cause is in +[Troubleshooting](docs/troubleshooting.md#citus-discovery-stopped). + Read [Troubleshooting](docs/troubleshooting.md) first. It covers Python installation, common errors, and managed database environments. @@ -241,6 +252,7 @@ superuser. Report the missing privilege instead. ## More documentation +- [Citus Extension Support](docs/citus.md) - discovery on databases with the Citus extension - [Advanced Usage](docs/advanced-usage.md) - CLI reference, focused analysis, automation - [Output Format](docs/output-format.md) - report structure - [Performance Considerations](docs/performance_considerations.md) - impact and timing diff --git a/README.md b/README.md index faab0ec..83364ac 100644 --- a/README.md +++ b/README.md @@ -88,6 +88,13 @@ The discovery tool supports two database engines. Set the engine in your config | **PostgreSQL** (default) | `engine: postgres` / `--engine postgres` | config, schema, performance, security, features, data_size | See [PostgreSQL Privileges](#required-postgresql-privileges) below | | **MySQL/Vitess** | `engine: mysql` / `--engine mysql` | config, schema, performance, replication, security, features | [MySQL Setup Guide](docs/mysql.md) | +PostgreSQL databases that use the Citus extension have initial, limited +support. No setting is needed: when the `citus` extension is installed, it is +detected, and regular tables are discovered next to the tables that Citus +manages. When the database holds Citus tables, every active primary node of the +cluster must be reachable during the run. See +[Citus Extension Support](docs/citus.md). + ### Quick Examples ```bash @@ -151,6 +158,17 @@ GRANT pg_read_all_settings TO planetscale_discovery; GRANT pg_monitor TO planetscale_discovery; ``` +On a database with the Citus extension, run the SQL above on the coordinator. +No additional SQL is needed by default. With +`citus.enable_create_role_propagation` on, which is the default, the role is +created on every node. The discovery role also reads the `pg_dist_*` metadata +tables, and on a database with Citus tables it calls +`run_command_on_all_nodes()`. Both are granted to `PUBLIC` by default. A missing +`SELECT` on a `pg_dist_*` table stops the run. On a database with Citus tables, +so does a missing `EXECUTE` on `run_command_on_all_nodes()`, or a node that does +not accept the discovery role. See +[Citus Extension Support](docs/citus.md#requirements). + Alternatively, you can use an existing superuser account for complete analysis. Query workload capture needs `pg_monitor` and a readable `pg_stat_statements`. @@ -403,6 +421,7 @@ All analysis runs locally — no data is sent to external services. ## Additional Documentation - [MySQL Setup Guide](docs/mysql.md) - MySQL/Vitess-specific setup, privileges, and PlanetScale configuration +- [Citus Extension Support](docs/citus.md) - Discovery on PostgreSQL databases that use the Citus extension - [Output Format](docs/output-format.md) - JSON report structure and markdown summary details - [Troubleshooting](docs/troubleshooting.md) - Python installation, common errors, managed database environments - [Advanced Usage](docs/advanced-usage.md) - CLI reference, focused analysis, automation, pipx installation diff --git a/docs/citus.md b/docs/citus.md new file mode 100644 index 0000000..c76756e --- /dev/null +++ b/docs/citus.md @@ -0,0 +1,195 @@ +# Citus Extension Support + +Discovery has initial, limited support for PostgreSQL databases that use the +Citus extension. This page describes what is collected from such a database, +how the values are computed, and what the cluster must provide for a run to +succeed. + +## Scope + +Support covers database discovery only. No setting is needed: when the +`citus` extension is installed in the database you connect to, it is detected, +and additional Citus analysis is run alongside the regular PostgreSQL analysis. +Regular tables are discovered next to the tables that Citus manages, in the same +report. + +A *Citus table* is a table recorded in Citus's metadata. For a Citus table, the +coordinator stores only an empty *shell*. The rows, the column statistics and +the activity counters are stored on the table's *shards*, and each shard is +stored on one or more nodes as a *placement*. Values read from the shell alone +do not describe the table. A 154 GB distributed table, for example, reports 0 +rows and 8 kB on the coordinator. + +## What is collected + +For every table, `citus.table_type` is reported with the label that Citus itself +uses: + +| `table_type` | Meaning | +| --- | --- | +| `distributed` | Rows are spread across shards by a distribution column. A table with one shard and no distribution column also has this type. | +| `reference` | One shard, with a placement on every node. | +| `local` | A Citus local table, whose data stays on the coordinator. | +| `schema` | A table in a distributed schema (schema-based sharding). | +| `regular` | A PostgreSQL table that Citus does not manage. | + +For Citus tables, the following are also collected: + +- The distribution column, the colocation group, the shard count, the + replication model and the number of placements per shard. +- Row estimates, table and index sizes, column statistics, and table and index + activity counters, read from every shard placement. +- The cluster's nodes, the database size on each node, the colocation groups, + the placements per node, distributed schemas and distributed functions. These + are reported under `analysis_results.citus`. + +No table rows are read. Every value comes from catalogs and statistics views, +read on the coordinator, and on each placement through +`run_command_on_all_nodes()`. Most common values and histograms are never read. + +Shards are hidden from discovery's own session, whatever +`citus.show_shards_for_app_name_prefixes` is set to for the discovery role, so +a shard placement on the coordinator is never reported as a table. The setting +is changed for that session only. + +## How values are computed + +The existing fields of the report keep their names, types and positions. For a +Citus table, each existing field holds the logical value for the table as a +whole, so a reader of the report needs no change. The coordinator's own values +are kept under `citus.coordinator`, and the detail for each shard under +`citus.shard_statistics`. See [Output Format](output-format.md#citus-tables). + +Values are combined per shard, following how Citus runs queries. A read runs on +one placement, and a write is applied to every placement: + +- **Row estimates, sizes and live tuples** are summed across shards, using one + placement per shard, the one analyzed most recently. A reference table on + three nodes is therefore counted once, not three times. +- **Write counters** (`n_tup_ins`, `n_tup_upd`, `n_tup_del`) are summed across + shards, using each shard's largest placement value, so counters reset on one + node do not lower the total. +- **Read counters** (`seq_scan`, `idx_scan` and related counters) are summed + over every placement. +- **`null_frac`** is a mean weighted by each shard's row estimate. The + combination adds no error of its own, so the result is as accurate as the + shards' statistics. +- **`correlation`** is a mean weighted by each shard's row estimate. The query + planner uses each shard's own correlation, so the value describes how ordered + a typical shard is. +- **`n_distinct`** keeps PostgreSQL's convention: a positive value is a count, + and a negative value is a fraction of the rows. The method used is recorded in + `n_distinct_method`. Each shard's value is itself an estimate from ANALYZE, + so the Result column describes the combination only: + + | Method | Used for | Result | + | --- | --- | --- | + | `single_placement` | A table with one shard | The shard's value, unchanged | + | `distribution_column_sum` | The distribution column of a hash or range distributed table, whose values are each stored in one shard | A sum, which adds no error | + | `weighted_fraction` | Every shard reports a fraction | An estimate, which assumes the fraction holds across shards | + | `max_absolute` | Every shard reports a count | A lower bound | + | `mixed_max_absolute` | Shards disagree | A lower bound | + + A combined count above 10% of the table's rows is stored as a negative + fraction, as ANALYZE does. + +A shard with no statistics, because it was never analyzed, is counted under +`shards_never_analyzed` and left out of the sums. A table whose every shard is +unanalyzed reports `estimated_rows: -1`, the same value PostgreSQL reports for +an unanalyzed regular table. + +The "Total Size" in the summary is the logical size: regular tables plus each +Citus table counted once. A shard that a move left on the coordinator, until +Citus's deferred cleanup drops it, is not counted. The summary also reports +cluster storage, the sum of the database size on every node, which includes +every placement and any shard not yet cleaned up. + +## Requirements + +A Citus run either reads every shard or stops. A partial result is never +written for the database. Before any analysis runs, the following are checked: + +- **Connect to the coordinator.** A connection to a worker stops the run. +- **Every active primary node must be reachable**, whether or not it holds a + placement, and must accept a connection from the discovery role, because + `run_command_on_all_nodes()` connects to each of them as that role. Inactive + nodes and secondaries are not contacted. A node added during the run stops + it, because its shards would be read under the wrong node. +- **Every shard must have an active placement on an active primary node.** A + shard whose placements are only on inactive nodes or secondaries stops the + run. +- **Every worker that holds placements must have the Citus metadata.** A + worker recognizes its shards from its own copy of the metadata. Citus 11 and + later sync it to every node, but a cluster upgraded from Citus 10 has none on + its workers until `citus_finalize_upgrade_to_citus11()` is run. +- **The discovery role needs the Citus privileges.** It needs `EXECUTE` on + `run_command_on_all_nodes`, `relation_is_a_known_shard`, `shard_name`, + `column_to_column_name` and `citus_version`, and `SELECT` on the `pg_dist_*` + metadata tables. All of them are granted to `PUBLIC` by default. +- **`citus.use_secondary_nodes` must be `never`** for the discovery role, which + is the default. + +The node checks, the `run_command_on_all_nodes` privilege and the +`citus.use_secondary_nodes` setting apply only when the database holds at least +one Citus table. A database where the extension is installed but no table is +distributed contacts no node, so it cannot stop the run for those reasons. The +coordinator check and the metadata privileges always apply. Without them, a +database with no Citus table cannot be told apart from one whose Citus tables +cannot be read. + +During the shard read, every valid index of a Citus table must also be found on +every shard of that table. A valid index missing on a shard stops the run, +because its size would be partial and its scan counts would drop it from the +list of unused indexes. There are two usual causes, and the message gives the +remedy for each: + +1. On a partition's index, created from its parent's index, the shard indexes + can lack their shard id after PostgreSQL named them automatically. `SELECT + fix_all_partition_shard_index_names();` on the coordinator renames them. +2. An index created with `citus.enable_ddl_propagation` off exists only on the + coordinator. Recreate it with DDL propagation on, or drop it. + +An invalid index never stops the run. A failed `CREATE INDEX CONCURRENTLY` +leaves the coordinator's index invalid, with shard indexes on only some shards. +A shard index that the build completed can be valid on its node and take +scans, whatever the coordinator records. An invalid index therefore reports +the values summed from the shards where it was found, and is flagged with +`citus.reason: "invalid_index"`. Compare `citus.shards_read` with +`citus.shard_count` to see how many shards hold it. + +When any check fails, database discovery stops, no database result is written, +and the message names each node or privilege at fault. Database discovery alone +then exits with code 1. A combined run with a cloud provider continues with +cloud discovery, writes the report with `database_results` set to `null`, and +exits with code 0, as for any other database failure. Unlike other permission errors, which are +reported as warnings, a missing Citus privilege stops the run: without it, +every Citus table would be reported with its shell's values. See +[Troubleshooting](troubleshooting.md#citus-discovery-stopped). + +Run the privilege SQL from the [README](../README.md#required-postgresql-privileges) +on the coordinator. With `citus.enable_create_role_propagation` on, which is the +default, `CREATE ROLE` and `DROP ROLE` on the coordinator are propagated to +every node. + +Column statistics follow the same rule as on a regular table: `pg_stats` +returns a column only to a role that may `SELECT` it, and on a node that rule +applies to the shard. If `statistics` is empty for a Citus table whose shards +were analyzed, confirm that the discovery role's `SELECT` grant reached the +shards. + +## Limitations + +- **Workload capture** reads only the coordinator's catalogs. On a Citus + database, its row counts and column statistics do not describe Citus tables. +- **Secondary nodes** are not read. +- **Foreign tables** added to Citus's metadata with + `citus_add_local_table_to_metadata()` are reported as foreign tables, as on + regular PostgreSQL. Their rows are stored on their own server, so they have + no shard values to read, and they are not counted among the Citus tables. +- **Data size analysis** skips Citus tables, with the reason `citus_table`. + Reading their rows would scan every shard in full. See + [Data Size Analysis](data_size_analysis.md). +- **Statement statistics** in the performance report come from the + coordinator's `pg_stat_statements` only. +- The support was tested with Citus 14.1 on PostgreSQL 18. A Citus version that + lacks a function the read needs is reported as unsupported. diff --git a/docs/cleanup.md b/docs/cleanup.md index 80c0da8..fc653ce 100644 --- a/docs/cleanup.md +++ b/docs/cleanup.md @@ -18,6 +18,18 @@ DROP USER IF EXISTS planetscale_discovery; -- DROP USER planetscale_discovery; ``` +On a database with the Citus extension, run the cleanup on the coordinator. +With `citus.enable_create_role_propagation` on, which is the default, the +dropped role is propagated to every node. Then confirm on every node that the +role is gone. Each row of the result should show `0`: + +```sql +SELECT nodeid, success, result +FROM run_command_on_all_nodes( + $$SELECT count(*) FROM pg_roles WHERE rolname = 'planetscale_discovery'$$ +); +``` + ## Verify No Temporary Objects Remain Check that no temporary objects were left behind: diff --git a/docs/data_size_analysis.md b/docs/data_size_analysis.md index a4fb26a..aa9d649 100644 --- a/docs/data_size_analysis.md +++ b/docs/data_size_analysis.md @@ -152,6 +152,15 @@ FROM public.emails TABLESAMPLE BERNOULLI (10); - **Statistical representation** of data - Only applied to tables with > 1,000 rows +### Citus Tables + +On a database with the Citus extension, Citus tables are not analyzed. Each one +is listed under `tables_skipped` with the reason `citus_table`, and a warning is +recorded. The coordinator holds only a Citus table's empty shell, so the shell's +size and row estimate would pass both safeguards above, and every column query +would scan every shard in full. Regular tables on the same database are +analyzed as usual. See [Citus Extension Support](citus.md). + ### Performance Example For a table with 500,000 rows: diff --git a/docs/output-format.md b/docs/output-format.md index 0e113b5..222428c 100644 --- a/docs/output-format.md +++ b/docs/output-format.md @@ -23,7 +23,8 @@ The `` is the discovery start time, for example "schema": { ... }, "performance": { ... }, "security": { ... }, - "features": { ... } + "features": { ... }, + "citus": { ... } } }, "cloud_results": { ... }, @@ -31,6 +32,86 @@ The `` is the discovery start time, for example } ``` +`analysis_results.citus` appears only when the database has the Citus +extension. See [Citus tables](#citus-tables). + +### Citus tables + +On a database with the Citus extension, the existing fields keep their names, +types and positions, and a `citus` object is added to the rows that describe +tables. On a database without Citus, no `citus` key appears anywhere. See +[Citus Extension Support](citus.md) for how each value is computed. + +Every entry in `schema.table_analysis` receives `citus.table_type`: one of +`distributed`, `reference`, `local`, `schema` or `regular`. For a regular +table, nothing else changes. For a Citus table, the existing fields hold the +logical value for the table as a whole, read from its shards: + +| Section | Fields that hold logical values | +| --- | --- | +| `schema.table_analysis` | `estimated_rows`, `table_size_bytes`, `total_size_bytes`, their `_human` forms, `statistics` | +| `schema.index_analysis` | `estimated_rows`, `index_size_bytes`, `index_size_human`, `usage_stats` | +| `schema.partition_analysis` | `partition_size_bytes`, `partition_size_human`, `estimated_rows` | +| `schema.database_catalog` (current database) | `size_bytes`, `size_human` | +| `performance.table_statistics` | Every counter and `last_*` time | +| `performance.index_usage` | `idx_scan`, `idx_tup_read`, `idx_tup_fetch`, `index_size_bytes`, `index_size_human`, and the unused index lists computed from them | +| `performance.resource_utilization` | `largest_tables`, ranked by logical size, and `database_sizes` for the current database | + +The `citus` object on those rows holds `source`, always `"shards"`, and +`coordinator`, the values the coordinator reported for the same fields. The +other keys depend on the section: + +- `schema.table_analysis`: `table_type`, `distribution_column`, + `colocation_id`, `shard_count`, `partmethod`, `replication_model`, + `autoconverted`, `access_method`, `placements_per_shard` (`min` and `max`), + and `shard_statistics`. `shard_statistics` holds `shards_read`, + `shards_never_analyzed`, `inactive_placements`, `placement_size_bytes`, and a + `columns` list. Each `columns` entry holds `attname`, `inherited`, + `n_distinct_method`, and a `shards` list with each shard's `reltuples`, + `n_distinct`, `null_frac` and `correlation`. +- `schema.index_analysis` and `performance.index_usage`: `shards_read`. +- `performance.table_statistics`: `table_type`, `shards_read`, and + `shards_without_event`, which counts, for each `last_*` time, the shards that + have not recorded that event. +- `performance.resource_utilization.largest_tables`: no other key. A Citus + table that the coordinator's own ranking left out, because its shell is + small, has `coordinator: null`. +- `schema.partition_analysis` and the current database's row: no other key. + +An invalid index of a Citus table, such as one left by a failed `CREATE INDEX +CONCURRENTLY`, can exist on only some shards. In `schema.index_analysis` and +`performance.index_usage`, its fields hold the values summed from the shards +where it was found, and its `citus` object also holds `reason: +"invalid_index"` and `shard_count`. A `shards_read` below `shard_count` means +the index is missing on the other shards. + +`n_distinct_method` is one of `single_placement`, `distribution_column_sum`, +`weighted_fraction`, `max_absolute` or `mixed_max_absolute`. The last two give +a lower bound rather than an exact count. + +`analysis_results.citus` reports the cluster: + +| Key | Contents | +| --- | --- | +| `version` | `extension_version`, `citus_version`, `node_role` | +| `nodes` | Each `pg_dist_node` entry, with `database_size_bytes` | +| `colocation_groups` | Each colocation group, with its shard count, replication factor, distribution column type and tables | +| `tables` | Each Citus table, with its type, distribution column, `estimated_rows`, `logical_size_bytes` (one placement per shard) and `placement_size_bytes` (every placement) | +| `shard_placement` | Placement counts and bytes `by_node` and `by_table` | +| `distributed_schemas` | Each distributed schema and its colocation id | +| `distributed_functions` | Each distributed function, with its distribution argument and colocation id | +| `summary` | `node_count`, `table_counts` by type (including `regular`), `citus_table_logical_size_bytes`, `cluster_storage_bytes` | + +On a database with the extension installed and no Citus table, no node is +contacted, so each node's `database_size_bytes` and `cluster_storage_bytes` are +`null`. + +`database_results.connection_info.citus` describes the connection only: the +extension version, the Citus version, and the node role (`coordinator` or +`worker`). Node and table counts are reported once, in +`analysis_results.citus.summary`, which `summary.database_summary.citus` copies +unchanged. There, `table_counts` holds every table type, `regular` included. + ### Cloud results `cloud_results.providers` holds one entry per provider that ran, keyed by @@ -43,7 +124,7 @@ regions analyzed, per-region resources, a summary, and complexity factors. When `--local-summary` is used, the markdown file includes: 1. **Overview**: Which modules ran and their status -2. **Database Summary**: PostgreSQL version, table count, size, extensions, users +2. **Database Summary**: PostgreSQL version, table count, size, extensions, users, and on a Citus database the Citus version, node count, tables by type and cluster storage 3. **Cloud Summary**: Per-provider instance and cluster counts 4. **Information Not Collected**: Any gaps due to permissions or missing extensions 5. **Errors**: Any errors encountered during discovery diff --git a/docs/performance_considerations.md b/docs/performance_considerations.md index ed2b8b8..eb29460 100644 --- a/docs/performance_considerations.md +++ b/docs/performance_considerations.md @@ -13,6 +13,7 @@ This document outlines the performance characteristics of each analyzer and prov | PerformanceAnalyzer | ✅ Low | Yes | Only reads `pg_stat_*` views | | SecurityAnalyzer | ✅ Low | Yes | Only queries system catalogs and roles | | FeatureAnalyzer | ✅ Low | Yes | Only queries system catalogs + one COUNT on static table | +| CitusAnalyzer | ✅ Low | Yes | Runs only with the Citus extension. Reads catalogs and statistics views on every node | | **DataSizeAnalyzer** | ⚠️ **HIGH** | **Use Caution** | **Requires table scans - OPT-IN ONLY** | ## Detailed Analysis @@ -132,6 +133,30 @@ LIMIT 20; --- +### ✅ CitusAnalyzer and the shard read - SAFE + +These run only when the database has the Citus extension. See +[Citus Extension Support](citus.md). + +**What they do:** +- Read the `pg_dist_*` metadata tables on the coordinator +- Read `pg_class`, `pg_stats`, `pg_stat_user_tables` and `pg_stat_user_indexes` + for the shard placements on every node, through `run_command_on_all_nodes()` +- Read the database size on every node + +**Performance characteristics:** +- No user table access and no table scans +- Each read runs one command on every active primary node, and the shards are + read in batches of at most 1,000 shards per call. Each command looks up its + batch's shards by name, so its work on a node follows the batch size, not the + number of relations on the node +- On a cluster with three nodes and 237 shards, the check before the run and + the shard read each took under one second + +**Production safety:** ✅ Always safe + +--- + ### ⚠️ DataSizeAnalyzer - USE CAUTION **What it does:** @@ -144,6 +169,10 @@ LIMIT 20; - Uses `pg_column_size()` which must read actual data - Default sampling (10%) reduces cost but still significant - Large tables can take minutes even with sampling +- On a database with the Citus extension, Citus tables are skipped with the + reason `citus_table`. The size limit and the sampling decision both read the + table's empty shell on the coordinator, so neither would stop a full scan of + every shard. See [Data Size Analysis](data_size_analysis.md#citus-tables). **Example query:** ```sql diff --git a/docs/providers/azure.md b/docs/providers/azure.md index 3f36d38..13622c1 100644 --- a/docs/providers/azure.md +++ b/docs/providers/azure.md @@ -469,7 +469,7 @@ Deliberately out of scope in this release: | Not collected | Why | | --- | --- | -| **Azure Cosmos DB for PostgreSQL** (Citus) | Planned as a follow-up. It uses a separate SDK and a different cluster/node model. | +| **Azure Cosmos DB for PostgreSQL** (Citus) | It uses a separate SDK and a different cluster and node model. | | **Actual storage *used*** | Only *provisioned* storage size is collected. Live usage would require Azure Monitor, an additional permission, and one metrics call per server. | | **Private endpoints** | Flexible Server uses VNet injection or public access with firewall rules, both of which are collected. | | **Database lists per server** | The database module already enumerates databases on the server it connects to. | diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index a1e324a..4153b1e 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -113,6 +113,39 @@ telnet your-postgres-host 5432 CREATE EXTENSION IF NOT EXISTS pg_stat_statements; ``` +### Citus Discovery Stopped + +On a database with the Citus extension, database discovery reads every shard +placement or stops. The message begins with `Citus discovery stopped` and names +each node or privilege at fault. No database result is written. See +[Citus Extension Support](citus.md#requirements). + +The exit code depends on what the run covers: + +- **Database discovery alone** exits with code 1, and writes no report. +- **A combined run**, with a cloud provider configured, continues with cloud + discovery, writes the report with `database_results` set to `null`, and exits + with code 0, as for any other database failure. A wrapper that checks the exit + code alone does not detect it. Check `database_results` in the report, or the + log for `Citus discovery stopped`. + +| Message | Remedy | +| --- | --- | +| `Discovery is connected to a Citus worker` | Connect to the coordinator. | +| `: unreachable` or `: returned no result` | Start the node, or restore the network path from the coordinator to it. | +| `: the node rejected the discovery role` | Allow the discovery role on that node: a `pg_hba.conf` entry, or a `.pgpass` file or `pg_dist_authinfo` entry on the coordinator. | +| `: Citus metadata is not synced to the node` | The worker cannot recognize its shards. After an upgrade from Citus 10, run `SELECT citus_finalize_upgrade_to_citus11();` on the coordinator. Otherwise run `SELECT start_metadata_sync_to_node('', );` on the coordinator for that node. | +| `node : not in pg_dist_node when discovery started` | A node was added to the cluster during the run. Run discovery again once the change is finished. | +| `no active placement on a node that discovery can read` | The shard's placements are on an inactive node or a secondary. Activate the node, or move the placement. | +| `missing EXECUTE privilege` or `missing SELECT privilege` | Run the `GRANT` printed in the message, on the coordinator. | +| `citus.use_secondary_nodes` | Set it to `never` for the discovery role. | +| `index : dropped during discovery` | The index was dropped during the run. Run discovery again. | +| `index : not found on N of M shards` | On a partition's index, created from its parent's index: the shard indexes can lack their shard id after PostgreSQL named them automatically. Run `SELECT fix_all_partition_shard_index_names();` on the coordinator, which renames them. On any other index, or when that function does not clear the error: the index was not created on those shards, typically because it was created with `citus.enable_ddl_propagation` off. Recreate it with DDL propagation on, or drop it. An invalid index, such as one left by a failed `CREATE INDEX CONCURRENTLY`, never stops the run. | +| `function not found` or `table not found; this Citus version is not supported` | Run discovery against a supported Citus version. | + +Do not work around these errors by connecting as a superuser. Report the node or +the privilege named in the message instead. + ### Missing Dependencies ```bash diff --git a/planetscale_discovery/cli.py b/planetscale_discovery/cli.py index e204a81..017a280 100644 --- a/planetscale_discovery/cli.py +++ b/planetscale_discovery/cli.py @@ -516,6 +516,22 @@ def _db_summary_lines_postgres(analysis: Dict, conn_info: Dict) -> List[str]: if total_size > 0: lines.append(f"- **Total Size:** {_format_bytes(total_size)}") + citus_results = analysis.get("citus", {}) + if "summary" in citus_results: + citus_summary = citus_results["summary"] + version = citus_results.get("version", {}).get("extension_version", "") + counts = citus_summary.get("table_counts", {}) + lines.append( + f"- **Citus:** {version}, {citus_summary.get('node_count', 0)} node(s)" + ) + lines.append( + "- **Citus Tables:** " + + ", ".join(f"{count} {kind}" for kind, count in sorted(counts.items())) + ) + storage = citus_summary.get("cluster_storage_bytes") + if storage is not None: + lines.append(f"- **Cluster Storage:** {_format_bytes(storage)}") + features_results = analysis.get("features", {}) if "extensions_analysis" in features_results: extensions = features_results["extensions_analysis"].get( @@ -999,9 +1015,20 @@ def _format_bytes(bytes_value: int) -> str: def _handle_database_error(e: Exception, logger: Any, engine: str = "postgres") -> None: """Print user-friendly database connection error messages.""" + from .database.citus import CitusClusterError + error_str = str(e) - if any( + # Checked first: a Citus remote error can contain "Connection refused", + # which the next branch would blame on the coordinator's host and port. + if isinstance(e, CitusClusterError): + logger.error( + f"Citus discovery stopped. Every node of the Citus cluster must be up " + f"and must accept the discovery role, and the role needs the Citus " + f"privileges listed in docs/citus.md. No database result was written.\n" + f" {error_str}" + ) + elif any( s in error_str for s in ["could not connect to server", "Connection refused", "Can't connect"] ): @@ -1286,6 +1313,9 @@ def generate_combined_summary(results: Dict[str, Any]) -> Dict[str, Any]: ), "analysis_gaps": len(analysis_gaps), } + citus_summary = analysis.get("citus", {}).get("summary") + if citus_summary: + summary["database_summary"]["citus"] = citus_summary # Extract cloud summary if results.get("cloud_results"): diff --git a/planetscale_discovery/database/analyzers/citus_analyzer.py b/planetscale_discovery/database/analyzers/citus_analyzer.py new file mode 100644 index 0000000..c70da32 --- /dev/null +++ b/planetscale_discovery/database/analyzers/citus_analyzer.py @@ -0,0 +1,219 @@ +""" +Citus Cluster Analysis Module +Reports the nodes, colocation groups, tables and shard placements of a Citus +cluster. It runs only when the citus extension is detected. +""" + +from typing import Any, Dict, List + +from planetscale_discovery.common.base_analyzer import DatabaseAnalyzer +from planetscale_discovery.database.citus import ( + DATABASE_SIZE_COMMAND, + TABLE_RELKINDS, + CitusContext, + citus_errors, + fetchall, + relation_filter, + run_on_all_nodes, +) + + +class CitusAnalyzer(DatabaseAnalyzer): + """Analyzes the topology and data placement of a Citus cluster.""" + + def __init__(self, connection, config=None, logger=None): + super().__init__(connection, config, logger) + ctx = self.config.get("citus") + if ctx is None: + raise ValueError("CitusAnalyzer needs the Citus context from detection") + self.citus: CitusContext = ctx + + def analyze(self) -> Dict[str, Any]: + """Run the Citus analysis. Any failure stops the discovery.""" + with citus_errors("The Citus cluster analysis failed.", "citus analysis"): + return self._analyze() + + def _analyze(self) -> Dict[str, Any]: + ctx = self.citus + stats = ctx.shard_statistics(self.connection) + nodes = self._get_nodes() + tables = self._get_tables(stats) + regular_count = self._count_regular_tables() + + type_counts = dict(ctx.table_counts()) + type_counts["regular"] = regular_count + logical = sum(t["logical_size_bytes"] for t in tables) + # None when no node was read, rather than a total that omits every node. + storage = ( + sum(n["database_size_bytes"] or 0 for n in nodes) if ctx.tables else None + ) + + return { + "version": { + "extension_version": ctx.extension_version, + "citus_version": ctx.citus_version, + "node_role": ctx.node_role, + }, + "nodes": nodes, + "colocation_groups": self._get_colocation_groups(), + "tables": tables, + "shard_placement": self._get_shard_placement(stats), + "distributed_schemas": self._get_distributed_schemas(), + "distributed_functions": self._get_distributed_functions(), + "summary": { + "node_count": len(nodes), + "table_counts": type_counts, + "citus_table_logical_size_bytes": logical, + "cluster_storage_bytes": storage, + }, + } + + def _get_nodes(self) -> List[Dict[str, Any]]: + """pg_dist_node, with each node's database size. + + The sum of the sizes is the cluster's storage, every placement included. + Without a Citus table, no node is contacted, and the sizes are None. + """ + ctx = self.citus + sizes = ( + run_on_all_nodes(self.connection, ctx, DATABASE_SIZE_COMMAND) + if ctx.tables + else {} + ) + nodes = [] + for node in ctx.nodes: + size = sizes.get(node["nodeid"]) + nodes.append(dict(node, database_size_bytes=int(size) if size else None)) + known = {node["nodeid"] for node in ctx.nodes} + for nodeid, size in sizes.items(): + # The coordinator, when it is not recorded in the metadata. + if nodeid not in known: + nodes.append( + { + "nodeid": nodeid, + "groupid": ctx.local_group_id, + "nodename": ctx.local_hostname, + "nodeport": ctx.port, + "noderole": "primary", + "nodecluster": ctx.cluster_name, + "isactive": True, + "shouldhaveshards": False, + "hasmetadata": True, + "database_size_bytes": int(size) if size else None, + } + ) + return nodes + + def _get_colocation_groups(self) -> List[Dict[str, Any]]: + groups = fetchall( + self.connection, + """ + SELECT colocationid AS colocation_id, + shardcount AS shard_count, + replicationfactor AS replication_factor, + distributioncolumntype::pg_catalog.regtype::text + AS distribution_column_type + FROM pg_catalog.pg_dist_colocation + ORDER BY colocationid + """, + ) + # Grouped once: schema-based sharding gives every schema its own group. + tables_by_group: Dict[int, List[str]] = {} + for t in self.citus.tables.values(): + tables_by_group.setdefault(t.colocation_id, []).append( + f"{t.schema_name}.{t.table_name}" + ) + for group in groups: + group["tables"] = tables_by_group.get(group["colocation_id"], []) + return groups + + def _get_tables(self, stats) -> List[Dict[str, Any]]: + tables = [] + for table in self.citus.tables.values(): + combined = stats.tables[table.oid] + tables.append( + { + "schema_name": table.schema_name, + "table_name": table.table_name, + **table.describe(), + "placements_per_shard": combined["placements_per_shard"], + "estimated_rows": combined["estimated_rows"], + "logical_size_bytes": combined["total_size_bytes"], + "placement_size_bytes": combined["placement_size_bytes"], + "shards_read": combined["shards_read"], + "shards_never_analyzed": combined["shards_never_analyzed"], + "inactive_placements": combined["inactive_placements"], + } + ) + return tables + + def _get_shard_placement(self, stats) -> Dict[str, Any]: + """Placement counts and bytes per node and per table, not per shard.""" + ctx = self.citus + labels: Dict[int, str] = {} + for node in ctx.nodes: + if node.get("isactive") and node.get("noderole") == "primary": + labels[node["groupid"]] = f"{node['nodename']}:{node['nodeport']}" + if ctx.local_group_id is not None and not ctx.coordinator_in_metadata(): + labels[ctx.local_group_id] = ctx.coordinator_label() + + by_node: Dict[str, Dict[str, int]] = {} + by_table: Dict[str, Dict[str, Dict[str, int]]] = {} + for placement in stats.placements: + node = labels.get(placement["groupid"], f"group {placement['groupid']}") + table = ctx.tables[placement["table_oid"]] + name = f"{table.schema_name}.{table.table_name}" + totals = by_node.setdefault(node, {"placements": 0, "bytes": 0}) + totals["placements"] += 1 + totals["bytes"] += placement["total_size_bytes"] + per_node = by_table.setdefault(name, {}).setdefault( + node, {"placements": 0, "bytes": 0} + ) + per_node["placements"] += 1 + per_node["bytes"] += placement["total_size_bytes"] + return {"by_node": by_node, "by_table": by_table} + + def _get_distributed_schemas(self) -> List[Dict[str, Any]]: + if not self.citus.has_pg_dist_schema: + return [] + return fetchall( + self.connection, + """ + SELECT schemaid::pg_catalog.regnamespace::text AS schema_name, + colocationid AS colocation_id + FROM pg_catalog.pg_dist_schema + ORDER BY 1 + """, + ) + + def _get_distributed_functions(self) -> List[Dict[str, Any]]: + return fetchall( + self.connection, + """ + SELECT objid::pg_catalog.regprocedure::text AS function, + distribution_argument_index, + colocationid AS colocation_id + FROM pg_catalog.pg_dist_object + WHERE classid = 'pg_catalog.pg_proc'::pg_catalog.regclass + AND (distribution_argument_index IS NOT NULL OR colocationid IS NOT NULL) + ORDER BY 1 + """, + ) + + def _count_regular_tables(self) -> int: + rows = fetchall( + self.connection, + f""" + SELECT pg_catalog.count(*) AS regular_count + FROM pg_catalog.pg_class c + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace + WHERE c.relkind IN {TABLE_RELKINDS} + AND n.nspname NOT IN ('information_schema', 'pg_catalog', 'pg_toast') + {relation_filter(self.citus, "c.oid", "n.oid")} + AND NOT EXISTS ( + SELECT 1 FROM pg_catalog.pg_dist_partition p + WHERE p.logicalrelid::pg_catalog.oid = c.oid + ) + """, + ) + return int(rows[0]["regular_count"]) if rows else 0 diff --git a/planetscale_discovery/database/analyzers/data_size_analyzer.py b/planetscale_discovery/database/analyzers/data_size_analyzer.py index 9925715..add508a 100644 --- a/planetscale_discovery/database/analyzers/data_size_analyzer.py +++ b/planetscale_discovery/database/analyzers/data_size_analyzer.py @@ -69,6 +69,7 @@ def __init__( "64kb": 65536, }, ) + self.citus = self.config.get("citus") def analyze(self) -> Dict[str, Any]: """ @@ -111,11 +112,27 @@ def analyze(self) -> Dict[str, Any]: # Get list of tables to analyze tables = self._get_tables_to_analyze() + citus_skipped = 0 for table in tables: schema_name = table["schema_name"] table_name = table["table_name"] table_size_gb = table["size_gb"] + # The coordinator holds only a Citus table's shell, so its size and + # row estimate would pass both safeguards, and every column query + # would scan every shard in the cluster without sampling. + if self.citus is not None and table.get("table_oid") in self.citus.tables: + results["tables_skipped"].append( + { + "schema": schema_name, + "table": table_name, + "reason": "citus_table", + } + ) + results["summary"]["total_tables_skipped"] += 1 + citus_skipped += 1 + continue + # Skip large tables if configured if table_size_gb > self.max_table_size_gb: self.logger.info( @@ -161,6 +178,12 @@ def analyze(self) -> Dict[str, Any]: ) results["summary"]["total_tables_skipped"] += 1 + if citus_skipped: + self.add_warning( + f"Skipped {citus_skipped} Citus table(s): data size analysis reads " + "rows, and on a Citus table it would read every shard in full" + ) + self.logger.info( f"Data size analysis complete: {results['summary']['total_tables_analyzed']} tables analyzed, " f"{results['summary']['total_tables_skipped']} skipped" diff --git a/planetscale_discovery/database/analyzers/feature_analyzer.py b/planetscale_discovery/database/analyzers/feature_analyzer.py index d65f099..9491066 100644 --- a/planetscale_discovery/database/analyzers/feature_analyzer.py +++ b/planetscale_discovery/database/analyzers/feature_analyzer.py @@ -6,6 +6,11 @@ from typing import Dict, Any, List from planetscale_discovery.common.base_analyzer import DatabaseAnalyzer +from planetscale_discovery.database.citus import ( + event_trigger_filter, + type_filter, + where, +) class FeatureAnalyzer(DatabaseAnalyzer): @@ -14,6 +19,7 @@ class FeatureAnalyzer(DatabaseAnalyzer): def __init__(self, connection, config=None, logger=None): super().__init__(connection, config, logger) self._version_num = None + self.citus = self.config.get("citus") def _get_server_version_num(self) -> int: """Get PostgreSQL version number (cached).""" @@ -178,9 +184,10 @@ def _analyze_postgis(self) -> Dict[str, Any]: def _get_custom_data_types(self) -> Dict[str, Any]: """Analyze custom data types.""" try: + types = type_filter(self.citus, "t.oid", "n.oid", "t.typrelid") with self.connection.cursor() as cursor: # Composite types - cursor.execute(""" + cursor.execute(f""" SELECT n.nspname as schema_name, t.typname as type_name, @@ -192,12 +199,13 @@ def _get_custom_data_types(self) -> Dict[str, Any]: JOIN pg_roles r ON t.typowner = r.oid WHERE t.typtype = 'c' -- composite types AND n.nspname NOT IN ('information_schema', 'pg_catalog', 'pg_toast') + {types} ORDER BY n.nspname, t.typname """) composite_types = [dict(row) for row in cursor.fetchall()] # Enum types - cursor.execute(""" + cursor.execute(f""" SELECT n.nspname as schema_name, t.typname as type_name, @@ -210,13 +218,14 @@ def _get_custom_data_types(self) -> Dict[str, Any]: LEFT JOIN pg_enum e ON t.oid = e.enumtypid WHERE t.typtype = 'e' -- enum types AND n.nspname NOT IN ('information_schema', 'pg_catalog', 'pg_toast') + {types} GROUP BY n.nspname, t.typname, t.oid, r.rolname ORDER BY n.nspname, t.typname """) enum_types = [dict(row) for row in cursor.fetchall()] # Domain types - cursor.execute(""" + cursor.execute(f""" SELECT n.nspname as schema_name, t.typname as type_name, @@ -232,12 +241,13 @@ def _get_custom_data_types(self) -> Dict[str, Any]: LEFT JOIN pg_constraint c ON c.contypid = t.oid WHERE t.typtype = 'd' -- domain types AND n.nspname NOT IN ('information_schema', 'pg_catalog', 'pg_toast') + {types} ORDER BY n.nspname, t.typname """) domain_types = [dict(row) for row in cursor.fetchall()] # Range types - cursor.execute(""" + cursor.execute(f""" SELECT n.nspname as schema_name, t.typname as type_name, @@ -249,6 +259,7 @@ def _get_custom_data_types(self) -> Dict[str, Any]: JOIN pg_roles ro ON t.typowner = ro.oid JOIN pg_range r ON t.oid = r.rngtypid WHERE n.nspname NOT IN ('information_schema', 'pg_catalog', 'pg_toast') + {types} ORDER BY n.nspname, t.typname """) range_types = [dict(row) for row in cursor.fetchall()] @@ -798,7 +809,10 @@ def _get_event_triggers(self) -> List[Dict[str, Any]]: """Analyze event triggers.""" try: with self.connection.cursor() as cursor: - cursor.execute(""" + event_filter = where( + event_trigger_filter(self.citus, "pg_event_trigger.oid") + ) + cursor.execute(f""" SELECT evtname as trigger_name, evtevent as event, @@ -807,6 +821,7 @@ def _get_event_triggers(self) -> List[Dict[str, Any]]: evtenabled as enabled, evttags as tags FROM pg_event_trigger + {event_filter} ORDER BY evtname """) diff --git a/planetscale_discovery/database/analyzers/performance_analyzer.py b/planetscale_discovery/database/analyzers/performance_analyzer.py index b3cb70e..6b5fc5a 100644 --- a/planetscale_discovery/database/analyzers/performance_analyzer.py +++ b/planetscale_discovery/database/analyzers/performance_analyzer.py @@ -3,10 +3,20 @@ Analyzes performance metrics, query statistics, and resource utilization. """ -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional from planetscale_discovery.common.base_analyzer import DatabaseAnalyzer from planetscale_discovery.common.sanitize import redact_sql +from planetscale_discovery.database.citus import ( + INDEX_READ_COUNTERS, + TABLE_COUNTER_KEYS, + CitusContext, + apply_logical_database_size, + citus_errors, + ShardStatistics, + invalid_index_flag, + size_pretty, +) def _redact_query_field(row: Dict[str, Any], field: str = "query") -> Dict[str, Any]: @@ -30,11 +40,28 @@ def _redact_statement_rows(rows: List[Dict[str, Any]]) -> List[Dict[str, Any]]: return [_redact_query_field(dict(row)) for row in rows] +def _index_usage(index_stats: List[Dict[str, Any]]) -> Dict[str, Any]: + """Summarize index rows, listing the indexes that record no scan.""" + unused_indexes = [idx for idx in index_stats if idx["idx_scan"] == 0] + return { + "all_indexes": index_stats, + "unused_indexes": unused_indexes, + "unused_count": len(unused_indexes), + "total_unused_size_bytes": sum( + idx["index_size_bytes"] for idx in unused_indexes + ), + } + + +LARGEST_TABLES_LIMIT = 20 + + class PerformanceAnalyzer(DatabaseAnalyzer): """Analyzes PostgreSQL performance metrics and usage patterns.""" def __init__(self, connection, config=None, logger=None): super().__init__(connection, config, logger) + self.citus: Optional[CitusContext] = self.config.get("citus") def analyze(self) -> Dict[str, Any]: """Run complete performance analysis.""" @@ -51,8 +78,140 @@ def analyze(self) -> Dict[str, Any]: "replication_lag": self._get_replication_lag(), } + if self.citus is not None: + with citus_errors( + "Citus activity counters could not be combined.", + "performance analysis", + ): + self._apply_citus(results, self.citus) + return results + def _apply_citus(self, results: Dict[str, Any], ctx: CitusContext) -> None: + """Fill the existing fields of Citus tables and indexes with logical values. + + Reads and writes on a Citus table run on its shards, so the coordinator's + counters for the shell stay at zero. + """ + stats = ctx.shard_statistics(self.connection) + indexes = { + (info["schema_name"], info["index_name"]): info + for info in stats.indexes.values() + } + + for row in results["table_statistics"]: + table = ctx.table_by_name(row["schemaname"], row["tablename"]) + if table is None: + continue + combined = stats.tables[table.oid] + coordinator = {key: row.get(key) for key in TABLE_COUNTER_KEYS} + row.update(combined["counters"]) + row["citus"] = { + "table_type": table.table_type, + "source": "shards", + "shards_read": combined["shards_read"], + "shards_without_event": combined["shards_without_event"], + "coordinator": coordinator, + } + + largest = results["resource_utilization"].get("largest_tables") + database_sizes = results["resource_utilization"].get("database_sizes") + + pretty = size_pretty( + self.connection, + [info["index_size_bytes"] for info in stats.indexes.values()] + + [ + value + for combined in stats.tables.values() + for value in ( + combined["table_size_bytes"], + combined["total_size_bytes"], + ) + ], + ) + + usage = results["index_usage"] + if "all_indexes" in usage: + for row in usage["all_indexes"]: + info = indexes.get((row["schemaname"], row["indexrelname"])) + if info is None: + continue + row["citus"] = { + "source": "shards", + "shards_read": info["shards_read"], + **invalid_index_flag(info), + "coordinator": { + key: row.get(key) + for key in INDEX_READ_COUNTERS + + ("index_size_bytes", "index_size_human") + }, + } + for key in INDEX_READ_COUNTERS: + row[key] = info[key] + row["index_size_bytes"] = info["index_size_bytes"] + row["index_size_human"] = pretty.get(info["index_size_bytes"]) + results["index_usage"] = _index_usage(usage["all_indexes"]) + + if largest is not None: + self._apply_citus_largest_tables(largest, stats, pretty, ctx) + + if database_sizes is not None: + apply_logical_database_size(self.connection, ctx, database_sizes, "datname") + database_sizes.sort(key=lambda e: e.get("size_bytes") or 0, reverse=True) + + def _apply_citus_largest_tables( + self, + largest: List[Dict[str, Any]], + stats: ShardStatistics, + pretty: Dict[int, str], + ctx: CitusContext, + ) -> None: + """Rank Citus tables by their logical size among the largest tables. + + The query ranks by the shell's size, so a large distributed table + would be missing from the list. + """ + listed = set() + for row in largest: + table = ctx.table_by_name(row["schemaname"], row["tablename"]) + if table is None: + continue + listed.add(table.oid) + combined = stats.tables[table.oid] + row["citus"] = { + "source": "shards", + "coordinator": { + key: row.get(key) + for key in ( + "total_size_bytes", + "total_size_human", + "table_size_bytes", + "table_size_human", + ) + }, + } + row["total_size_bytes"] = combined["total_size_bytes"] + row["total_size_human"] = pretty.get(combined["total_size_bytes"]) + row["table_size_bytes"] = combined["table_size_bytes"] + row["table_size_human"] = pretty.get(combined["table_size_bytes"]) + for oid, table in ctx.tables.items(): + if oid in listed: + continue + combined = stats.tables[oid] + largest.append( + { + "schemaname": table.schema_name, + "tablename": table.table_name, + "total_size_bytes": combined["total_size_bytes"], + "total_size_human": pretty.get(combined["total_size_bytes"]), + "table_size_bytes": combined["table_size_bytes"], + "table_size_human": pretty.get(combined["table_size_bytes"]), + "citus": {"source": "shards", "coordinator": None}, + } + ) + largest.sort(key=lambda r: r.get("total_size_bytes") or 0, reverse=True) + del largest[LARGEST_TABLES_LIMIT:] + def _get_connection_analysis(self) -> Dict[str, Any]: """Analyze current connections and connection patterns.""" try: @@ -328,17 +487,7 @@ def _get_index_usage(self) -> List[Dict[str, Any]]: index_stats = [dict(row) for row in cursor.fetchall()] - # Find unused indexes - unused_indexes = [idx for idx in index_stats if idx["idx_scan"] == 0] - - return { - "all_indexes": index_stats, - "unused_indexes": unused_indexes, - "unused_count": len(unused_indexes), - "total_unused_size_bytes": sum( - idx["index_size_bytes"] for idx in unused_indexes - ), - } + return _index_usage(index_stats) except Exception as e: self.logger.error(f"Failed to get index usage: {e}") diff --git a/planetscale_discovery/database/analyzers/schema_analyzer.py b/planetscale_discovery/database/analyzers/schema_analyzer.py index 960b6b2..a827624 100644 --- a/planetscale_discovery/database/analyzers/schema_analyzer.py +++ b/planetscale_discovery/database/analyzers/schema_analyzer.py @@ -3,9 +3,32 @@ Analyzes database structure, schemas, tables, and database objects. """ -from typing import Dict, Any, List +from typing import Dict, Any, List, Optional from planetscale_discovery.common.base_analyzer import DatabaseAnalyzer +from planetscale_discovery.database.citus import ( + CitusContext, + apply_logical_database_size, + citus_errors, + function_filter, + invalid_index_flag, + namespace_filter, + namespace_name_filter, + relation_filter, + size_pretty, + where, +) + +# Fields of a table entry that hold logical values for a Citus table. The +# coordinator's own values are kept under citus.coordinator. +CITUS_TABLE_FIELDS = ( + "estimated_rows", + "table_size_bytes", + "table_size_human", + "total_size_bytes", + "total_size_human", + "statistics", +) class SchemaAnalyzer(DatabaseAnalyzer): @@ -14,6 +37,7 @@ class SchemaAnalyzer(DatabaseAnalyzer): def __init__(self, connection, config=None, logger=None): super().__init__(connection, config, logger) self._version_num = None + self.citus: Optional[CitusContext] = self.config.get("citus") def _get_server_version_num(self) -> int: """Get PostgreSQL version number (cached).""" @@ -44,8 +68,127 @@ def analyze(self) -> Dict[str, Any]: "inheritance_analysis": self._get_inheritance_analysis(), } + if self.citus is not None: + self._apply_citus(results, self.citus) + return results + def _apply_citus(self, results: Dict[str, Any], ctx: CitusContext) -> None: + """Fill the existing fields of Citus tables with logical values. + + Regular tables keep their values and receive only a table type. + """ + with citus_errors( + "Citus table values could not be combined.", "schema analysis" + ): + self._apply_citus_values(results, ctx) + + def _apply_citus_values(self, results: Dict[str, Any], ctx: CitusContext) -> None: + stats = ctx.shard_statistics(self.connection) + indexes = { + (info["schema_name"], info["index_name"]): info + for info in stats.indexes.values() + } + sizes = [ + value + for combined in stats.tables.values() + for value in (combined["table_size_bytes"], combined["total_size_bytes"]) + ] + [info["index_size_bytes"] for info in stats.indexes.values()] + + for table_info in results["table_analysis"]: + # By oid: a table renamed between the classification and this + # listing would otherwise be labeled regular, with its shell's values. + table = ctx.tables.get(table_info.get("table_oid")) + if table is None: + table_info["citus"] = {"table_type": "regular"} + continue + combined = stats.tables[table.oid] + coordinator = {key: table_info.get(key) for key in CITUS_TABLE_FIELDS} + table_info["estimated_rows"] = combined["estimated_rows"] + table_info["table_size_bytes"] = combined["table_size_bytes"] + table_info["total_size_bytes"] = combined["total_size_bytes"] + table_info["statistics"] = combined["statistics"] + table_info["citus"] = { + **table.describe(), + "placements_per_shard": combined["placements_per_shard"], + "source": "shards", + "coordinator": coordinator, + "shard_statistics": { + "shards_read": combined["shards_read"], + "shards_never_analyzed": combined["shards_never_analyzed"], + "inactive_placements": combined["inactive_placements"], + "placement_size_bytes": combined["placement_size_bytes"], + "columns": combined["columns"], + }, + } + + pretty = size_pretty(self.connection, sizes) + + for table_info in results["table_analysis"]: + if table_info["citus"].get("source") == "shards": + table_info["table_size_human"] = pretty.get( + table_info["table_size_bytes"] + ) + table_info["total_size_human"] = pretty.get( + table_info["total_size_bytes"] + ) + + apply_logical_database_size( + self.connection, ctx, results["database_catalog"], "database_name" + ) + + for index_info in results["index_analysis"]: + info = indexes.get((index_info["schema_name"], index_info["index_name"])) + if info is None: + continue + index_info["citus"] = { + "source": "shards", + "shards_read": info["shards_read"], + **invalid_index_flag(info), + "coordinator": { + key: index_info.get(key) + for key in ( + "estimated_rows", + "index_size_bytes", + "index_size_human", + "usage_stats", + ) + }, + } + index_info["estimated_rows"] = info["estimated_rows"] + index_info["index_size_bytes"] = info["index_size_bytes"] + index_info["index_size_human"] = pretty.get(info["index_size_bytes"]) + # PostgreSQL keeps no usage counters for a partitioned index, so a + # Citus parent's index has no usage_stats, as a regular one has none. + if info["idx_scan"] is not None: + index_info["usage_stats"] = { + "idx_scan": info["idx_scan"], + "idx_tup_read": info["idx_tup_read"], + "idx_tup_fetch": info["idx_tup_fetch"], + } + + for partition in results["partition_analysis"]: + table = ctx.table_by_name( + partition["partition_schema"], partition["partition_name"] + ) + if table is None: + continue + combined = stats.tables[table.oid] + partition["citus"] = { + "source": "shards", + "coordinator": { + key: partition.get(key) + for key in ( + "partition_size_bytes", + "partition_size_human", + "estimated_rows", + ) + }, + } + partition["partition_size_bytes"] = combined["table_size_bytes"] + partition["partition_size_human"] = pretty.get(combined["table_size_bytes"]) + partition["estimated_rows"] = combined["estimated_rows"] + def _get_temp_schema_counts(self) -> Dict[str, int]: """Count per-backend temp namespaces. Not migration-relevant; surfaced only as a workload signal (high counts indicate unpooled connection @@ -142,7 +285,7 @@ def _get_schema_inventory(self) -> List[Dict[str, Any]]: # They have no migration value and can number in the tens of # thousands on databases with high max_connections + unpooled # workloads — see _get_temp_schema_counts for a count signal. - cursor.execute(""" + cursor.execute(f""" SELECT n.nspname as schema_name, r.rolname as owner, @@ -154,6 +297,7 @@ def _get_schema_inventory(self) -> List[Dict[str, Any]]: JOIN pg_roles r ON n.nspowner = r.oid WHERE n.nspname NOT LIKE 'pg_temp_%' AND n.nspname NOT LIKE 'pg_toast_temp_%' + {namespace_filter(self.citus, "n.oid")} ORDER BY is_system_schema, n.nspname """) @@ -310,6 +454,7 @@ def _get_table_analysis(self) -> List[Dict[str, Any]]: JOIN pg_roles r ON c.relowner = r.oid WHERE c.relkind IN {relkind_filter} AND n.nspname NOT IN ('information_schema', 'pg_catalog', 'pg_toast') + {relation_filter(self.citus, "c.oid", "n.oid")} ORDER BY n.nspname, c.relname """) @@ -481,7 +626,7 @@ def _get_index_analysis(self) -> List[Dict[str, Any]]: """Get detailed information about all indexes.""" try: with self.connection.cursor() as cursor: - cursor.execute(""" + cursor.execute(f""" SELECT n.nspname as schema_name, i.relname as index_name, @@ -509,6 +654,7 @@ def _get_index_analysis(self) -> List[Dict[str, Any]]: JOIN pg_namespace n ON i.relnamespace = n.oid JOIN pg_am am ON i.relam = am.oid WHERE n.nspname NOT IN ('information_schema', 'pg_catalog', 'pg_toast') + {relation_filter(self.citus, "t.oid", "n.oid")} ORDER BY n.nspname, t.relname, i.relname """) @@ -554,7 +700,7 @@ def _get_constraint_analysis(self) -> List[Dict[str, Any]]: """Get detailed information about all constraints.""" try: with self.connection.cursor() as cursor: - cursor.execute(""" + cursor.execute(f""" SELECT n.nspname as schema_name, c.relname as table_name, @@ -578,6 +724,7 @@ def _get_constraint_analysis(self) -> List[Dict[str, Any]]: LEFT JOIN pg_class f_c ON con.confrelid = f_c.oid LEFT JOIN pg_namespace f_n ON f_c.relnamespace = f_n.oid WHERE n.nspname NOT IN ('information_schema', 'pg_catalog', 'pg_toast') + {relation_filter(self.citus, "con.conrelid", "n.oid")} ORDER BY n.nspname, c.relname, con.conname """) @@ -591,7 +738,7 @@ def _get_function_analysis(self) -> List[Dict[str, Any]]: """Get detailed information about all functions and procedures.""" try: with self.connection.cursor() as cursor: - cursor.execute(""" + cursor.execute(f""" SELECT n.nspname as schema_name, p.proname as function_name, @@ -613,6 +760,7 @@ def _get_function_analysis(self) -> List[Dict[str, Any]]: JOIN pg_language l ON p.prolang = l.oid JOIN pg_roles r ON p.proowner = r.oid WHERE n.nspname NOT IN ('information_schema', 'pg_catalog', 'pg_toast') + {function_filter(self.citus, "p.oid", "n.oid")} ORDER BY n.nspname, p.proname """) @@ -626,7 +774,7 @@ def _get_trigger_analysis(self) -> List[Dict[str, Any]]: """Get detailed information about all triggers.""" try: with self.connection.cursor() as cursor: - cursor.execute(""" + cursor.execute(f""" SELECT n.nspname as schema_name, c.relname as table_name, @@ -642,6 +790,8 @@ def _get_trigger_analysis(self) -> List[Dict[str, Any]]: JOIN pg_proc p ON t.tgfoid = p.oid WHERE n.nspname NOT IN ('information_schema', 'pg_catalog', 'pg_toast') AND NOT t.tgisinternal + {relation_filter(self.citus, "c.oid", "n.oid")} + {function_filter(self.citus, "p.oid", "p.pronamespace")} ORDER BY n.nspname, c.relname, t.tgname """) @@ -655,7 +805,7 @@ def _get_view_analysis(self) -> List[Dict[str, Any]]: """Get detailed information about views and materialized views.""" try: with self.connection.cursor() as cursor: - cursor.execute(""" + cursor.execute(f""" SELECT n.nspname as schema_name, c.relname as view_name, @@ -676,6 +826,7 @@ def _get_view_analysis(self) -> List[Dict[str, Any]]: JOIN pg_roles r ON c.relowner = r.oid WHERE c.relkind IN ('v', 'm') -- views and materialized views AND n.nspname NOT IN ('information_schema', 'pg_catalog', 'pg_toast') + {relation_filter(self.citus, "c.oid", "n.oid")} ORDER BY n.nspname, c.relname """) @@ -689,7 +840,7 @@ def _get_sequence_analysis(self) -> List[Dict[str, Any]]: """Get detailed information about sequences.""" try: with self.connection.cursor() as cursor: - cursor.execute(""" + cursor.execute(f""" SELECT schemaname as schema_name, sequencename as sequence_name, @@ -704,6 +855,7 @@ def _get_sequence_analysis(self) -> List[Dict[str, Any]]: NULL::boolean as is_called FROM pg_sequences WHERE schemaname NOT IN ('information_schema', 'pg_catalog', 'pg_toast') + {namespace_name_filter(self.citus, "schemaname")} ORDER BY schemaname, sequencename """) @@ -727,7 +879,7 @@ def _get_partition_analysis(self) -> List[Dict[str, Any]]: return [] with self.connection.cursor() as cursor: - cursor.execute(""" + cursor.execute(f""" SELECT parent_ns.nspname as parent_schema, parent_class.relname as parent_table, @@ -743,6 +895,7 @@ def _get_partition_analysis(self) -> List[Dict[str, Any]]: JOIN pg_class child_class ON i.inhrelid = child_class.oid JOIN pg_namespace child_ns ON child_class.relnamespace = child_ns.oid WHERE child_class.relispartition = true + {relation_filter(self.citus, "child_class.oid", "child_ns.oid")} ORDER BY parent_ns.nspname, parent_class.relname, child_ns.nspname, child_class.relname """) @@ -763,6 +916,9 @@ def _get_inheritance_analysis(self) -> List[Dict[str, Any]]: if has_declarative_partitioning else "false as is_partition" ) + inheritance_filter = where( + relation_filter(self.citus, "child_class.oid", "child_ns.oid") + ) with self.connection.cursor() as cursor: cursor.execute(f""" @@ -778,6 +934,7 @@ def _get_inheritance_analysis(self) -> List[Dict[str, Any]]: JOIN pg_namespace parent_ns ON parent_class.relnamespace = parent_ns.oid JOIN pg_class child_class ON i.inhrelid = child_class.oid JOIN pg_namespace child_ns ON child_class.relnamespace = child_ns.oid + {inheritance_filter} ORDER BY parent_ns.nspname, parent_class.relname, i.inhseqno """) diff --git a/planetscale_discovery/database/analyzers/security_analyzer.py b/planetscale_discovery/database/analyzers/security_analyzer.py index de1848b..26d214e 100644 --- a/planetscale_discovery/database/analyzers/security_analyzer.py +++ b/planetscale_discovery/database/analyzers/security_analyzer.py @@ -6,6 +6,11 @@ from typing import Dict, Any, List from planetscale_discovery.common.base_analyzer import DatabaseAnalyzer +from planetscale_discovery.database.citus import ( + function_filter, + namespace_filter, + relation_filter, +) class SecurityAnalyzer(DatabaseAnalyzer): @@ -13,6 +18,7 @@ class SecurityAnalyzer(DatabaseAnalyzer): def __init__(self, connection, config=None, logger=None): super().__init__(connection, config, logger) + self.citus = self.config.get("citus") self._version_num = None def _get_server_version_num(self) -> int: @@ -147,7 +153,7 @@ def _get_permission_analysis(self) -> Dict[str, Any]: database_privileges = [dict(row) for row in cursor.fetchall()] # Schema privileges - cursor.execute(""" + cursor.execute(f""" SELECT n.nspname as schema_name, n.nspacl as schema_acl, @@ -158,11 +164,12 @@ def _get_permission_analysis(self) -> Dict[str, Any]: FROM pg_namespace n WHERE n.nspname NOT IN ('information_schema', 'pg_catalog', 'pg_toast') AND n.nspacl IS NOT NULL + {namespace_filter(self.citus, "n.oid")} """) schema_privileges = [dict(row) for row in cursor.fetchall()] # Table privileges - cursor.execute(""" + cursor.execute(f""" SELECT n.nspname as schema_name, c.relname as table_name, @@ -176,12 +183,13 @@ def _get_permission_analysis(self) -> Dict[str, Any]: WHERE c.relkind IN ('r', 'v', 'm', 'S', 'f') AND n.nspname NOT IN ('information_schema', 'pg_catalog', 'pg_toast') AND c.relacl IS NOT NULL + {relation_filter(self.citus, "c.oid", "n.oid")} ORDER BY n.nspname, c.relname """) table_privileges = [dict(row) for row in cursor.fetchall()] # Function privileges - cursor.execute(""" + cursor.execute(f""" SELECT n.nspname as schema_name, p.proname as function_name, @@ -194,6 +202,7 @@ def _get_permission_analysis(self) -> Dict[str, Any]: JOIN pg_namespace n ON p.pronamespace = n.oid WHERE n.nspname NOT IN ('information_schema', 'pg_catalog', 'pg_toast') AND p.proacl IS NOT NULL + {function_filter(self.citus, "p.oid", "n.oid")} ORDER BY n.nspname, p.proname """) function_privileges = [dict(row) for row in cursor.fetchall()] @@ -303,7 +312,7 @@ def _get_privilege_escalation(self) -> List[Dict[str, Any]]: risks = [] # Functions with SECURITY DEFINER - cursor.execute(""" + cursor.execute(f""" SELECT n.nspname as schema_name, p.proname as function_name, @@ -317,6 +326,7 @@ def _get_privilege_escalation(self) -> List[Dict[str, Any]]: JOIN pg_language l ON p.prolang = l.oid WHERE p.prosecdef = true AND n.nspname NOT IN ('information_schema', 'pg_catalog', 'pg_toast') + {function_filter(self.citus, "p.oid", "n.oid")} ORDER BY n.nspname, p.proname """) diff --git a/planetscale_discovery/database/citus.py b/planetscale_discovery/database/citus.py new file mode 100644 index 0000000..5b71f58 --- /dev/null +++ b/planetscale_discovery/database/citus.py @@ -0,0 +1,1613 @@ +""" +Citus support for PostgreSQL discovery. + +A database with the citus extension holds regular tables next to tables that +Citus manages. For a Citus table, the coordinator stores an empty shell; the +rows, column statistics and activity counters are stored on the shard +placements across the nodes. This module detects the extension, checks that +every placement can be read, reads the shards' catalogs through +run_command_on_all_nodes(), and combines the values per logical table. + +A failure that would leave a Citus table described by its shell raises +CitusClusterError, which stops the discovery. A partial report is never written. +""" + +import json +import re +from contextlib import contextmanager +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import Any, Dict, Iterable, Iterator, List, Optional, Sequence, Tuple + +CITUS_EXTENSIONS = ("citus", "citus_columnar") + +REQUIRED_FUNCTIONS = ( + "citus_version", + "column_to_column_name", + "relation_is_a_known_shard", + "run_command_on_all_nodes", + "shard_name", +) + +REQUIRED_TABLES = ( + "pg_dist_colocation", + "pg_dist_local_group", + "pg_dist_node", + "pg_dist_object", + "pg_dist_partition", + "pg_dist_placement", + "pg_dist_shard", +) + +OPTIONAL_TABLES = ("pg_dist_schema",) + +# Needed only to read the nodes, which happens only when a Citus table exists. +REMOTE_FUNCTIONS = ("run_command_on_all_nodes",) + +# run_command_on_all_nodes() holds every node's result in memory on the +# coordinator, so the shards are read in batches. +SHARD_BATCH_SIZE = 1000 + +ACTIVE_SHARD_STATE = 1 + +# The node id run_command_on_all_nodes() reports for the current node when +# pg_dist_node does not list it. +UNLISTED_NODE_ID = 0 + +# Problems listed in one error message before the rest are only counted. +MAX_LISTED_PROBLEMS = 20 + +READ_COUNTERS = ("seq_scan", "seq_tup_read", "idx_scan", "idx_tup_fetch") +WRITE_COUNTERS = ("n_tup_ins", "n_tup_upd", "n_tup_del", "n_tup_hot_upd") +PLACEMENT_COUNTERS = ("n_live_tup", "n_dead_tup", "n_mod_since_analyze") +RUN_COUNTERS = ( + "vacuum_count", + "autovacuum_count", + "analyze_count", + "autoanalyze_count", +) +EVENT_TIMES = ("last_vacuum", "last_autovacuum", "last_analyze", "last_autoanalyze") +INDEX_READ_COUNTERS = ("idx_scan", "idx_tup_read", "idx_tup_fetch") + +TABLE_COUNTER_KEYS = ( + READ_COUNTERS + WRITE_COUNTERS + PLACEMENT_COUNTERS + EVENT_TIMES + RUN_COUNTERS +) + +# The relations whose shards store data: tables and partitioned tables. The +# classification and every remote command share this filter, so a Citus table +# is never expected whose shard the remote commands do not read. A foreign +# table added to the metadata stores its rows on its own server, so it has no +# shard values to read, and it is reported as a foreign table, as on regular +# PostgreSQL. +TABLE_RELKINDS = "('r', 'p')" + +# The remote commands run in a session whose search_path this tool does not +# control, so every relation and function in them is qualified. +# +# Each command looks up its batch's shards by schema and name through pg_class's +# (relname, relnamespace) index, so a node's work follows the batch size rather +# than the number of relations on the node. A view is filtered once in its own +# CTE: pg_stat_user_tables groups every relation, so a join condition is not +# pushed into it, and a filter repeated per shard would grow with the square of +# the batch. +_SHARDS = ( + "WITH shard AS MATERIALIZED (" + " SELECT c.oid, n.nspname, c.relname, c.relkind, c.reltuples" + " FROM ROWS FROM (pg_catalog.unnest(%L::pg_catalog.name[])," + " pg_catalog.unnest(%L::pg_catalog.name[])) AS w(nspname, relname)" + " JOIN pg_catalog.pg_namespace n ON n.nspname = w.nspname" + " JOIN pg_catalog.pg_class c ON c.relnamespace = n.oid AND c.relname = w.relname" + f" WHERE c.relkind IN {TABLE_RELKINDS}" + " AND pg_catalog.relation_is_a_known_shard(c.oid))" +) + +TABLE_COMMAND = ( + _SHARDS + ", stat AS MATERIALIZED (" + " SELECT * FROM pg_catalog.pg_stat_user_tables" + " WHERE relid = ANY (ARRAY(SELECT oid FROM shard)))" + " SELECT COALESCE(pg_catalog.json_agg(t), '[]')::text FROM (" + " SELECT c.nspname AS schema_name, c.relname, c.relkind::text AS relkind," + " c.reltuples," + " pg_catalog.pg_relation_size(c.oid) AS table_size_bytes," + " pg_catalog.pg_total_relation_size(c.oid) AS total_size_bytes," + " s.seq_scan, s.seq_tup_read, s.idx_scan, s.idx_tup_fetch," + " s.n_tup_ins, s.n_tup_upd, s.n_tup_del, s.n_tup_hot_upd," + " s.n_live_tup, s.n_dead_tup, s.n_mod_since_analyze," + " s.last_vacuum, s.last_autovacuum, s.last_analyze, s.last_autoanalyze," + " s.vacuum_count, s.autovacuum_count, s.analyze_count, s.autoanalyze_count" + " FROM shard c LEFT JOIN stat s ON s.relid = c.oid" + ") t" +) + +# most_common_vals, most_common_freqs and histogram_bounds are never selected: +# they hold values from the customer's data. +COLUMN_COMMAND = ( + _SHARDS + ", stat AS MATERIALIZED (" + " SELECT schemaname, tablename, attname, inherited, null_frac, n_distinct," + " correlation FROM pg_catalog.pg_stats" + " WHERE tablename = ANY (ARRAY(SELECT relname FROM shard))" + " AND schemaname = ANY (ARRAY(SELECT nspname FROM shard)))" + " SELECT COALESCE(pg_catalog.json_agg(t), '[]')::text FROM (" + " SELECT s.schemaname AS schema_name, s.tablename AS relname, s.attname," + " s.inherited, s.null_frac, s.n_distinct, s.correlation" + " FROM shard c JOIN stat s" + " ON s.schemaname = c.nspname AND s.tablename = c.relname" + ") t" +) + +INDEX_COMMAND = ( + _SHARDS + " SELECT COALESCE(pg_catalog.json_agg(t), '[]')::text FROM (" + " SELECT n.nspname AS schema_name, i.relname AS index_relname," + " i.reltuples, s.idx_scan, s.idx_tup_read, s.idx_tup_fetch," + " pg_catalog.pg_relation_size(i.oid) AS index_size_bytes" + " FROM shard c" + " JOIN pg_catalog.pg_index x ON x.indrelid = c.oid" + " JOIN pg_catalog.pg_class i ON i.oid = x.indexrelid" + " JOIN pg_catalog.pg_namespace n ON n.oid = i.relnamespace" + " LEFT JOIN pg_catalog.pg_stat_user_indexes s ON s.indexrelid = i.oid" + ") t" +) + +DATABASE_SIZE_COMMAND = ( + "SELECT pg_catalog.pg_database_size(pg_catalog.current_database())::text" +) + +CONNECTIVITY_COMMAND = "SELECT 1" + +REMOTE_COMMANDS = ( + TABLE_COMMAND, + COLUMN_COMMAND, + INDEX_COMMAND, + DATABASE_SIZE_COMMAND, + CONNECTIVITY_COMMAND, +) + + +class CitusClusterError(ConnectionError): + """A Citus node, shard or privilege that discovery needs is unavailable. + + It subclasses ConnectionError so that the CLI treats it like a coordinator + that cannot be reached: the discovery stops, and no database result is + written. + """ + + def __init__(self, summary: str, problems: Optional[List[Dict[str, str]]] = None): + self.summary = summary + self.problems = problems or [] + super().__init__(self._render()) + + def _render(self) -> str: + lines = [self.summary] + for problem in self.problems[:MAX_LISTED_PROBLEMS]: + line = f" - {problem['subject']}: {problem['cause']}" + if problem.get("detail"): + line += f" ({problem['detail']})" + lines.append(line) + hidden = len(self.problems) - MAX_LISTED_PROBLEMS + if hidden > 0: + lines.append(f" - and {hidden} more") + return "\n".join(lines) + + +@contextmanager +def citus_errors(summary: str, subject: str, cause: str = "failed") -> Iterator[None]: + """Raise any failure inside the block as a CitusClusterError, which stops the + discovery. A CitusClusterError raised inside passes through unchanged.""" + try: + yield + except CitusClusterError: + raise + except Exception as e: + raise CitusClusterError( + summary, [{"subject": subject, "cause": cause, "detail": str(e)}] + ) from e + + +# --- SQL fragments used by the analyzers ------------------------------------ + + +def _extension_member(catalog: str, oid_expr: str) -> str: + # oid_expr must be qualified: inside this subquery a bare "oid" resolves + # to pg_extension.oid. + return ( + "EXISTS (SELECT 1 FROM pg_catalog.pg_depend citus_dep" + " JOIN pg_catalog.pg_extension citus_ext" + " ON citus_ext.oid = citus_dep.refobjid" + f" WHERE citus_dep.classid = '{catalog}'::pg_catalog.regclass" + f" AND citus_dep.objid = {oid_expr}" + " AND citus_dep.refclassid = 'pg_catalog.pg_extension'::pg_catalog.regclass" + " AND citus_dep.deptype = 'e'" + " AND citus_ext.extname IN ('citus', 'citus_columnar'))" + ) + + +def _not_known_shard(oid_expr: str) -> str: + return f"NOT pg_catalog.relation_is_a_known_shard({oid_expr})" + + +def _conditions(parts: Iterable[str]) -> str: + return "".join(f" AND {part}" for part in parts) + + +def where(conditions: str) -> str: + """Turn a filter's " AND ..." conditions into a WHERE clause. + + For a query that has no WHERE clause of its own. Returns '' without Citus, + so the query's text stays exactly as it is on a server without Citus. + """ + return "WHERE " + conditions.removeprefix(" AND ") if conditions else "" + + +def relation_filter( + ctx: Optional["CitusContext"], rel_oid: str, ns_oid: Optional[str] = None +) -> str: + """Exclude shard relations and relations that belong to the extension.""" + if ctx is None: + return "" + parts = [ + _not_known_shard(rel_oid), + "NOT " + _extension_member("pg_catalog.pg_class", rel_oid), + ] + if ns_oid: + parts.append("NOT " + _extension_member("pg_catalog.pg_namespace", ns_oid)) + return _conditions(parts) + + +def namespace_filter(ctx: Optional["CitusContext"], ns_oid: str) -> str: + """Exclude the schemas that belong to the extension (citus, citus_internal).""" + if ctx is None: + return "" + return _conditions(["NOT " + _extension_member("pg_catalog.pg_namespace", ns_oid)]) + + +def namespace_name_filter(ctx: Optional["CitusContext"], nspname_expr: str) -> str: + """namespace_filter() for a view that exposes the schema only by name.""" + if ctx is None: + return "" + return _conditions( + [ + "NOT EXISTS (SELECT 1 FROM pg_catalog.pg_namespace citus_ns" + f" WHERE citus_ns.nspname = {nspname_expr} AND " + + _extension_member("pg_catalog.pg_namespace", "citus_ns.oid") + + ")" + ] + ) + + +def function_filter(ctx: Optional["CitusContext"], proc_oid: str, ns_oid: str) -> str: + """Exclude the functions that belong to the extension.""" + if ctx is None: + return "" + return _conditions( + [ + "NOT " + _extension_member("pg_catalog.pg_proc", proc_oid), + "NOT " + _extension_member("pg_catalog.pg_namespace", ns_oid), + ] + ) + + +def type_filter( + ctx: Optional["CitusContext"], type_oid: str, ns_oid: str, typrelid: str +) -> str: + """Exclude the extension's types and the row types of shards. + + pg_type is not hidden the way pg_class is, so the row type of a shard + placement on this node is visible. + """ + if ctx is None: + return "" + return _conditions( + [ + "NOT " + _extension_member("pg_catalog.pg_type", type_oid), + "NOT " + _extension_member("pg_catalog.pg_namespace", ns_oid), + f"({typrelid} = 0 OR ({_not_known_shard(typrelid)} AND NOT " + + _extension_member("pg_catalog.pg_class", typrelid) + + "))", + ] + ) + + +def event_trigger_filter(ctx: Optional["CitusContext"], evt_oid: str) -> str: + """Exclude the event triggers that belong to the extension.""" + if ctx is None: + return "" + return _conditions( + ["NOT " + _extension_member("pg_catalog.pg_event_trigger", evt_oid)] + ) + + +# --- Detection and metadata -------------------------------------------------- + + +@dataclass +class CitusTable: + """A table recorded in pg_dist_partition.""" + + oid: int + schema_name: str + table_name: str + table_type: str + partmethod: str + repmodel: str + colocation_id: int + distribution_column: Optional[str] + autoconverted: Optional[bool] + shard_count: int + access_method: Optional[str] + + def describe(self) -> Dict[str, Any]: + return { + "table_type": self.table_type, + "distribution_column": self.distribution_column, + "colocation_id": self.colocation_id, + "shard_count": self.shard_count, + "partmethod": self.partmethod, + "replication_model": self.repmodel, + "autoconverted": self.autoconverted, + "access_method": self.access_method, + } + + +@dataclass +class ShardStatistics: + """Values combined per logical table and index from every shard placement.""" + + tables: Dict[int, Dict[str, Any]] = field(default_factory=dict) + indexes: Dict[int, Dict[str, Any]] = field(default_factory=dict) + placements: List[Dict[str, Any]] = field(default_factory=list) + # By node group: shard relations a move left behind, before cleanup drops them. + orphaned_size_bytes: Dict[int, int] = field(default_factory=dict) + + +@dataclass +class CitusContext: + """What discovery knows about the Citus cluster it is connected to.""" + + extension_version: str + citus_version: Optional[str] = None + local_group_id: Optional[int] = None + current_database: Optional[str] = None + has_pg_dist_schema: bool = False + has_autoconverted: bool = False + has_connection_error_warning: bool = False + use_secondary_nodes: Optional[str] = None + cluster_name: Optional[str] = None + # How run_command_on_all_nodes() reaches the current node when pg_dist_node + # does not list it: citus.local_hostname and port, reported as node id 0. + local_hostname: Optional[str] = None + port: Optional[int] = None + nodes: List[Dict[str, Any]] = field(default_factory=list) + tables: Dict[int, CitusTable] = field(default_factory=dict) + shards: Dict[int, Tuple[int, str]] = field(default_factory=dict) + placements: List[Dict[str, Any]] = field(default_factory=list) + _shard_statistics: Optional[ShardStatistics] = None + _tables_by_name: Optional[Dict[Tuple[str, str], CitusTable]] = None + _shell_size_bytes: Optional[int] = None + + @property + def node_role(self) -> str: + return "coordinator" if self.local_group_id == 0 else "worker" + + def table_by_name(self, schema_name: str, table_name: str) -> Optional[CitusTable]: + """For rows that carry no oid. Look a row up in self.tables by its oid + when it has one: a name can change between two queries.""" + if self._tables_by_name is None: + self._tables_by_name = { + (t.schema_name, t.table_name): t for t in self.tables.values() + } + return self._tables_by_name.get((schema_name, table_name)) + + def readable_node_ids(self) -> List[int]: + """The nodes run_command_on_all_nodes() reads, from its definition.""" + return [ + n["nodeid"] + for n in self.nodes + if n.get("isactive") + and n.get("noderole") == "primary" + and (self.cluster_name is None or n.get("nodecluster") == self.cluster_name) + ] + + def expected_node_ids(self) -> List[int]: + """Every node id run_command_on_all_nodes() can report: the nodes read in + the preflight, and the current node when pg_dist_node does not list it.""" + expected = self.readable_node_ids() + if not self.coordinator_in_metadata(): + expected.append(UNLISTED_NODE_ID) + return expected + + def readable_groups(self) -> set: + readable = set(self.readable_node_ids()) + groups = {n["groupid"] for n in self.nodes if n["nodeid"] in readable} + # The function adds the current node when the metadata does not hold it. + if self.local_group_id is not None and not self.coordinator_in_metadata(): + groups.add(self.local_group_id) + return groups + + def coordinator_in_metadata(self) -> bool: + return any(n["groupid"] == self.local_group_id for n in self.nodes) + + def coordinator_label(self) -> str: + """The current node, named as run_command_on_all_nodes() reaches it when + pg_dist_node does not list it.""" + return f"{self.local_hostname}:{self.port}" + + def node_label(self, nodeid: Optional[int]) -> str: + for node in self.nodes: + if node["nodeid"] == nodeid: + return f"{node['nodename']}:{node['nodeport']}" + if nodeid == UNLISTED_NODE_ID and not self.coordinator_in_metadata(): + return self.coordinator_label() + return f"node {nodeid}" + + def table_counts(self) -> Dict[str, int]: + counts: Dict[str, int] = {} + for table in self.tables.values(): + counts[table.table_type] = counts.get(table.table_type, 0) + 1 + return counts + + def connection_info(self) -> Dict[str, Any]: + """What describes the connection itself, for connection_info.citus. + + Counts of the cluster's tables and nodes are reported once, in the + citus module's summary, so that two copies cannot disagree. + """ + return { + "extension_version": self.extension_version, + "citus_version": self.citus_version, + "node_role": self.node_role, + } + + def shard_statistics(self, connection) -> ShardStatistics: + """Read the shards on the first call, and return the cached result.""" + if self._shard_statistics is None: + self._shard_statistics = collect_shard_statistics(connection, self) + return self._shard_statistics + + +def fetchall(connection, sql: str, params: Optional[Sequence[Any]] = None): + """Every row of a query, as dicts.""" + with connection.cursor() as cursor: + cursor.execute(sql, params) + return [dict(row) for row in cursor.fetchall()] + + +def _fetchone(connection, sql: str, params: Optional[Sequence[Any]] = None): + rows = fetchall(connection, sql, params) + return rows[0] if rows else None + + +def detect_citus(connection) -> Optional[CitusContext]: + """Return a context when the citus extension is installed, else None.""" + row = _fetchone( + connection, + "SELECT extversion FROM pg_catalog.pg_extension WHERE extname = 'citus'", + ) + if not row: + return None + return CitusContext(extension_version=row["extversion"]) + + +def preflight(connection, ctx: CitusContext) -> None: + """Confirm that every Citus node, shard and privilege is available. + + Raises CitusClusterError naming each problem. On success, the context holds + the cluster's nodes, the Citus tables and their placements. + + The checks that only reading the nodes needs run only when the database + holds a Citus table. Without one, no node is contacted, so an unused + extension cannot stop the run. + """ + with citus_errors( + "Citus metadata could not be read on the coordinator.", + "coordinator", + cause="query failed", + ): + _hide_shards(connection) + remote_problems = _check_privileges(connection, ctx) + _read_settings(connection, ctx) + _read_metadata(connection, ctx) + if not ctx.tables: + return + if remote_problems: + raise CitusClusterError( + "The discovery role cannot read the Citus shards.", remote_problems + ) + _check_secondary_nodes(ctx) + _check_metadata_sync(ctx) + _check_placements(ctx) + run_on_all_nodes(connection, ctx, CONNECTIVITY_COMMAND) + + +def _hide_shards(connection) -> None: + """Hide shards from pg_class for the rest of this session. + + A role or database default for citus.show_shards_for_app_name_prefixes + would otherwise list every shard placement on the coordinator as a table + in each analyzer that relies on Citus hiding them. The setting applies to + this connection only. + """ + fetchall( + connection, + "SELECT pg_catalog.set_config(" + "'citus.show_shards_for_app_name_prefixes', '', false) AS hidden", + ) + + +def _check_privileges(connection, ctx: CitusContext) -> List[Dict[str, str]]: + """Raise for what reading the coordinator's metadata needs. + + Returns the problems with run_command_on_all_nodes(), which only reading + the nodes needs, so that the caller raises them only when there is a Citus + table to read. + """ + problems = [] + remote_problems = [] + # Quoted, so that a GRANT for a role such as app-reader or AppUser can be run. + role = _fetchone(connection, "SELECT pg_catalog.quote_ident(current_user) AS role")[ + "role" + ] + + functions = fetchall( + connection, + """ + SELECT p.proname, + pg_catalog.pg_get_function_identity_arguments(p.oid) AS args, + pg_catalog.has_function_privilege(p.oid, 'EXECUTE') AS can_execute + FROM pg_catalog.pg_proc p + WHERE p.pronamespace = 'pg_catalog'::pg_catalog.regnamespace + AND p.proname = ANY(%s) + """, + (list(REQUIRED_FUNCTIONS),), + ) + + def bucket(name: str) -> List[Dict[str, str]]: + return remote_problems if name in REMOTE_FUNCTIONS else problems + + found = {f["proname"] for f in functions} + for name in REQUIRED_FUNCTIONS: + if name not in found: + bucket(name).append( + { + "subject": f"pg_catalog.{name}()", + "cause": "function not found; this Citus version is not supported", + } + ) + for function in functions: + if not function["can_execute"]: + bucket(function["proname"]).append( + { + "subject": f"pg_catalog.{function['proname']}()", + "cause": "missing EXECUTE privilege", + "detail": f"GRANT EXECUTE ON FUNCTION pg_catalog.{function['proname']}" + f"({function['args']}) TO {role}", + } + ) + if function["proname"] == "run_command_on_all_nodes": + if "give_warning_for_connection_errors" in (function["args"] or ""): + ctx.has_connection_error_warning = True + + tables = fetchall( + connection, + """ + SELECT t.name, + pg_catalog.to_regclass('pg_catalog.' || t.name) IS NOT NULL AS present, + CASE WHEN pg_catalog.to_regclass('pg_catalog.' || t.name) IS NULL THEN NULL + ELSE pg_catalog.has_table_privilege( + pg_catalog.to_regclass('pg_catalog.' || t.name), 'SELECT') + END AS can_select + FROM pg_catalog.unnest(%s::text[]) AS t(name) + """, + (list(REQUIRED_TABLES + OPTIONAL_TABLES),), + ) + for table in tables: + if not table["present"]: + if table["name"] in REQUIRED_TABLES: + problems.append( + { + "subject": f"pg_catalog.{table['name']}", + "cause": "table not found; this Citus version is not supported", + } + ) + continue + if table["name"] == "pg_dist_schema": + ctx.has_pg_dist_schema = True + if not table["can_select"]: + problems.append( + { + "subject": f"pg_catalog.{table['name']}", + "cause": "missing SELECT privilege", + "detail": f"GRANT SELECT ON pg_catalog.{table['name']} TO {role}", + } + ) + + if problems: + raise CitusClusterError( + "The discovery role cannot read the Citus metadata it needs.", problems + ) + return remote_problems + + +def _read_settings(connection, ctx: CitusContext) -> None: + row = _fetchone( + connection, + """ + SELECT (SELECT groupid FROM pg_catalog.pg_dist_local_group) AS local_group_id, + pg_catalog.citus_version() AS citus_version, + pg_catalog.current_database() AS current_database, + pg_catalog.current_setting('citus.use_secondary_nodes', true) + AS use_secondary_nodes, + pg_catalog.current_setting('citus.cluster_name', true) AS cluster_name, + pg_catalog.current_setting('citus.local_hostname', true) + AS local_hostname, + pg_catalog.current_setting('port') AS port, + EXISTS ( + SELECT 1 FROM pg_catalog.pg_attribute + WHERE attrelid = 'pg_catalog.pg_dist_partition'::pg_catalog.regclass + AND attname = 'autoconverted' AND NOT attisdropped + ) AS has_autoconverted + """, + ) + ctx.local_group_id = row["local_group_id"] + ctx.citus_version = row["citus_version"] + ctx.current_database = row["current_database"] + ctx.use_secondary_nodes = row["use_secondary_nodes"] + ctx.cluster_name = row["cluster_name"] + ctx.local_hostname = row.get("local_hostname") + ctx.port = int(row["port"]) if row.get("port") else None + ctx.has_autoconverted = bool(row["has_autoconverted"]) + + # Stops the run even when pg_dist_partition is empty: a worker whose + # metadata is not synced has an empty pg_dist_partition while it holds + # shards, which relation_is_a_known_shard() cannot then recognize. + if ctx.local_group_id != 0: + raise CitusClusterError( + "Discovery is connected to a Citus worker. Connect to the coordinator.", + [ + { + "subject": "connection", + "cause": f"pg_dist_local_group.groupid is {ctx.local_group_id}, " + "and only the coordinator has group 0", + } + ], + ) + + +def _check_secondary_nodes(ctx: CitusContext) -> None: + if ctx.use_secondary_nodes not in (None, "never"): + raise CitusClusterError( + "Discovery reads only primary nodes.", + [ + { + "subject": "citus.use_secondary_nodes", + "cause": f"set to '{ctx.use_secondary_nodes}' for this session; " + "set it to 'never' for the discovery role", + } + ], + ) + + +def _classification_sql(ctx: CitusContext) -> str: + # The CASE expression is the one Citus uses for citus_table_type in its + # citus_tables and citus_shards views. + schema_case = ( + "WHEN p.colocationid IN (SELECT colocationid FROM pg_catalog.pg_dist_schema)" + " THEN 'schema' " + if ctx.has_pg_dist_schema + else "" + ) + autoconverted = "p.autoconverted" if ctx.has_autoconverted else "NULL::boolean" + return f""" + SELECT p.logicalrelid::pg_catalog.oid AS table_oid, + n.nspname AS schema_name, + c.relname AS table_name, + CASE {schema_case} + WHEN p.partkey IS NOT NULL THEN 'distributed' + WHEN p.repmodel = 't' THEN 'reference' + WHEN p.colocationid = 0 THEN 'local' + ELSE 'distributed' + END AS table_type, + p.partmethod::text AS partmethod, + p.repmodel::text AS repmodel, + p.colocationid AS colocation_id, + pg_catalog.column_to_column_name(p.logicalrelid, p.partkey) + AS distribution_column, + {autoconverted} AS autoconverted, + (SELECT pg_catalog.count(*) FROM pg_catalog.pg_dist_shard s + WHERE s.logicalrelid = p.logicalrelid) AS shard_count, + a.amname AS access_method + FROM pg_catalog.pg_dist_partition p + JOIN pg_catalog.pg_class c ON c.oid = p.logicalrelid + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace + LEFT JOIN pg_catalog.pg_am a ON a.oid = c.relam + WHERE c.relkind IN {TABLE_RELKINDS} + AND NOT {_extension_member('pg_catalog.pg_class', 'p.logicalrelid')} + ORDER BY n.nspname, c.relname + """ + + +def _read_metadata(connection, ctx: CitusContext) -> None: + ctx.nodes = fetchall( + connection, + """ + SELECT nodeid, groupid, nodename, nodeport, noderole::text AS noderole, + nodecluster::text AS nodecluster, isactive, shouldhaveshards, hasmetadata + FROM pg_catalog.pg_dist_node + ORDER BY nodeid + """, + ) + ctx.tables = { + row["table_oid"]: CitusTable( + oid=row["table_oid"], + schema_name=row["schema_name"], + table_name=row["table_name"], + table_type=row["table_type"], + partmethod=row["partmethod"], + repmodel=row["repmodel"], + colocation_id=row["colocation_id"], + distribution_column=row["distribution_column"], + autoconverted=row["autoconverted"], + shard_count=int(row["shard_count"]), + access_method=row["access_method"], + ) + for row in fetchall(connection, _classification_sql(ctx)) + } + ctx._tables_by_name = None + ctx.shards = { + row["shardid"]: (row["table_oid"], unqualified(row["shard_relname"])) + for row in fetchall( + connection, + """ + SELECT s.shardid, s.logicalrelid::pg_catalog.oid AS table_oid, + pg_catalog.shard_name(s.logicalrelid, s.shardid) AS shard_relname + FROM pg_catalog.pg_dist_shard s + """, + ) + if row["table_oid"] in ctx.tables + } + ctx.placements = [ + row + for row in fetchall( + connection, + """ + SELECT shardid, groupid, shardstate + FROM pg_catalog.pg_dist_placement + ORDER BY shardid, groupid + """, + ) + if row["shardid"] in ctx.shards + ] + + +def _check_metadata_sync(ctx: CitusContext) -> None: + """Stop when a worker that holds placements has no Citus metadata. + + relation_is_a_known_shard() reads the node's own metadata, so such a node + returns none of its shards, and the shard read would stop later naming + shards instead of the node. This follows an upgrade from Citus 10 that was + not finalized. + """ + holding = { + p["groupid"] for p in ctx.placements if p["shardstate"] == ACTIVE_SHARD_STATE + } + readable = set(ctx.readable_node_ids()) + problems = [] + for node in ctx.nodes: + if ( + node["nodeid"] in readable + and node["groupid"] != ctx.local_group_id + and node["groupid"] in holding + and node.get("hasmetadata") is False + ): + problems.append( + { + "subject": f"{node['nodename']}:{node['nodeport']}", + "cause": "Citus metadata is not synced to the node", + "detail": ( + "on the coordinator, run SELECT " + "citus_finalize_upgrade_to_citus11(); after an upgrade from " + "Citus 10, or SELECT start_metadata_sync_to_node(" + f"'{node['nodename']}', {node['nodeport']});" + ), + } + ) + if problems: + raise CitusClusterError("A Citus node cannot recognize its shards.", problems) + + +def _check_placements(ctx: CitusContext) -> None: + readable = ctx.readable_groups() + active: Dict[int, List[int]] = {} + for placement in ctx.placements: + if placement["shardstate"] == ACTIVE_SHARD_STATE: + active.setdefault(placement["shardid"], []).append(placement["groupid"]) + + groups_to_nodes: Dict[int, List[str]] = {} + for node in ctx.nodes: + state = "active" if node.get("isactive") else "inactive" + groups_to_nodes.setdefault(node["groupid"], []).append( + f"{node['nodename']}:{node['nodeport']} ({node['noderole']}, {state})" + ) + + problems = [] + for shardid, (table_oid, _) in sorted(ctx.shards.items()): + groups = active.get(shardid, []) + if any(group in readable for group in groups): + continue + table = ctx.tables[table_oid] + holders = [label for g in groups for label in groups_to_nodes.get(g, [])] + problems.append( + { + "subject": f"shard {shardid} of {table.schema_name}.{table.table_name}", + "cause": "no active placement on a node that discovery can read", + "detail": "placements: " + (", ".join(holders) or "none"), + } + ) + if problems: + raise CitusClusterError( + "Some shards cannot be read. Every Citus node that holds a " + "placement must be active and primary.", + problems, + ) + + +def unqualified(name: str) -> str: + """The last identifier of a possibly qualified, possibly quoted name. + + shard_name() qualifies the name with its schema outside the search path, + and quotes either part when needed. + """ + parts = [] + current = [] + quoted = False + i = 0 + while i < len(name): + ch = name[i] + if quoted: + if ch == '"': + if i + 1 < len(name) and name[i + 1] == '"': + current.append('"') + i += 1 + else: + quoted = False + else: + current.append(ch) + elif ch == '"': + quoted = True + elif ch == ".": + parts.append("".join(current)) + current = [] + else: + current.append(ch) + i += 1 + parts.append("".join(current)) + return parts[-1] + + +# --- Running commands on every node ------------------------------------------ + + +def run_on_all_nodes( + connection, + ctx: CitusContext, + command: str, + shards: Optional[Sequence[Tuple[str, str]]] = None, +) -> Dict[int, str]: + """Run a constant command on every readable node; return each node's result. + + A batch's shard schemas and names are passed as query parameters and placed + into the command on the server with format(), so no command is built in + Python. Any node that fails raises CitusClusterError. + """ + if shards is None: + command_sql = "%s" + params: List[Any] = [command] + else: + command_sql = ( + "pg_catalog.format(%s, %s::pg_catalog.name[]::text," + " %s::pg_catalog.name[]::text)" + ) + params = [command, [s for s, _ in shards], [r for _, r in shards]] + warning_arg = ( + ", give_warning_for_connection_errors := true" + if ctx.has_connection_error_warning + else "" + ) + sql = ( + "SELECT nodeid, success, result FROM pg_catalog.run_command_on_all_nodes(" + f"{command_sql}{warning_arg})" + ) + try: + rows = fetchall(connection, sql, params) + except Exception as e: + raise CitusClusterError( + "A Citus node could not be read.", + [ + { + "subject": "run_command_on_all_nodes", + "cause": "failed", + "detail": str(e), + } + ], + ) from e + + returned = {row["nodeid"] for row in rows} + # The function reads pg_dist_node on every call, and the preflight read it + # once. A node it did not see would be read as the coordinator's group. + unexpected = sorted(returned - set(ctx.expected_node_ids())) + if unexpected: + raise CitusClusterError( + "The Citus cluster's nodes changed during discovery. Run discovery " + "again.", + [ + { + "subject": f"node {nodeid}", + "cause": "not in pg_dist_node when discovery started", + } + for nodeid in unexpected + ], + ) + + problems = [] + results: Dict[int, str] = {} + for row in rows: + if not row["success"]: + problems.append( + { + "subject": ctx.node_label(row["nodeid"]), + "cause": _failure_cause(row["result"]), + "detail": (row["result"] or "").strip()[:300], + } + ) + else: + results[row["nodeid"]] = row["result"] + for nodeid in ctx.readable_node_ids(): + if nodeid not in returned: + problems.append( + {"subject": ctx.node_label(nodeid), "cause": "returned no result"} + ) + if problems: + raise CitusClusterError( + "A Citus node could not be read. Every node must be up and must " + "accept the discovery role.", + problems, + ) + return results + + +def _failure_cause(message: Optional[str]) -> str: + text = (message or "").lower() + if "password authentication failed" in text or "no pg_hba.conf entry" in text: + return "the node rejected the discovery role" + if "permission denied" in text: + return "permission denied on the node" + if ( + "connection" in text + or "could not connect" in text + or "timeout" in text + or "timed out" in text + ): + return "unreachable" + return "command failed" + + +# --- Reading and combining shard statistics ---------------------------------- + + +def _batches(ctx: CitusContext) -> List[List[int]]: + batches: List[List[int]] = [] + current: List[int] = [] + by_table: Dict[int, List[int]] = {} + for shardid, (table_oid, _) in sorted(ctx.shards.items()): + by_table.setdefault(table_oid, []).append(shardid) + for table_oid in sorted(by_table): + for shardid in by_table[table_oid]: + current.append(shardid) + if len(current) >= SHARD_BATCH_SIZE: + batches.append(current) + current = [] + if current: + batches.append(current) + return batches + + +def _rows(results: Dict[int, str]) -> Iterable[Tuple[int, Dict[str, Any]]]: + for nodeid, text in results.items(): + for row in json.loads(text or "[]"): + yield nodeid, row + + +def collect_shard_statistics(connection, ctx: CitusContext) -> ShardStatistics: + """Read every shard placement's catalogs and combine them per table.""" + with citus_errors("Citus shard statistics could not be read.", "shard read"): + return _collect_shard_statistics(connection, ctx) + + +def _collect_shard_statistics(connection, ctx: CitusContext) -> ShardStatistics: + active_groups: Dict[int, set] = {} + for placement in ctx.placements: + if placement["shardstate"] == ACTIVE_SHARD_STATE: + active_groups.setdefault(placement["shardid"], set()).add( + placement["groupid"] + ) + + # (schema, shard relname) -> shardid, so a remote row maps to its shard. + shard_by_name = { + (ctx.tables[table_oid].schema_name, relname): shardid + for shardid, (table_oid, relname) in ctx.shards.items() + } + + indexes = _read_index_names(connection, ctx) + index_shard_by_name: Dict[Tuple[str, str], Tuple[int, int]] = {} + group_by_node = {node["nodeid"]: node["groupid"] for node in ctx.nodes} + + def group_of(nodeid: int) -> Optional[int]: + if nodeid in group_by_node: + return group_by_node[nodeid] + # run_command_on_all_nodes() adds the current node when the metadata + # does not hold it. Any other node is rejected in run_on_all_nodes(). + return ctx.local_group_id if nodeid == UNLISTED_NODE_ID else None + + table_rows: Dict[int, List[Tuple[int, Dict[str, Any]]]] = {} + column_rows: Dict[Tuple[int, int], List[Dict[str, Any]]] = {} + index_rows: Dict[Tuple[int, int], List[Tuple[int, Dict[str, Any]]]] = {} + orphaned: Dict[int, int] = {} + + def placement_of(nodeid: int, row: Dict[str, Any]) -> Optional[Tuple[int, int]]: + shardid = shard_by_name.get((row["schema_name"], row["relname"])) + if shardid is None: + return None + group = group_of(nodeid) + # A relation left on a node after a shard move has no active placement. + if group not in active_groups.get(shardid, set()): + return None + return shardid, group + + for batch in _batches(ctx): + for shard_index in _read_index_shard_names(connection, batch): + # An index created since the index list was read is not reported, + # as if it had been created after the run. + if shard_index["index_oid"] not in indexes: + continue + index_shard_by_name[ + ( + indexes[shard_index["index_oid"]]["schema_name"], + unqualified(shard_index["shard_index_name"]), + ) + ] = (shard_index["index_oid"], shard_index["shardid"]) + + names = [ + (ctx.tables[ctx.shards[shardid][0]].schema_name, ctx.shards[shardid][1]) + for shardid in batch + ] + for nodeid, row in _rows( + run_on_all_nodes(connection, ctx, TABLE_COMMAND, names) + ): + where = placement_of(nodeid, row) + if where: + table_rows.setdefault(where[0], []).append((where[1], row)) + elif (row["schema_name"], row["relname"]) in shard_by_name: + # A shard relation that a move left on the node, until Citus's + # deferred cleanup drops it. pg_database_size() still counts it. + group = group_of(nodeid) + if group is not None: + orphaned[group] = orphaned.get(group, 0) + ( + row.get("total_size_bytes") or 0 + ) + + for nodeid, row in _rows( + run_on_all_nodes(connection, ctx, COLUMN_COMMAND, names) + ): + where = placement_of(nodeid, row) + if where: + column_rows.setdefault(where, []).append(row) + + for nodeid, row in _rows( + run_on_all_nodes(connection, ctx, INDEX_COMMAND, names) + ): + key = index_shard_by_name.get((row["schema_name"], row["index_relname"])) + if key is None: + continue + index_oid, shardid = key + group = group_of(nodeid) + if group in active_groups.get(shardid, set()): + index_rows.setdefault((index_oid, shardid), []).append((group, row)) + + # Grouped once, so that combining stays linear in the number of shards. A + # cluster with schema-based sharding can hold 100k tables, each with its own + # shard. + shards_by_table: Dict[int, List[int]] = {} + for shardid, (table_oid, _) in ctx.shards.items(): + shards_by_table.setdefault(table_oid, []).append(shardid) + inactive_by_shard: Dict[int, int] = {} + for placement in ctx.placements: + if placement["shardstate"] != ACTIVE_SHARD_STATE: + shardid = placement["shardid"] + inactive_by_shard[shardid] = inactive_by_shard.get(shardid, 0) + 1 + rows_by_index: Dict[int, Dict[int, List[Tuple[int, Dict[str, Any]]]]] = {} + for (index_oid, shardid), rows in index_rows.items(): + rows_by_index.setdefault(index_oid, {})[shardid] = rows + + stats = ShardStatistics(orphaned_size_bytes=orphaned) + missing = [] + for table_oid, table in ctx.tables.items(): + by_shard: Dict[int, List[Tuple[int, Dict[str, Any]]]] = {} + for shardid in shards_by_table.get(table_oid, []): + rows = table_rows.get(shardid) + if not rows: + missing.append( + { + "subject": f"shard {shardid} of {table.schema_name}.{table.table_name}", + "cause": "not found on any node that holds an active placement", + } + ) + continue + by_shard[shardid] = rows + if missing: + continue + inactive = sum(inactive_by_shard.get(shardid, 0) for shardid in by_shard) + combined = combine_table(table, by_shard) + combined["inactive_placements"] = inactive + statistics, detail = combine_columns( + table, combined.pop("_chosen"), column_rows + ) + combined["statistics"] = statistics + combined["columns"] = detail + stats.tables[table_oid] = combined + for shardid, rows in by_shard.items(): + for group, row in rows: + stats.placements.append( + { + "table_oid": table_oid, + "shardid": shardid, + "groupid": group, + "total_size_bytes": row.get("total_size_bytes") or 0, + } + ) + if missing: + raise CitusClusterError("Some shards could not be read.", missing) + + unmatched = [] + for index_oid, info in indexes.items(): + table_stats = stats.tables.get(info["table_oid"]) + if table_stats is None: + continue + found = rows_by_index.get(index_oid, {}) + table_shards = shards_by_table.get(info["table_oid"], []) + invalid = info.get("is_valid") is False + missing_shards = set(table_shards) - set(found) + # A failed CREATE INDEX CONCURRENTLY leaves the coordinator's index + # invalid, with shard indexes on only some shards. Those can be valid on + # their nodes and take scans, so the shards found are reported, flagged, + # instead of stopping a read-only run until a DDL. + if missing_shards and not invalid: + unmatched.append( + (index_oid, _unmatched_index(info, len(missing_shards), len(found))) + ) + continue + combined = combine_index(found, table_stats["chosen_groups"]) + if invalid and not found: + # No shard index exists, so none was scanned. + combined.update({key: 0 for key in INDEX_READ_COUNTERS}) + stats.indexes[index_oid] = dict(info, **combined, shard_count=len(table_shards)) + # An index missing on some shards would report a partial size, and its scan + # counts would drop it from unused_indexes, so the run stops instead. + if unmatched: + # An index dropped since the index list was read has no shard index + # left, which the remedies below would misdiagnose. + current = _read_index_names(connection, ctx) + dropped = [ + { + "subject": problem["subject"], + "cause": "dropped during discovery", + } + for index_oid, problem in unmatched + if index_oid not in current + ] + if dropped: + raise CitusClusterError( + "The Citus indexes changed during discovery. Run discovery again.", + dropped, + ) + raise CitusClusterError( + "Some shard indexes could not be matched to their index.", + [problem for _, problem in unmatched], + ) + return stats + + +NOT_PROPAGATED_REMEDY = ( + "the index was not created on those shards, typically because it was " + "created with citus.enable_ddl_propagation off; recreate it with DDL " + "propagation on, or drop it" +) + + +def _unmatched_index(info: Dict[str, Any], missing: int, found: int) -> Dict[str, str]: + # Only an index that PostgreSQL created on a partition from its parent's + # index is named automatically, which is what the Citus function repairs. + if info.get("is_partition_index"): + remedy = ( + "a partition's shard index can lack its shard id after automatic " + "naming; run SELECT fix_all_partition_shard_index_names(); on the " + "coordinator. If the index is still not found, " + NOT_PROPAGATED_REMEDY + ) + else: + remedy = NOT_PROPAGATED_REMEDY + return { + "subject": f"index {info['schema_name']}.{info['index_name']}", + "cause": f"not found on {missing} of {missing + found} shards", + "detail": remedy, + } + + +def invalid_index_flag(info: Dict[str, Any]) -> Dict[str, Any]: + """The citus keys that flag an invalid index, whose values can cover only the + shards its build reached.""" + if info.get("is_valid") is not False: + return {} + return {"reason": "invalid_index", "shard_count": info["shard_count"]} + + +def _read_index_names(connection, ctx: CitusContext) -> Dict[int, Dict[str, Any]]: + if not ctx.tables: + return {} + rows = fetchall( + connection, + """ + SELECT x.indexrelid::pg_catalog.oid AS index_oid, + x.indrelid::pg_catalog.oid AS table_oid, + n.nspname AS schema_name, + i.relname AS index_name, + x.indisvalid AS is_valid, + i.relispartition AS is_partition_index + FROM pg_catalog.pg_index x + JOIN pg_catalog.pg_class i ON i.oid = x.indexrelid + JOIN pg_catalog.pg_namespace n ON n.oid = i.relnamespace + WHERE x.indrelid = ANY(%s::pg_catalog.oid[]) + """, + (list(ctx.tables),), + ) + return {row["index_oid"]: row for row in rows} + + +def _read_index_shard_names(connection, shard_ids: List[int]) -> List[Dict[str, Any]]: + return fetchall( + connection, + """ + SELECT x.indexrelid::pg_catalog.oid AS index_oid, s.shardid, + pg_catalog.shard_name(x.indexrelid, s.shardid) AS shard_index_name + FROM pg_catalog.pg_index x + JOIN pg_catalog.pg_dist_shard s ON s.logicalrelid::pg_catalog.oid = x.indrelid + WHERE s.shardid = ANY(%s::pg_catalog.int8[]) + """, + ([int(s) for s in shard_ids],), + ) + + +def parse_timestamp(value: Any) -> Optional[datetime]: + """Parse a timestamptz as json_agg() writes it.""" + if value is None or isinstance(value, datetime): + return value + text = str(value) + # Python 3.10 accepts only 3 or 6 fractional digits, and PostgreSQL drops + # trailing zeros. + match = re.match(r"^(.*T\d{2}:\d{2}:\d{2})(?:\.(\d+))?(.*)$", text) + if match: + head, fraction, tail = match.groups() + text = head + ("." + (fraction + "000000")[:6] if fraction else "") + tail + parsed = datetime.fromisoformat(text) + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=timezone.utc) + return parsed + + +def _sum(values: Iterable[Any]) -> Optional[int]: + present = [v for v in values if v is not None] + return sum(present) if present else None + + +def _analyzed_at(row: Dict[str, Any]) -> Optional[datetime]: + times = [parse_timestamp(row.get(k)) for k in ("last_analyze", "last_autoanalyze")] + times = [t for t in times if t is not None] + return max(times) if times else None + + +def _chosen(rows: List[Tuple[int, Dict[str, Any]]]) -> Tuple[int, Dict[str, Any]]: + """The placement analyzed most recently, then the lowest group id.""" + + def key(item): + group, row = item + at = _analyzed_at(row) + return (0 if at else 1, -(at.timestamp()) if at else 0, group) + + return sorted(rows, key=key)[0] + + +def _reltuples(row: Dict[str, Any]) -> Optional[float]: + value = row.get("reltuples") + return None if value is None else float(value) + + +def _analyzed(reltuples: Optional[float]) -> bool: + # reltuples is -1 until the first VACUUM or ANALYZE, and 0 for an analyzed + # empty shard. + return reltuples is not None and reltuples >= 0 + + +def combine_table( + table: CitusTable, by_shard: Dict[int, List[Tuple[int, Dict[str, Any]]]] +) -> Dict[str, Any]: + """Combine one table's placements into logical values. + + A read runs on one placement and a write is applied to every placement, so + read counters are summed over every placement, write counters take each + shard's largest placement, and row estimates take one placement per shard. + """ + chosen = {shardid: _chosen(rows) for shardid, rows in by_shard.items()} + all_rows = [row for rows in by_shard.values() for _, row in rows] + chosen_rows = [row for _, row in chosen.values()] + + analyzed = [r for r in chosen_rows if _analyzed(_reltuples(r))] + never_analyzed = len(chosen_rows) - len(analyzed) + estimated_rows = ( + int(round(sum(_reltuples(r) or 0 for r in analyzed))) if analyzed else -1 + ) + + counters: Dict[str, Any] = {} + for key in READ_COUNTERS + RUN_COUNTERS: + counters[key] = _sum(r.get(key) for r in all_rows) + for key in WRITE_COUNTERS: + counters[key] = _sum( + _max(row.get(key) for _, row in rows) for rows in by_shard.values() + ) + for key in PLACEMENT_COUNTERS: + counters[key] = _sum(r.get(key) for r in chosen_rows) + + shards_without_event: Dict[str, int] = {} + for key in EVENT_TIMES: + per_shard = [] + for rows in by_shard.values(): + times = [parse_timestamp(row.get(key)) for _, row in rows] + times = [t for t in times if t is not None] + per_shard.append(max(times) if times else None) + recorded = [t for t in per_shard if t is not None] + counters[key] = min(recorded) if recorded else None + shards_without_event[key] = len(per_shard) - len(recorded) + + return { + "estimated_rows": estimated_rows, + "table_size_bytes": sum(r.get("table_size_bytes") or 0 for r in chosen_rows), + "total_size_bytes": sum(r.get("total_size_bytes") or 0 for r in chosen_rows), + "placement_size_bytes": sum(r.get("total_size_bytes") or 0 for r in all_rows), + "counters": counters, + "shards_read": len(by_shard), + "shards_never_analyzed": never_analyzed, + "shards_without_event": shards_without_event, + "placements_per_shard": { + "min": min(len(rows) for rows in by_shard.values()) if by_shard else 0, + "max": max(len(rows) for rows in by_shard.values()) if by_shard else 0, + }, + "chosen_groups": {shardid: group for shardid, (group, _) in chosen.items()}, + "_chosen": {shardid: (group, row) for shardid, (group, row) in chosen.items()}, + } + + +def _max(values: Iterable[Any]) -> Optional[int]: + present = [v for v in values if v is not None] + return max(present) if present else None + + +def _absolute_distinct(n_distinct: float, reltuples: float) -> float: + return n_distinct if n_distinct >= 0 else -n_distinct * reltuples + + +def _postgres_convention(distinct: float, rows: float) -> float: + # ANALYZE stores a negative fraction when the distinct count exceeds 10% of + # the rows, so the count is assumed to grow with the table. + if rows > 0 and distinct > 0.1 * rows: + return -min(distinct / rows, 1.0) + return distinct + + +def combine_n_distinct( + table: CitusTable, attname: str, shards: List[Dict[str, Any]] +) -> Tuple[Optional[float], Optional[str]]: + """Combine the shards' n_distinct values, and name the method used.""" + values = [ + (float(s["n_distinct"]), float(s["reltuples"])) + for s in shards + if s.get("n_distinct") is not None and _analyzed(s.get("reltuples")) + ] + if not values: + return None, None + total_rows = sum(rows for _, rows in values) + + if table.shard_count == 1: + return values[0][0], "single_placement" + + if attname == table.distribution_column and table.partmethod in ("h", "r"): + distinct = sum(_absolute_distinct(nd, rows) for nd, rows in values) + return _postgres_convention(distinct, total_rows), "distribution_column_sum" + + if all(nd < 0 for nd, _ in values): + if total_rows > 0: + return ( + sum(nd * rows for nd, rows in values) / total_rows, + "weighted_fraction", + ) + return sum(nd for nd, _ in values) / len(values), "weighted_fraction" + + if all(nd >= 0 for nd, _ in values): + return max(nd for nd, _ in values), "max_absolute" + + distinct = max(_absolute_distinct(nd, rows) for nd, rows in values) + return _postgres_convention(distinct, total_rows), "mixed_max_absolute" + + +def _weighted_mean(values: List[Tuple[float, float]]) -> Optional[float]: + """Mean of (value, weight) pairs; a plain mean when no weight is positive.""" + if not values: + return None + weight = sum(w for _, w in values if w > 0) + if weight > 0: + return sum(v * w for v, w in values if w > 0) / weight + return sum(v for v, _ in values) / len(values) + + +def combine_columns( + table: CitusTable, + chosen: Dict[int, Tuple[int, Dict[str, Any]]], + column_rows: Dict[Tuple[int, int], List[Dict[str, Any]]], +) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + """Combine pg_stats rows into one row per column, as pg_stats reports them. + + Each shard contributes the rows of its chosen placement. Returns the rows + for the existing statistics field and the detail for each shard. + """ + grouped: Dict[Tuple[str, bool], List[Dict[str, Any]]] = {} + for shardid, (group, table_row) in sorted(chosen.items()): + reltuples = _reltuples(table_row) + for row in column_rows.get((shardid, group), []): + grouped.setdefault((row["attname"], bool(row.get("inherited"))), []).append( + { + "shardid": shardid, + "reltuples": reltuples, + "n_distinct": row.get("n_distinct"), + "null_frac": row.get("null_frac"), + "correlation": row.get("correlation"), + } + ) + + statistics = [] + detail = [] + for (attname, inherited), shards in sorted(grouped.items()): + analyzed = [s for s in shards if _analyzed(s["reltuples"])] + null_frac = _weighted_mean( + [ + (float(s["null_frac"]), s["reltuples"]) + for s in analyzed + if s["null_frac"] is not None + ] + ) + correlation = _weighted_mean( + [ + (float(s["correlation"]), s["reltuples"]) + for s in analyzed + if s["correlation"] is not None + ] + ) + n_distinct, method = combine_n_distinct(table, attname, shards) + statistics.append( + { + "schemaname": table.schema_name, + "tablename": table.table_name, + "attname": attname, + "n_distinct": n_distinct, + "null_frac": null_frac, + "correlation": correlation, + } + ) + detail.append( + { + "attname": attname, + "inherited": inherited, + "n_distinct_method": method, + "shards": shards, + } + ) + return statistics, detail + + +def combine_index( + by_shard: Dict[int, List[Tuple[int, Dict[str, Any]]]], + chosen_groups: Dict[int, int], +) -> Dict[str, Any]: + """Combine one index's placements: reads over every placement, size once.""" + all_rows = [row for rows in by_shard.values() for _, row in rows] + one_per_shard = [] + for shardid, rows in by_shard.items(): + group = chosen_groups.get(shardid) + match = [row for g, row in rows if g == group] or [rows[0][1]] + one_per_shard.append(match[0]) + analyzed = [r for r in one_per_shard if _analyzed(_reltuples(r))] + combined: Dict[str, Any] = { + key: _sum(r.get(key) for r in all_rows) for key in INDEX_READ_COUNTERS + } + combined["index_size_bytes"] = sum( + r.get("index_size_bytes") or 0 for r in one_per_shard + ) + combined["estimated_rows"] = ( + int(round(sum(_reltuples(r) or 0 for r in analyzed))) if analyzed else -1 + ) + combined["shards_read"] = len(by_shard) + return combined + + +def logical_database_size(connection, ctx: CitusContext, coordinator_size: int) -> int: + """The current database's size with each Citus table counted once. + + pg_database_size() on the coordinator counts the shells, the placements + stored on the coordinator and the shard relations that a move left there, + and nothing stored on the workers. + """ + stats = ctx.shard_statistics(connection) + if ctx._shell_size_bytes is None: + row = _fetchone( + connection, + "SELECT COALESCE(pg_catalog.sum(pg_catalog.pg_total_relation_size(t)), 0)" + "::pg_catalog.int8 AS shells" + " FROM pg_catalog.unnest(%s::pg_catalog.oid[]) AS t", + (list(ctx.tables),), + ) + ctx._shell_size_bytes = int(row["shells"]) if row else 0 + on_coordinator = sum( + p["total_size_bytes"] + for p in stats.placements + if p["groupid"] == ctx.local_group_id + ) + if ctx.local_group_id is not None: + on_coordinator += stats.orphaned_size_bytes.get(ctx.local_group_id, 0) + logical = sum(t["total_size_bytes"] for t in stats.tables.values()) + return max(coordinator_size - on_coordinator - ctx._shell_size_bytes, 0) + logical + + +def apply_logical_database_size( + connection, ctx: CitusContext, rows: List[Dict[str, Any]], name_key: str +) -> None: + """Rewrite the current database's row with its logical size. + + The coordinator's values are kept under citus.coordinator. The schema and + performance reports both use this, so that they report the same size. + """ + current = [ + (row, logical_database_size(connection, ctx, row.get("size_bytes") or 0)) + for row in rows + if row.get(name_key) == ctx.current_database + ] + if not current: + return + pretty = size_pretty(connection, [size for _, size in current]) + for row, size in current: + row["citus"] = { + "source": "shards", + "coordinator": { + "size_bytes": row.get("size_bytes"), + "size_human": row.get("size_human"), + }, + } + row["size_bytes"] = size + row["size_human"] = pretty.get(size) + + +def size_pretty(connection, values: Iterable[Optional[int]]) -> Dict[int, str]: + """pg_size_pretty() for many values in one round trip, so the text matches.""" + distinct = sorted({int(v) for v in values if v is not None}) + if not distinct: + return {} + rows = fetchall( + connection, + "SELECT v, pg_catalog.pg_size_pretty(v) AS pretty" + " FROM pg_catalog.unnest(%s::pg_catalog.int8[]) AS v", + (distinct,), + ) + return {int(row["v"]): row["pretty"] for row in rows} diff --git a/planetscale_discovery/database/discovery.py b/planetscale_discovery/database/discovery.py index 8c31daf..bdb35ca 100644 --- a/planetscale_discovery/database/discovery.py +++ b/planetscale_discovery/database/discovery.py @@ -20,6 +20,8 @@ from .analyzers.security_analyzer import SecurityAnalyzer from .analyzers.feature_analyzer import FeatureAnalyzer from .analyzers.data_size_analyzer import DataSizeAnalyzer +from .analyzers.citus_analyzer import CitusAnalyzer +from .citus import CitusClusterError, detect_citus, preflight class DatabaseDiscoveryTool: @@ -47,6 +49,7 @@ def discover(self): data_size_config=self.config.database.data_size, excluded_databases=self.config.database.excluded_databases, statement_timeout=self.config.database.statement_timeout, + citus_detection=True, ) if not discovery.connect(): raise ConnectionError("Failed to connect to PostgreSQL database") @@ -77,6 +80,7 @@ def __init__( data_size_config=None, excluded_databases=None, statement_timeout: str = "300s", + citus_detection: bool = False, ): # statement_timeout is a server GUC, not a libpq connection parameter, so # it is held separately. Leaving it inside connection_params would make @@ -89,6 +93,10 @@ def __init__( ) self.data_size_config = data_size_config self.excluded_databases = excluded_databases or [] + # Only callers that can stop on CitusClusterError opt in. Workload + # capture does not: its connection runs every query in one transaction. + self.citus_detection = citus_detection + self.citus = None self.connection = None self.results = { "timestamp": datetime.now(timezone.utc).isoformat(), @@ -148,11 +156,13 @@ def run_analysis( "security": SecurityAnalyzer, "features": FeatureAnalyzer, "data_size": DataSizeAnalyzer, + "citus": CitusAnalyzer, } - # Run all modules if none specified + # Run all modules if none specified. The citus module is added only + # when the extension is detected. if modules is None: - modules = list(available_modules.keys()) + modules = [name for name in available_modules if name != "citus"] self.logger.info(f"Running analysis modules: {modules}") @@ -188,10 +198,22 @@ def run_analysis( return self.results try: + # Raises CitusClusterError before any analyzer runs, so a cluster + # that cannot be read fails in seconds, not after a full run. + citus_ctx = self._prepare_citus(analyzer_connection) + if citus_ctx is not None and "citus" not in modules: + modules = list(modules) + ["citus"] + citus_config = {"citus": citus_ctx} if citus_ctx is not None else {} + for module_name in modules: if module_name not in available_modules: self.logger.warning(f"Unknown module: {module_name}") continue + if module_name == "citus" and citus_ctx is None: + self.logger.warning( + "Skipping citus analysis: the citus extension was not detected" + ) + continue try: self.logger.info(f"Running {module_name} analysis...") @@ -201,7 +223,10 @@ def run_analysis( if module_name == "schema": analyzer = available_modules[module_name]( analyzer_connection, - config={"excluded_databases": self.excluded_databases}, + config={ + "excluded_databases": self.excluded_databases, + **citus_config, + }, ) elif module_name == "data_size" and self.data_size_config: # Pass configuration to data_size analyzer @@ -215,8 +240,13 @@ def run_analysis( "target_schemas": self.data_size_config.target_schemas, "check_column_types": self.data_size_config.check_column_types, "size_thresholds": self.data_size_config.size_thresholds, + **citus_config, }, ) + elif citus_config: + analyzer = available_modules[module_name]( + analyzer_connection, config=citus_config + ) else: analyzer = available_modules[module_name](analyzer_connection) @@ -232,6 +262,10 @@ def run_analysis( f"Completed {module_name} analysis in {elapsed:.1f}s" ) + except CitusClusterError: + # Recorded as a module failure, it would let a report with + # shell values for every Citus table be written. + raise except Exception as e: self.logger.error(f"Failed to run {module_name} analysis: {e}") self.results["analysis_results"][module_name] = { @@ -252,6 +286,26 @@ def run_analysis( return self.results + def _prepare_citus(self, connection): + """Detect the citus extension and confirm the cluster can be read. + + Returns the context when the extension is installed, else None. Raises + CitusClusterError when a node, shard or privilege is unavailable. + """ + if not self.citus_detection: + return None + ctx = detect_citus(connection) + if ctx is None: + return None + self.logger.info( + f"Detected the citus extension {ctx.extension_version}; " + "checking every node and shard" + ) + preflight(connection, ctx) + self.citus = ctx + self.results["connection_info"]["citus"] = ctx.connection_info() + return ctx + def _add_analysis_gap( self, module: str, diff --git a/test-data/citus-coverage.sql b/test-data/citus-coverage.sql new file mode 100644 index 0000000..dfcd909 --- /dev/null +++ b/test-data/citus-coverage.sql @@ -0,0 +1,669 @@ +-- Citus coverage dataset for discovery testing. +-- +-- Creates every table kind that discovery labels, the regular PostgreSQL +-- variants it must leave unchanged, and the Citus cases it handles specially, +-- with enough data (about 6 GB logical) for sizes, row estimates, column +-- statistics and index counters to be measurable. It needs no other table on +-- the cluster. +-- +-- Citus tables: +-- cov.orders_hash hash distributed table, 8 shards +-- cov.orders_replicated hash distributed table with 2 placements per shard +-- (citus.shard_replication_factor = 2), with updates +-- and deletes, so write counters differ from reads +-- cov.currencies reference table (a foreign key target) +-- cov.invoices Citus local table, converted automatically by a +-- foreign key to a reference table (autoconverted) +-- cov.app_settings Citus local table, added explicitly +-- cov.audit_log single-shard distributed table, no distribution column +-- cov.events partitioned distributed table, 3 monthly partitions +-- cov_tenant_acme.* tables in a distributed schema (schema-based sharding) +-- cov_tenant_globex.* a second distributed schema, smaller +-- +-- Statistics edge states: +-- cov.empty_analyzed distributed, analyzed while empty: 0 rows +-- cov.never_analyzed distributed, autovacuum off, never analyzed: -1 rows +-- cov.partly_analyzed distributed, autovacuum off, one shard analyzed +-- +-- Identifiers that Citus must quote or shorten in shard names: +-- "Cov Mixed"."Order Lines" quoted, mixed case, with a space, and a quoted +-- distribution column and index +-- cov.measurement_rollups_by_region_and_hour_for_discovery_tests +-- a 58-character name, and a 62-character index +-- name, which exceed 63 characters with a shard id +-- +-- Other Citus objects: +-- cov.tenant_order_count() a distributed function +-- +-- Regular PostgreSQL tables, which discovery must report unchanged: +-- cov.regular_notes a regular table in a schema with Citus tables +-- cov.audit_log_999999 a regular table named like a shard of cov.audit_log +-- cov.log_20261029101112123456 +-- a regular table whose digit suffix exceeds bigint +-- cov_plain.metrics a regular partitioned table, 2 partitions and a default +-- cov_plain.legacy_child a regular table inheriting from cov_plain.legacy_parent +-- +-- Optional sections at the end, commented out: an invalid index from a failed +-- CREATE INDEX CONCURRENTLY, an index created with DDL propagation off (which +-- stops discovery on purpose), a foreign table added to the Citus metadata +-- (which needs postgres_fdw on every node), and a columnar table. +-- +-- Run it on the coordinator as a superuser. It needs no extension besides +-- citus, and the replicated table needs at least two worker nodes: +-- +-- psql "" -f test-data/citus-coverage.sql +-- +-- It is idempotent: the schemas it creates, and the server of the optional +-- foreign table, are dropped first. +-- The teardown at the end of the file removes everything it created. + +\set ON_ERROR_STOP on +\timing on + +DROP SCHEMA IF EXISTS cov CASCADE; +DROP SCHEMA IF EXISTS cov_tenant_acme CASCADE; +DROP SCHEMA IF EXISTS cov_tenant_globex CASCADE; +DROP SCHEMA IF EXISTS cov_plain CASCADE; +DROP SCHEMA IF EXISTS "Cov Mixed" CASCADE; +DROP SERVER IF EXISTS cov_loopback CASCADE; + +-- Every schema below is a regular schema unless it is created with schema-based +-- sharding turned on explicitly. +SET citus.enable_schema_based_sharding TO off; +CREATE SCHEMA cov; +CREATE SCHEMA cov_plain; +CREATE SCHEMA "Cov Mixed"; + +-- --------------------------------------------------------------------------- +-- Hash distributed table: the common case, 8 shards spread across the workers. +-- About 2,000,000 rows. +-- --------------------------------------------------------------------------- + +CREATE TABLE cov.orders_hash ( + tenant_id int NOT NULL, + order_id bigint NOT NULL, + status text NOT NULL, + total numeric(12, 2) NOT NULL, + created_at timestamptz NOT NULL, + PRIMARY KEY (tenant_id, order_id) +); +SELECT create_distributed_table('cov.orders_hash', 'tenant_id', + shard_count => 8, colocate_with => 'none'); + +INSERT INTO cov.orders_hash +SELECT 1 + (g % 5000), g, + (ARRAY['open', 'paid', 'shipped', 'cancelled'])[1 + g % 4], + round((random() * 1000)::numeric, 2), + timestamptz '2026-01-01' + g * interval '12 seconds' +FROM generate_series(1, 2000000) AS g; + +CREATE INDEX orders_hash_status_idx ON cov.orders_hash (status); + +-- --------------------------------------------------------------------------- +-- Replicated distributed table: 2 placements per shard. A write is applied to +-- both placements and a read runs on one, so discovery takes each shard's +-- largest write counter, sums read counters over both placements, and takes +-- rows and sizes from one placement per shard. About 1,000,000 rows, then +-- updates and deletes. +-- --------------------------------------------------------------------------- + +SET citus.shard_replication_factor TO 2; +CREATE TABLE cov.orders_replicated ( + tenant_id int NOT NULL, + order_id bigint NOT NULL, + status text NOT NULL, + total numeric(12, 2) NOT NULL, + PRIMARY KEY (tenant_id, order_id) +); +SELECT create_distributed_table('cov.orders_replicated', 'tenant_id', + shard_count => 8, colocate_with => 'none'); +RESET citus.shard_replication_factor; + +INSERT INTO cov.orders_replicated +SELECT 1 + (g % 2000), g, + (ARRAY['open', 'paid', 'shipped'])[1 + g % 3], + round((random() * 500)::numeric, 2) +FROM generate_series(1, 1000000) AS g; + +UPDATE cov.orders_replicated SET status = 'refunded' WHERE order_id % 10 = 0; +DELETE FROM cov.orders_replicated WHERE order_id % 50 = 0; + +CREATE INDEX orders_replicated_status_idx ON cov.orders_replicated (status); + +-- --------------------------------------------------------------------------- +-- Reference table: the target of the foreign key that converts cov.invoices. +-- --------------------------------------------------------------------------- + +CREATE TABLE cov.currencies ( + code char(3) PRIMARY KEY, + name text NOT NULL +); +SELECT create_reference_table('cov.currencies'); + +INSERT INTO cov.currencies +SELECT upper(substr(md5(g::text), 1, 3)), 'Currency ' || g +FROM generate_series(1, 400) AS g +ON CONFLICT DO NOTHING; + +-- --------------------------------------------------------------------------- +-- Citus local table, converted automatically (autoconverted = true): a +-- foreign key from a regular table to a reference table converts it, because +-- citus.enable_local_reference_table_foreign_keys is on by default. About +-- 2,000,000 rows, 20% NULL memo. +-- --------------------------------------------------------------------------- + +CREATE TABLE cov.invoices ( + invoice_id bigint PRIMARY KEY, + currency char(3) NOT NULL, + amount numeric(12, 2) NOT NULL, + issued_at timestamptz NOT NULL, + memo text +); + +INSERT INTO cov.invoices +SELECT g, + c.codes[1 + (g % array_length(c.codes, 1))], + round((random() * 10000)::numeric, 2), + timestamptz '2026-01-01' + g * interval '15 seconds', + CASE WHEN g % 5 = 0 THEN NULL ELSE repeat(md5(g::text), 3) END +FROM generate_series(1, 2000000) AS g, + (SELECT array_agg(code ORDER BY code) AS codes FROM cov.currencies) AS c; + +CREATE INDEX invoices_currency_idx ON cov.invoices (currency); + +ALTER TABLE cov.invoices + ADD CONSTRAINT invoices_currency_fk + FOREIGN KEY (currency) REFERENCES cov.currencies (code); + +-- --------------------------------------------------------------------------- +-- Citus local table, added explicitly (autoconverted = false). About +-- 1,000,000 rows. Its data stays on the coordinator. +-- --------------------------------------------------------------------------- + +CREATE TABLE cov.app_settings ( + setting_id bigint PRIMARY KEY, + scope text NOT NULL, + setting_key text NOT NULL, + setting_value text +); + +INSERT INTO cov.app_settings +SELECT g, + (ARRAY['global', 'team', 'user'])[1 + g % 3], + 'key_' || (g % 5000), + repeat(md5(g::text), 2) +FROM generate_series(1, 1000000) AS g; + +CREATE INDEX app_settings_key_idx ON cov.app_settings (setting_key); + +SELECT citus_add_local_table_to_metadata('cov.app_settings'); + +-- --------------------------------------------------------------------------- +-- Single-shard distributed table: no distribution column, one shard on a +-- worker. Labeled 'distributed', with distribution_column null and +-- shard_count 1. About 3,000,000 rows. +-- --------------------------------------------------------------------------- + +CREATE TABLE cov.audit_log ( + log_id bigint PRIMARY KEY, + actor text NOT NULL, + action text NOT NULL, + logged_at timestamptz NOT NULL, + detail text +); +SELECT create_distributed_table('cov.audit_log', NULL, colocate_with => 'none'); + +INSERT INTO cov.audit_log +SELECT g, + 'user_' || (g % 1000), + (ARRAY['create', 'update', 'delete', 'login'])[1 + g % 4], + timestamptz '2026-01-01' + g * interval '10 seconds', + repeat(md5(g::text), 2) +FROM generate_series(1, 3000000) AS g; + +CREATE INDEX audit_log_actor_idx ON cov.audit_log (actor); + +-- --------------------------------------------------------------------------- +-- Partitioned distributed table: distributed on device_id, range partitioned +-- by month. Each partition is a distributed table of its own, with 16 shards. +-- About 12,000,000 rows, loaded in four batches so that the coordinator +-- routes each batch without holding all rows at once. +-- +-- Columns are chosen to exercise each n_distinct method: +-- device_id distribution column -> distribution_column_sum +-- event_id unique per shard -> weighted_fraction +-- kind 5 values in every shard -> max_absolute +-- score 10% NULL -> weighted null_frac +-- --------------------------------------------------------------------------- + +CREATE TABLE cov.events ( + event_id bigint NOT NULL, + device_id int NOT NULL, + created_at timestamptz NOT NULL, + kind text NOT NULL, + score double precision, + payload text, + PRIMARY KEY (device_id, event_id, created_at) +) PARTITION BY RANGE (created_at); + +CREATE TABLE cov.events_2026_08 PARTITION OF cov.events + FOR VALUES FROM ('2026-08-01 00:00:00+00') TO ('2026-09-01 00:00:00+00'); +CREATE TABLE cov.events_2026_09 PARTITION OF cov.events + FOR VALUES FROM ('2026-09-01 00:00:00+00') TO ('2026-10-01 00:00:00+00'); +CREATE TABLE cov.events_2026_10 PARTITION OF cov.events + FOR VALUES FROM ('2026-10-01 00:00:00+00') TO ('2026-11-01 00:00:00+00'); + +SELECT create_distributed_table('cov.events', 'device_id', shard_count => 16); + +CREATE INDEX events_kind_idx ON cov.events (kind); + +-- 92 days from 2026-08-01 00:00 UTC ends at 2026-11-01 00:00 UTC, so every +-- generated time falls in a partition. +INSERT INTO cov.events +SELECT g, + 1 + (g % 10000), + timestamptz '2026-08-01 00:00:00+00' + + ((g::bigint * 7919) % (92 * 86400)) * interval '1 second', + (ARRAY['click', 'view', 'purchase', 'error', 'scroll'])[1 + g % 5], + CASE WHEN g % 10 = 0 THEN NULL ELSE random() * 100 END, + repeat(md5(g::text), 2) +FROM generate_series(1, 3000000) AS g; + +INSERT INTO cov.events +SELECT g, + 1 + (g % 10000), + timestamptz '2026-08-01 00:00:00+00' + + ((g::bigint * 7919) % (92 * 86400)) * interval '1 second', + (ARRAY['click', 'view', 'purchase', 'error', 'scroll'])[1 + g % 5], + CASE WHEN g % 10 = 0 THEN NULL ELSE random() * 100 END, + repeat(md5(g::text), 2) +FROM generate_series(3000001, 6000000) AS g; + +INSERT INTO cov.events +SELECT g, + 1 + (g % 10000), + timestamptz '2026-08-01 00:00:00+00' + + ((g::bigint * 7919) % (92 * 86400)) * interval '1 second', + (ARRAY['click', 'view', 'purchase', 'error', 'scroll'])[1 + g % 5], + CASE WHEN g % 10 = 0 THEN NULL ELSE random() * 100 END, + repeat(md5(g::text), 2) +FROM generate_series(6000001, 9000000) AS g; + +INSERT INTO cov.events +SELECT g, + 1 + (g % 10000), + timestamptz '2026-08-01 00:00:00+00' + + ((g::bigint * 7919) % (92 * 86400)) * interval '1 second', + (ARRAY['click', 'view', 'purchase', 'error', 'scroll'])[1 + g % 5], + CASE WHEN g % 10 = 0 THEN NULL ELSE random() * 100 END, + repeat(md5(g::text), 2) +FROM generate_series(9000001, 12000000) AS g; + +-- --------------------------------------------------------------------------- +-- Schema-based sharding: every table created in a distributed schema is a +-- single-shard table, colocated with the schema's other tables. Labeled +-- 'schema'. Two tenants of different sizes, so their shards differ. +-- --------------------------------------------------------------------------- + +SET citus.enable_schema_based_sharding TO on; +CREATE SCHEMA cov_tenant_acme; +CREATE SCHEMA cov_tenant_globex; +SET citus.enable_schema_based_sharding TO off; + +CREATE TABLE cov_tenant_acme.orders ( + order_id bigint PRIMARY KEY, + customer_id int NOT NULL, + status text NOT NULL, + total numeric(12, 2) NOT NULL, + placed_at timestamptz NOT NULL, + notes text +); +CREATE TABLE cov_tenant_acme.order_items ( + order_id bigint NOT NULL REFERENCES cov_tenant_acme.orders (order_id), + line_no int NOT NULL, + sku text NOT NULL, + quantity int NOT NULL, + price numeric(10, 2) NOT NULL, + PRIMARY KEY (order_id, line_no) +); + +CREATE TABLE cov_tenant_globex.orders (LIKE cov_tenant_acme.orders INCLUDING ALL); +CREATE TABLE cov_tenant_globex.order_items ( + LIKE cov_tenant_acme.order_items INCLUDING ALL, + FOREIGN KEY (order_id) REFERENCES cov_tenant_globex.orders (order_id) +); + +INSERT INTO cov_tenant_acme.orders +SELECT g, + 1 + (g % 50000), + (ARRAY['open', 'paid', 'shipped', 'cancelled'])[1 + g % 4], + round((random() * 500)::numeric, 2), + timestamptz '2026-01-01' + g * interval '20 seconds', + CASE WHEN g % 3 = 0 THEN repeat(md5(g::text), 2) END +FROM generate_series(1, 1500000) AS g; + +INSERT INTO cov_tenant_acme.order_items +SELECT o, l, 'SKU-' || ((o * 31 + l) % 20000), 1 + (o + l) % 5, + round((random() * 100)::numeric, 2) +FROM generate_series(1, 1500000) AS o, generate_series(1, 3) AS l; + +INSERT INTO cov_tenant_globex.orders +SELECT g, + 1 + (g % 8000), + (ARRAY['open', 'paid', 'shipped', 'cancelled'])[1 + g % 4], + round((random() * 500)::numeric, 2), + timestamptz '2026-01-01' + g * interval '60 seconds', + CASE WHEN g % 3 = 0 THEN repeat(md5(g::text), 2) END +FROM generate_series(1, 500000) AS g; + +INSERT INTO cov_tenant_globex.order_items +SELECT o, l, 'SKU-' || ((o * 31 + l) % 20000), 1 + (o + l) % 5, + round((random() * 100)::numeric, 2) +FROM generate_series(1, 500000) AS o, generate_series(1, 3) AS l; + +-- --------------------------------------------------------------------------- +-- Statistics edge states. reltuples is -1 until a shard is first analyzed or +-- vacuumed, and 0 once an empty shard is analyzed. autovacuum is turned off on +-- the last two tables, and Citus copies that setting to their shards, so +-- nothing analyzes them behind this file's back. They have no index, because +-- building an index also sets reltuples. +-- --------------------------------------------------------------------------- + +CREATE TABLE cov.empty_analyzed (tenant_id int NOT NULL, note text); +SELECT create_distributed_table('cov.empty_analyzed', 'tenant_id', + shard_count => 4, colocate_with => 'none'); + +CREATE TABLE cov.never_analyzed (tenant_id int NOT NULL, note text) + WITH (autovacuum_enabled = false); +SELECT create_distributed_table('cov.never_analyzed', 'tenant_id', + shard_count => 4, colocate_with => 'none'); +INSERT INTO cov.never_analyzed +SELECT 1 + (g % 1000), md5(g::text) FROM generate_series(1, 200000) AS g; + +CREATE TABLE cov.partly_analyzed (tenant_id int NOT NULL, note text) + WITH (autovacuum_enabled = false); +SELECT create_distributed_table('cov.partly_analyzed', 'tenant_id', + shard_count => 4, colocate_with => 'none'); +INSERT INTO cov.partly_analyzed +SELECT 1 + (g % 1000), md5(g::text) FROM generate_series(1, 200000) AS g; + +-- Analyze one shard only. ANALYZE of a relation that a node does not hold +-- fails on that node, which run_command_on_workers() reports instead of +-- raising. +SELECT nodename, success +FROM run_command_on_workers( + (SELECT format('ANALYZE %s', shard_name('cov.partly_analyzed'::regclass, min(shardid))) + FROM pg_dist_shard WHERE logicalrelid = 'cov.partly_analyzed'::regclass)); + +-- --------------------------------------------------------------------------- +-- Identifiers that need quoting or shortening. A shard name is the table's +-- name followed by _; past 63 characters Citus shortens it and adds a +-- hash, and a name with capitals or spaces is quoted. +-- --------------------------------------------------------------------------- + +CREATE TABLE "Cov Mixed"."Order Lines" ( + "Tenant Id" int NOT NULL, + "Line No" int NOT NULL, + "Amount" numeric(10, 2) NOT NULL, + PRIMARY KEY ("Tenant Id", "Line No") +); +SELECT create_distributed_table('"Cov Mixed"."Order Lines"', 'Tenant Id', + shard_count => 4, colocate_with => 'none'); +INSERT INTO "Cov Mixed"."Order Lines" +SELECT 1 + (g % 3000), g, round((random() * 100)::numeric, 2) +FROM generate_series(1, 600000) AS g; +CREATE INDEX "Amount Lookup" ON "Cov Mixed"."Order Lines" ("Amount"); + +CREATE TABLE cov.measurement_rollups_by_region_and_hour_for_discovery_tests ( + region_id int NOT NULL, + hour timestamptz NOT NULL, + reading double precision, + PRIMARY KEY (region_id, hour) +); +SELECT create_distributed_table( + 'cov.measurement_rollups_by_region_and_hour_for_discovery_tests', 'region_id', + shard_count => 4, colocate_with => 'none'); +INSERT INTO cov.measurement_rollups_by_region_and_hour_for_discovery_tests +SELECT 1 + (g % 500), timestamptz '2026-01-01' + (g / 500) * interval '1 hour', + random() * 100 +FROM generate_series(0, 599999) AS g; +CREATE INDEX measurement_rollups_by_region_and_hour_for_discovery_tests_idx + ON cov.measurement_rollups_by_region_and_hour_for_discovery_tests (reading); + +-- --------------------------------------------------------------------------- +-- A distributed function, colocated with cov.orders_hash and routed by its +-- argument. +-- --------------------------------------------------------------------------- + +CREATE FUNCTION cov.tenant_order_count(p_tenant int) RETURNS bigint + LANGUAGE sql STABLE + AS $$ SELECT count(*) FROM cov.orders_hash WHERE tenant_id = p_tenant $$; +SELECT create_distributed_function('cov.tenant_order_count(int)', 'p_tenant', + colocate_with => 'cov.orders_hash'); + +-- --------------------------------------------------------------------------- +-- Regular PostgreSQL tables. None references a Citus table, so none is +-- converted, and discovery must report each with its own values. +-- --------------------------------------------------------------------------- + +-- In a schema that also holds Citus tables. +CREATE TABLE cov.regular_notes ( + note_id bigint PRIMARY KEY, + author text NOT NULL, + body text +); +INSERT INTO cov.regular_notes +SELECT g, 'author_' || (g % 300), repeat(md5(g::text), 4) +FROM generate_series(1, 500000) AS g; + +-- Named like a shard of cov.audit_log. 999999 is no shard id, so Citus does +-- not treat it as a shard, and neither may discovery. +CREATE TABLE cov.audit_log_999999 (log_id bigint PRIMARY KEY, actor text NOT NULL); +INSERT INTO cov.audit_log_999999 +SELECT g, 'user_' || (g % 100) FROM generate_series(1, 100000) AS g; + +-- A name whose 20-digit suffix exceeds bigint, so a shard id parsed from every +-- relation name on a node would fail the run. +CREATE TABLE cov.log_20261029101112123456 (entry_id bigint PRIMARY KEY, line text); +INSERT INTO cov.log_20261029101112123456 +SELECT g, md5(g::text) FROM generate_series(1, 10000) AS g; + +-- A regular partitioned table, with an index on the parent. +CREATE TABLE cov_plain.metrics ( + metric_id bigint NOT NULL, + measured_on date NOT NULL, + name text NOT NULL, + value double precision, + PRIMARY KEY (metric_id, measured_on) +) PARTITION BY RANGE (measured_on); +CREATE TABLE cov_plain.metrics_2026_h1 PARTITION OF cov_plain.metrics + FOR VALUES FROM ('2026-01-01') TO ('2026-07-01'); +CREATE TABLE cov_plain.metrics_2026_h2 PARTITION OF cov_plain.metrics + FOR VALUES FROM ('2026-07-01') TO ('2027-01-01'); +CREATE TABLE cov_plain.metrics_default PARTITION OF cov_plain.metrics DEFAULT; +CREATE INDEX metrics_name_idx ON cov_plain.metrics (name); +INSERT INTO cov_plain.metrics +SELECT g, date '2026-01-01' + (g % 500), 'metric_' || (g % 40), random() * 100 +FROM generate_series(1, 1000000) AS g; + +-- Table inheritance, which Citus cannot distribute. +CREATE TABLE cov_plain.legacy_parent (item_id bigint NOT NULL, label text); +CREATE TABLE cov_plain.legacy_child (extra text) INHERITS (cov_plain.legacy_parent); +INSERT INTO cov_plain.legacy_parent +SELECT g, 'parent_' || g FROM generate_series(1, 100000) AS g; +INSERT INTO cov_plain.legacy_child +SELECT g, 'child_' || g, md5(g::text) FROM generate_series(1, 200000) AS g; + +-- --------------------------------------------------------------------------- +-- Statistics and activity. ANALYZE on the coordinator propagates to every +-- shard placement. The loop then makes index and sequential scans, so that +-- the scan counters are not all zero. cov.never_analyzed and +-- cov.partly_analyzed are left out on purpose. +-- --------------------------------------------------------------------------- + +ANALYZE cov.orders_hash; +ANALYZE cov.orders_replicated; +ANALYZE cov.currencies; +ANALYZE cov.invoices; +ANALYZE cov.app_settings; +ANALYZE cov.audit_log; +ANALYZE cov.events; +ANALYZE cov_tenant_acme.orders; +ANALYZE cov_tenant_acme.order_items; +ANALYZE cov_tenant_globex.orders; +ANALYZE cov_tenant_globex.order_items; +ANALYZE cov.empty_analyzed; +ANALYZE "Cov Mixed"."Order Lines"; +ANALYZE cov.measurement_rollups_by_region_and_hour_for_discovery_tests; +ANALYZE cov.regular_notes; +ANALYZE cov.audit_log_999999; +ANALYZE cov.log_20261029101112123456; +ANALYZE cov_plain.metrics; +ANALYZE cov_plain.legacy_parent; +ANALYZE cov_plain.legacy_child; + +DO $$ +DECLARE + i int; +BEGIN + FOR i IN 1..200 LOOP + PERFORM count(*) FROM cov.events WHERE device_id = 1 + (i * 37) % 10000; + PERFORM 1 FROM cov.audit_log WHERE log_id = i * 1000; + PERFORM 1 FROM cov.invoices WHERE invoice_id = i * 500; + PERFORM 1 FROM cov.app_settings WHERE setting_id = i * 250; + PERFORM 1 FROM cov_tenant_acme.orders WHERE order_id = i * 700; + PERFORM 1 FROM cov_tenant_globex.order_items WHERE order_id = i * 200; + PERFORM 1 FROM cov.orders_hash WHERE tenant_id = 1 + i % 5000 AND order_id = i; + PERFORM 1 FROM cov.orders_replicated + WHERE tenant_id = 1 + i % 2000 AND order_id = i; + PERFORM 1 FROM "Cov Mixed"."Order Lines" + WHERE "Tenant Id" = 1 + i % 3000 AND "Line No" = i; + PERFORM 1 FROM cov.regular_notes WHERE note_id = i * 100; + END LOOP; + PERFORM count(*) FROM cov.events WHERE kind = 'error'; + PERFORM count(*) FROM cov.audit_log WHERE actor = 'user_7'; + PERFORM count(*) FROM cov.orders_replicated WHERE status = 'refunded'; + PERFORM cov.tenant_order_count(42); +END +$$; + +-- --------------------------------------------------------------------------- +-- What discovery should now see. Compare with the report. +-- --------------------------------------------------------------------------- + +SELECT table_name, citus_table_type, distribution_column, colocation_id, + shard_count, table_size +FROM citus_tables +WHERE table_name::text ILIKE '%cov%' +ORDER BY table_name::text; + +SELECT logicalrelid, autoconverted +FROM pg_dist_partition +WHERE logicalrelid::text IN ('cov.invoices', 'cov.app_settings', 'cov.remote_currencies'); + +SELECT logicalrelid, shardid, count(*) AS placements +FROM pg_dist_shard JOIN pg_dist_placement USING (shardid) +WHERE logicalrelid = 'cov.orders_replicated'::regclass +GROUP BY logicalrelid, shardid +ORDER BY shardid; + +SELECT objid::regprocedure AS function, distribution_argument_index, colocationid +FROM pg_dist_object +WHERE classid = 'pg_proc'::regclass AND objid = 'cov.tenant_order_count(int)'::regprocedure; + +SELECT n.nspname, c.relname, c.relkind, c.relispartition +FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace +WHERE n.nspname IN ('cov', 'cov_plain') AND c.relkind IN ('r', 'p', 'f') + AND NOT EXISTS (SELECT 1 FROM pg_dist_partition p WHERE p.logicalrelid = c.oid) +ORDER BY 1, 2; + +-- --------------------------------------------------------------------------- +-- Optional: an invalid index. A unique index fails on the one shard that holds +-- a duplicate, which leaves the coordinator's index invalid and shard indexes +-- on only some shards. Discovery reports it with the shards where it was +-- found, flagged invalid_index, and does not stop. Uncomment to run. +-- --------------------------------------------------------------------------- + +-- CREATE TABLE cov.cic_probe (tenant_id int NOT NULL, k int NOT NULL, payload text); +-- SELECT create_distributed_table('cov.cic_probe', 'tenant_id', +-- shard_count => 8, colocate_with => 'none'); +-- INSERT INTO cov.cic_probe +-- SELECT t, k, md5((t * 100 + k)::text) +-- FROM generate_series(1, 20000) AS t, generate_series(1, 10) AS k; +-- INSERT INTO cov.cic_probe VALUES (1, 1, 'duplicate'); +-- ANALYZE cov.cic_probe; +-- \set ON_ERROR_STOP off +-- CREATE UNIQUE INDEX CONCURRENTLY cic_probe_tenant_k ON cov.cic_probe (tenant_id, k); +-- \set ON_ERROR_STOP on + +-- --------------------------------------------------------------------------- +-- Optional: an index created with DDL propagation off. It exists only on the +-- coordinator, so discovery stops with exit code 1 and the propagation +-- remedy. Run only test_an_unpropagated_index_stops_with_its_remedy while it +-- exists, then drop it with propagation off, because its shards have no index +-- to drop. Uncomment to run. +-- --------------------------------------------------------------------------- + +-- SET citus.enable_ddl_propagation TO off; +-- CREATE INDEX orders_hash_unpropagated_idx ON cov.orders_hash (created_at); +-- RESET citus.enable_ddl_propagation; +-- +-- To remove it: +-- SET citus.enable_ddl_propagation TO off; +-- DROP INDEX cov.orders_hash_unpropagated_idx; +-- RESET citus.enable_ddl_propagation; + +-- --------------------------------------------------------------------------- +-- Optional: a foreign table added to the Citus metadata. Its rows stay on its +-- server, so discovery reports it as a foreign table, not as a Citus table. +-- Creating it opens no connection to the server. Citus copies the extension, +-- the server and the table to every node, so postgres_fdw must be available on +-- each of them. Uncomment to run. +-- --------------------------------------------------------------------------- + +-- CREATE EXTENSION IF NOT EXISTS postgres_fdw; +-- DO $$ +-- BEGIN +-- EXECUTE format( +-- 'CREATE SERVER cov_loopback FOREIGN DATA WRAPPER postgres_fdw' +-- ' OPTIONS (host %L, dbname %L)', 'localhost', current_database()); +-- END +-- $$; +-- CREATE FOREIGN TABLE cov.remote_currencies (code char(3), name text) +-- SERVER cov_loopback OPTIONS (schema_name 'cov', table_name 'currencies'); +-- SELECT citus_add_local_table_to_metadata('cov.remote_currencies'); + +-- --------------------------------------------------------------------------- +-- Optional: a columnar table. Installs the citus_columnar extension, which +-- also adds its own schemas and objects for discovery to exclude. Uncomment +-- to run. +-- --------------------------------------------------------------------------- + +-- CREATE EXTENSION IF NOT EXISTS citus_columnar; +-- CREATE TABLE cov.metrics_archive ( +-- device_id int NOT NULL, +-- measured_at timestamptz NOT NULL, +-- metric text NOT NULL, +-- value double precision +-- ) USING columnar; +-- SELECT create_distributed_table('cov.metrics_archive', 'device_id', +-- colocate_with => 'cov.events'); +-- INSERT INTO cov.metrics_archive +-- SELECT 1 + (g % 10000), timestamptz '2026-08-01' + g * interval '1 second', +-- (ARRAY['cpu', 'mem', 'disk'])[1 + g % 3], random() * 100 +-- FROM generate_series(1, 3000000) AS g; +-- ANALYZE cov.metrics_archive; + +-- --------------------------------------------------------------------------- +-- Teardown. Uncomment to remove everything this file created. Drop the +-- postgres_fdw extension only if nothing else on the cluster uses it. +-- --------------------------------------------------------------------------- + +-- DROP SCHEMA IF EXISTS cov CASCADE; +-- DROP SCHEMA IF EXISTS cov_tenant_acme CASCADE; +-- DROP SCHEMA IF EXISTS cov_tenant_globex CASCADE; +-- DROP SCHEMA IF EXISTS cov_plain CASCADE; +-- DROP SCHEMA IF EXISTS "Cov Mixed" CASCADE; +-- DROP SERVER IF EXISTS cov_loopback CASCADE; +-- DROP EXTENSION IF EXISTS citus_columnar; +-- DROP EXTENSION IF EXISTS postgres_fdw; diff --git a/tests/fixtures/database_responses.py b/tests/fixtures/database_responses.py index f6e2152..6236ccb 100644 --- a/tests/fixtures/database_responses.py +++ b/tests/fixtures/database_responses.py @@ -2,6 +2,9 @@ Database query response fixtures for testing """ +import json +import re + POSTGRES_VERSION_RESPONSE = [ { "version": "PostgreSQL 17.4 on aarch64-unknown-linux-gnu, compiled by gcc (GCC) 12.4.0, 64-bit" @@ -148,6 +151,515 @@ }, ] +# +# Citus. The shapes and values are copied from a test cluster (Citus 14.1 on +# PostgreSQL 18: one coordinator and two workers, pgbench schema), reduced to +# two shards of one distributed table and one reference table. +# + +CITUS_BRANCHES_OID = 17507 +CITUS_REFERENCE_OID = 17549 +CITUS_REGULAR_OID = 17572 +CITUS_BRANCHES_PKEY_OID = 17508 +CITUS_REFERENCE_PKEY_OID = 17552 + +CITUS_NODES = [ + { + "nodeid": 1, + "groupid": 0, + "nodename": "citus-coordinator", + "nodeport": 5432, + "noderole": "primary", + "nodecluster": "default", + "isactive": True, + "shouldhaveshards": False, + "hasmetadata": True, + }, + { + "nodeid": 2, + "groupid": 1, + "nodename": "citus-worker-1", + "nodeport": 5432, + "noderole": "primary", + "nodecluster": "default", + "isactive": True, + "shouldhaveshards": True, + "hasmetadata": True, + }, + { + "nodeid": 3, + "groupid": 2, + "nodename": "citus-worker-2", + "nodeport": 5432, + "noderole": "primary", + "nodecluster": "default", + "isactive": True, + "shouldhaveshards": True, + "hasmetadata": True, + }, +] + +CITUS_TABLES = [ + { + "table_oid": CITUS_BRANCHES_OID, + "schema_name": "public", + "table_name": "pgbench_branches", + "table_type": "distributed", + "partmethod": "h", + "repmodel": "s", + "colocation_id": 3, + "distribution_column": "bid", + "autoconverted": False, + "shard_count": 2, + "access_method": "heap", + }, + { + "table_oid": CITUS_REFERENCE_OID, + "schema_name": "public", + "table_name": "reference_t1", + "table_type": "reference", + "partmethod": "n", + "repmodel": "t", + "colocation_id": 4, + "distribution_column": None, + "autoconverted": False, + "shard_count": 1, + "access_method": "heap", + }, +] + +CITUS_SHARDS = [ + { + "shardid": 102137, + "table_oid": CITUS_BRANCHES_OID, + "shard_relname": "pgbench_branches_102137", + }, + { + "shardid": 102138, + "table_oid": CITUS_BRANCHES_OID, + "shard_relname": "pgbench_branches_102138", + }, + { + "shardid": 102265, + "table_oid": CITUS_REFERENCE_OID, + "shard_relname": "reference_t1_102265", + }, +] + +CITUS_PLACEMENTS = [ + {"shardid": 102137, "groupid": 1, "shardstate": 1}, + {"shardid": 102138, "groupid": 2, "shardstate": 1}, + {"shardid": 102265, "groupid": 0, "shardstate": 1}, + {"shardid": 102265, "groupid": 1, "shardstate": 1}, + {"shardid": 102265, "groupid": 2, "shardstate": 1}, +] + + +def citus_table_row(relname, reltuples, **overrides): + """One row of the remote table command, as json_agg() writes it.""" + row = { + "schema_name": "public", + "relname": relname, + "relkind": "r", + "reltuples": reltuples, + "table_size_bytes": 16384, + "total_size_bytes": 65536, + "seq_scan": 1, + "seq_tup_read": 0, + "idx_scan": 0, + "idx_tup_fetch": 0, + "n_tup_ins": int(reltuples) if reltuples > 0 else 0, + "n_tup_upd": 0, + "n_tup_del": 0, + "n_tup_hot_upd": 0, + "n_live_tup": int(reltuples) if reltuples > 0 else 0, + "n_dead_tup": 0, + "n_mod_since_analyze": 0, + "last_vacuum": None, + "last_autovacuum": None, + "last_analyze": "2026-10-27T22:17:09.167203+00:00", + "last_autoanalyze": None, + "vacuum_count": 0, + "autovacuum_count": 0, + "analyze_count": 1, + "autoanalyze_count": 0, + } + row.update(overrides) + return row + + +def _reference_row(reltuples, analyzed_at, **overrides): + return citus_table_row( + "reference_t1_102265", + reltuples, + table_size_bytes=4222812160, + total_size_bytes=5801107456, + n_tup_ins=50000000, + last_analyze=analyzed_at, + **overrides, + ) + + +CITUS_REMOTE_TABLE_ROWS = { + 1: [_reference_row(4.999656e07, "2026-10-27T23:27:52.237928+00:00")], + 2: [ + citus_table_row("pgbench_branches_102137", 317), + _reference_row(4.9999892e07, "2026-10-27T23:27:40.1+00:00"), + ], + 3: [ + citus_table_row("pgbench_branches_102138", 313), + _reference_row(5.000156e07, "2026-10-27T23:27:45+00:00"), + ], +} + + +def citus_column_row(relname, attname, n_distinct, null_frac=0.0, correlation=1.0): + return { + "schema_name": "public", + "relname": relname, + "attname": attname, + "inherited": False, + "null_frac": null_frac, + "n_distinct": n_distinct, + "correlation": correlation, + } + + +CITUS_REMOTE_COLUMN_ROWS = { + 1: [citus_column_row("reference_t1_102265", "id", -1.0)], + 2: [ + citus_column_row("pgbench_branches_102137", "bid", -1.0), + citus_column_row("reference_t1_102265", "id", -1.0), + ], + 3: [ + citus_column_row("pgbench_branches_102138", "bid", -1.0), + citus_column_row("reference_t1_102265", "id", -1.0), + ], +} + + +def citus_index_row(index_relname, reltuples, idx_scan=0, size=16384): + return { + "schema_name": "public", + "index_relname": index_relname, + "reltuples": reltuples, + "idx_scan": idx_scan, + "idx_tup_read": idx_scan, + "idx_tup_fetch": idx_scan, + "index_size_bytes": size, + } + + +CITUS_REMOTE_INDEX_ROWS = { + 1: [citus_index_row("reference_t1_pkey_102265", 4.999656e07, size=1577099264)], + 2: [ + citus_index_row("pgbench_branches_pkey_102137", 317), + citus_index_row("reference_t1_pkey_102265", 4.9999892e07, size=1577099264), + ], + 3: [ + citus_index_row("pgbench_branches_pkey_102138", 313), + citus_index_row("reference_t1_pkey_102265", 5.000156e07, size=1577099264), + ], +} + +CITUS_INDEXES = [ + { + "index_oid": CITUS_BRANCHES_PKEY_OID, + "table_oid": CITUS_BRANCHES_OID, + "schema_name": "public", + "index_name": "pgbench_branches_pkey", + }, + { + "index_oid": CITUS_REFERENCE_PKEY_OID, + "table_oid": CITUS_REFERENCE_OID, + "schema_name": "public", + "index_name": "reference_t1_pkey", + }, +] + +CITUS_INDEX_SHARD_NAMES = [ + { + "index_oid": CITUS_BRANCHES_PKEY_OID, + "shardid": 102137, + "shard_index_name": "pgbench_branches_pkey_102137", + }, + { + "index_oid": CITUS_BRANCHES_PKEY_OID, + "shardid": 102138, + "shard_index_name": "pgbench_branches_pkey_102138", + }, + { + "index_oid": CITUS_REFERENCE_PKEY_OID, + "shardid": 102265, + "shard_index_name": "reference_t1_pkey_102265", + }, +] + +CITUS_FUNCTION_ARGS = { + "citus_version": "", + "column_to_column_name": "table_name regclass, column_var_text text", + "relation_is_a_known_shard": "regclass", + "run_command_on_all_nodes": ( + "command text, parallel boolean, give_warning_for_connection_errors boolean," + " OUT nodeid integer, OUT success boolean, OUT result text" + ), + "shard_name": "object_name regclass, shard_id bigint", +} + +CITUS_DIST_TABLES = ( + "pg_dist_colocation", + "pg_dist_local_group", + "pg_dist_node", + "pg_dist_object", + "pg_dist_partition", + "pg_dist_placement", + "pg_dist_shard", + "pg_dist_schema", +) + +CITUS_DATABASE_SIZES = {1: "19327047359", 2: "88322569919", 3: "89151944383"} + + +class FakeCursor: + def __init__(self, connection): + self.connection = connection + self.rows = [] + + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + def execute(self, sql, params=None): + self.connection.executed.append((sql, params)) + self.rows = self.connection.respond(sql, params) + + def fetchall(self): + return list(self.rows) + + def fetchone(self): + return self.rows[0] if self.rows else None + + def close(self): + pass + + +class FakeConnection: + """Answers each query with the rows of the first handler whose marker it holds. + + A handler's response is a list of rows, a callable taking (sql, params), or + an exception to raise. Queries that match no handler return no rows. + """ + + def __init__(self, handlers): + self.handlers = list(handlers) + self.executed = [] + self.closed = 0 + + def cursor(self, *args, **kwargs): + return FakeCursor(self) + + def respond(self, sql, params): + for marker, response in self.handlers: + if marker in sql: + if isinstance(response, Exception): + raise response + if callable(response): + return response(sql, params) + return [dict(row) for row in response] + return [] + + def statements(self): + return [sql for sql, _ in self.executed] + + +def citus_handlers( + role="discovery", + local_group_id=0, + use_secondary_nodes="never", + missing_functions=(), + no_execute=(), + missing_tables=(), + no_select=(), + give_warning=True, + nodes=None, + tables=None, + shards=None, + placements=None, + node_failures=None, + remote_tables=None, + remote_columns=None, + remote_indexes=None, + indexes=None, + index_shard_names=None, +): + """Handlers that make a FakeConnection behave like the test cluster.""" + from planetscale_discovery.database import citus + + indexes = CITUS_INDEXES if indexes is None else indexes + index_shard_names = ( + CITUS_INDEX_SHARD_NAMES if index_shard_names is None else index_shard_names + ) + + nodes = CITUS_NODES if nodes is None else nodes + tables = CITUS_TABLES if tables is None else tables + shards = CITUS_SHARDS if shards is None else shards + placements = CITUS_PLACEMENTS if placements is None else placements + node_failures = node_failures or {} + remote = { + citus.TABLE_COMMAND: ( + CITUS_REMOTE_TABLE_ROWS if remote_tables is None else remote_tables + ), + citus.COLUMN_COMMAND: ( + CITUS_REMOTE_COLUMN_ROWS if remote_columns is None else remote_columns + ), + citus.INDEX_COMMAND: ( + CITUS_REMOTE_INDEX_ROWS if remote_indexes is None else remote_indexes + ), + } + + def functions(sql, params): + args = dict(CITUS_FUNCTION_ARGS) + if not give_warning: + args["run_command_on_all_nodes"] = "command text, parallel boolean" + return [ + {"proname": name, "args": a, "can_execute": name not in no_execute} + for name, a in args.items() + if name not in missing_functions + ] + + def dist_tables(sql, params): + return [ + { + "name": name, + "present": name not in missing_tables, + "can_select": None if name in missing_tables else name not in no_select, + } + for name in CITUS_DIST_TABLES + ] + + def run_command(sql, params): + command = params[0] + # params[1] and params[2] hold the batch's shard schemas and names. + batch = ( + {int(name.rsplit("_", 1)[1]) for name in params[2]} + if len(params) > 2 + else None + ) + # (node id reported, node id whose fixture data it returns) + readable = [ + (n["nodeid"], n["nodeid"]) + for n in nodes + if n["isactive"] and n["noderole"] == "primary" + ] + # Like the real function, the current node is added as node id 0 when + # pg_dist_node does not list it. It returns the coordinator's data. + if not any(n["groupid"] == local_group_id for n in nodes): + readable.append((citus.UNLISTED_NODE_ID, 1)) + rows = [] + for nodeid, data in readable: + if nodeid in node_failures: + rows.append( + { + "nodeid": nodeid, + "success": False, + "result": node_failures[nodeid], + } + ) + continue + if command in remote: + node_rows = [ + r + for r in remote[command].get(data, []) + if batch is None + # The server selects by the table shard's name; an index row + # may name its table shard when its own name lacks a shard id. + or int( + ( + r.get("relname") + or r.get("table_relname") + or r.get("index_relname") + ).rsplit("_", 1)[1] + ) + in batch + ] + result = json.dumps(node_rows) + elif command == citus.DATABASE_SIZE_COMMAND: + result = CITUS_DATABASE_SIZES.get(data, "0") + else: + result = "1" + rows.append({"nodeid": nodeid, "success": True, "result": result}) + return rows + + shard_names_by_shard = {} + for row in index_shard_names: + shard_names_by_shard.setdefault(row["shardid"], []).append(row) + + def index_shard_names_in(sql, params): + return [ + row + for shardid in params[0] + for row in shard_names_by_shard.get(shardid, []) + ] + + def size_pretty(sql, params): + return [{"v": v, "pretty": f"{v} bytes"} for v in params[0]] + + def classified(sql, params): + # Applies the relkind filter as the server would, so a fixture table + # may carry a relkind; one without it is an ordinary table. + filtered = f"c.relkind IN {citus.TABLE_RELKINDS}" in sql + return [ + {k: v for k, v in row.items() if k != "relkind"} + for row in tables + if not filtered or row.get("relkind", "r") in ("r", "p") + ] + + return [ + ("WHERE extname = 'citus'", [{"extversion": "14.1-1"}]), + # Quoted as quote_ident() does for a name that is not lowercase and plain. + ( + "pg_catalog.quote_ident(current_user) AS role", + [ + { + "role": ( + role + if re.fullmatch(r"[a-z_][a-z0-9_$]*", role) + else '"' + role.replace('"', '""') + '"' + ) + } + ], + ), + ("pg_get_function_identity_arguments(p.oid) AS args", functions), + ("pg_catalog.to_regclass('pg_catalog.' || t.name)", dist_tables), + ( + "AS local_group_id", + [ + { + "local_group_id": local_group_id, + "citus_version": "Citus 14.1.0 on x86_64-pc-linux-gnu", + "current_database": "postgres", + "use_secondary_nodes": use_secondary_nodes, + "cluster_name": "default", + "local_hostname": "citus-coordinator", + "port": "5432", + "has_autoconverted": True, + } + ], + ), + ("nodecluster::text AS nodecluster", nodes), + ("END AS table_type", classified), + ("AS shard_relname", shards), + ("SELECT shardid, groupid, shardstate", placements), + ("run_command_on_all_nodes(", run_command), + ("AS shard_index_name", index_shard_names_in), + ("i.relname AS index_name", indexes), + ("pg_size_pretty(v) AS pretty", size_pretty), + ("AS shells", [{"shells": 24576}]), + ] + + POSTGRES_CONNECTIONS_RESPONSE = [ { "pid": 12345, diff --git a/tests/integration/test_citus_discovery.py b/tests/integration/test_citus_discovery.py new file mode 100644 index 0000000..5ba7cc1 --- /dev/null +++ b/tests/integration/test_citus_discovery.py @@ -0,0 +1,902 @@ +""" +Discovery against a live Citus cluster. + +Every check compares the report with what the cluster itself says, so the test +holds for any Citus coordinator rather than one seeded schema. Set the +coordinator's connection string to run it; without it, every test skips. + + CITUS_E2E_DSN="host=... port=5432 dbname=postgres user=... password=..." + +The run is read-only: it reads catalogs and statistics views on the coordinator +and, through run_command_on_shards(), run_command_on_placements() and +run_command_on_all_nodes(), on each shard placement. A test whose case the +cluster does not hold skips; test-data/citus-coverage.sql creates every case. +""" + +import json +import os + +import pytest + +psycopg2 = pytest.importorskip("psycopg2") + +from psycopg2.extensions import parse_dsn # noqa: E402 +from psycopg2.extras import RealDictCursor # noqa: E402 + +from planetscale_discovery.database import citus # noqa: E402 +from planetscale_discovery.database.citus import CitusClusterError # noqa: E402 +from planetscale_discovery.database.discovery import PostgreSQLDiscovery # noqa: E402 + +pytestmark = [pytest.mark.integration, pytest.mark.db] + +DSN = os.environ.get("CITUS_E2E_DSN") + + +@pytest.fixture(scope="module") +def dsn(): + if not DSN: + pytest.skip("CITUS_E2E_DSN is not set") + return DSN + + +@pytest.fixture(scope="module") +def connection(dsn): + conn = psycopg2.connect( + dsn, + cursor_factory=RealDictCursor, + options="-c default_transaction_read_only=on", + ) + conn.autocommit = True + yield conn + conn.close() + + +# Shards are shown to the discovery connection, as a role default would show +# them, so the tests below exercise discovery's own hiding and filters rather +# than Citus's default. +SHOW_SHARDS = "-c citus.show_shards_for_app_name_prefixes=*" + + +@pytest.fixture(scope="module") +def report(dsn): + params = parse_dsn(dsn) + params["options"] = f"-c default_transaction_read_only=on {SHOW_SHARDS}" + discovery = PostgreSQLDiscovery(params, citus_detection=True) + assert discovery.connect() + try: + results = discovery.run_analysis() + finally: + discovery.close() + if "citus" not in results["analysis_results"]: + pytest.skip("the citus extension is not installed on CITUS_E2E_DSN") + return results + + +def fetch(connection, sql): + with connection.cursor() as cursor: + cursor.execute(sql) + return [dict(row) for row in cursor.fetchall()] + + +def tables_by_name(report): + return { + (t["schema_name"], t["table_name"]): t + for t in report["analysis_results"]["schema"]["table_analysis"] + } + + +def test_every_user_table_is_discovered(report, connection): + """Catches a regular or Citus table dropped from table_analysis.""" + expected = { + (row["nspname"], row["relname"]) + for row in fetch( + connection, + """ + SELECT n.nspname, c.relname + FROM pg_catalog.pg_class c + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace + WHERE c.relkind IN ('r', 'p') + AND n.nspname NOT IN ('information_schema', 'pg_catalog', 'pg_toast') + AND NOT pg_catalog.relation_is_a_known_shard(c.oid) + AND NOT EXISTS ( + SELECT 1 FROM pg_catalog.pg_depend d + WHERE d.classid = 'pg_catalog.pg_class'::pg_catalog.regclass + AND d.objid = c.oid AND d.deptype = 'e') + """, + ) + } + assert set(tables_by_name(report)) == expected + + +def test_labels_match_citus_own_view(report, connection): + """Catches a table type that disagrees with Citus's citus_tables view.""" + view = fetch( + connection, + """ + SELECT c.oid::pg_catalog.regclass::text AS name + FROM pg_catalog.pg_class c + JOIN pg_catalog.pg_depend d ON d.objid = c.oid AND d.deptype = 'e' + JOIN pg_catalog.pg_extension e ON e.oid = d.refobjid AND e.extname = 'citus' + WHERE c.relname = 'citus_tables' + """, + )[0]["name"] + citus_types = { + (row["nspname"], row["relname"]): row["citus_table_type"] + for row in fetch( + connection, + f""" + SELECT n.nspname, c.relname, t.citus_table_type + FROM {view} t + JOIN pg_catalog.pg_class c ON c.oid = t.table_name::pg_catalog.oid + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace + """, + ) + } + for key, table in tables_by_name(report).items(): + assert table["citus"]["table_type"] == citus_types.get(key, "regular"), key + + +def test_regular_tables_keep_their_coordinator_values(report, connection): + """Catches a regular table on a Citus database whose row estimate is replaced, or + that receives coordinator values meant for Citus tables. + """ + reltuples = { + (row["nspname"], row["relname"]): row["reltuples"] + for row in fetch( + connection, + """ + SELECT n.nspname, c.relname, c.reltuples::bigint AS reltuples + FROM pg_catalog.pg_class c + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace + """, + ) + } + for key, table in tables_by_name(report).items(): + if table["citus"]["table_type"] == "regular": + assert "coordinator" not in table["citus"] + assert table["estimated_rows"] == reltuples[key], key + + +def single_placement_tables(report): + """Citus tables with one placement per shard and every shard analyzed. + + Covers distributed (including single-shard), schema and local tables, whose + values each shard placement holds once. Reference tables, with a placement + on every node, are covered by the unit tests. + """ + for key, table in tables_by_name(report).items(): + citus_info = table["citus"] + if citus_info["table_type"] == "regular": + continue + if citus_info["placements_per_shard"]["max"] != 1: + continue + if citus_info["shard_statistics"]["shards_never_analyzed"]: + continue + yield key, table + + +def on_shards(connection, schema, name, command): + """run_command_on_shards() substitutes the shard name for one %s.""" + return fetch( + connection, + f""" + SELECT result FROM pg_catalog.run_command_on_shards( + '"{schema}"."{name}"'::pg_catalog.regclass, $cmd${command}$cmd$) + """, + ) + + +def test_rows_and_sizes_match_the_shards(report, connection): + """Catches a row estimate or size that disagrees with the shards' own values. + + The row estimate follows the Citus docs' example of summing each shard's + reltuples through run_command_on_shards(). + """ + checked = set() + for (schema, name), table in single_placement_tables(report): + shards = on_shards( + connection, + schema, + name, + "SELECT c.reltuples::bigint || ',' || pg_catalog.pg_relation_size(c.oid)" + " || ',' || pg_catalog.pg_total_relation_size(c.oid)" + " FROM pg_catalog.pg_class c WHERE c.oid = '%s'::pg_catalog.regclass", + ) + values = [tuple(map(int, row["result"].split(","))) for row in shards] + expected = tuple(sum(column) for column in zip(*values)) + reported = ( + table["estimated_rows"], + table["table_size_bytes"], + table["total_size_bytes"], + ) + assert reported == expected, (schema, name) + checked.add(table["citus"]["table_type"]) + assert checked, "no Citus table with one placement per shard was found" + + +def test_single_shard_statistics_match_the_shard(report, connection): + """Catches column statistics of a single-shard table (schema, local, or + distributed without a distribution column) that differ from its shard's.""" + for (schema, name), table in single_placement_tables(report): + if table["citus"]["shard_count"] != 1: + continue + shard = on_shards( + connection, + schema, + name, + "SELECT COALESCE(pg_catalog.json_agg(pg_catalog.json_build_array(" + "s.attname, s.inherited, s.null_frac, s.n_distinct, s.correlation)), '[]')" + " FROM pg_catalog.pg_stats s" + " WHERE pg_catalog.to_regclass(pg_catalog.format(" + "'%%I.%%I', s.schemaname, s.tablename)) = '%s'::pg_catalog.regclass", + )[0]["result"] + expected = { + (attname, inherited): (null_frac, n_distinct, correlation) + for attname, inherited, null_frac, n_distinct, correlation in json.loads( + shard + ) + } + detail = table["citus"]["shard_statistics"]["columns"] + reported = { + (column["attname"], column["inherited"]): ( + row["null_frac"], + row["n_distinct"], + row["correlation"], + ) + for column, row in zip(detail, table["statistics"]) + } + assert reported.keys() == expected.keys(), (schema, name) + for key, values in reported.items(): + assert values == pytest.approx(expected[key], abs=1e-6), (schema, name, key) + + +def strings(value): + """Every string in the report, at any depth.""" + if isinstance(value, dict): + for item in value.values(): + yield from strings(item) + elif isinstance(value, list): + for item in value: + yield from strings(item) + elif isinstance(value, str): + yield value + + +def test_no_shard_object_is_reported_anywhere(report, connection, dsn): + """Catches a shard placement, its index, constraint or row type reported as + a customer object, with shards shown to the discovery connection.""" + shard_names = { + row["name"] + for row in fetch( + connection, + """ + SELECT pg_catalog.shard_name(s.logicalrelid, s.shardid) AS name + FROM pg_catalog.pg_dist_shard s + UNION + SELECT pg_catalog.shard_name(x.indexrelid, s.shardid) + FROM pg_catalog.pg_dist_shard s + JOIN pg_catalog.pg_index x ON x.indrelid = s.logicalrelid + UNION + SELECT c.conname || '_' || s.shardid + FROM pg_catalog.pg_dist_shard s + JOIN pg_catalog.pg_constraint c ON c.conrelid = s.logicalrelid + """, + ) + } + with psycopg2.connect(dsn, options=SHOW_SHARDS) as shown: + with shown.cursor() as cursor: + cursor.execute( + "SELECT pg_catalog.count(*) FROM pg_catalog.pg_class" + " WHERE pg_catalog.relation_is_a_known_shard(oid)" + ) + visible = cursor.fetchone()[0] + if not visible: + pytest.skip("the coordinator holds no shard placement, so none can leak") + leaked = {s for s in strings(report["analysis_results"]) if s in shard_names} + assert not leaked + + +def test_no_extension_object_is_reported(report, connection): + """Catches the extension's schemas, views, functions, types or event + triggers, or grants on them, reported as customer objects.""" + members = fetch( + connection, + """ + SELECT d.classid::pg_catalog.regclass::text AS catalog, + COALESCE(cn.nspname, pn.nspname, tn.nspname) AS schema_name, + COALESCE(n.nspname, c.relname, p.proname, t.typname, e.evtname) AS name + FROM pg_catalog.pg_depend d + JOIN pg_catalog.pg_extension x ON x.oid = d.refobjid AND x.extname = 'citus' + LEFT JOIN pg_catalog.pg_namespace n + ON d.classid = 'pg_catalog.pg_namespace'::pg_catalog.regclass + AND n.oid = d.objid + LEFT JOIN pg_catalog.pg_class c + ON d.classid = 'pg_catalog.pg_class'::pg_catalog.regclass + AND c.oid = d.objid + LEFT JOIN pg_catalog.pg_namespace cn ON cn.oid = c.relnamespace + LEFT JOIN pg_catalog.pg_proc p + ON d.classid = 'pg_catalog.pg_proc'::pg_catalog.regclass + AND p.oid = d.objid + LEFT JOIN pg_catalog.pg_namespace pn ON pn.oid = p.pronamespace + LEFT JOIN pg_catalog.pg_type t + ON d.classid = 'pg_catalog.pg_type'::pg_catalog.regclass + AND t.oid = d.objid + LEFT JOIN pg_catalog.pg_namespace tn ON tn.oid = t.typnamespace + LEFT JOIN pg_catalog.pg_event_trigger e + ON d.classid = 'pg_catalog.pg_event_trigger'::pg_catalog.regclass + AND e.oid = d.objid + WHERE d.deptype = 'e' + """, + ) + + def named(catalog): + return { + (m["schema_name"], m["name"]) for m in members if m["catalog"] == catalog + } + + schemas = {m["name"] for m in members if m["catalog"] == "pg_namespace"} + relations = named("pg_class") + functions = named("pg_proc") + types = named("pg_type") + events = {m["name"] for m in members if m["catalog"] == "pg_event_trigger"} + assert schemas and relations and functions, "no extension member was found" + + results = report["analysis_results"] + schema = results["schema"] + features = results["features"] + privileges = results["security"]["permission_analysis"] + + assert not {s["schema_name"] for s in schema["schema_inventory"]} & schemas + assert ( + not {(v["schema_name"], v["view_name"]) for v in schema["view_analysis"]} + & relations + ) + assert ( + not { + (f["schema_name"], f["function_name"]) for f in schema["function_analysis"] + } + & functions + ) + reported_types = { + (t["schema_name"], t["type_name"]) + for kind in ("composite_types", "enum_types", "domain_types", "range_types") + for t in features["custom_data_types"][kind] + } + # A view's row type is not an extension member itself; it shares the view's name. + assert not reported_types & (types | relations) + assert not {t["trigger_name"] for t in features["event_triggers"]} & events + assert not {p["schema_name"] for p in privileges["schema_privileges"]} & schemas + assert ( + not { + (p["schema_name"], p["table_name"]) for p in privileges["table_privileges"] + } + & relations + ) + assert ( + not { + (p["schema_name"], p["function_name"]) + for p in privileges["function_privileges"] + } + & functions + ) + + +def test_nodes_and_cluster_storage_match_the_cluster(report, connection): + """Catches a node, the coordinator included, missing from the report, or a + cluster storage that is not the sum of every node's database size.""" + sizes = { + row["nodeid"]: int(row["result"]) + for row in fetch( + connection, + """ + SELECT nodeid, result FROM pg_catalog.run_command_on_all_nodes( + 'SELECT pg_catalog.pg_database_size(pg_catalog.current_database())') + """, + ) + } + citus_report = report["analysis_results"]["citus"] + if citus_report["summary"]["cluster_storage_bytes"] is None: + pytest.skip("the database holds no Citus table, so no node is read") + reported = {n["nodeid"]: n for n in citus_report["nodes"]} + assert set(sizes) <= set(reported) + for nodeid, size in sizes.items(): + assert reported[nodeid]["database_size_bytes"] == pytest.approx(size, rel=0.01) + assert reported[nodeid]["nodename"] and reported[nodeid]["nodeport"] + assert citus_report["summary"]["node_count"] == len(reported) + assert citus_report["summary"]["cluster_storage_bytes"] == pytest.approx( + sum(sizes.values()), rel=0.01 + ) + + +def citus_tables(report): + for key, table in tables_by_name(report).items(): + if table["citus"]["table_type"] != "regular": + yield key, table + + +def on_placements(connection, schema, name, command): + """Each placement's result, with the shard and the node group it is on.""" + rows = fetch( + connection, + f""" + SELECT p.shardid, n.groupid, p.success, p.result + FROM pg_catalog.run_command_on_placements( + '"{schema}"."{name}"'::pg_catalog.regclass, $cmd${command}$cmd$) p + JOIN pg_catalog.pg_dist_node n + ON n.nodename = p.nodename AND n.nodeport = p.nodeport + """, + ) + failed = [row["result"] for row in rows if not row["success"]] + assert not failed, (schema, name, failed[:1]) + return rows + + +PLACEMENT_VALUES = ( + "SELECT pg_catalog.json_build_object(" + "'reltuples', c.reltuples," + " 'table_size', pg_catalog.pg_relation_size(c.oid)," + " 'total_size', pg_catalog.pg_total_relation_size(c.oid)," + " 'analyzed', pg_catalog.date_part('epoch'," + " GREATEST(s.last_analyze, s.last_autoanalyze))," + " 'n_tup_ins', s.n_tup_ins, 'n_tup_upd', s.n_tup_upd, 'n_tup_del', s.n_tup_del," + " 'seq_scan', s.seq_scan, 'idx_scan', s.idx_scan)::text" + " FROM pg_catalog.pg_class c" + " LEFT JOIN pg_catalog.pg_stat_user_tables s ON s.relid = c.oid" + " WHERE c.oid = '%s'::pg_catalog.regclass" +) + + +def summed(values): + present = [v for v in values if v is not None] + return sum(present) if present else None + + +def largest(values): + present = [v for v in values if v is not None] + return max(present) if present else None + + +def test_every_citus_table_follows_the_placement_rules(report, connection): + """Catches a replicated shard counted once per placement for rows, sizes or + writes, a read on one placement left out, or rows taken from a placement + other than the one analyzed most recently. + + Covers every Citus table, so reference tables and tables with + citus.shard_replication_factor above 1 are checked against the live + placements rather than only against fakes. + """ + statistics = { + (row["schemaname"], row["tablename"]): row + for row in report["analysis_results"]["performance"]["table_statistics"] + } + replicated = 0 + for (schema, name), table in citus_tables(report): + by_shard = {} + for row in on_placements(connection, schema, name, PLACEMENT_VALUES): + by_shard.setdefault(row["shardid"], []).append( + (row["groupid"], json.loads(row["result"])) + ) + if max(len(rows) for rows in by_shard.values()) > 1: + replicated += 1 + + chosen = [ + sorted( + rows, + key=lambda item: ( + item[1]["analyzed"] is None, + -(item[1]["analyzed"] or 0), + item[0], + ), + )[0][1] + for rows in by_shard.values() + ] + analyzed = [r["reltuples"] for r in chosen if r["reltuples"] >= 0] + assert table["estimated_rows"] == ( + int(round(sum(analyzed))) if analyzed else -1 + ), (schema, name) + assert table["table_size_bytes"] == sum(r["table_size"] for r in chosen) + assert table["total_size_bytes"] == sum(r["total_size"] for r in chosen) + + counters = statistics.get((schema, name)) + if counters is None: + continue + for key in ("n_tup_ins", "n_tup_upd", "n_tup_del"): + expected = summed( + largest(r[key] for _, r in rows) for rows in by_shard.values() + ) + assert counters[key] == expected, (schema, name, key) + for key in ("seq_scan", "idx_scan"): + expected = summed(r[key] for rows in by_shard.values() for _, r in rows) + assert counters[key] == expected, (schema, name, key) + assert replicated, "no Citus table with more than one placement per shard" + + +def test_unanalyzed_shards_are_counted_and_left_out(report, connection): + """Catches a shard that was never analyzed (reltuples -1) summed into the rows, or an + analyzed empty shard (reltuples 0) counted as never analyzed.""" + checked = 0 + for (schema, name), table in citus_tables(report): + never = table["citus"]["shard_statistics"]["shards_never_analyzed"] + if not never: + continue + reltuples = [ + float(row["result"]) + for row in on_shards( + connection, + schema, + name, + "SELECT c.reltuples FROM pg_catalog.pg_class c" + " WHERE c.oid = '%s'::pg_catalog.regclass", + ) + ] + assert never == sum(1 for r in reltuples if r < 0), (schema, name) + analyzed = [r for r in reltuples if r >= 0] + assert table["estimated_rows"] == ( + int(round(sum(analyzed))) if analyzed else -1 + ), (schema, name) + checked += 1 + if not checked: + pytest.skip("no Citus table has a shard that was never analyzed") + + +def shard_index_names(connection): + """(table schema, table, index schema, index) -> {shardid: shard index name}. + + shard_name() applies Citus's own quoting and shortening; parse_ident() + removes the quoting and the schema. + """ + names = {} + for row in fetch( + connection, + """ + SELECT tn.nspname AS table_schema, t.relname AS table_name, + n.nspname AS index_schema, i.relname AS index_name, s.shardid, + (pg_catalog.parse_ident( + pg_catalog.shard_name(x.indexrelid, s.shardid)))[ + pg_catalog.array_length(pg_catalog.parse_ident( + pg_catalog.shard_name(x.indexrelid, s.shardid)), 1)] + AS shard_index + FROM pg_catalog.pg_dist_shard s + JOIN pg_catalog.pg_index x ON x.indrelid = s.logicalrelid + JOIN pg_catalog.pg_class t ON t.oid = x.indrelid + JOIN pg_catalog.pg_namespace tn ON tn.oid = t.relnamespace + JOIN pg_catalog.pg_class i ON i.oid = x.indexrelid + JOIN pg_catalog.pg_namespace n ON n.oid = i.relnamespace + """, + ): + key = ( + row["table_schema"], + row["table_name"], + row["index_schema"], + row["index_name"], + ) + names.setdefault(key, {})[row["shardid"]] = row["shard_index"] + return names + + +SHARD_INDEX_SIZES = ( + "SELECT COALESCE(pg_catalog.json_object_agg(" + "i.relname, pg_catalog.pg_relation_size(i.oid)), '{}')::text" + " FROM pg_catalog.pg_index x" + " JOIN pg_catalog.pg_class i ON i.oid = x.indexrelid" + " WHERE x.indrelid = '%s'::pg_catalog.regclass" +) + + +def shard_index_sizes(connection, schema, name): + """{shardid: {shard index name: size}} for a table with one placement per shard.""" + return { + row["shardid"]: json.loads(row["result"]) + for row in fetch( + connection, + f""" + SELECT shardid, result FROM pg_catalog.run_command_on_shards( + '"{schema}"."{name}"'::pg_catalog.regclass, $cmd${SHARD_INDEX_SIZES}$cmd$) + """, + ) + } + + +def test_citus_indexes_hold_their_shard_indexes_values(report, connection): + """Catches a shard index missed because Citus quoted or shortened its name, a + size other than the sum of its shard indexes, or an invalid index that is not + flagged or is reported with shards where it was never built.""" + names = shard_index_names(connection) + indexes = { + (i["schema_name"], i["index_name"]): i + for i in report["analysis_results"]["schema"]["index_analysis"] + } + checked = 0 + for (schema, name), table in citus_tables(report): + if table["citus"]["placements_per_shard"]["max"] != 1: + continue + sizes = shard_index_sizes(connection, schema, name) + for (t_schema, t_name, i_schema, i_name), by_shard in names.items(): + if (t_schema, t_name) != (schema, name): + continue + entry = indexes[(i_schema, i_name)] + found = { + shardid: sizes[shardid][shard_index] + for shardid, shard_index in by_shard.items() + if shard_index in sizes.get(shardid, {}) + } + assert entry["citus"]["shards_read"] == len(found), (i_schema, i_name) + assert entry["index_size_bytes"] == sum(found.values()), (i_schema, i_name) + if entry["is_valid"]: + assert len(found) == len(by_shard), (i_schema, i_name) + assert "reason" not in entry["citus"] + else: + assert entry["citus"]["reason"] == "invalid_index" + assert entry["citus"]["shard_count"] == len(by_shard) + checked += 1 + assert checked, "no index on a Citus table with one placement per shard" + + +def test_regular_partitions_keep_their_coordinator_values(report, connection): + """Catches a regular partition on a Citus database given a citus key or + values other than its own.""" + own = { + (row["nspname"], row["relname"]): row + for row in fetch( + connection, + """ + SELECT n.nspname, c.relname, c.reltuples::bigint AS reltuples, + pg_catalog.pg_relation_size(c.oid) AS size_bytes + FROM pg_catalog.pg_class c + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace + WHERE c.relispartition AND c.relkind IN ('r', 'p') + AND NOT EXISTS ( + SELECT 1 FROM pg_catalog.pg_dist_partition p + WHERE p.logicalrelid::pg_catalog.oid = c.oid) + """, + ) + } + checked = 0 + for partition in report["analysis_results"]["schema"]["partition_analysis"]: + key = (partition["partition_schema"], partition["partition_name"]) + if key not in own: + continue + assert "citus" not in partition, key + assert partition["partition_size_bytes"] == own[key]["size_bytes"], key + assert partition["estimated_rows"] == own[key]["reltuples"], key + checked += 1 + if not checked: + pytest.skip("the cluster holds no regular partitioned table") + + +def test_regular_inheritance_is_reported(report, connection): + """Catches a regular inheritance child dropped from inheritance_analysis on a + Citus database.""" + expected = { + (row["parent_schema"], row["parent_table"], row["child_schema"], row["child"]) + for row in fetch( + connection, + """ + SELECT pn.nspname AS parent_schema, p.relname AS parent_table, + cn.nspname AS child_schema, c.relname AS child + FROM pg_catalog.pg_inherits i + JOIN pg_catalog.pg_class p ON p.oid = i.inhparent + JOIN pg_catalog.pg_namespace pn ON pn.oid = p.relnamespace + JOIN pg_catalog.pg_class c ON c.oid = i.inhrelid + JOIN pg_catalog.pg_namespace cn ON cn.oid = c.relnamespace + WHERE NOT c.relispartition + AND NOT pg_catalog.relation_is_a_known_shard(c.oid) + """, + ) + } + if not expected: + pytest.skip("the cluster holds no table inheritance") + reported = { + (r["parent_schema"], r["parent_table"], r["child_schema"], r["child_table"]) + for r in report["analysis_results"]["schema"]["inheritance_analysis"] + } + assert expected <= reported + + +def test_foreign_tables_in_citus_metadata_are_reported_as_foreign_tables( + report, connection +): + """Catches a foreign table added to the Citus metadata stopping the run, being + labeled a Citus table, or missing from the foreign tables.""" + managed = { + (row["nspname"], row["relname"]) + for row in fetch( + connection, + """ + SELECT n.nspname, c.relname + FROM pg_catalog.pg_dist_partition p + JOIN pg_catalog.pg_class c ON c.oid = p.logicalrelid::pg_catalog.oid + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace + WHERE c.relkind = 'f' + """, + ) + } + if not managed: + pytest.skip("no foreign table is in the Citus metadata") + features = report["analysis_results"]["features"] + foreign = { + (t["schema_name"], t["table_name"]) + for t in features["foreign_data_wrappers"]["foreign_tables"] + } + assert managed <= foreign + assert not managed & set(tables_by_name(report)) + + +def test_distributed_functions_match_the_metadata(report, connection): + """Catches a distributed function missing from the citus section, or reported + with another distribution argument or colocation group.""" + expected = { + (row["function"], row["distribution_argument_index"], row["colocation_id"]) + for row in fetch( + connection, + """ + SELECT objid::pg_catalog.regprocedure::text AS function, + distribution_argument_index, colocationid AS colocation_id + FROM pg_catalog.pg_dist_object + WHERE classid = 'pg_catalog.pg_proc'::pg_catalog.regclass + AND (distribution_argument_index IS NOT NULL + OR colocationid IS NOT NULL) + """, + ) + } + if not expected: + pytest.skip("the cluster holds no distributed function") + reported = { + (f["function"], f["distribution_argument_index"], f["colocation_id"]) + for f in report["analysis_results"]["citus"]["distributed_functions"] + } + assert reported == expected + + +def rows_read(plan, relations): + """Rows each catalog in relations yields or discards, over every loop of an + EXPLAIN ANALYZE plan, whether it is read by a scan of the table or its index. + """ + total = 0 + if plan.get("Relation Name") in relations: + per_loop = ( + plan.get("Actual Rows", 0) + + plan.get("Rows Removed by Filter", 0) + + plan.get("Rows Removed by Index Recheck", 0) + ) + total += per_loop * plan.get("Actual Loops", 1) + for child in plan.get("Plans", []): + total += rows_read(child, relations) + return total + + +@pytest.mark.parametrize( + "command", [citus.TABLE_COMMAND, citus.COLUMN_COMMAND, citus.INDEX_COMMAND] +) +def test_shard_commands_read_only_their_batch(connection, command): + """Catches a shard command that reads every relation, index or column on a + node for each batch, which grows with the square of the shard count. + + A one-shard batch reads a few catalog rows for each of the table's columns and + indexes; reading a node's whole pg_class, or one schema's, takes hundreds + even on a small cluster. pg_statistic and pg_namespace are not counted: on a + small node a scan of them is cheaper than their index, and the planner + switches as the catalog grows. + """ + shard = fetch( + connection, + """ + SELECT n.nspname, (pg_catalog.parse_ident( + pg_catalog.shard_name(s.logicalrelid, s.shardid)))[ + pg_catalog.array_length(pg_catalog.parse_ident( + pg_catalog.shard_name(s.logicalrelid, s.shardid)), 1)] AS relname, + (SELECT pg_catalog.count(*) FROM pg_catalog.pg_attribute a + WHERE a.attrelid = c.oid AND a.attnum > 0) + + (SELECT pg_catalog.count(*) FROM pg_catalog.pg_index x + WHERE x.indrelid = c.oid) AS width + FROM pg_catalog.pg_dist_shard s + JOIN pg_catalog.pg_class c ON c.oid = s.logicalrelid::pg_catalog.oid + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace + WHERE c.relkind = 'r' + ORDER BY s.shardid + LIMIT 1 + """, + ) + if not shard: + pytest.skip("the cluster holds no shard") + with connection.cursor() as cursor: + cursor.execute( + "SELECT nodeid, success, result FROM pg_catalog.run_command_on_all_nodes(" + "'EXPLAIN (ANALYZE, FORMAT JSON) ' || pg_catalog.format(%s," + " %s::pg_catalog.name[]::text, %s::pg_catalog.name[]::text))", + (command, [shard[0]["nspname"]], [shard[0]["relname"]]), + ) + plans = cursor.fetchall() + for row in plans: + assert row["success"], row["result"] + plan = json.loads(row["result"])[0]["Plan"] + read = rows_read(plan, {"pg_class", "pg_index", "pg_attribute"}) + assert read <= 10 + 3 * shard[0]["width"], (row["nodeid"], read) + + +def test_distributed_schemas_match_the_metadata(report, connection): + """Catches the distributed schemas of a cluster with schema-based sharding + missing from the citus section, or reported with another colocation group.""" + expected = { + (row["schema_name"], row["colocation_id"]) + for row in fetch( + connection, + """ + SELECT schemaid::pg_catalog.regnamespace::text AS schema_name, + colocationid AS colocation_id + FROM pg_catalog.pg_dist_schema + """, + ) + } + if not expected: + pytest.skip("the cluster holds no distributed schema") + reported = { + (s["schema_name"], s["colocation_id"]) + for s in report["analysis_results"]["citus"]["distributed_schemas"] + } + assert reported == expected + + +def test_an_unpropagated_index_stops_with_its_remedy(connection, dsn): + """Catches a valid index that exists on no shard, as after CREATE INDEX with + citus.enable_ddl_propagation off, reported with a partial size instead of + stopping, or stopping with the partition naming remedy. + + Runs only while such an index exists; the optional section of + test-data/citus-coverage.sql creates one. + """ + tables = { + (row["nspname"], row["relname"]) + for row in fetch( + connection, + """ + SELECT n.nspname, c.relname + FROM pg_catalog.pg_dist_partition p + JOIN pg_catalog.pg_class c ON c.oid = p.logicalrelid::pg_catalog.oid + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace + WHERE c.relkind IN ('r', 'p') AND p.repmodel <> 't' + """, + ) + } + valid = { + (row["nspname"], row["relname"]) + for row in fetch( + connection, + """ + SELECT n.nspname, i.relname + FROM pg_catalog.pg_index x + JOIN pg_catalog.pg_class i ON i.oid = x.indexrelid + JOIN pg_catalog.pg_namespace n ON n.oid = i.relnamespace + WHERE x.indisvalid AND NOT i.relispartition + """, + ) + } + unpropagated = set() + sizes = {} + for (t_schema, t_name, i_schema, i_name), by_shard in shard_index_names( + connection + ).items(): + if (t_schema, t_name) not in tables or (i_schema, i_name) not in valid: + continue + if (t_schema, t_name) not in sizes: + sizes[(t_schema, t_name)] = shard_index_sizes(connection, t_schema, t_name) + found = sizes[(t_schema, t_name)] + if not any(name in found.get(s, {}) for s, name in by_shard.items()): + unpropagated.add(f"index {i_schema}.{i_name}") + if not unpropagated: + pytest.skip("no valid index on a Citus table is missing from all its shards") + + params = parse_dsn(dsn) + params["options"] = "-c default_transaction_read_only=on" + discovery = PostgreSQLDiscovery(params, citus_detection=True) + assert discovery.connect() + try: + with pytest.raises(CitusClusterError) as error: + discovery.run_analysis() + finally: + discovery.close() + message = str(error.value) + for subject in unpropagated: + assert subject in message + assert "citus.enable_ddl_propagation off" in message + assert "fix_all_partition_shard_index_names" not in message diff --git a/tests/unit/test_citus_analyzer.py b/tests/unit/test_citus_analyzer.py new file mode 100644 index 0000000..4a0571f --- /dev/null +++ b/tests/unit/test_citus_analyzer.py @@ -0,0 +1,176 @@ +"""The Citus module: nodes, colocation groups, tables and placements.""" + +import pytest + +from planetscale_discovery.database import citus +from planetscale_discovery.database.analyzers.citus_analyzer import CitusAnalyzer +from planetscale_discovery.database.citus import CitusClusterError +from tests.fixtures.database_responses import FakeConnection, citus_handlers + +COLOCATION = ( + "FROM pg_catalog.pg_dist_colocation", + [ + { + "colocation_id": 3, + "shard_count": 2, + "replication_factor": 1, + "distribution_column_type": "integer", + }, + { + "colocation_id": 4, + "shard_count": 1, + "replication_factor": -1, + "distribution_column_type": "-", + }, + ], +) +REGULAR = ("AS regular_count", [{"regular_count": 2}]) + + +def analyzer(*extra): + connection = FakeConnection(list(extra) + citus_handlers() + [COLOCATION, REGULAR]) + ctx = citus.detect_citus(connection) + citus.preflight(connection, ctx) + return CitusAnalyzer(connection, config={"citus": ctx}) + + +class TestCitusAnalyzer: + def test_the_summary_counts_regular_tables_beside_citus_tables(self): + """Catches regular tables left out of the Citus summary's table counts.""" + summary = analyzer().analyze()["summary"] + assert summary["table_counts"] == { + "distributed": 1, + "reference": 1, + "regular": 2, + } + assert summary["node_count"] == 3 + assert ( + summary["cluster_storage_bytes"] == 19327047359 + 88322569919 + 89151944383 + ) + + def test_a_coordinator_missing_from_the_metadata_is_still_reported(self): + """Catches the coordinator dropped from nodes, node_count and + cluster_storage_bytes when pg_dist_node lists the workers only, or its + placements labeled without its host and port.""" + from tests.fixtures.database_responses import CITUS_NODES + + workers = [n for n in CITUS_NODES if n["groupid"] != 0] + connection = FakeConnection( + citus_handlers(nodes=workers) + [COLOCATION, REGULAR] + ) + ctx = citus.detect_citus(connection) + citus.preflight(connection, ctx) + result = CitusAnalyzer(connection, config={"citus": ctx}).analyze() + + coordinator = [n for n in result["nodes"] if n["groupid"] == 0] + assert len(coordinator) == 1 + assert coordinator[0]["nodeid"] == citus.UNLISTED_NODE_ID + assert (coordinator[0]["nodename"], coordinator[0]["nodeport"]) == ( + "citus-coordinator", + 5432, + ) + assert coordinator[0]["database_size_bytes"] == 19327047359 + assert result["summary"]["node_count"] == 3 + assert ( + result["summary"]["cluster_storage_bytes"] + == 19327047359 + 88322569919 + 89151944383 + ) + by_node = result["shard_placement"]["by_node"] + assert "citus-coordinator:5432" in by_node + assert not [label for label in by_node if label.startswith(("group ", "None"))] + + def test_tables_report_logical_and_placement_sizes(self): + """Catches a reference table's logical size counted once for every placement.""" + tables = {t["table_name"]: t for t in analyzer().analyze()["tables"]} + reference = tables["reference_t1"] + assert reference["logical_size_bytes"] == 5801107456 + assert reference["placement_size_bytes"] == 3 * 5801107456 + assert reference["placements_per_shard"] == {"min": 3, "max": 3} + assert tables["pgbench_branches"]["distribution_column"] == "bid" + + def test_placements_are_summed_per_node(self): + """Catches a placement on the coordinator left out of the totals for each node.""" + by_node = analyzer().analyze()["shard_placement"]["by_node"] + assert by_node["citus-coordinator:5432"] == { + "placements": 1, + "bytes": 5801107456, + } + assert by_node["citus-worker-1:5432"]["placements"] == 2 + + def test_colocation_groups_list_their_tables(self): + """Catches a colocation group listed without its tables.""" + groups = analyzer().analyze()["colocation_groups"] + assert groups[0]["tables"] == ["public.pgbench_branches"] + + def test_tables_are_grouped_by_colocation_once(self): + """Catches a scan of every table for each colocation group: schema-based + sharding gives every schema its own group, so 10k schemas of 100k + tables make that 1e9 iterations.""" + instance = analyzer() + + class CountingDict(dict): + scans = 0 + + def values(self): + CountingDict.scans += 1 + return super().values() + + instance.citus.tables = CountingDict(instance.citus.tables) + instance._get_colocation_groups() + assert CountingDict.scans == 1 + + def test_without_a_citus_table_no_node_is_read(self): + """Catches the node size read contacting nodes for an unused extension, + or a cluster storage total that silently omits every node.""" + connection = FakeConnection( + citus_handlers(tables=[], shards=[], placements=[], node_failures={1: "x"}) + + [COLOCATION, REGULAR] + ) + ctx = citus.detect_citus(connection) + citus.preflight(connection, ctx) + result = CitusAnalyzer(connection, config={"citus": ctx}).analyze() + assert not any( + "run_command_on_all_nodes(" in sql for sql in connection.statements() + ) + assert all(n["database_size_bytes"] is None for n in result["nodes"]) + assert result["summary"]["cluster_storage_bytes"] is None + assert result["summary"]["table_counts"] == {"regular": 2} + + def test_distributed_schemas_are_reported(self): + """Catches the distributed schemas of a cluster with schema-based sharding + reported as an empty list.""" + schemas = ( + "FROM pg_catalog.pg_dist_schema\n", + [ + {"schema_name": "tenant_acme", "colocation_id": 7}, + {"schema_name": "tenant_globex", "colocation_id": 8}, + ], + ) + assert analyzer(schemas).analyze()["distributed_schemas"] == schemas[1] + + def test_a_citus_version_without_distributed_schemas_reads_none(self): + """Catches a query on pg_dist_schema on a Citus version that lacks it, which + would stop the run as a failed section.""" + connection = FakeConnection( + citus_handlers(missing_tables=("pg_dist_schema",)) + [COLOCATION, REGULAR] + ) + ctx = citus.detect_citus(connection) + citus.preflight(connection, ctx) + result = CitusAnalyzer(connection, config={"citus": ctx}).analyze() + assert result["distributed_schemas"] == [] + assert not any( + "FROM pg_catalog.pg_dist_schema\n" in sql for sql in connection.statements() + ) + + def test_it_requires_the_context_from_detection(self): + """Catches a CitusAnalyzer created without a context, where every section would + fail on None. Discovery creates the module only after detection. + """ + with pytest.raises(ValueError, match="Citus context"): + CitusAnalyzer(FakeConnection([]), config={}) + + def test_a_failing_section_stops_instead_of_going_missing(self): + """Catches a Citus report written with a section silently absent.""" + failing = ("FROM pg_catalog.pg_dist_object", RuntimeError("permission denied")) + with pytest.raises(CitusClusterError, match="cluster analysis failed"): + analyzer(failing).analyze() diff --git a/tests/unit/test_citus_detection.py b/tests/unit/test_citus_detection.py new file mode 100644 index 0000000..e80933b --- /dev/null +++ b/tests/unit/test_citus_detection.py @@ -0,0 +1,648 @@ +"""Citus detection and the preflight check that stops a discovery early.""" + +import re +from pathlib import Path + +import pytest + +from planetscale_discovery.database import citus +from planetscale_discovery.database.citus import CitusClusterError +from tests.fixtures.database_responses import ( + CITUS_NODES, + CITUS_PLACEMENTS, + FakeConnection, + citus_handlers, +) + + +def prepared(**kwargs): + connection = FakeConnection(citus_handlers(**kwargs)) + ctx = citus.detect_citus(connection) + return connection, ctx + + +class TestDetection: + def test_a_server_without_the_extension_runs_no_other_citus_query(self): + """Catches Citus queries reaching a server that does not run Citus.""" + connection = FakeConnection([]) + assert citus.detect_citus(connection) is None + assert len(connection.executed) == 1 + + def test_the_extension_is_detected(self): + """Catches a database with the citus extension not detected as one.""" + _, ctx = prepared() + assert ctx.extension_version == "14.1-1" + + +class TestPreflight: + def test_a_readable_cluster_passes_and_fills_the_context(self): + """Catches a healthy cluster stopped by the preflight, or a context left without + its nodes and tables. + """ + connection, ctx = prepared() + citus.preflight(connection, ctx) + assert ctx.node_role == "coordinator" + assert ctx.table_counts() == {"distributed": 1, "reference": 1} + assert ctx.readable_groups() == {0, 1, 2} + + def test_connection_info_holds_no_count_the_summary_also_holds(self): + """Catches a second copy of the table or node counts, which disagreed + with analysis_results.citus.summary: Citus types only against + regular included, and pg_dist_node rows against nodes read.""" + connection, ctx = prepared() + citus.preflight(connection, ctx) + assert ctx.connection_info() == { + "extension_version": "14.1-1", + "citus_version": "Citus 14.1.0 on x86_64-pc-linux-gnu", + "node_role": "coordinator", + } + + def test_a_connection_to_a_worker_stops(self): + """Catches discovery reporting a worker's view as the cluster's.""" + connection, ctx = prepared(local_group_id=1) + with pytest.raises(CitusClusterError, match="Connect to the coordinator"): + citus.preflight(connection, ctx) + + def test_a_missing_function_stops_as_unsupported(self): + """Catches a Citus version that lacks a function the read needs being run + anyway. + """ + connection, ctx = prepared(missing_functions=("relation_is_a_known_shard",)) + with pytest.raises(CitusClusterError) as error: + citus.preflight(connection, ctx) + assert "relation_is_a_known_shard" in str(error.value) + assert "not supported" in str(error.value) + + def test_a_missing_metadata_table_stops_as_unsupported(self): + """Catches a Citus version that lacks a metadata table being run anyway.""" + connection, ctx = prepared(missing_tables=("pg_dist_placement",)) + with pytest.raises(CitusClusterError) as error: + citus.preflight(connection, ctx) + assert "pg_catalog.pg_dist_placement: table not found" in str(error.value) + assert "not supported" in str(error.value) + + def test_shards_are_hidden_for_the_session_first(self): + """Catches shards listed as tables when a role default shows them.""" + connection, ctx = prepared() + connection.executed.clear() + citus.preflight(connection, ctx) + first = connection.statements()[0] + assert ( + "set_config('citus.show_shards_for_app_name_prefixes', '', false)" in first + ) + + def test_the_coordinator_is_read_when_the_metadata_omits_it(self): + """Catches the coordinator's placements left unread when pg_dist_node does not + list it. run_command_on_all_nodes() adds the current node in that case. + """ + workers = [n for n in CITUS_NODES if n["groupid"] != 0] + connection, ctx = prepared(nodes=workers) + citus.preflight(connection, ctx) + assert ctx.readable_groups() == {0, 1, 2} + + def test_a_missing_execute_privilege_names_the_grant(self): + """Catches a permission error that is reported as a warning.""" + connection, ctx = prepared(no_execute=("run_command_on_all_nodes",)) + with pytest.raises(CitusClusterError) as error: + citus.preflight(connection, ctx) + assert "GRANT EXECUTE ON FUNCTION pg_catalog.run_command_on_all_nodes(" in str( + error.value + ) + assert "TO discovery" in str(error.value) + + def test_a_missing_select_privilege_names_the_grant(self): + """Catches a missing SELECT on a pg_dist_* table reported without the GRANT that + fixes it. + """ + connection, ctx = prepared(no_select=("pg_dist_placement",)) + with pytest.raises(CitusClusterError) as error: + citus.preflight(connection, ctx) + assert "GRANT SELECT ON pg_catalog.pg_dist_placement TO discovery" in str( + error.value + ) + + @pytest.mark.parametrize("role", ["app-reader", "AppUser"]) + def test_a_grant_names_a_role_that_needs_quoting_quoted(self, role): + """Catches a GRANT remedy that is a syntax error for app-reader, or grants to + appuser, the name folded to lowercase, instead of AppUser.""" + # A missing metadata privilege stops the preflight before the remote one + # is checked, so each remedy is read from its own run. + connection, ctx = prepared(role=role, no_select=("pg_dist_placement",)) + with pytest.raises(CitusClusterError) as error: + citus.preflight(connection, ctx) + assert f'GRANT SELECT ON pg_catalog.pg_dist_placement TO "{role}"' in str( + error.value + ) + assert any( + "pg_catalog.quote_ident(current_user)" in sql + for sql in connection.statements() + ) + + connection, ctx = prepared(role=role, no_execute=("run_command_on_all_nodes",)) + with pytest.raises(CitusClusterError) as error: + citus.preflight(connection, ctx) + assert f') TO "{role}"' in str(error.value) + + def test_reading_from_secondaries_stops(self): + """Catches write counters read from a secondary, where they are 0.""" + connection, ctx = prepared(use_secondary_nodes="always") + with pytest.raises(CitusClusterError, match="use_secondary_nodes"): + citus.preflight(connection, ctx) + + def test_a_shard_only_on_an_inactive_node_stops(self): + """Catches a shard that no node read, summed as if it held no rows.""" + nodes = [dict(n) for n in CITUS_NODES] + nodes[2]["isactive"] = False + connection, ctx = prepared(nodes=nodes) + with pytest.raises(CitusClusterError) as error: + citus.preflight(connection, ctx) + assert "shard 102138 of public.pgbench_branches" in str(error.value) + assert "citus-worker-2:5432 (primary, inactive)" in str(error.value) + + def test_a_shard_only_on_a_secondary_stops(self): + """Catches a shard whose only placement is on a secondary, which + run_command_on_all_nodes() does not read, summed as if it held no rows. + """ + nodes = [dict(n) for n in CITUS_NODES] + [ + dict(CITUS_NODES[2], nodeid=4, groupid=3, noderole="secondary") + ] + placements = [dict(p) for p in CITUS_PLACEMENTS] + placements[1]["groupid"] = 3 + connection, ctx = prepared(nodes=nodes, placements=placements) + with pytest.raises(CitusClusterError, match="shard 102138"): + citus.preflight(connection, ctx) + + def test_inactive_nodes_and_secondaries_are_not_expected_to_answer(self): + """Catches a run that stops because an inactive node or a secondary, which + run_command_on_all_nodes() does not contact, returned no result.""" + nodes = [dict(n) for n in CITUS_NODES] + [ + dict(CITUS_NODES[2], nodeid=4, groupid=3, isactive=False), + dict(CITUS_NODES[2], nodeid=5, groupid=2, noderole="secondary"), + ] + connection, ctx = prepared(nodes=nodes) + citus.preflight(connection, ctx) + assert 4 not in ctx.readable_node_ids() + assert 5 not in ctx.readable_node_ids() + + def test_a_worker_without_citus_metadata_stops_up_front(self): + """Catches a worker whose metadata was never synced, as after an upgrade from + Citus 10 that was not finalized, stopping the run later with a list of + unreadable shards instead of naming the node and the remedy.""" + nodes = [dict(n) for n in CITUS_NODES] + nodes[1]["hasmetadata"] = False + connection, ctx = prepared(nodes=nodes) + with pytest.raises(CitusClusterError) as error: + citus.preflight(connection, ctx) + message = str(error.value) + assert ( + "citus-worker-1:5432: Citus metadata is not synced to the node" in message + ) + assert "citus_finalize_upgrade_to_citus11()" in message + assert not any( + "run_command_on_all_nodes(" in sql for sql in connection.statements() + ), "the check runs before any node is read" + + def test_a_worker_without_metadata_or_placements_passes(self): + """Catches a worker without metadata that holds no placement stopping the + run, although none of its shards could be misread.""" + nodes = [dict(n) for n in CITUS_NODES] + [ + dict(CITUS_NODES[2], nodeid=4, groupid=3, hasmetadata=False) + ] + connection, ctx = prepared(nodes=nodes) + citus.preflight(connection, ctx) + + def test_the_coordinator_row_never_needs_hasmetadata(self): + """Catches the coordinator stopped for hasmetadata false on its own + pg_dist_node row, although it holds the metadata it would sync.""" + nodes = [dict(n) for n in CITUS_NODES] + nodes[0]["hasmetadata"] = False + connection, ctx = prepared(nodes=nodes) + citus.preflight(connection, ctx) + + def test_a_node_added_during_the_run_stops_it(self): + """Catches the rows of a node added to pg_dist_node after the preflight + credited to the coordinator, where a reference table's copy is added to + the coordinator's and other shards are taken for leftovers of a move.""" + message = a_node_added_during_the_run() + assert "nodes changed during discovery" in message + assert "node 9: not in pg_dist_node when discovery started" in message + + def test_node_id_zero_is_rejected_when_the_coordinator_is_listed(self): + """Catches node id 0 accepted as the coordinator while pg_dist_node lists + it under its own id, so a second copy of its rows would be read.""" + connection, ctx = prepared() + citus.preflight(connection, ctx) + with_node_answering(connection, citus.UNLISTED_NODE_ID) + with pytest.raises(CitusClusterError, match="node 0: not in pg_dist_node"): + ctx.shard_statistics(connection) + + def test_an_unlisted_coordinator_is_named_by_its_host(self): + """Catches a failure on a coordinator that pg_dist_node does not list named + "node 0", which tells the customer nothing.""" + workers = [n for n in CITUS_NODES if n["groupid"] != 0] + message = stopped_preflight( + nodes=workers, + node_failures={citus.UNLISTED_NODE_ID: "could not connect: timeout"}, + ) + assert "citus-coordinator:5432: unreachable" in message + + def test_a_primary_without_placements_must_answer(self): + """Catches an active primary that holds no placement left out of the check, + although run_command_on_all_nodes() connects to it on every read.""" + nodes = [dict(n) for n in CITUS_NODES] + [ + dict(CITUS_NODES[2], nodeid=4, groupid=3, nodename="citus-worker-3") + ] + message = stopped_preflight( + nodes=nodes, + node_failures={4: "connection to the remote node failed: timeout"}, + ) + assert "citus-worker-3:5432: unreachable" in message + + def test_an_inactive_placement_beside_an_active_one_passes(self): + """Catches an inactive placement stopping a run in which every shard has an + active placement. + """ + placements = CITUS_PLACEMENTS + [ + {"shardid": 102137, "groupid": 2, "shardstate": 3} + ] + connection, ctx = prepared(placements=placements) + citus.preflight(connection, ctx) + + def test_an_unreachable_node_stops_and_is_named(self): + """Catches an unreachable node that does not stop the run, or that the message + does not name. + """ + connection, ctx = prepared( + node_failures={ + 3: "connection to the remote node failed: Connection refused" + } + ) + with pytest.raises(CitusClusterError) as error: + citus.preflight(connection, ctx) + assert "citus-worker-2:5432: unreachable" in str(error.value) + + def test_a_rejected_login_is_named(self): + """Catches a node that rejects the discovery role reported as unreachable.""" + connection, ctx = prepared( + node_failures={2: 'FATAL: password authentication failed for user "d"'} + ) + with pytest.raises(CitusClusterError, match="rejected the discovery role"): + citus.preflight(connection, ctx) + + def test_a_failed_classification_stops(self): + """Catches every Citus table reported with its shell's values.""" + handlers = [("END AS table_type", RuntimeError("permission denied"))] + connection = FakeConnection(handlers + citus_handlers()) + ctx = citus.detect_citus(connection) + with pytest.raises(CitusClusterError, match="metadata could not be read"): + citus.preflight(connection, ctx) + + def test_it_is_a_connection_error(self): + """Catches a Citus failure that bypasses the CLI's fatal path, which catches + ConnectionError. + """ + assert issubclass(CitusClusterError, ConnectionError) + + +class TestCitusErrors: + def test_a_failure_is_raised_as_a_citus_cluster_error(self): + """Catches a failure in a Citus read escaping as a plain exception, which + run_analysis() records as a module failure and writes a report with shell + values.""" + with pytest.raises(CitusClusterError) as error: + with citus.citus_errors("Values could not be read.", "schema analysis"): + raise KeyError("pgbench_branches") + assert error.value.summary == "Values could not be read." + assert error.value.problems == [ + { + "subject": "schema analysis", + "cause": "failed", + "detail": "'pgbench_branches'", + } + ] + assert isinstance(error.value.__cause__, KeyError) + + def test_a_citus_cluster_error_passes_through_unchanged(self): + """Catches a node or privilege problem replaced by a generic summary that no + longer names the node or the GRANT.""" + original = CitusClusterError( + "A node is down.", [{"subject": "n", "cause": "x"}] + ) + with pytest.raises(CitusClusterError) as error: + with citus.citus_errors("Values could not be read.", "schema analysis"): + raise original + assert error.value is original + + +class TestUnusedExtension: + """The extension is installed, and no table was ever distributed.""" + + NO_TABLES = dict(tables=[], shards=[], placements=[]) + + def test_no_node_is_contacted(self): + """Catches a single node whose loopback login fails stopping a run that + has no Citus table to misreport.""" + connection, ctx = prepared( + node_failures={1: "connection to the remote node localhost:5432 failed"}, + **self.NO_TABLES, + ) + citus.preflight(connection, ctx) + assert not any( + "run_command_on_all_nodes(" in sql for sql in connection.statements() + ) + + def test_checks_that_only_the_node_read_needs_are_skipped(self): + """Catches the secondary node setting or the run_command_on_all_nodes privilege + stopping a run that reads no node. + """ + connection, ctx = prepared( + use_secondary_nodes="always", + no_execute=("run_command_on_all_nodes",), + **self.NO_TABLES, + ) + citus.preflight(connection, ctx) + + def test_unreadable_metadata_still_stops(self): + """Catches a run that continues without SELECT on pg_dist_partition. Without it, + no Citus table cannot be told apart from a Citus table that cannot be read. + """ + connection, ctx = prepared(no_select=("pg_dist_partition",), **self.NO_TABLES) + with pytest.raises(CitusClusterError, match="pg_dist_partition"): + citus.preflight(connection, ctx) + + def test_a_worker_connection_still_stops(self): + """Catches a worker without synced metadata, whose pg_dist_partition is + empty while it holds shards, being read as a regular database.""" + connection, ctx = prepared(local_group_id=1, **self.NO_TABLES) + with pytest.raises(CitusClusterError, match="Connect to the coordinator"): + citus.preflight(connection, ctx) + + def test_the_node_read_checks_apply_once_a_citus_table_exists(self): + """Catches the node read checks skipped on a database that holds Citus tables.""" + connection, ctx = prepared(no_execute=("run_command_on_all_nodes",)) + with pytest.raises(CitusClusterError, match="cannot read the Citus shards"): + citus.preflight(connection, ctx) + + +class TestOptionalFeatures: + def test_features_are_detected_rather_than_assumed(self): + """Catches SQL that names a column or argument an older Citus lacks.""" + connection, ctx = prepared( + missing_tables=("pg_dist_schema",), give_warning=False + ) + citus.preflight(connection, ctx) + assert not ctx.has_pg_dist_schema + assert not ctx.has_connection_error_warning + sql = " ".join(connection.statements()) + assert "pg_dist_schema)" not in citus._classification_sql(ctx) + assert "give_warning_for_connection_errors :=" not in sql + + def test_the_classification_follows_citus_own_order(self): + """Catches a reordered CASE, which mislabels tables that match two + branches (a reference table also has colocationid != 0). + + Agreement with Citus's own citus_tables view is checked against a live + cluster in tests/integration/test_citus_discovery.py. + """ + _, ctx = prepared() + ctx.has_pg_dist_schema = True + sql = citus._classification_sql(ctx) + order = [ + sql.index(fragment) + for fragment in ( + "THEN 'schema'", + "WHEN p.partkey IS NOT NULL THEN 'distributed'", + "WHEN p.repmodel = 't' THEN 'reference'", + "WHEN p.colocationid = 0 THEN 'local'", + "ELSE 'distributed'", + ) + ] + assert order == sorted(order) + + +class TestNames: + @pytest.mark.parametrize( + "name, expected", + [ + ("orders_102008", "orders_102008"), + ("sales.orders_102008", "orders_102008"), + ('"Sales"."Orders_102008"', "Orders_102008"), + ('"a.b"."c""d_1"', 'c"d_1'), + ], + ) + def test_shard_names_lose_their_schema(self, name, expected): + """Catches a shard in a quoted schema, or one other than public, that maps to + nothing.""" + assert citus.unqualified(name) == expected + + +class TestSqlFragments: + def test_fragments_are_empty_without_citus(self): + """Catches a Citus predicate added to an analyzer's SQL on a server without + Citus, where every analyzer must run exactly today's SQL. + """ + assert citus.relation_filter(None, "c.oid", "n.oid") == "" + assert citus.namespace_filter(None, "n.oid") == "" + assert citus.function_filter(None, "p.oid", "n.oid") == "" + assert citus.type_filter(None, "t.oid", "n.oid", "t.typrelid") == "" + assert citus.event_trigger_filter(None, "pg_event_trigger.oid") == "" + assert citus.namespace_name_filter(None, "schemaname") == "" + + def test_where_turns_conditions_into_a_clause_or_nothing(self): + """Catches a WHERE clause added without Citus, which changes the SQL that + every non-Citus run, workload capture's schema read included, executes.""" + _, ctx = prepared() + assert citus.where(citus.relation_filter(None, "c.oid")) == "" + clause = citus.where(citus.relation_filter(ctx, "c.oid")) + assert clause.startswith( + "WHERE NOT pg_catalog.relation_is_a_known_shard(c.oid)" + ) + assert " AND NOT EXISTS" in clause + + def test_extension_membership_is_matched_by_identity(self): + """Catches a customer schema named like a Citus schema being hidden.""" + _, ctx = prepared() + fragment = citus.namespace_filter(ctx, "n.oid") + assert "deptype = 'e'" in fragment + assert not re.search(r"nspname\s+(NOT\s+)?IN", fragment) + + +TROUBLESHOOTING = Path(__file__).parents[2] / "docs" / "troubleshooting.md" + + +def documented_messages(): + """The messages quoted in the troubleshooting table for a stopped Citus run.""" + text = TROUBLESHOOTING.read_text() + section = text.split("### Citus Discovery Stopped", 1)[1].split("\n### ", 1)[0] + rows = [line for line in section.splitlines() if line.startswith("| `")] + return {m for row in rows for m in re.findall(r"`([^`]+)`", row.split("|")[1])} + + +def stopped_preflight(**kwargs): + connection, ctx = prepared(**kwargs) + with pytest.raises(CitusClusterError) as error: + citus.preflight(connection, ctx) + return str(error.value) + + +def a_node_returns_nothing(): + connection, ctx = prepared() + connection.handlers.insert( + 0, + ( + "run_command_on_all_nodes(", + lambda sql, params: [ + {"nodeid": 1, "success": True, "result": "1"}, + {"nodeid": 2, "success": True, "result": "1"}, + ], + ), + ) + with pytest.raises(CitusClusterError) as error: + citus.preflight(connection, ctx) + return str(error.value) + + +def an_index_on_no_shard(): + from tests.fixtures.database_responses import CITUS_REMOTE_INDEX_ROWS + + remote = { + nodeid: [ + r + for r in rows + if not r["index_relname"].startswith("pgbench_branches_pkey_") + ] + for nodeid, rows in CITUS_REMOTE_INDEX_ROWS.items() + } + connection, ctx = prepared(remote_indexes=remote) + citus.preflight(connection, ctx) + with pytest.raises(CitusClusterError) as error: + ctx.shard_statistics(connection) + return str(error.value) + + +def with_node_answering(connection, nodeid): + """Add a node that answers every run_command_on_all_nodes() call, as one + added to pg_dist_node after the preflight read it would.""" + original = next( + r for m, r in connection.handlers if m == "run_command_on_all_nodes(" + ) + + def respond(sql, params): + rows = original(sql, params) + return rows + [dict(rows[-1], nodeid=nodeid)] + + connection.handlers.insert(0, ("run_command_on_all_nodes(", respond)) + + +def a_node_added_during_the_run(): + connection, ctx = prepared() + citus.preflight(connection, ctx) + with_node_answering(connection, 9) + with pytest.raises(CitusClusterError) as error: + ctx.shard_statistics(connection) + return str(error.value) + + +def an_index_dropped_during_the_read(): + from tests.fixtures.database_responses import ( + CITUS_INDEX_SHARD_NAMES, + CITUS_INDEXES, + CITUS_REMOTE_INDEX_ROWS, + ) + + kept = [r for r in CITUS_INDEXES if r["index_name"] != "pgbench_branches_pkey"] + kept_oids = {r["index_oid"] for r in kept} + remote = { + nodeid: [ + r + for r in rows + if not r["index_relname"].startswith("pgbench_branches_pkey_") + ] + for nodeid, rows in CITUS_REMOTE_INDEX_ROWS.items() + } + connection, ctx = prepared( + remote_indexes=remote, + index_shard_names=[ + r for r in CITUS_INDEX_SHARD_NAMES if r["index_oid"] in kept_oids + ], + ) + citus.preflight(connection, ctx) + reads = [] + + def index_list(sql, params): + reads.append(sql) + return CITUS_INDEXES if len(reads) == 1 else kept + + connection.handlers.insert(0, ("i.relname AS index_name", index_list)) + with pytest.raises(CitusClusterError) as error: + ctx.shard_statistics(connection) + return str(error.value) + + +def a_shard_only_on_an_inactive_node(): + nodes = [dict(n) for n in CITUS_NODES] + nodes[2]["isactive"] = False + return stopped_preflight(nodes=nodes) + + +# Each message the troubleshooting table quotes, and the situation it names. +SITUATIONS = { + "Discovery is connected to a Citus worker": lambda: stopped_preflight( + local_group_id=1 + ), + ": unreachable": lambda: stopped_preflight( + node_failures={3: "connection to the remote node failed: Connection refused"} + ), + ": returned no result": a_node_returns_nothing, + ": the node rejected the discovery role": lambda: stopped_preflight( + node_failures={2: 'FATAL: password authentication failed for user "d"'} + ), + "no active placement on a node that discovery can read": ( + a_shard_only_on_an_inactive_node + ), + "node : not in pg_dist_node when discovery started": ( + a_node_added_during_the_run + ), + ": Citus metadata is not synced to the node": lambda: stopped_preflight( + nodes=[dict(n, hasmetadata=n["groupid"] == 0) for n in CITUS_NODES] + ), + "missing EXECUTE privilege": lambda: stopped_preflight( + no_execute=("run_command_on_all_nodes",) + ), + "missing SELECT privilege": lambda: stopped_preflight( + no_select=("pg_dist_placement",) + ), + "citus.use_secondary_nodes": lambda: stopped_preflight( + use_secondary_nodes="always" + ), + "index : not found on N of M shards": an_index_on_no_shard, + "index : dropped during discovery": an_index_dropped_during_the_read, + "function not found": lambda: stopped_preflight( + missing_functions=("relation_is_a_known_shard",) + ), + "table not found; this Citus version is not supported": lambda: stopped_preflight( + missing_tables=("pg_dist_placement",) + ), +} + + +def as_pattern(message): + pattern = re.escape(message) + pattern = pattern.replace(re.escape(""), r"[\w.-]+:\d+") + pattern = pattern.replace(re.escape(""), r"\S+") + pattern = pattern.replace(re.escape(""), r"\d+") + return pattern.replace(re.escape("N of M"), r"\d+ of \d+") + + +class TestTroubleshootingMessages: + def test_every_documented_message_has_a_situation(self): + """Catches a row added to the troubleshooting table, or a situation removed, + without the other, so a message is documented that nothing checks.""" + assert documented_messages() == set(SITUATIONS) + + @pytest.mark.parametrize("message", sorted(SITUATIONS)) + def test_the_run_prints_the_documented_message(self, message): + """Catches a stop message reworded in the code, so the customer cannot find + its remedy in the troubleshooting table.""" + printed = SITUATIONS[message]() + assert re.search(as_pattern(message), printed), (message, printed) diff --git a/tests/unit/test_citus_shard_statistics.py b/tests/unit/test_citus_shard_statistics.py new file mode 100644 index 0000000..379df94 --- /dev/null +++ b/tests/unit/test_citus_shard_statistics.py @@ -0,0 +1,1456 @@ +"""Reading every shard placement, and combining the values per logical table.""" + +import re +from datetime import datetime, timezone + +import pytest + +from planetscale_discovery.database import citus +from planetscale_discovery.database.analyzers.performance_analyzer import ( + PerformanceAnalyzer, +) +from planetscale_discovery.database.analyzers.schema_analyzer import SchemaAnalyzer +from planetscale_discovery.database.citus import CitusClusterError, CitusTable +from tests.fixtures.database_responses import ( + CITUS_BRANCHES_OID, + CITUS_PLACEMENTS, + CITUS_REFERENCE_OID, + CITUS_REFERENCE_PKEY_OID, + CITUS_REMOTE_INDEX_ROWS, + CITUS_REMOTE_TABLE_ROWS, + CITUS_SHARDS, + CITUS_TABLES, + FakeConnection, + citus_handlers, + citus_index_row, + citus_table_row, +) + + +def ready(**kwargs): + connection = FakeConnection(citus_handlers(**kwargs)) + ctx = citus.detect_citus(connection) + citus.preflight(connection, ctx) + return connection, ctx + + +def keys_named(value, name, path="$"): + """The path of every dict key called name, at any depth.""" + if isinstance(value, dict): + for key, item in value.items(): + if key == name: + yield f"{path}.{key}" + yield from keys_named(item, name, f"{path}.{key}") + elif isinstance(value, list): + for i, item in enumerate(value): + yield from keys_named(item, name, f"{path}[{i}]") + + +def table(**overrides): + values = dict( + oid=1, + schema_name="public", + table_name="orders", + table_type="distributed", + partmethod="h", + repmodel="s", + colocation_id=3, + distribution_column="tenant_id", + autoconverted=False, + shard_count=2, + access_method="heap", + ) + values.update(overrides) + return CitusTable(**values) + + +class TestPlacementRules: + def test_a_reference_table_counts_its_rows_and_writes_once(self): + """Catches a reference table on 3 nodes reported with 3x its rows.""" + connection, ctx = ready() + combined = ctx.shard_statistics(connection).tables[CITUS_REFERENCE_OID] + # The coordinator's placement was analyzed most recently. + assert combined["estimated_rows"] == 49996560 + assert combined["counters"]["n_tup_ins"] == 50000000 + assert combined["total_size_bytes"] == 5801107456 + assert combined["placement_size_bytes"] == 3 * 5801107456 + assert combined["placements_per_shard"] == {"min": 3, "max": 3} + + def test_reads_are_summed_over_every_placement(self): + """Catches a reference table's reads counted on one placement only, while each + read runs on any one of them. + """ + connection, ctx = ready() + combined = ctx.shard_statistics(connection).tables[CITUS_REFERENCE_OID] + assert combined["counters"]["seq_scan"] == 3 + + def test_live_tuples_take_one_placement_per_shard(self): + """Catches a reference table's live tuples counted on every node.""" + connection, ctx = ready() + combined = ctx.shard_statistics(connection).tables[CITUS_REFERENCE_OID] + assert combined["counters"]["n_live_tup"] == 49996560 + + def test_an_inactive_placement_is_not_read_but_is_counted(self): + """Catches a stale inactive placement added to the table's values.""" + placements = CITUS_PLACEMENTS + [ + {"shardid": 102137, "groupid": 2, "shardstate": 3} + ] + remote = {k: list(v) for k, v in CITUS_REMOTE_TABLE_ROWS.items()} + remote[3] = remote[3] + [citus_table_row("pgbench_branches_102137", 9999)] + connection, ctx = ready(placements=placements, remote_tables=remote) + combined = ctx.shard_statistics(connection).tables[CITUS_BRANCHES_OID] + assert combined["estimated_rows"] == 630 + assert combined["inactive_placements"] == 1 + assert combined["placements_per_shard"] == {"min": 1, "max": 1} + assert combined["counters"]["seq_scan"] == 2 + + def test_index_reads_sum_every_placement_and_size_counts_once(self): + """Catches a reference index's size counted on every node, or its + scans counted on one.""" + remote_indexes = { + nodeid: [ + r + for r in CITUS_REMOTE_INDEX_ROWS[nodeid] + if not r["index_relname"].startswith("reference_t1") + ] + + [ + citus_index_row( + "reference_t1_pkey_102265", 5e7, idx_scan=nodeid, size=1577099264 + ) + ] + for nodeid in (1, 2, 3) + } + connection, ctx = ready(remote_indexes=remote_indexes) + index = ctx.shard_statistics(connection).indexes[CITUS_REFERENCE_PKEY_OID] + assert index["idx_scan"] == 1 + 2 + 3 + assert index["index_size_bytes"] == 1577099264 + + def test_a_distributed_table_sums_its_shards(self): + """Catches a distributed table's rows, writes or size taken from one shard + instead of all of them. + """ + connection, ctx = ready() + combined = ctx.shard_statistics(connection).tables[CITUS_BRANCHES_OID] + assert combined["estimated_rows"] == 630 + assert combined["counters"]["n_tup_ins"] == 630 + assert combined["total_size_bytes"] == 2 * 65536 + + def test_statement_replication_takes_one_placement_per_shard(self): + """Catches writes counted once per placement under repmodel 'c'.""" + replicated = table(repmodel="c") + by_shard = { + 1: [ + (1, citus_table_row("orders_1", 100, n_tup_ins=100, seq_scan=2)), + (2, citus_table_row("orders_1", 100, n_tup_ins=100, seq_scan=3)), + ], + 2: [ + (1, citus_table_row("orders_2", 50, n_tup_ins=50, seq_scan=1)), + (2, citus_table_row("orders_2", 50, n_tup_ins=50, seq_scan=0)), + ], + } + combined = citus.combine_table(replicated, by_shard) + assert combined["estimated_rows"] == 150 + assert combined["counters"]["n_tup_ins"] == 150 + assert combined["counters"]["seq_scan"] == 6 + + def test_a_placement_with_reset_counters_does_not_undercount_writes(self): + """Catches pg_stat_reset() on one node lowering a shard's writes.""" + by_shard = { + 1: [ + (1, citus_table_row("orders_1", 100, n_tup_ins=0)), + (2, citus_table_row("orders_1", 100, n_tup_ins=100)), + ] + } + combined = citus.combine_table(table(shard_count=1), by_shard) + assert combined["counters"]["n_tup_ins"] == 100 + + def test_an_analyzed_empty_shard_is_not_never_analyzed(self): + """Catches reltuples 0 being read as -1 (never analyzed).""" + by_shard = { + 1: [(1, citus_table_row("orders_1", 0))], + 2: [(2, citus_table_row("orders_2", 10))], + } + combined = citus.combine_table(table(), by_shard) + assert combined["shards_never_analyzed"] == 0 + assert combined["estimated_rows"] == 10 + + def test_a_table_whose_shards_were_never_analyzed_reports_minus_one(self): + """Catches a table with no analyzed shard reported with 0 rows, which reads as + an empty table. + """ + by_shard = { + 1: [(1, citus_table_row("orders_1", -1, last_analyze=None))], + 2: [(2, citus_table_row("orders_2", -1, last_analyze=None))], + } + combined = citus.combine_table(table(), by_shard) + assert combined["estimated_rows"] == -1 + assert combined["shards_never_analyzed"] == 2 + assert combined["counters"]["last_analyze"] is None + + def test_the_oldest_shard_analyze_time_is_reported(self): + """Catches the newest analyze time reported, which hides a shard with stale + statistics. + """ + by_shard = { + 1: [ + ( + 1, + citus_table_row( + "orders_1", 1, last_analyze="2026-10-01T00:00:00+00:00" + ), + ) + ], + 2: [ + ( + 2, + citus_table_row( + "orders_2", 1, last_analyze="2026-10-05T00:00:00.5+00:00" + ), + ) + ], + } + combined = citus.combine_table(table(), by_shard) + assert combined["counters"]["last_analyze"] == datetime( + 2026, 10, 1, tzinfo=timezone.utc + ) + + +class TestColumnStatistics: + def shards(self, *values): + return [ + {"shardid": i, "reltuples": rows, "n_distinct": nd} + for i, (rows, nd) in enumerate(values) + ] + + def test_a_single_shard_table_uses_its_value(self): + """Catches the n_distinct of a table with one shard changed by a combination + method. + """ + value, method = citus.combine_n_distinct( + table(shard_count=1, table_type="reference"), "id", self.shards((100, -1.0)) + ) + assert (value, method) == (-1.0, "single_placement") + + def test_the_distribution_column_adds_up_across_shards(self): + """Catches the distinct values of the distribution column not added up. Each + value is stored in one shard, so the counts add up exactly. + """ + value, method = citus.combine_n_distinct( + table(), "tenant_id", self.shards((1000, 5.0), (1000, 7.0)) + ) + assert (value, method) == (12.0, "distribution_column_sum") + + def test_an_append_table_does_not_add_up(self): + """Catches append shards, which can share values, summed as disjoint.""" + value, method = citus.combine_n_distinct( + table(partmethod="a"), "tenant_id", self.shards((1000, 5.0), (1000, 7.0)) + ) + assert (value, method) == (7.0, "max_absolute") + + def test_fractions_are_weighted_by_rows(self): + """Catches a plain mean of n_distinct fractions, which lets a small shard skew + the table. + """ + value, method = citus.combine_n_distinct( + table(), "email", self.shards((300, -1.0), (100, -0.5)) + ) + assert method == "weighted_fraction" + assert value == pytest.approx(-0.875) + + def test_mixed_signs_give_a_lower_bound(self): + """Catches shards that disagree in sign combined into a count above the largest + one. + """ + value, method = citus.combine_n_distinct( + table(), "city", self.shards((1000, 40.0), (100, -0.9)) + ) + assert (value, method) == (90.0, "mixed_max_absolute") + + def test_a_summed_count_above_ten_percent_becomes_a_fraction(self): + """Catches a combined count left absolute where ANALYZE would store a + fraction, as for a unique distribution column.""" + value, method = citus.combine_n_distinct( + table(), "tenant_id", self.shards((100, -1.0), (300, -1.0)) + ) + assert (value, method) == (-1.0, "distribution_column_sum") + + def test_null_frac_and_correlation_are_weighted_by_rows(self): + """Catches a plain mean, which lets a small shard skew the table.""" + chosen = { + 1: (1, citus_table_row("orders_1", 300)), + 2: (2, citus_table_row("orders_2", 100)), + } + column_rows = { + (1, 1): [ + { + "attname": "note", + "inherited": False, + "null_frac": 0.0, + "n_distinct": 5, + "correlation": 1.0, + } + ], + (2, 2): [ + { + "attname": "note", + "inherited": False, + "null_frac": 0.4, + "n_distinct": 5, + "correlation": -1.0, + } + ], + } + statistics, _ = citus.combine_columns(table(), chosen, column_rows) + assert statistics[0]["null_frac"] == pytest.approx(0.1) + assert statistics[0]["correlation"] == pytest.approx(0.5) + + def test_a_never_analyzed_shard_does_not_weigh_null_frac(self): + """Catches a -1 reltuples entering a weighted mean.""" + ctable = table(shard_count=2) + chosen = { + 1: (1, citus_table_row("orders_1", 100)), + 2: (2, citus_table_row("orders_2", -1)), + } + column_rows = { + (1, 1): [ + { + "attname": "note", + "inherited": False, + "null_frac": 0.2, + "n_distinct": 5, + "correlation": 0.5, + } + ], + (2, 2): [ + { + "attname": "note", + "inherited": False, + "null_frac": 0.9, + "n_distinct": 5, + "correlation": -0.5, + } + ], + } + statistics, detail = citus.combine_columns(ctable, chosen, column_rows) + assert statistics[0]["null_frac"] == pytest.approx(0.2) + assert statistics[0]["correlation"] == pytest.approx(0.5) + assert len(detail[0]["shards"]) == 2 + + @pytest.mark.parametrize("shard_count", [1, 2]) + def test_a_column_with_no_analyzed_shard_reports_unknown(self, shard_count): + """Catches a 0 reported for n_distinct, which reads as no distinct values, + when no shard counts as analyzed. TRUNCATE sets reltuples back to -1 and + keeps the old pg_stats rows, so a column can have rows in this state.""" + ctable = table(shard_count=shard_count) + chosen = { + i: (i, citus_table_row(f"orders_{i}", -1)) + for i in range(1, shard_count + 1) + } + column_rows = { + (i, i): [ + { + "attname": "note", + "inherited": False, + "null_frac": 0.2, + "n_distinct": 5, + "correlation": 0.5, + } + ] + for i in range(1, shard_count + 1) + } + statistics, detail = citus.combine_columns(ctable, chosen, column_rows) + assert statistics[0]["n_distinct"] is None + assert statistics[0]["null_frac"] is None + assert statistics[0]["correlation"] is None + assert detail[0]["n_distinct_method"] is None + + +class TestShardRead: + def test_a_relation_without_an_active_placement_is_ignored(self): + """Catches an orphaned shard left after a move counted twice.""" + remote = {k: list(v) for k, v in CITUS_REMOTE_TABLE_ROWS.items()} + remote[3] = remote[3] + [citus_table_row("pgbench_branches_102137", 9999)] + connection, ctx = ready(remote_tables=remote) + combined = ctx.shard_statistics(connection).tables[CITUS_BRANCHES_OID] + assert combined["estimated_rows"] == 630 + + def test_an_orphaned_shard_on_the_coordinator_is_left_out_of_its_size(self): + """Catches a shard relation that a move left on the coordinator, before the + deferred cleanup drops it, counted in the logical database size on top of + the shard's active placement.""" + coordinator_size = 19327047359 + connection, ctx = ready() + baseline = citus.logical_database_size(connection, ctx, coordinator_size) + + remote = {k: list(v) for k, v in CITUS_REMOTE_TABLE_ROWS.items()} + remote[1] = remote[1] + [ + citus_table_row("pgbench_branches_102137", 317, total_size_bytes=7000000) + ] + connection, ctx = ready(remote_tables=remote) + stats = ctx.shard_statistics(connection) + assert stats.orphaned_size_bytes == {0: 7000000} + assert stats.tables[CITUS_BRANCHES_OID]["estimated_rows"] == 630 + size = citus.logical_database_size(connection, ctx, coordinator_size) + assert size == baseline - 7000000 + + def test_both_reports_rewrite_the_current_database_the_same_way(self): + """Catches the schema catalog and the performance report disagreeing on the + current database's size, or another database's row being rewritten.""" + connection, ctx = ready() + catalog = [ + { + "database_name": "postgres", + "size_bytes": 19327047359, + "size_human": "18 GB", + }, + {"database_name": "other", "size_bytes": 8192, "size_human": "8192 bytes"}, + ] + sizes = [ + {"datname": "postgres", "size_bytes": 19327047359, "size_human": "18 GB"}, + {"datname": "other", "size_bytes": 8192, "size_human": "8192 bytes"}, + ] + citus.apply_logical_database_size(connection, ctx, catalog, "database_name") + citus.apply_logical_database_size(connection, ctx, sizes, "datname") + + def without_name(row): + return { + k: v for k, v in row.items() if k not in ("database_name", "datname") + } + + assert without_name(catalog[0]) == without_name(sizes[0]) + assert catalog[0]["citus"]["coordinator"] == { + "size_bytes": 19327047359, + "size_human": "18 GB", + } + assert catalog[0]["size_bytes"] != 19327047359 + assert "citus" not in catalog[1] and "citus" not in sizes[1] + + def test_a_node_that_fails_during_the_read_stops(self): + """Catches the other nodes' values written as the table's whole.""" + connection, ctx = ready() + connection.handlers.insert( + 0, + ( + "run_command_on_all_nodes(", + lambda sql, params: [ + {"nodeid": 1, "success": True, "result": "[]"}, + {"nodeid": 2, "success": True, "result": "[]"}, + {"nodeid": 3, "success": False, "result": "connection lost"}, + ], + ), + ) + with pytest.raises(CitusClusterError, match="citus-worker-2"): + ctx.shard_statistics(connection) + + def test_a_node_missing_from_the_results_stops(self): + """Catches a node that returned nothing treated as holding nothing.""" + connection, ctx = ready() + connection.handlers.insert( + 0, + ( + "run_command_on_all_nodes(", + lambda sql, params: [ + {"nodeid": 1, "success": True, "result": "[]"}, + {"nodeid": 2, "success": True, "result": "[]"}, + ], + ), + ) + with pytest.raises( + CitusClusterError, match="citus-worker-2:5432: returned no result" + ): + ctx.shard_statistics(connection) + + def test_a_citus_foreign_table_does_not_stop_the_run(self): + """Catches a foreign table added with citus_add_local_table_to_metadata() + entering the shard read: its shard has relkind 'f', which no remote + command returns, so the healthy run stopped as a missing shard.""" + foreign = { + "table_oid": 17600, + "schema_name": "public", + "table_name": "remote_events", + "relkind": "f", + "table_type": "local", + "partmethod": "n", + "repmodel": "s", + "colocation_id": 0, + "distribution_column": None, + "autoconverted": False, + "shard_count": 1, + "access_method": None, + } + connection, ctx = ready( + tables=CITUS_TABLES + [foreign], + shards=CITUS_SHARDS + + [ + { + "shardid": 102300, + "table_oid": 17600, + "shard_relname": "remote_events_102300", + } + ], + placements=CITUS_PLACEMENTS + + [{"shardid": 102300, "groupid": 0, "shardstate": 1}], + ) + stats = ctx.shard_statistics(connection) + assert 17600 not in ctx.tables + assert set(stats.tables) == {CITUS_BRANCHES_OID, CITUS_REFERENCE_OID} + + def test_the_classification_and_the_shard_read_share_one_relkind_filter(self): + """Catches the two drifting apart, which makes a healthy run stop on a + Citus table whose shard the remote commands never return.""" + _, ctx = ready() + expected = f"c.relkind IN {citus.TABLE_RELKINDS}" + assert expected in citus._classification_sql(ctx) + for command in (citus.TABLE_COMMAND, citus.COLUMN_COMMAND, citus.INDEX_COMMAND): + assert expected in command + + def test_the_regular_count_shares_the_relkind_filter(self, monkeypatch): + """Catches the count of regular tables and the Citus classification drifting + apart, which counts a relation twice or not at all in table_counts. + + The filter is changed for the test, because a copy of today's value would + match it. + """ + from planetscale_discovery.database.analyzers import citus_analyzer + + connection, ctx = ready() + monkeypatch.setattr(citus_analyzer, "TABLE_RELKINDS", "('r', 'p', 'f')") + citus_analyzer.CitusAnalyzer(connection, config={"citus": ctx}).analyze() + count = [s for s in connection.statements() if "AS regular_count" in s] + assert count, "the regular tables were not counted" + assert "c.relkind IN ('r', 'p', 'f')" in count[0] + + def misnamed_index(self): + """pgbench_branches_pkey on shard 102138 lost its shard id when + PostgreSQL named it automatically, so shard_name() does not match it.""" + remote = {k: list(v) for k, v in CITUS_REMOTE_INDEX_ROWS.items()} + remote[3] = [ + ( + dict( + citus_index_row("pgbench_branches_pkey_idx", 313), + table_relname="pgbench_branches_102138", + ) + if r["index_relname"] == "pgbench_branches_pkey_102138" + else r + ) + for r in remote[3] + ] + return remote + + def test_an_index_missing_on_a_shard_stops_with_the_remedy(self): + """Catches an index reported with a partial size and, once matched on no + shard, idx_scan None, which drops it from unused_indexes.""" + from tests.fixtures.database_responses import CITUS_INDEXES + + indexes = [dict(i, is_partition_index=True) for i in CITUS_INDEXES] + connection, ctx = ready(remote_indexes=self.misnamed_index(), indexes=indexes) + with pytest.raises(CitusClusterError) as error: + ctx.shard_statistics(connection) + message = str(error.value) + assert ( + "index public.pgbench_branches_pkey: not found on 1 of 2 shards" in message + ) + assert "fix_all_partition_shard_index_names()" in message + assert "If the index is still not found" in message + + def test_an_index_dropped_during_the_read_is_named_as_dropped(self): + """Catches an index dropped between the two index reads reported as one + created with DDL propagation off, a remedy that does not apply.""" + from tests.fixtures.database_responses import ( + CITUS_BRANCHES_PKEY_OID, + CITUS_INDEX_SHARD_NAMES, + CITUS_INDEXES, + ) + + def gone(row): + return row["index_oid"] == CITUS_BRANCHES_PKEY_OID + + remote = { + nodeid: [ + r + for r in rows + if not r["index_relname"].startswith("pgbench_branches_pkey_") + ] + for nodeid, rows in CITUS_REMOTE_INDEX_ROWS.items() + } + connection, ctx = ready( + remote_indexes=remote, + index_shard_names=[r for r in CITUS_INDEX_SHARD_NAMES if not gone(r)], + ) + reads = [] + + def index_list(sql, params): + reads.append(sql) + return ( + CITUS_INDEXES + if len(reads) == 1 + else [r for r in CITUS_INDEXES if not gone(r)] + ) + + connection.handlers.insert(0, ("i.relname AS index_name", index_list)) + with pytest.raises(CitusClusterError) as error: + ctx.shard_statistics(connection) + message = str(error.value) + assert "indexes changed during discovery" in message + assert "index public.pgbench_branches_pkey: dropped during discovery" in message + assert "enable_ddl_propagation" not in message + + def test_an_index_created_during_the_read_is_left_out(self): + """Catches an index created between the two index reads stopping the run + with a bare KeyError, although leaving it out misreports nothing.""" + from tests.fixtures.database_responses import ( + CITUS_BRANCHES_PKEY_OID, + CITUS_INDEX_SHARD_NAMES, + ) + + created = { + "index_oid": 99999, + "shardid": 102137, + "shard_index_name": "pgbench_branches_created_idx_102137", + } + connection, ctx = ready(index_shard_names=CITUS_INDEX_SHARD_NAMES + [created]) + stats = ctx.shard_statistics(connection) + assert 99999 not in stats.indexes + assert CITUS_BRANCHES_PKEY_OID in stats.indexes + + def test_an_index_not_on_a_partition_is_not_sent_to_the_partition_fix(self): + """Catches an index created with DDL propagation off, on no shard of a table + that is not partitioned, sent to a function that renames nothing. + """ + remote = { + nodeid: [ + r + for r in rows + if not r["index_relname"].startswith("pgbench_branches_pkey_") + ] + for nodeid, rows in CITUS_REMOTE_INDEX_ROWS.items() + } + connection, ctx = ready(remote_indexes=remote) + with pytest.raises(CitusClusterError) as error: + ctx.shard_statistics(connection) + message = str(error.value) + assert "not found on 2 of 2 shards" in message + assert "fix_all_partition_shard_index_names" not in message + index_names = [ + sql for sql, _ in connection.executed if "x.indisvalid AS is_valid" in sql + ] + assert "i.relispartition AS is_partition_index" in index_names[0] + assert "citus.enable_ddl_propagation off" in message + + @staticmethod + def invalid_branches_pkey(): + from tests.fixtures.database_responses import CITUS_INDEXES + + return [ + dict(i, is_valid=i["index_name"] != "pgbench_branches_pkey") + for i in CITUS_INDEXES + ] + + def test_an_invalid_index_missing_on_a_shard_does_not_stop(self): + """Catches a failed CREATE INDEX CONCURRENTLY, which builds shard indexes on + only some shards, stopping a read-only run until the index is rebuilt. + """ + from tests.fixtures.database_responses import CITUS_BRANCHES_PKEY_OID + + connection, ctx = ready( + remote_indexes=self.misnamed_index(), indexes=self.invalid_branches_pkey() + ) + stats = ctx.shard_statistics(connection) + invalid = stats.indexes[CITUS_BRANCHES_PKEY_OID] + assert (invalid["shards_read"], invalid["shard_count"]) == (1, 2) + assert invalid["index_size_bytes"] == 16384 + + def test_an_invalid_index_on_no_shard_records_no_scan(self): + """Catches an invalid index that exists on no shard reported with idx_scan + None, which drops it from unused_indexes. + """ + from tests.fixtures.database_responses import CITUS_BRANCHES_PKEY_OID + + remote = { + nodeid: [ + r + for r in rows + if not r["index_relname"].startswith("pgbench_branches_pkey_") + ] + for nodeid, rows in CITUS_REMOTE_INDEX_ROWS.items() + } + connection, ctx = ready( + remote_indexes=remote, indexes=self.invalid_branches_pkey() + ) + invalid = ctx.shard_statistics(connection).indexes[CITUS_BRANCHES_PKEY_OID] + assert (invalid["shards_read"], invalid["idx_scan"]) == (0, 0) + assert invalid["index_size_bytes"] == 0 + + def test_a_shard_found_on_no_node_stops(self): + """Catches a shard that no node returned summed as if it held no rows.""" + remote = {k: list(v) for k, v in CITUS_REMOTE_TABLE_ROWS.items()} + remote[3] = [r for r in remote[3] if r["relname"] != "pgbench_branches_102138"] + connection, ctx = ready(remote_tables=remote) + with pytest.raises(CitusClusterError, match="shard 102138"): + ctx.shard_statistics(connection) + + def test_shards_are_read_in_batches(self, monkeypatch): + """Catches one call holding every shard of a large cluster in memory.""" + monkeypatch.setattr(citus, "SHARD_BATCH_SIZE", 2) + connection, ctx = ready() + ctx.shard_statistics(connection) + batches = [ + params[1] + for sql, params in connection.executed + if "run_command_on_all_nodes(" in sql + and params + and params[0] == citus.TABLE_COMMAND + ] + assert sorted(len(b) for b in batches) == [1, 2] + + +class TestScale: + """Combining stays linear in the number of tables, shards and indexes.""" + + TABLES = 20000 + + def cluster(self): + tables, shards, placements, indexes, shard_names = [], [], [], [], [] + remote_tables = {2: [], 3: []} + remote_columns = {2: [], 3: []} + remote_indexes = {2: [], 3: []} + for i in range(self.TABLES): + schema, oid, shardid = f"tenant_{i}", 100000 + i, 200000 + i + node, relname = 2 + i % 2, f"orders_{shardid}" + tables.append( + { + "table_oid": oid, + "schema_name": schema, + "table_name": "orders", + "table_type": "schema", + "partmethod": "n", + "repmodel": "s", + "colocation_id": 10 + i, + "distribution_column": None, + "autoconverted": False, + "shard_count": 1, + "access_method": "heap", + } + ) + shards.append( + {"shardid": shardid, "table_oid": oid, "shard_relname": relname} + ) + placements.append( + {"shardid": shardid, "groupid": node - 1, "shardstate": 1} + ) + remote_tables[node].append( + citus_table_row(relname, 100, schema_name=schema) + ) + remote_columns[node].append( + { + "schema_name": schema, + "relname": relname, + "attname": "order_id", + "inherited": False, + "null_frac": 0.0, + "n_distinct": -1.0, + "correlation": 1.0, + } + ) + for offset, name in enumerate(("orders_pkey", "orders_status_idx")): + index_oid = 300000 + 2 * i + offset + indexes.append( + { + "index_oid": index_oid, + "table_oid": oid, + "schema_name": schema, + "index_name": name, + } + ) + shard_names.append( + { + "index_oid": index_oid, + "shardid": shardid, + "shard_index_name": f"{name}_{shardid}", + } + ) + remote_indexes[node].append( + dict(citus_index_row(f"{name}_{shardid}", 100), schema_name=schema) + ) + return citus_handlers( + tables=tables, + shards=shards, + placements=placements, + remote_tables=remote_tables, + remote_columns=remote_columns, + remote_indexes=remote_indexes, + indexes=indexes, + index_shard_names=shard_names, + ) + + def test_shards_and_placements_are_scanned_a_fixed_number_of_times(self): + """Catches a loop that rescans every shard or placement for each table: + 100k schema tables make that 1e10 iterations. Counting scans makes the + test independent of the machine's speed.""" + + class CountingDict(dict): + scans = 0 + + def items(self): + CountingDict.scans += 1 + return super().items() + + class CountingList(list): + scans = 0 + + def __iter__(self): + CountingList.scans += 1 + return super().__iter__() + + connection = FakeConnection(self.cluster()) + ctx = citus.detect_citus(connection) + citus.preflight(connection, ctx) + ctx.shards = CountingDict(ctx.shards) + ctx.placements = CountingList(ctx.placements) + + stats = citus.collect_shard_statistics(connection, ctx) + + assert len(stats.tables) == self.TABLES + assert CountingDict.scans <= 5 + assert CountingList.scans <= 5 + + def test_a_large_schema_sharded_cluster_is_combined_quickly(self): + """Catches a loop that rescans every index's rows for each index, which + takes minutes at this size and hours at 100k schemas.""" + import time + + connection = FakeConnection(self.cluster()) + ctx = citus.detect_citus(connection) + citus.preflight(connection, ctx) + started = time.monotonic() + stats = citus.collect_shard_statistics(connection, ctx) + elapsed = time.monotonic() - started + + assert len(stats.indexes) == 2 * self.TABLES + assert stats.tables[100000 + self.TABLES - 1]["estimated_rows"] == 100 + assert elapsed < 10, f"combining {self.TABLES} tables took {elapsed:.1f}s" + + +class TestRemoteCommands: + FORBIDDEN = ( + "most_common_vals", + "most_common_freqs", + "histogram_bounds", + "pg_dist_authinfo", + "pg_dist_poolinfo", + "shardminvalue", + "shardmaxvalue", + ) + + @pytest.mark.parametrize("command", citus.REMOTE_COMMANDS) + def test_commands_are_selects(self, command): + """Catches a remote command that writes on a node, including through a + data-modifying CTE.""" + assert command.startswith(("SELECT ", "WITH ")) + writes = re.findall( + r"\b(INSERT|UPDATE|DELETE|MERGE|TRUNCATE|COPY|CREATE|DROP|ALTER|GRANT" + r"|ANALYZE|VACUUM)\b", + command, + ) + assert writes == [] + + @pytest.mark.parametrize("command", citus.REMOTE_COMMANDS) + def test_commands_name_no_value_or_credential(self, command): + """Catches customer values or credentials read from a node.""" + for name in self.FORBIDDEN: + assert name not in command + + @pytest.mark.parametrize("command", citus.REMOTE_COMMANDS) + def test_commands_qualify_every_catalog_object(self, command): + """Catches a planted relation or function shadowing a catalog on a node. + + A remote session does not inherit this tool's search_path pin. + """ + unqualified = re.findall( + r"(?