Skip to main content

alibi_sqlx/store/
device_codes.rs

1use super::SqlxStore;
2use super::entities::device_code::Model;
3use crate::error::record_not_updated;
4use crate::model::{self, ActiveRow, SqlxModel};
5use crate::pool::{Engine, Exec};
6use crate::schema::AuthSchema;
7use crate::sql::Sql;
8use crate::value::SqlValue;
9use alibi_core::error::{AuthError, AuthResult};
10use alibi_core::store::DeviceCodeStore;
11use alibi_core::types::{CreateDeviceCode, DeviceCode, UpdateDeviceCode};
12use async_trait::async_trait;
13use uuid::Uuid;
14
15impl<S: AuthSchema + Send + Sync> SqlxStore<S> {
16    async fn create_device_code_with_connection(
17        &self,
18        exec: Exec<'_>,
19        input: CreateDeviceCode,
20        fields: serde_json::Map<String, serde_json::Value>,
21    ) -> AuthResult<DeviceCode> {
22        let mut active = ActiveRow::new();
23        active.set(
24            "id",
25            self.generated_id(self.exec(), "deviceCode", "deviceCode", "id")
26                .await?
27                .unwrap_or_else(|| Uuid::new_v4().to_string()),
28        );
29        active.set("deviceCode", input.device_code);
30        active.set("userCode", input.user_code);
31        active.set("userId", input.user_id);
32        active.set("expiresAt", input.expires_at);
33        active.set("status", input.status);
34        active.set("lastPolledAt", input.last_polled_at);
35        active.set("pollingInterval", input.polling_interval);
36        active.set("clientId", input.client_id);
37        active.set("scope", input.scope);
38        if fields.is_empty() {
39            return model::insert::<Model>(exec, &active)
40                .await
41                .map(|model| DeviceCode::from(&model));
42        }
43        if fields
44            .keys()
45            .any(|name| Model::COLUMN_NAMES.contains(&name.as_str()))
46        {
47            return Err(AuthError::bad_request(
48                "Application device fields cannot replace core columns",
49            ));
50        }
51        let present = active.present().collect::<Vec<_>>();
52        let mut sql = Sql::with(exec.engine(), "INSERT INTO ");
53        sql.ident(Model::TABLE);
54        sql.push(" (");
55        for (index, name) in present
56            .iter()
57            .map(|(name, _)| *name)
58            .chain(fields.keys().map(String::as_str))
59            .enumerate()
60        {
61            if index > 0 {
62                sql.push(", ");
63            }
64            sql.ident(name);
65        }
66        sql.push(") VALUES (");
67        for (index, (name, value)) in present.into_iter().enumerate() {
68            if index > 0 {
69                sql.push(", ");
70            }
71            model::bind_model_value::<Model>(&mut sql, name, value.clone());
72        }
73        for value in fields.into_values() {
74            sql.push(", ");
75            if value.is_null() {
76                sql.push("NULL");
77            } else {
78                sql.bind(device_field_value(value));
79            }
80        }
81        sql.push(")");
82        model::returning::<Model>(&mut sql);
83        exec.fetch_optional::<Model>(sql)
84            .await?
85            .map(|model| DeviceCode::from(&model))
86            .ok_or_else(|| AuthError::internal("Device code insertion returned no row"))
87    }
88
89    async fn find_device_code(&self, column: &str, value: &str) -> AuthResult<Option<DeviceCode>> {
90        let mut sql = model::select_model::<Model>(self.exec());
91        sql.push(" WHERE ");
92        sql.compare(Model::TABLE, column, " = ", value);
93        model::limit_one(&mut sql);
94        Ok(self
95            .exec()
96            .fetch_optional::<Model>(sql)
97            .await?
98            .map(|model| DeviceCode::from(&model)))
99    }
100
101    /// `UPDATE device_code SET ... WHERE id = ? AND ...`, returning the affected count.
102    async fn update_device_codes(
103        &self,
104        sets: Vec<(&'static str, SqlValue)>,
105        guards: Vec<(&'static str, Option<SqlValue>)>,
106    ) -> AuthResult<u64> {
107        let mut sql = Sql::with(self.exec().engine(), "UPDATE ");
108        sql.ident(Model::TABLE);
109        sql.push(" SET ");
110        for (index, (column, value)) in sets.into_iter().enumerate() {
111            if index > 0 {
112                sql.push(", ");
113            }
114            sql.assign(column, value);
115        }
116        for (index, (column, value)) in guards.into_iter().enumerate() {
117            sql.push(if index == 0 { " WHERE " } else { " AND " });
118            sql.column(Model::TABLE, column);
119            match value {
120                Some(value) => {
121                    sql.push(" = ");
122                    sql.bind(value);
123                }
124                None => {
125                    sql.push(" IS NULL");
126                }
127            }
128        }
129        self.exec().execute(sql).await
130    }
131}
132
133#[async_trait]
134impl<S> DeviceCodeStore for SqlxStore<S>
135where
136    S: AuthSchema + Send + Sync,
137{
138    async fn create_device_code(&self, input: CreateDeviceCode) -> AuthResult<DeviceCode> {
139        self.create_device_code_with_connection(self.exec(), input, serde_json::Map::new())
140            .await
141    }
142    async fn create_device_code_with_fields(
143        &self,
144        input: CreateDeviceCode,
145        fields: serde_json::Map<String, serde_json::Value>,
146    ) -> AuthResult<DeviceCode> {
147        self.create_device_code_with_connection(self.exec(), input, fields)
148            .await
149    }
150    async fn device_code_fields(
151        &self,
152        id: &str,
153    ) -> AuthResult<serde_json::Map<String, serde_json::Value>> {
154        let mut query = if self.exec().engine() == Engine::Postgres {
155            Sql::with(
156                Engine::Postgres,
157                "SELECT to_jsonb(t)::text FROM \"deviceCode\" t WHERE \"id\" = ",
158            )
159        } else {
160            #[derive(sqlx::FromRow)]
161            struct ColumnName {
162                name: String,
163            }
164            let columns = self
165                .exec()
166                .fetch_all::<ColumnName>(Sql::with(
167                    Engine::Sqlite,
168                    "SELECT name FROM pragma_table_info('deviceCode')",
169                ))
170                .await?;
171            let mut query = Sql::with(Engine::Sqlite, "SELECT json_object(");
172            for (index, column) in columns.into_iter().enumerate() {
173                if index > 0 {
174                    query.push(", ");
175                }
176                query.bind(column.name.clone());
177                query.push(", ");
178                query.ident(&column.name);
179            }
180            query.push(") FROM \"deviceCode\" WHERE \"id\" = ");
181            query
182        };
183        query.bind(id);
184        let json = self.exec().fetch_scalar::<String>(query).await?;
185        let mut fields = json
186            .map(|json| serde_json::from_str::<serde_json::Map<String, serde_json::Value>>(&json))
187            .transpose()?
188            .unwrap_or_default();
189        fields.retain(|name, _| !Model::COLUMN_NAMES.contains(&name.as_str()));
190        Ok(fields)
191    }
192    async fn consume_device_code(
193        &self,
194        id: &str,
195        status: &str,
196        ownership: &serde_json::Map<String, serde_json::Value>,
197    ) -> AuthResult<Option<DeviceCode>> {
198        let mut sql = Sql::with(
199            self.exec().engine(),
200            "DELETE FROM \"deviceCode\" WHERE id = ",
201        );
202        sql.bind(id);
203        sql.push(" AND status = ");
204        sql.bind(status);
205        for (field, value) in ownership {
206            sql.push(" AND ");
207            sql.ident(field);
208            if value.is_null() {
209                sql.push(" IS NULL");
210            } else {
211                sql.push(" = ");
212                sql.bind(device_field_value(value.clone()));
213            }
214        }
215        model::returning::<Model>(&mut sql);
216        self.exec()
217            .fetch_optional::<Model>(sql)
218            .await
219            .map(|row| row.as_ref().map(DeviceCode::from))
220    }
221
222    async fn get_device_code_by_device_code(
223        &self,
224        device_code: &str,
225    ) -> AuthResult<Option<DeviceCode>> {
226        self.find_device_code("deviceCode", device_code).await
227    }
228
229    async fn get_device_code_by_user_code(
230        &self,
231        user_code: &str,
232    ) -> AuthResult<Option<DeviceCode>> {
233        self.find_device_code("userCode", user_code).await
234    }
235
236    async fn update_device_code(
237        &self,
238        id: &str,
239        update: UpdateDeviceCode,
240    ) -> AuthResult<DeviceCode> {
241        let mut sql = model::by_id::<Model>(self.exec(), id);
242        model::limit_one(&mut sql);
243        let Some(model) = self.exec().fetch_optional::<Model>(sql).await? else {
244            return Err(AuthError::not_found("Device code not found"));
245        };
246
247        let mut active = model.into_active();
248        if let Some(status) = update.status {
249            active.set("status", status);
250        }
251        if let Some(user_id) = update.user_id {
252            active.set("userId", user_id);
253        }
254        if let Some(last_polled_at) = update.last_polled_at {
255            active.set("lastPolledAt", last_polled_at);
256        }
257
258        model::update::<Model>(self.exec(), &active)
259            .await?
260            .map(|model| DeviceCode::from(&model))
261            .ok_or_else(record_not_updated)
262    }
263
264    async fn update_device_code_if_status(
265        &self,
266        id: &str,
267        current_status: &str,
268        update: UpdateDeviceCode,
269    ) -> AuthResult<bool> {
270        let mut sets = Vec::new();
271        if let Some(status) = update.status {
272            sets.push(("status", status.into()));
273        }
274        if let Some(user_id) = update.user_id {
275            sets.push(("userId", user_id.into()));
276        }
277        if let Some(last_polled_at) = update.last_polled_at {
278            sets.push(("lastPolledAt", last_polled_at.into()));
279        }
280        self.update_device_codes(
281            sets,
282            vec![
283                ("id", Some(id.into())),
284                ("status", Some(current_status.into())),
285            ],
286        )
287        .await
288        .map(|affected| affected == 1)
289    }
290
291    async fn claim_device_code(&self, id: &str, user_id: &str) -> AuthResult<bool> {
292        self.update_device_codes(
293            vec![("userId", user_id.into())],
294            vec![
295                ("id", Some(id.into())),
296                ("status", Some("pending".into())),
297                ("userId", None),
298            ],
299        )
300        .await
301        .map(|affected| affected == 1)
302    }
303
304    async fn delete_device_code(&self, id: &str) -> AuthResult<()> {
305        self.exec()
306            .execute(model::delete_by_id::<Model>(self.exec(), id))
307            .await
308            .map(|_| ())
309    }
310
311    async fn delete_device_code_if_status(&self, id: &str, status: &str) -> AuthResult<bool> {
312        let mut sql = model::delete_by_id::<Model>(self.exec(), id);
313        sql.push(" AND ");
314        sql.compare(Model::TABLE, "status", " = ", status);
315        self.exec().execute(sql).await.map(|affected| affected == 1)
316    }
317}
318
319fn device_field_value(value: serde_json::Value) -> SqlValue {
320    match value {
321        serde_json::Value::String(value) => value.into(),
322        serde_json::Value::Bool(value) => value.into(),
323        serde_json::Value::Number(value) => value
324            .as_i64()
325            .map_or_else(|| value.as_f64().into(), SqlValue::from),
326        serde_json::Value::Null => SqlValue::Text(None),
327        value => SqlValue::Json(Some(Box::new(value))),
328    }
329}