Skip to main content

ironflow_ops_postgres/admin/
size.rs

1//! Database and table size operations.
2
3use async_trait::async_trait;
4use ironflow_core::error::OperationError;
5use ironflow_core::operation::{Operation, OperationContext, TypedOperation};
6use serde::{Deserialize, Serialize};
7use serde_json::Value;
8use sqlx::{PgPool, Row};
9
10use crate::helpers::{pg_error, to_value};
11
12/// Output of [`DatabaseSize`].
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct DatabaseSizeOutput {
15    /// Database name.
16    pub database: String,
17    /// Size in bytes.
18    pub size_bytes: i64,
19}
20
21/// Get the size of a database in bytes.
22///
23/// # Examples
24///
25/// ```no_run
26/// use ironflow_ops_postgres::admin::size::DatabaseSize;
27/// use ironflow_core::operation::Operation;
28///
29/// # fn example(pool: sqlx::PgPool) {
30/// let op = DatabaseSize::new(pool, "mydb");
31/// assert_eq!(op.kind(), "postgres");
32/// # }
33/// ```
34pub struct DatabaseSize {
35    pool: PgPool,
36    database: String,
37}
38
39impl DatabaseSize {
40    /// Create a new database-size operation.
41    pub fn new(pool: PgPool, database: impl Into<String>) -> Self {
42        Self {
43            pool,
44            database: database.into(),
45        }
46    }
47
48    /// Execute and return a typed result.
49    ///
50    /// # Errors
51    ///
52    /// Returns [`OperationError::External`] on connection errors.
53    pub async fn run(&self, _ctx: &OperationContext) -> Result<DatabaseSizeOutput, OperationError> {
54        let row = sqlx::query("SELECT pg_database_size($1) AS size")
55            .bind(&self.database)
56            .fetch_one(&self.pool)
57            .await
58            .map_err(pg_error)?;
59        let size: i64 = row.try_get("size").map_err(pg_error)?;
60        Ok(DatabaseSizeOutput {
61            database: self.database.clone(),
62            size_bytes: size,
63        })
64    }
65}
66
67#[async_trait]
68impl Operation for DatabaseSize {
69    fn kind(&self) -> &str {
70        "postgres"
71    }
72    async fn execute(&self, ctx: &OperationContext) -> Result<Value, OperationError> {
73        to_value(&self.run(ctx).await?)
74    }
75    fn input(&self) -> Option<Value> {
76        Some(serde_json::json!({ "database": self.database }))
77    }
78}
79
80impl TypedOperation for DatabaseSize {
81    type Output = DatabaseSizeOutput;
82}
83
84/// Output of [`TableSize`].
85#[derive(Debug, Clone, Serialize, Deserialize)]
86pub struct TableSizeOutput {
87    /// Fully qualified table name.
88    pub table: String,
89    /// Total size in bytes (table + indexes + toast).
90    pub total_bytes: i64,
91    /// Table-only size in bytes.
92    pub table_bytes: i64,
93    /// Index size in bytes.
94    pub index_bytes: i64,
95}
96
97/// Get the size of a table including indexes.
98///
99/// # Examples
100///
101/// ```no_run
102/// use ironflow_ops_postgres::admin::size::TableSize;
103/// use ironflow_core::operation::Operation;
104///
105/// # fn example(pool: sqlx::PgPool) {
106/// let op = TableSize::new(pool, "public", "users");
107/// assert_eq!(op.kind(), "postgres");
108/// # }
109/// ```
110pub struct TableSize {
111    pool: PgPool,
112    schema: String,
113    table: String,
114}
115
116impl TableSize {
117    /// Create a new table-size operation.
118    pub fn new(pool: PgPool, schema: impl Into<String>, table: impl Into<String>) -> Self {
119        Self {
120            pool,
121            schema: schema.into(),
122            table: table.into(),
123        }
124    }
125
126    /// Execute and return a typed result.
127    ///
128    /// # Errors
129    ///
130    /// Returns [`OperationError::External`] on connection errors.
131    pub async fn run(&self, _ctx: &OperationContext) -> Result<TableSizeOutput, OperationError> {
132        let qualified = format!("{}.{}", self.schema, self.table);
133        let row = sqlx::query(
134            "SELECT pg_total_relation_size($1::regclass) AS total, \
135                    pg_table_size($1::regclass) AS tbl, \
136                    pg_indexes_size($1::regclass) AS idx",
137        )
138        .bind(&qualified)
139        .fetch_one(&self.pool)
140        .await
141        .map_err(pg_error)?;
142        Ok(TableSizeOutput {
143            table: qualified,
144            total_bytes: row.try_get("total").map_err(pg_error)?,
145            table_bytes: row.try_get("tbl").map_err(pg_error)?,
146            index_bytes: row.try_get("idx").map_err(pg_error)?,
147        })
148    }
149}
150
151#[async_trait]
152impl Operation for TableSize {
153    fn kind(&self) -> &str {
154        "postgres"
155    }
156    async fn execute(&self, ctx: &OperationContext) -> Result<Value, OperationError> {
157        to_value(&self.run(ctx).await?)
158    }
159    fn input(&self) -> Option<Value> {
160        Some(serde_json::json!({ "schema": self.schema, "table": self.table }))
161    }
162}
163
164impl TypedOperation for TableSize {
165    type Output = TableSizeOutput;
166}
167
168#[cfg(test)]
169mod tests {
170    use super::*;
171
172    #[tokio::test]
173    async fn database_size_kind() {
174        let pool = PgPool::connect_lazy("postgres://localhost/test").unwrap();
175        let op = DatabaseSize::new(pool, "test");
176        assert_eq!(op.kind(), "postgres");
177    }
178
179    #[tokio::test]
180    async fn table_size_kind() {
181        let pool = PgPool::connect_lazy("postgres://localhost/test").unwrap();
182        let op = TableSize::new(pool, "public", "users");
183        assert_eq!(op.kind(), "postgres");
184    }
185}