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    ) -> AuthResult<DeviceCode> {
21        let mut active = ActiveRow::new();
22        active.set(
23            "id",
24            self.generated_id(self.exec(), "deviceCode", "device_code", "id")
25                .await?
26                .unwrap_or_else(|| Uuid::new_v4().to_string()),
27        );
28        active.set("device_code", input.device_code);
29        active.set("user_code", input.user_code);
30        active.set("user_id", input.user_id);
31        active.set("expires_at", input.expires_at);
32        active.set("status", input.status);
33        active.set("last_polled_at", input.last_polled_at);
34        active.set("polling_interval", input.polling_interval);
35        active.set("client_id", input.client_id);
36        active.set("scope", input.scope);
37        model::insert::<Model>(exec, &active)
38            .await
39            .map(|model| DeviceCode::from(&model))
40    }
41
42    async fn find_device_code(&self, column: &str, value: &str) -> AuthResult<Option<DeviceCode>> {
43        let mut sql = model::select_model::<Model>(self.exec());
44        sql.push(" WHERE ");
45        sql.compare(Model::TABLE, column, " = ", value);
46        model::limit_one(&mut sql);
47        Ok(self
48            .exec()
49            .fetch_optional::<Model>(sql)
50            .await?
51            .map(|model| DeviceCode::from(&model)))
52    }
53
54    /// `UPDATE device_code SET ... WHERE id = ? AND ...`, returning the affected count.
55    async fn update_device_codes(
56        &self,
57        sets: Vec<(&'static str, SqlValue)>,
58        guards: Vec<(&'static str, Option<SqlValue>)>,
59    ) -> AuthResult<u64> {
60        let mut sql = Sql::with(self.exec().engine(), "UPDATE ");
61        sql.ident(Model::TABLE);
62        sql.push(" SET ");
63        for (index, (column, value)) in sets.into_iter().enumerate() {
64            if index > 0 {
65                sql.push(", ");
66            }
67            sql.assign(column, value);
68        }
69        for (index, (column, value)) in guards.into_iter().enumerate() {
70            sql.push(if index == 0 { " WHERE " } else { " AND " });
71            sql.column(Model::TABLE, column);
72            match value {
73                Some(value) => {
74                    sql.push(" = ");
75                    sql.bind(value);
76                }
77                None => {
78                    sql.push(" IS NULL");
79                }
80            }
81        }
82        self.exec().execute(sql).await
83    }
84}
85
86#[async_trait]
87impl<S> DeviceCodeStore for SqlxStore<S>
88where
89    S: AuthSchema + Send + Sync,
90{
91    async fn create_device_code(&self, input: CreateDeviceCode) -> AuthResult<DeviceCode> {
92        self.create_device_code_with_connection(self.exec(), input)
93            .await
94    }
95    async fn create_device_code_with_fields(
96        &self,
97        input: CreateDeviceCode,
98        fields: serde_json::Map<String, serde_json::Value>,
99    ) -> AuthResult<DeviceCode> {
100        if fields.is_empty() {
101            return self.create_device_code(input).await;
102        }
103        self.in_transaction(true, async move |tx| {
104            let exec = crate::pool::Exec::tx(tx);
105            let create = Sql::with(exec.engine(), "CREATE TABLE IF NOT EXISTS device_code_fields (device_code_id TEXT PRIMARY KEY REFERENCES device_code(id) ON DELETE CASCADE, fields TEXT NOT NULL)");
106            _ = exec.execute(create).await?;
107            let row = self.create_device_code_with_connection(exec, input).await?;
108            let mut insert = Sql::with(exec.engine(), "INSERT INTO device_code_fields (device_code_id, fields) VALUES (");
109            insert.bind(row.id.clone());insert.push(", ");insert.bind(serde_json::to_string(&fields)?);insert.push(")");
110            _ = exec.execute(insert).await?;
111            Ok(row)
112        }).await
113    }
114    async fn device_code_fields(
115        &self,
116        id: &str,
117    ) -> AuthResult<serde_json::Map<String, serde_json::Value>> {
118        if !super::migrator::has_table(self.exec(), "device_code_fields").await? {
119            return Ok(serde_json::Map::new());
120        }
121        let mut sql = Sql::with(
122            self.exec().engine(),
123            "SELECT fields FROM device_code_fields WHERE device_code_id = ",
124        );
125        sql.bind(id);
126        self.exec()
127            .fetch_scalar::<String>(sql)
128            .await?
129            .map(|fields| serde_json::from_str(&fields).map_err(Into::into))
130            .transpose()
131            .map(Option::unwrap_or_default)
132    }
133    async fn consume_device_code(
134        &self,
135        id: &str,
136        status: &str,
137        ownership: &serde_json::Map<String, serde_json::Value>,
138    ) -> AuthResult<Option<DeviceCode>> {
139        let mut sql = Sql::with(self.exec().engine(), "DELETE FROM device_code WHERE id = ");
140        sql.bind(id);
141        sql.push(" AND status = ");
142        sql.bind(status);
143        for (field, value) in ownership {
144            sql.push(" AND EXISTS (SELECT 1 FROM device_code_fields WHERE device_code_id = device_code.id AND ");
145            if self.exec().engine() == Engine::Postgres {
146                sql.push("CAST(fields AS JSONB) -> ");
147                sql.bind(field);
148                sql.push(" = CAST(");
149                sql.bind(value.to_string());
150                sql.push(" AS JSONB)");
151            } else {
152                sql.push("json_extract(fields, ");
153                sql.bind(format!("$.{}", serde_json::to_string(field)?));
154                sql.push(") IS json_extract(");
155                sql.bind(value.to_string());
156                sql.push(", '$')");
157            }
158            sql.push(")");
159        }
160        sql.push(" RETURNING ");
161        for (index, column) in Model::COLUMN_NAMES.iter().enumerate() {
162            if index > 0 {
163                sql.push(", ");
164            }
165            sql.ident(column);
166        }
167        self.exec()
168            .fetch_optional::<Model>(sql)
169            .await
170            .map(|row| row.as_ref().map(DeviceCode::from))
171    }
172
173    async fn get_device_code_by_device_code(
174        &self,
175        device_code: &str,
176    ) -> AuthResult<Option<DeviceCode>> {
177        self.find_device_code("device_code", device_code).await
178    }
179
180    async fn get_device_code_by_user_code(
181        &self,
182        user_code: &str,
183    ) -> AuthResult<Option<DeviceCode>> {
184        self.find_device_code("user_code", user_code).await
185    }
186
187    async fn update_device_code(
188        &self,
189        id: &str,
190        update: UpdateDeviceCode,
191    ) -> AuthResult<DeviceCode> {
192        let mut sql = model::by_id::<Model>(self.exec(), id);
193        model::limit_one(&mut sql);
194        let Some(model) = self.exec().fetch_optional::<Model>(sql).await? else {
195            return Err(AuthError::not_found("Device code not found"));
196        };
197
198        let mut active = model.into_active();
199        if let Some(status) = update.status {
200            active.set("status", status);
201        }
202        if let Some(user_id) = update.user_id {
203            active.set("user_id", user_id);
204        }
205        if let Some(last_polled_at) = update.last_polled_at {
206            active.set("last_polled_at", last_polled_at);
207        }
208
209        model::update::<Model>(self.exec(), &active)
210            .await?
211            .map(|model| DeviceCode::from(&model))
212            .ok_or_else(record_not_updated)
213    }
214
215    async fn update_device_code_if_status(
216        &self,
217        id: &str,
218        current_status: &str,
219        update: UpdateDeviceCode,
220    ) -> AuthResult<bool> {
221        let mut sets = Vec::new();
222        if let Some(status) = update.status {
223            sets.push(("status", status.into()));
224        }
225        if let Some(user_id) = update.user_id {
226            sets.push(("user_id", user_id.into()));
227        }
228        if let Some(last_polled_at) = update.last_polled_at {
229            sets.push(("last_polled_at", last_polled_at.into()));
230        }
231        self.update_device_codes(
232            sets,
233            vec![
234                ("id", Some(id.into())),
235                ("status", Some(current_status.into())),
236            ],
237        )
238        .await
239        .map(|affected| affected == 1)
240    }
241
242    async fn claim_device_code(&self, id: &str, user_id: &str) -> AuthResult<bool> {
243        self.update_device_codes(
244            vec![("user_id", user_id.into())],
245            vec![
246                ("id", Some(id.into())),
247                ("status", Some("pending".into())),
248                ("user_id", None),
249            ],
250        )
251        .await
252        .map(|affected| affected == 1)
253    }
254
255    async fn delete_device_code(&self, id: &str) -> AuthResult<()> {
256        self.exec()
257            .execute(model::delete_by_id::<Model>(self.exec(), id))
258            .await
259            .map(|_| ())
260    }
261
262    async fn delete_device_code_if_status(&self, id: &str, status: &str) -> AuthResult<bool> {
263        let mut sql = model::delete_by_id::<Model>(self.exec(), id);
264        sql.push(" AND ");
265        sql.compare(Model::TABLE, "status", " = ", status);
266        self.exec().execute(sql).await.map(|affected| affected == 1)
267    }
268}
alibi_sqlx/store/device_codes.rs

alibi_sqlx/store/
device_codes.rs
