sui_rpc_benchmark/direct/
query_template_generator.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4/// This module generates SQL query templates for benchmarking, including
5/// query templates based on primary key columns and indexed columns.
6///
7/// The primary key query templates ("pk queries") select a row by each PK,
8/// while the "index queries" filter by indexed columns. Instead
9/// of returning just a list of tables and indexes, this module
10/// returns a vector of QueryTemplate objects, each of which is
11/// ready to be executed. This approach streamlines the pipeline
12/// so we can directly run these queries as part of the benchmark.
13use tokio_postgres::NoTls;
14use tracing::{debug, info};
15use url::Url;
16
17#[derive(Debug, Clone)]
18pub struct QueryTemplate {
19    pub query_template: String,
20    pub table_name: String,
21    pub needed_columns: Vec<String>,
22}
23
24pub struct QueryTemplateGenerator {
25    db_url: Url,
26}
27
28impl QueryTemplateGenerator {
29    pub fn new(db_url: Url) -> Self {
30        Self { db_url }
31    }
32
33    pub async fn generate_query_templates(&self) -> Result<Vec<QueryTemplate>, anyhow::Error> {
34        let (client, connection) = tokio_postgres::connect(self.db_url.as_str(), NoTls).await?;
35        tokio::spawn(connection);
36
37        let pk_query = r#"
38            SELECT 
39                tc.table_name,  
40                array_agg(kcu.column_name ORDER BY kcu.ordinal_position)::text[] as primary_key_columns  
41            FROM information_schema.table_constraints tc  
42            JOIN information_schema.key_column_usage kcu 
43                ON tc.constraint_name = kcu.constraint_name  
44            WHERE tc.constraint_type = 'PRIMARY KEY'
45                AND tc.table_schema = 'public'
46                AND tc.table_name != '__diesel_schema_migrations'
47            GROUP BY tc.table_name
48            ORDER BY tc.table_name;
49        "#;
50
51        let mut queries = Vec::new();
52        let rows = client.query(pk_query, &[]).await?;
53        let tables: Vec<String> = rows
54            .iter()
55            .map(|row| row.get::<_, String>("table_name"))
56            .collect::<std::collections::HashSet<_>>()
57            .into_iter()
58            .collect();
59        info!(
60            "Found {} active tables in database: {:?}",
61            tables.len(),
62            tables
63        );
64
65        // Process primary key queries - now each row has all columns for a table
66        for row in rows {
67            let table: String = row.get("table_name");
68            let pk_columns: Vec<String> = row.get("primary_key_columns");
69            queries.push(self.create_pk_benchmark_query(&table, &pk_columns));
70        }
71
72        let idx_query = r#"
73            SELECT 
74                t.relname AS table_name,
75                i.relname AS index_name,
76                array_agg(a.attname ORDER BY k.i) AS column_names
77            FROM pg_class t
78            JOIN pg_index ix ON t.oid = ix.indrelid
79            JOIN pg_class i ON ix.indexrelid = i.oid
80            JOIN pg_attribute a ON t.oid = a.attrelid
81            JOIN generate_subscripts(ix.indkey, 1) k(i) ON a.attnum = ix.indkey[k.i]
82            WHERE t.relkind = 'r'
83                AND t.relnamespace = (SELECT oid FROM pg_namespace WHERE nspname = 'public')
84                AND NOT ix.indisprimary
85                AND t.relname != '__diesel_schema_migrations'
86            GROUP BY t.relname, i.relname
87            ORDER BY t.relname, i.relname;
88        "#;
89
90        let rows = client.query(idx_query, &[]).await?;
91        for row in rows {
92            let table: String = row.get("table_name");
93            let columns: Vec<String> = row.get("column_names");
94            queries.push(self.create_index_benchmark_query(&table, &columns));
95        }
96
97        debug!("Generated {} queries:", queries.len());
98        for (i, query) in queries.iter().enumerate() {
99            debug!(
100                "  {}. Table: {}, Template: {}",
101                i + 1,
102                query.table_name,
103                query.query_template
104            );
105        }
106
107        Ok(queries)
108    }
109
110    /// An example query template:
111    /// SELECT * FROM tx_kinds WHERE tx_kind = $1 AND tx_sequence_number = $2 LIMIT 1
112    fn create_pk_benchmark_query(&self, table: &str, columns: &[String]) -> QueryTemplate {
113        let conditions = columns
114            .iter()
115            .enumerate()
116            .map(|(i, col)| format!("{} = ${}", col, i + 1))
117            .collect::<Vec<_>>()
118            .join(" AND ");
119
120        QueryTemplate {
121            query_template: format!("SELECT * FROM {} WHERE {} LIMIT 1", table, conditions),
122            table_name: table.to_string(),
123            needed_columns: columns.to_vec(),
124        }
125    }
126
127    fn create_index_benchmark_query(&self, table: &str, columns: &[String]) -> QueryTemplate {
128        let conditions = columns
129            .iter()
130            .enumerate()
131            .map(|(i, col)| format!("{} = ${}", col, i + 1))
132            .collect::<Vec<_>>()
133            .join(" AND ");
134
135        QueryTemplate {
136            query_template: format!("SELECT * FROM {} WHERE {} LIMIT 50", table, conditions),
137            table_name: table.to_string(),
138            needed_columns: columns.to_vec(),
139        }
140    }
141}
sui_rpc_benchmark/direct/query_template_generator.rs

sui_rpc_benchmark/direct/
query_template_generator.rs
