Fix the DB table rows metric

crates/autopilot/src/database/mod.rs

@@ -41,22 +41,18 @@ impl Postgres {
         let metrics = Metrics::get();
 
         // update table row metrics
-        for &table in database::TABLES {
-            let mut ex = self.pool.acquire().await?;
-            let count = count_rows_in_table(&mut ex, table).await?;
-            metrics.table_rows.with_label_values(&[table]).set(count);
-        }
-
-        // update table row metrics
-        for &table in database::LARGE_TABLES {
-            let mut ex = self.pool.acquire().await?;
-            let count = estimate_rows_in_table(&mut ex, table).await?;
+        let mut ex = self.pool.acquire().await?;
+        for table in database::all_tables(&mut ex).await? {
+            let count = if database::LARGE_TABLES.iter().any(|t| *t == table) {
+                estimate_rows_in_table(&mut ex, table).await?
+            } else {
+                count_rows_in_table(&mut ex, table).await?
+            };
             metrics.table_rows.with_label_values(&[table]).set(count);
         }
 
         // update unused app data metric
         {
-            let mut ex = self.pool.acquire().await?;
             let count = count_unused_app_data(&mut ex).await?;
             metrics.unused_app_data.set(count);
         }
@@ -80,7 +76,7 @@ async fn count_rows_in_table(ex: &mut PgConnection, table: &str) -> sqlx::Result
 }
 
 async fn estimate_rows_in_table(ex: &mut PgConnection, table: &str) -> sqlx::Result<i64> {
-    let query = format!("SELECT reltuples::bigint FROM pg_class WHERE relname='{table}';");
+    let query = format!("SELECT reltuples::bigint FROM pg_class WHERE oid = '{table}'::regclass;");
     sqlx::query_scalar(&query).fetch_one(ex).await
 }
 

crates/database/Cargo.toml

@@ -14,6 +14,7 @@ hex = { workspace = true }
 sqlx = { workspace = true }
 strum = { workspace = true }
 serde_json = { workspace = true }
+tokio = { workspace = true, features = ["sync"] }
 
 [dev-dependencies]
 maplit = { workspace = true }

crates/database/src/lib.rs

@@ -25,7 +25,8 @@ pub mod trades;
 
 use {
     byte_array::ByteArray,
-    sqlx::{Executor, PgPool},
+    sqlx::{Executor, PgConnection, PgPool},
+    tokio::sync::OnceCell,
 };
 
 // Design:
@@ -48,43 +49,42 @@ use {
 
 pub type PgTransaction<'a> = sqlx::Transaction<'a, sqlx::Postgres>;
 
-/// The names of tables we use in the db.
-pub const TABLES: &[&str] = &[
-    "orders",
-    "trades",
-    "invalidations",
-    "last_indexed_blocks",
-    "quotes",
-    "settlements",
-    "presignature_events",
-    "order_quotes",
-    "solver_competitions",
-    "auctions",
-    "competition_auctions",
-    "onchain_placed_orders",
-    "ethflow_orders",
-    "order_execution",
-    "interactions",
-    "ethflow_refunds",
-    "settlement_scores",
-    "settlement_observations",
-    "auction_prices",
+/// The names of potentially big volume tables we use in the db.
+pub const LARGE_TABLES: &[&str] = &[
     "auction_participants",
-    "app_data",
-    "jit_orders",
+    "auction_prices",
+    "order_events",
+    "proposed_trade_executions",
 ];
 
-/// The names of potentially big volume tables we use in the db.
-pub const LARGE_TABLES: &[&str] = &["order_events"];
+pub async fn all_tables(ex: &mut PgConnection) -> sqlx::Result<&'static Vec<String>> {
+    static TABLES: OnceCell<Vec<String>> = OnceCell::const_new();
+
+    TABLES
+        .get_or_try_init(|| async {
+            #[derive(sqlx::FromRow, Debug)]
+            struct TableName(String);
+
+            const QUERY: &str = "SELECT tablename FROM pg_tables WHERE schemaname = 'public' AND \
+                                 tablename NOT LIKE '%flyway%' AND tableowner NOT IN ('rdsadmin', \
+                                 'cow_protocol_admin')";
+
+            let table_names: Vec<String> = sqlx::query_as(QUERY)
+                .fetch_all(ex)
+                .await?
+                .into_iter()
+                .map(|TableName(name)| name)
+                .collect();
 
-pub fn all_tables() -> impl Iterator<Item = &'static str> {
-    TABLES.iter().copied().chain(LARGE_TABLES.iter().copied())
+            Ok(table_names)
+        })
+        .await
 }
 
 /// Delete all data in the database. Only used by tests.
 #[allow(non_snake_case)]
 pub async fn clear_DANGER_(ex: &mut PgTransaction<'_>) -> sqlx::Result<()> {
-    for table in all_tables() {
+    for table in all_tables(ex).await? {
         ex.execute(format!("TRUNCATE {table};").as_str()).await?;
     }
     Ok(())