claw_spawn/infrastructure/
postgres_droplet_repo.rs1use crate::domain::{Droplet, DropletStatus};
2use crate::infrastructure::{DropletRepository, RepositoryError};
3use async_trait::async_trait;
4use sqlx::{PgPool, Row};
5use uuid::Uuid;
6
7pub struct PostgresDropletRepository {
8 pool: PgPool,
9}
10
11impl PostgresDropletRepository {
12 pub fn new(pool: PgPool) -> Self {
13 Self { pool }
14 }
15}
16
17#[async_trait]
18impl DropletRepository for PostgresDropletRepository {
19 async fn create(&self, droplet: &Droplet) -> Result<(), RepositoryError> {
20 let status_str = droplet_status_to_string(&droplet.status);
21
22 sqlx::query(
23 r#"
24 INSERT INTO droplets (id, name, region, size, image, status, ip_address, bot_id, created_at, destroyed_at)
25 VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
26 "#,
27 )
28 .bind(droplet.id)
29 .bind(&droplet.name)
30 .bind(&droplet.region)
31 .bind(&droplet.size)
32 .bind(&droplet.image)
33 .bind(status_str)
34 .bind(&droplet.ip_address)
35 .bind(droplet.bot_id)
36 .bind(droplet.created_at)
37 .bind(droplet.destroyed_at)
38 .execute(&self.pool)
39 .await?;
40
41 Ok(())
42 }
43
44 async fn get_by_id(&self, id: i64) -> Result<Droplet, RepositoryError> {
45 let row = sqlx::query(
46 r#"
47 SELECT id, name, region, size, image, status, ip_address, bot_id, created_at, destroyed_at
48 FROM droplets
49 WHERE id = $1
50 "#,
51 )
52 .bind(id)
53 .fetch_one(&self.pool)
54 .await
55 .map_err(|e| match e {
56 sqlx::Error::RowNotFound => RepositoryError::NotFound(format!("Droplet {}", id)),
57 _ => RepositoryError::DatabaseError(e),
58 })?;
59
60 Ok(row_to_droplet(&row)?)
61 }
62
63 async fn update_bot_assignment(
64 &self,
65 droplet_id: i64,
66 bot_id: Option<Uuid>,
67 ) -> Result<(), RepositoryError> {
68 sqlx::query(
69 r#"
70 UPDATE droplets
71 SET bot_id = $1
72 WHERE id = $2
73 "#,
74 )
75 .bind(bot_id)
76 .bind(droplet_id)
77 .execute(&self.pool)
78 .await?;
79
80 Ok(())
81 }
82
83 async fn update_status(&self, droplet_id: i64, status: &str) -> Result<(), RepositoryError> {
84 sqlx::query(
85 r#"
86 UPDATE droplets
87 SET status = $1
88 WHERE id = $2
89 "#,
90 )
91 .bind(status)
92 .bind(droplet_id)
93 .execute(&self.pool)
94 .await?;
95
96 Ok(())
97 }
98
99 async fn update_ip(&self, droplet_id: i64, ip: Option<String>) -> Result<(), RepositoryError> {
100 sqlx::query(
101 r#"
102 UPDATE droplets
103 SET ip_address = $1
104 WHERE id = $2
105 "#,
106 )
107 .bind(ip)
108 .bind(droplet_id)
109 .execute(&self.pool)
110 .await?;
111
112 Ok(())
113 }
114
115 async fn mark_destroyed(&self, droplet_id: i64) -> Result<(), RepositoryError> {
116 sqlx::query(
117 r#"
118 UPDATE droplets
119 SET status = 'destroyed', destroyed_at = $1
120 WHERE id = $2
121 "#,
122 )
123 .bind(chrono::Utc::now())
124 .bind(droplet_id)
125 .execute(&self.pool)
126 .await?;
127
128 Ok(())
129 }
130}
131
132fn droplet_status_to_string(status: &DropletStatus) -> String {
133 match status {
134 DropletStatus::New => "new".to_string(),
135 DropletStatus::Active => "active".to_string(),
136 DropletStatus::Off => "off".to_string(),
137 DropletStatus::Destroyed => "destroyed".to_string(),
138 DropletStatus::Error => "error".to_string(),
139 }
140}
141
142fn string_to_droplet_status(status: &str) -> Result<DropletStatus, RepositoryError> {
143 match status {
144 "new" => Ok(DropletStatus::New),
145 "active" => Ok(DropletStatus::Active),
146 "off" => Ok(DropletStatus::Off),
147 "destroyed" => Ok(DropletStatus::Destroyed),
148 "error" => Ok(DropletStatus::Error),
149 _ => Err(RepositoryError::InvalidData(format!(
150 "Unknown droplet status: {}",
151 status
152 ))),
153 }
154}
155
156fn row_to_droplet(row: &sqlx::postgres::PgRow) -> Result<Droplet, RepositoryError> {
157 let status_str: String = row.try_get("status")?;
158
159 Ok(Droplet {
160 id: row.try_get("id")?,
161 name: row.try_get("name")?,
162 region: row.try_get("region")?,
163 size: row.try_get("size")?,
164 image: row.try_get("image")?,
165 status: string_to_droplet_status(&status_str)?,
166 ip_address: row.try_get("ip_address")?,
167 bot_id: row.try_get("bot_id")?,
168 created_at: row.try_get("created_at")?,
169 destroyed_at: row.try_get("destroyed_at")?,
170 })
171}