conservator/builder/
insert.rs

1use crate::{Creatable, Domain};
2use std::marker::PhantomData;
3
4/// INSERT 查询构建器
5///
6/// 用于构建 INSERT 语句并支持不同的返回类型
7///
8/// # Example
9/// ```ignore
10/// // 返回主键
11/// let pk = CreateUser { name: "test".into(), email: "a@b.com".into() }
12///     .insert::<User>()
13///     .returning_pk(db)
14///     .await?;
15///
16/// // 返回完整实体
17/// let user = CreateUser { name: "test".into(), email: "a@b.com".into() }
18///     .insert::<User>()
19///     .returning_entity(db)
20///     .await?;
21/// ```
22/// 构建返回列名的 SQL 片段
23fn returning_columns<T: Domain>() -> String {
24    T::COLUMN_NAMES
25        .iter()
26        .map(|name| format!("\"{}\"", name))
27        .collect::<Vec<_>>()
28        .join(", ")
29}
30
31pub struct InsertBuilder<T: Domain, C: Creatable> {
32    data: C,
33    _phantom: PhantomData<T>,
34}
35
36impl<T: Domain, C: Creatable> InsertBuilder<T, C> {
37    pub fn new(data: C) -> Self {
38        Self {
39            data,
40            _phantom: PhantomData,
41        }
42    }
43
44    /// 执行 INSERT 并返回主键
45    ///
46    /// 生成: `INSERT INTO table (cols) VALUES (vals) RETURNING "pk_field"`
47    pub async fn returning_pk<'e, 'c: 'e, E: 'e + sqlx::Executor<'c, Database = sqlx::Postgres>>(
48        self,
49        executor: E,
50    ) -> Result<T::PrimaryKey, sqlx::Error>
51    where
52        T::PrimaryKey:
53            for<'r> sqlx::Decode<'r, sqlx::Postgres> + sqlx::Type<sqlx::Postgres> + Unpin,
54    {
55        let sql = format!(
56            "INSERT INTO {} {} VALUES {} RETURNING \"{}\"",
57            T::TABLE_NAME,
58            self.data.get_columns(),
59            self.data.get_insert_sql(),
60            T::PK_FIELD_NAME
61        );
62
63        let query = sqlx::query_scalar(&sql);
64        let query = self.data.bind_to_query_scalar(query);
65        query.fetch_one(executor).await
66    }
67
68    /// 执行 INSERT 并返回完整实体
69    ///
70    /// 生成: `INSERT INTO table (cols) VALUES (vals) RETURNING "col1", "col2", ...`
71    pub async fn returning_entity<
72        'e,
73        'c: 'e,
74        E: 'e + sqlx::Executor<'c, Database = sqlx::Postgres>,
75    >(
76        self,
77        executor: E,
78    ) -> Result<T, sqlx::Error> {
79        let sql = format!(
80            "INSERT INTO {} {} VALUES {} RETURNING {}",
81            T::TABLE_NAME,
82            self.data.get_columns(),
83            self.data.get_insert_sql(),
84            returning_columns::<T>()
85        );
86
87        let query = sqlx::query_as(&sql);
88        let query = self.data.build_for_query_as(query);
89        query.fetch_one(executor).await
90    }
91}
92
93pub struct InsertManyBuilder<T: Domain, C: Creatable> {
94    data: Vec<C>,
95    _phantom: PhantomData<T>,
96}
97
98impl<T: Domain, C: Creatable> InsertManyBuilder<T, C> {
99    pub fn new(data: Vec<C>) -> Self {
100        Self {
101            data,
102            _phantom: PhantomData,
103        }
104    }
105
106    /// 构建批量 VALUES 子句
107    fn build_values_sql(&self) -> String {
108        self.data
109            .iter()
110            .enumerate()
111            .map(|(idx, item)| item.get_batch_insert_sql(idx))
112            .collect::<Vec<_>>()
113            .join(", ")
114    }
115
116    /// 执行 INSERT 并返回主键列表
117    ///
118    /// 生成: `INSERT INTO table (cols) VALUES (v1), (v2), ... RETURNING "pk_field"`
119    pub async fn returning_pk<'e, 'c: 'e, E: 'e + sqlx::Executor<'c, Database = sqlx::Postgres>>(
120        self,
121        executor: E,
122    ) -> Result<Vec<T::PrimaryKey>, sqlx::Error>
123    where
124        T::PrimaryKey:
125            for<'r> sqlx::Decode<'r, sqlx::Postgres> + sqlx::Type<sqlx::Postgres> + Unpin,
126    {
127        if self.data.is_empty() {
128            return Ok(Vec::new());
129        }
130
131        let columns = self.data[0].get_columns();
132        let values_sql = self.build_values_sql();
133
134        let sql = format!(
135            "INSERT INTO {} {} VALUES {} RETURNING \"{}\"",
136            T::TABLE_NAME,
137            columns,
138            values_sql,
139            T::PK_FIELD_NAME
140        );
141
142        let mut query = sqlx::query_scalar(&sql);
143        for item in self.data {
144            query = item.bind_to_query_scalar(query);
145        }
146        query.fetch_all(executor).await
147    }
148
149    /// 执行 INSERT 并返回完整实体列表
150    ///
151    /// 生成: `INSERT INTO table (cols) VALUES (v1), (v2), ... RETURNING "col1", "col2", ...`
152    pub async fn returning_entity<
153        'e,
154        'c: 'e,
155        E: 'e + sqlx::Executor<'c, Database = sqlx::Postgres>,
156    >(
157        self,
158        executor: E,
159    ) -> Result<Vec<T>, sqlx::Error> {
160        if self.data.is_empty() {
161            return Ok(Vec::new());
162        }
163
164        let columns = self.data[0].get_columns();
165        let values_sql = self.build_values_sql();
166
167        let sql = format!(
168            "INSERT INTO {} {} VALUES {} RETURNING {}",
169            T::TABLE_NAME,
170            columns,
171            values_sql,
172            returning_columns::<T>()
173        );
174
175        let mut query = sqlx::query_as(&sql);
176        for item in self.data {
177            query = item.build_for_query_as(query);
178        }
179        query.fetch_all(executor).await
180    }
181
182    /// 执行 INSERT(不返回数据)
183    ///
184    /// 生成: `INSERT INTO table (cols) VALUES (v1), (v2), ...`
185    pub async fn execute<'e, 'c: 'e, E: 'e + sqlx::Executor<'c, Database = sqlx::Postgres>>(
186        self,
187        executor: E,
188    ) -> Result<u64, sqlx::Error> {
189        if self.data.is_empty() {
190            return Ok(0);
191        }
192
193        let columns = self.data[0].get_columns();
194        let values_sql = self.build_values_sql();
195
196        let sql = format!(
197            "INSERT INTO {} {} VALUES {}",
198            T::TABLE_NAME,
199            columns,
200            values_sql
201        );
202
203        let mut query = sqlx::query(&sql);
204        for item in self.data {
205            query = item.build_for_query(query);
206        }
207        let result = query.execute(executor).await?;
208        Ok(result.rows_affected())
209    }
210}