Skip to main content

claw_spawn/infrastructure/
postgres_droplet_repo.rs

1use 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}