ironflow_ops_postgres/admin/
connections.rs

1//! Connection and query monitoring 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 [`ActiveConnections`].
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct ActiveConnectionsOutput {
15    /// Number of active connections.
16    pub count: i64,
17}
18
19/// Count the number of active connections to the current database.
20///
21/// # Examples
22///
23/// ```no_run
24/// use ironflow_ops_postgres::admin::connections::ActiveConnections;
25/// use ironflow_core::operation::Operation;
26///
27/// # fn example(pool: sqlx::PgPool) {
28/// let op = ActiveConnections::new(pool);
29/// assert_eq!(op.kind(), "postgres");
30/// # }
31/// ```
32pub struct ActiveConnections {
33    pool: PgPool,
34}
35
36impl ActiveConnections {
37    /// Create a new active-connections operation.
38    pub fn new(pool: PgPool) -> Self {
39        Self { pool }
40    }
41
42    /// Execute and return a typed result.
43    ///
44    /// # Errors
45    ///
46    /// Returns [`OperationError::External`] on connection errors.
47    pub async fn run(
48        &self,
49        _ctx: &OperationContext,
50    ) -> Result<ActiveConnectionsOutput, OperationError> {
51        let row = sqlx::query(
52            "SELECT count(*) AS cnt FROM pg_stat_activity \
53             WHERE datname = current_database()",
54        )
55        .fetch_one(&self.pool)
56        .await
57        .map_err(pg_error)?;
58        let count: i64 = row.try_get("cnt").map_err(pg_error)?;
59        Ok(ActiveConnectionsOutput { count })
60    }
61}
62
63#[async_trait]
64impl Operation for ActiveConnections {
65    fn kind(&self) -> &str {
66        "postgres"
67    }
68    async fn execute(&self, ctx: &OperationContext) -> Result<Value, OperationError> {
69        to_value(&self.run(ctx).await?)
70    }
71}
72
73impl TypedOperation for ActiveConnections {
74    type Output = ActiveConnectionsOutput;
75}
76
77/// A running query description.
78#[derive(Debug, Clone, Serialize, Deserialize)]
79pub struct RunningQueryInfo {
80    /// Backend process ID.
81    pub pid: i32,
82    /// The SQL query text.
83    pub query: String,
84    /// Current state (e.g. `active`, `idle`).
85    pub state: String,
86    /// How long the query has been running, as a human-readable string.
87    pub duration: String,
88}
89
90/// Output of [`RunningQueries`].
91#[derive(Debug, Clone, Serialize, Deserialize)]
92pub struct RunningQueriesOutput {
93    /// Currently running queries.
94    pub queries: Vec<RunningQueryInfo>,
95}
96
97/// List currently running queries on the database.
98///
99/// # Examples
100///
101/// ```no_run
102/// use ironflow_ops_postgres::admin::connections::RunningQueries;
103/// use ironflow_core::operation::Operation;
104///
105/// # fn example(pool: sqlx::PgPool) {
106/// let op = RunningQueries::new(pool);
107/// assert_eq!(op.kind(), "postgres");
108/// # }
109/// ```
110pub struct RunningQueries {
111    pool: PgPool,
112}
113
114impl RunningQueries {
115    /// Create a new running-queries operation.
116    pub fn new(pool: PgPool) -> Self {
117        Self { pool }
118    }
119
120    /// Execute and return a typed result.
121    ///
122    /// # Errors
123    ///
124    /// Returns [`OperationError::External`] on connection errors.
125    pub async fn run(
126        &self,
127        _ctx: &OperationContext,
128    ) -> Result<RunningQueriesOutput, OperationError> {
129        let rows = sqlx::query(
130            "SELECT pid, query, state, \
131                    extract(epoch from (now() - query_start))::bigint AS duration_secs \
132             FROM pg_stat_activity \
133             WHERE datname = current_database() AND state = 'active' \
134             ORDER BY query_start",
135        )
136        .fetch_all(&self.pool)
137        .await
138        .map_err(pg_error)?;
139        let queries = rows
140            .iter()
141            .map(|r| {
142                let secs: i64 = r.try_get("duration_secs").unwrap_or(0);
143                Ok(RunningQueryInfo {
144                    pid: r.try_get::<i32, _>("pid").map_err(pg_error)?,
145                    query: r.try_get::<String, _>("query").map_err(pg_error)?,
146                    state: r.try_get::<String, _>("state").map_err(pg_error)?,
147                    duration: format!("{secs}s"),
148                })
149            })
150            .collect::<Result<Vec<_>, OperationError>>()?;
151        Ok(RunningQueriesOutput { queries })
152    }
153}
154
155#[async_trait]
156impl Operation for RunningQueries {
157    fn kind(&self) -> &str {
158        "postgres"
159    }
160    async fn execute(&self, ctx: &OperationContext) -> Result<Value, OperationError> {
161        to_value(&self.run(ctx).await?)
162    }
163}
164
165impl TypedOperation for RunningQueries {
166    type Output = RunningQueriesOutput;
167}
168
169#[cfg(test)]
170mod tests {
171    use super::*;
172
173    #[tokio::test]
174    async fn active_connections_kind() {
175        let pool = PgPool::connect_lazy("postgres://localhost/test").unwrap();
176        let op = ActiveConnections::new(pool);
177        assert_eq!(op.kind(), "postgres");
178    }
179
180    #[tokio::test]
181    async fn running_queries_kind() {
182        let pool = PgPool::connect_lazy("postgres://localhost/test").unwrap();
183        let op = RunningQueries::new(pool);
184        assert_eq!(op.kind(), "postgres");
185    }
186}
ironflow_ops_postgres/admin/connections.rs

ironflow_ops_postgres/admin/
connections.rs
