|
| 1 | +// This Source Code Form is subject to the terms of the Mozilla Public |
| 2 | +// License, v. 2.0. If a copy of the MPL was not distributed with this |
| 3 | +// file, You can obtain one at https://mozilla.org/MPL/2.0/. |
| 4 | + |
| 5 | +//! `omdb db whatis` subcommand |
| 6 | +//! |
| 7 | +//! Heuristically determine what type of object a given UUID refers to |
| 8 | +
|
| 9 | +use anyhow::Context; |
| 10 | +use async_bb8_diesel::AsyncRunQueryDsl; |
| 11 | +use clap::Args; |
| 12 | +use nexus_db_queries::db::DataStore; |
| 13 | +use uuid::Uuid; |
| 14 | + |
| 15 | +#[derive(Debug, Args, Clone)] |
| 16 | +pub(super) struct WhatisArgs { |
| 17 | + /// The UUID(s) to look up |
| 18 | + uuids: Vec<Uuid>, |
| 19 | + |
| 20 | + /// Show all tables that were checked |
| 21 | + #[clap(long)] |
| 22 | + debug: bool, |
| 23 | +} |
| 24 | + |
| 25 | +// `omdb db whatis UUID...` heuristically determines what type of object a UUID |
| 26 | +// refers to by first searching the database for unique UUID columns and then |
| 27 | +// searching those tables for a matching row. |
| 28 | +pub(super) async fn cmd_db_whatis( |
| 29 | + datastore: &DataStore, |
| 30 | + args: &WhatisArgs, |
| 31 | +) -> Result<(), anyhow::Error> { |
| 32 | + let conn = datastore.pool_connection_for_tests().await?; |
| 33 | + |
| 34 | + // Query CockroachDB's information schema to find UUID columns that are the |
| 35 | + // sole column in a unique constraint or primary key with no associated |
| 36 | + // WHERE clause (i.e., not a partial index) . This ensures we only query |
| 37 | + // columns that have their own index (avoiding table scans). |
| 38 | + // |
| 39 | + // We use this approach rather than hardcoding all the kinds of objects that |
| 40 | + // we know about so that this stays up-to-date with other changes to the |
| 41 | + // system. |
| 42 | + let query = " |
| 43 | + WITH constraint_column_counts AS ( |
| 44 | + SELECT |
| 45 | + constraint_catalog, |
| 46 | + constraint_schema, |
| 47 | + constraint_name, |
| 48 | + table_catalog, |
| 49 | + table_schema, |
| 50 | + table_name, |
| 51 | + COUNT(*) as column_count |
| 52 | + FROM |
| 53 | + information_schema.key_column_usage |
| 54 | + GROUP BY |
| 55 | + constraint_catalog, |
| 56 | + constraint_schema, |
| 57 | + constraint_name, |
| 58 | + table_catalog, |
| 59 | + table_schema, |
| 60 | + table_name |
| 61 | + ) |
| 62 | + SELECT DISTINCT |
| 63 | + tc.table_name, |
| 64 | + kcu.column_name |
| 65 | + FROM |
| 66 | + information_schema.table_constraints tc |
| 67 | + JOIN information_schema.key_column_usage kcu |
| 68 | + ON tc.constraint_catalog = kcu.constraint_catalog |
| 69 | + AND tc.constraint_schema = kcu.constraint_schema |
| 70 | + AND tc.constraint_name = kcu.constraint_name |
| 71 | + AND tc.table_catalog = kcu.table_catalog |
| 72 | + AND tc.table_schema = kcu.table_schema |
| 73 | + AND tc.table_name = kcu.table_name |
| 74 | + JOIN information_schema.columns c |
| 75 | + ON kcu.table_catalog = c.table_catalog |
| 76 | + AND kcu.table_schema = c.table_schema |
| 77 | + AND kcu.table_name = c.table_name |
| 78 | + AND kcu.column_name = c.column_name |
| 79 | + JOIN constraint_column_counts ccc |
| 80 | + ON tc.constraint_catalog = ccc.constraint_catalog |
| 81 | + AND tc.constraint_schema = ccc.constraint_schema |
| 82 | + AND tc.constraint_name = ccc.constraint_name |
| 83 | + AND tc.table_catalog = ccc.table_catalog |
| 84 | + AND tc.table_schema = ccc.table_schema |
| 85 | + AND tc.table_name = ccc.table_name |
| 86 | + LEFT JOIN pg_indexes idx |
| 87 | + ON idx.tablename = tc.table_name |
| 88 | + AND idx.indexname = tc.constraint_name |
| 89 | + AND idx.schemaname = tc.table_schema |
| 90 | + WHERE |
| 91 | + tc.constraint_type IN ('PRIMARY KEY', 'UNIQUE') |
| 92 | + AND tc.table_schema = 'public' |
| 93 | + AND c.data_type = 'uuid' |
| 94 | + AND ccc.column_count = 1 |
| 95 | + AND (idx.indexdef IS NULL OR idx.indexdef NOT LIKE '% WHERE %') |
| 96 | + ORDER BY |
| 97 | + tc.table_name, kcu.column_name |
| 98 | + "; |
| 99 | + |
| 100 | + #[derive(Debug, diesel::QueryableByName)] |
| 101 | + struct UniqueUuidColumn { |
| 102 | + #[diesel(sql_type = diesel::sql_types::Text)] |
| 103 | + table_name: String, |
| 104 | + #[diesel(sql_type = diesel::sql_types::Text)] |
| 105 | + column_name: String, |
| 106 | + } |
| 107 | + |
| 108 | + let unique_columns: Vec<UniqueUuidColumn> = diesel::sql_query(query) |
| 109 | + .load_async(&*conn) |
| 110 | + .await |
| 111 | + .context("querying information_schema for unique UUID columns")?; |
| 112 | + |
| 113 | + // Search separately for each UUID provided on the command line. |
| 114 | + for &uuid in &args.uuids { |
| 115 | + let mut found = false; |
| 116 | + |
| 117 | + // For each unique UUID column that we found above, see if there's a row |
| 118 | + // in the corresponding table for the given UUID value. |
| 119 | + for col in &unique_columns { |
| 120 | + let check_query = format!( |
| 121 | + "SELECT EXISTS(SELECT 1 FROM {} WHERE {} = $1) as exists", |
| 122 | + col.table_name, col.column_name |
| 123 | + ); |
| 124 | + |
| 125 | + let table_column = |
| 126 | + format!("{}.{}", col.table_name, col.column_name); |
| 127 | + if args.debug { |
| 128 | + eprintln!("checking table {table_column:?}"); |
| 129 | + } |
| 130 | + |
| 131 | + // `sql_query()` requires extracting the results with a struct. |
| 132 | + #[derive(Debug, diesel::QueryableByName)] |
| 133 | + struct ExistsResult { |
| 134 | + #[diesel(sql_type = diesel::sql_types::Bool)] |
| 135 | + exists: bool, |
| 136 | + } |
| 137 | + |
| 138 | + let exists_result: ExistsResult = diesel::sql_query(&check_query) |
| 139 | + .bind::<diesel::sql_types::Uuid, _>(uuid) |
| 140 | + .get_result_async(&*conn) |
| 141 | + .await |
| 142 | + .with_context(|| { |
| 143 | + format!("checking {table_column} for UUID {uuid}") |
| 144 | + })?; |
| 145 | + |
| 146 | + if exists_result.exists { |
| 147 | + println!("{uuid} found in {table_column}"); |
| 148 | + found = true; |
| 149 | + } |
| 150 | + } |
| 151 | + |
| 152 | + if !found { |
| 153 | + println!("{uuid} not found in any unique UUID column"); |
| 154 | + } |
| 155 | + } |
| 156 | + |
| 157 | + Ok(()) |
| 158 | +} |
0 commit comments