Skip to main content

rust_ef_postgres/
introspection.rs

1//! Database introspection  - ?reads schema from PostgreSQL information_schema.
2
3use crate::tls::PgTlsMode;
4use rust_ef::error::EFResult;
5use tokio_postgres::NoTls;
6
7/// Column information from database introspection.
8#[derive(Debug, Clone)]
9pub struct DbColumn {
10    pub name: String,
11    pub data_type: String,
12    pub is_nullable: bool,
13    pub is_primary_key: bool,
14    pub max_length: Option<usize>,
15}
16
17/// Table information from database introspection.
18#[derive(Debug, Clone)]
19pub struct DbTable {
20    pub name: String,
21    pub columns: Vec<DbColumn>,
22}
23
24/// Reads all tables and their columns from a PostgreSQL database.
25///
26/// The `tls` parameter controls transport encryption. For local dev
27/// (plaintext), pass [`PgTlsMode::Disable`]; for production, use
28/// [`PgTlsMode::Require`] with a configured `TlsConnector`.
29pub async fn introspect_postgres(
30    connection_string: &str,
31    tls: PgTlsMode,
32) -> EFResult<Vec<DbTable>> {
33    let config: tokio_postgres::Config = connection_string.parse().map_err(|e| {
34        rust_ef::error::EFError::connection(format!("Invalid connection string: {}", e))
35    })?;
36
37    let client = match tls {
38        PgTlsMode::Disable => {
39            let (client, connection) = config.connect(NoTls).await.map_err(|e| {
40                rust_ef::error::EFError::connection(format!("Connection failed: {}", e))
41            })?;
42            tokio::spawn(async move {
43                if let Err(e) = connection.await {
44                    eprintln!("PostgreSQL connection error: {}", e);
45                }
46            });
47            client
48        }
49        PgTlsMode::Require(connector) => {
50            let tls_connector = postgres_native_tls::MakeTlsConnector::new(connector);
51            let (client, connection) = config.connect(tls_connector).await.map_err(|e| {
52                rust_ef::error::EFError::connection(format!("Connection failed: {}", e))
53            })?;
54            tokio::spawn(async move {
55                if let Err(e) = connection.await {
56                    eprintln!("PostgreSQL connection error: {}", e);
57                }
58            });
59            client
60        }
61    };
62
63    // Query for tables
64    let table_rows = client
65        .query(
66            "SELECT table_name FROM information_schema.tables \
67             WHERE table_schema = 'public' AND table_type = 'BASE TABLE' \
68             AND table_name NOT LIKE '__ef_%' \
69             ORDER BY table_name",
70            &[],
71        )
72        .await
73        .map_err(|e| rust_ef::error::EFError::query(format!("Table query error: {}", e)))?;
74
75    let mut tables = Vec::new();
76    for table_row in &table_rows {
77        let table_name: String = table_row.get(0);
78
79        // Query for columns
80        let col_rows = client
81            .query(
82                "SELECT c.column_name, c.data_type, c.is_nullable, c.character_maximum_length, \
83                 CASE WHEN pk.column_name IS NOT NULL THEN true ELSE false END AS is_pk \
84                 FROM information_schema.columns c \
85                 LEFT JOIN ( \
86                   SELECT ku.column_name \
87                   FROM information_schema.table_constraints tc \
88                   JOIN information_schema.key_column_usage ku \
89                     ON tc.constraint_name = ku.constraint_name \
90                   WHERE tc.constraint_type = 'PRIMARY KEY' AND tc.table_name = $1 \
91                 ) pk ON c.column_name = pk.column_name \
92                 WHERE c.table_name = $1 \
93                 ORDER BY c.ordinal_position",
94                &[&table_name],
95            )
96            .await
97            .map_err(|e| rust_ef::error::EFError::query(format!("Column query error: {}", e)))?;
98
99        let mut columns = Vec::new();
100        for col_row in &col_rows {
101            let col_name: String = col_row.get(0);
102            let data_type: String = col_row.get(1);
103            let is_nullable_str: String = col_row.get(2);
104            let max_length: Option<i32> = col_row.get(3);
105            let is_pk: bool = col_row.get(4);
106
107            columns.push(DbColumn {
108                name: col_name,
109                data_type,
110                is_nullable: is_nullable_str == "YES",
111                is_primary_key: is_pk,
112                max_length: max_length.map(|n| n as usize),
113            });
114        }
115
116        tables.push(DbTable {
117            name: table_name,
118            columns,
119        });
120    }
121
122    Ok(tables)
123}