Skip to main content

wasi_pg_client/query/
result.rs

1//! Query result types.
2//!
3//! This module defines [`CommandTag`], [`QueryResult`], and [`ExecuteResult`].
4
5use std::sync::Arc;
6
7use crate::query::row::{FieldDescription, Row};
8
9// ---------------------------------------------------------------------------
10// CommandTag
11// ---------------------------------------------------------------------------
12
13/// The command tag returned by PostgreSQL for a completed command
14/// (e.g. `"SELECT 3"`, `"INSERT 0 1"`, `"UPDATE 4"`).
15#[derive(Debug, Clone, PartialEq, Eq, Default)]
16#[non_exhaustive]
17pub struct CommandTag {
18    tag: String,
19}
20
21impl CommandTag {
22    /// Create a new command tag from the raw string.
23    pub fn new(tag: String) -> Self {
24        Self { tag }
25    }
26
27    /// The raw tag string.
28    pub fn as_str(&self) -> &str {
29        &self.tag
30    }
31
32    /// Parse the number of rows affected from the tag.
33    ///
34    /// Returns `None` for commands that do not have a row count
35    /// (e.g. `CREATE TABLE`).
36    pub fn rows_affected(&self) -> Option<u64> {
37        // Tags have the form:
38        //   INSERT 0 <rows>
39        //   UPDATE <rows>
40        //   DELETE <rows>
41        //   SELECT <rows>
42        //   MOVE <rows>
43        //   FETCH <rows>
44        //   COPY <rows>
45        let parts: Vec<&str> = self.tag.split_whitespace().collect();
46        match parts.as_slice() {
47            ["INSERT", _, n] => n.parse().ok(),
48            ["UPDATE", n] => n.parse().ok(),
49            ["DELETE", n] => n.parse().ok(),
50            ["SELECT", n] => n.parse().ok(),
51            ["MOVE", n] => n.parse().ok(),
52            ["FETCH", n] => n.parse().ok(),
53            ["COPY", n] => n.parse().ok(),
54            _ => None,
55        }
56    }
57}
58
59// ---------------------------------------------------------------------------
60// QueryResult
61// ---------------------------------------------------------------------------
62
63/// Result of a query that returns rows.
64#[derive(Debug, Clone)]
65#[non_exhaustive]
66pub struct QueryResult {
67    rows: Vec<Row>,
68    command_tag: CommandTag,
69    columns: Arc<Vec<FieldDescription>>,
70}
71
72impl QueryResult {
73    pub(crate) fn new(
74        rows: Vec<Row>,
75        command_tag: CommandTag,
76        columns: Arc<Vec<FieldDescription>>,
77    ) -> Self {
78        Self {
79            rows,
80            command_tag,
81            columns,
82        }
83    }
84
85    /// Returns the number of rows in the result.
86    pub fn len(&self) -> usize {
87        self.rows.len()
88    }
89
90    /// Returns `true` if the result contains no rows.
91    pub fn is_empty(&self) -> bool {
92        self.rows.is_empty()
93    }
94
95    /// Returns the rows.
96    pub fn rows(&self) -> &[Row] {
97        &self.rows
98    }
99
100    /// Consumes the result and returns the rows vector.
101    pub fn into_rows(self) -> Vec<Row> {
102        self.rows
103    }
104
105    /// Returns the command tag.
106    pub fn command_tag(&self) -> &CommandTag {
107        &self.command_tag
108    }
109
110    /// Returns the column descriptions.
111    pub fn columns(&self) -> &[FieldDescription] {
112        &self.columns
113    }
114
115    /// Returns the number of rows affected (if parseable from the tag).
116    pub fn rows_affected(&self) -> Option<u64> {
117        self.command_tag.rows_affected()
118    }
119
120    /// Returns an iterator over the rows.
121    pub fn iter(&self) -> impl Iterator<Item = &Row> {
122        self.rows.iter()
123    }
124}
125
126// ---------------------------------------------------------------------------
127// ExecuteResult
128// ---------------------------------------------------------------------------
129
130/// Result for statements that do not return rows (INSERT, UPDATE, DELETE, DDL).
131#[derive(Debug, Clone, PartialEq, Eq)]
132#[non_exhaustive]
133pub struct ExecuteResult {
134    command_tag: CommandTag,
135}
136
137impl ExecuteResult {
138    pub(crate) fn new(command_tag: CommandTag) -> Self {
139        Self { command_tag }
140    }
141
142    /// Returns the command tag.
143    pub fn command_tag(&self) -> &CommandTag {
144        &self.command_tag
145    }
146
147    /// Returns the number of rows affected (if parseable from the tag).
148    pub fn rows_affected(&self) -> Option<u64> {
149        self.command_tag.rows_affected()
150    }
151}
152
153// ---------------------------------------------------------------------------
154// Tests
155// ---------------------------------------------------------------------------
156
157#[cfg(test)]
158mod tests {
159    use super::*;
160
161    #[test]
162    fn test_command_tag_rows_affected() {
163        assert_eq!(CommandTag::new("SELECT 3".into()).rows_affected(), Some(3));
164        assert_eq!(
165            CommandTag::new("INSERT 0 1".into()).rows_affected(),
166            Some(1)
167        );
168        assert_eq!(CommandTag::new("UPDATE 4".into()).rows_affected(), Some(4));
169        assert_eq!(CommandTag::new("DELETE 0".into()).rows_affected(), Some(0));
170        assert_eq!(CommandTag::new("CREATE TABLE".into()).rows_affected(), None);
171        assert_eq!(CommandTag::new("".into()).rows_affected(), None);
172    }
173
174    #[test]
175    fn test_query_result_empty() {
176        let qr = QueryResult::new(
177            Vec::new(),
178            CommandTag::new("SELECT 0".into()),
179            Arc::new(Vec::new()),
180        );
181        assert!(qr.is_empty());
182        assert_eq!(qr.len(), 0);
183        assert_eq!(qr.rows_affected(), Some(0));
184    }
185
186    #[test]
187    fn test_execute_result_rows_affected() {
188        let er = ExecuteResult::new(CommandTag::new("UPDATE 5".into()));
189        assert_eq!(er.rows_affected(), Some(5));
190    }
191}