1use crate::prelude::*;
2use serde::{Deserialize, Serialize};
3use sqlx::{Row, postgres::PgRow};
4use std::{
5 collections::BTreeMap,
6 sync::{Arc, LazyLock},
7};
8use utoipa::ToSchema;
9
10#[derive(Serialize, Deserialize, Clone)]
11pub struct ServerAllocation {
12 pub uuid: uuid::Uuid,
13 pub allocation: super::node_allocation::NodeAllocation,
14
15 pub notes: Option<compact_str::CompactString>,
16
17 pub created: chrono::NaiveDateTime,
18}
19
20impl BaseModel for ServerAllocation {
21 const NAME: &'static str = "server_allocation";
22
23 #[inline]
24 fn columns(prefix: Option<&str>) -> BTreeMap<&'static str, compact_str::CompactString> {
25 let prefix = prefix.unwrap_or_default();
26
27 let mut columns = BTreeMap::from([
28 (
29 "server_allocations.uuid",
30 compact_str::format_compact!("{prefix}uuid"),
31 ),
32 (
33 "server_allocations.notes",
34 compact_str::format_compact!("{prefix}notes"),
35 ),
36 (
37 "server_allocations.created",
38 compact_str::format_compact!("{prefix}created"),
39 ),
40 ]);
41
42 columns.extend(super::node_allocation::NodeAllocation::columns(Some(
43 "allocation_",
44 )));
45
46 columns
47 }
48
49 #[inline]
50 fn map(prefix: Option<&str>, row: &PgRow) -> Result<Self, crate::database::DatabaseError> {
51 let prefix = prefix.unwrap_or_default();
52
53 Ok(Self {
54 uuid: row.try_get(compact_str::format_compact!("{prefix}uuid").as_str())?,
55 allocation: super::node_allocation::NodeAllocation::map(Some("allocation_"), row)?,
56 notes: row.try_get(compact_str::format_compact!("{prefix}notes").as_str())?,
57 created: row.try_get(compact_str::format_compact!("{prefix}created").as_str())?,
58 })
59 }
60}
61
62impl ServerAllocation {
63 pub async fn create(
64 database: &crate::database::Database,
65 server_uuid: uuid::Uuid,
66 allocation_uuid: uuid::Uuid,
67 ) -> Result<uuid::Uuid, crate::database::DatabaseError> {
68 let row = sqlx::query(
69 r#"
70 INSERT INTO server_allocations (server_uuid, allocation_uuid)
71 VALUES ($1, $2)
72 RETURNING uuid
73 "#,
74 )
75 .bind(server_uuid)
76 .bind(allocation_uuid)
77 .fetch_one(database.write())
78 .await?;
79
80 Ok(row.try_get("uuid")?)
81 }
82
83 pub async fn create_random(
84 database: &crate::database::Database,
85 server: &super::server::Server,
86 ) -> Result<uuid::Uuid, crate::database::DatabaseError> {
87 let row = sqlx::query(
88 r#"
89 INSERT INTO server_allocations (server_uuid, allocation_uuid)
90 VALUES ($1, (
91 SELECT node_allocations.uuid FROM node_allocations
92 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
93 WHERE
94 node_allocations.node_uuid = $2
95 AND ($3 IS NULL OR node_allocations.ip = $3)
96 AND node_allocations.port BETWEEN $4 AND $5
97 AND server_allocations.uuid IS NULL
98 ORDER BY RANDOM()
99 LIMIT 1
100 ))
101 RETURNING uuid
102 "#,
103 )
104 .bind(server.uuid)
105 .bind(server.node.uuid)
106 .bind(server.allocation.as_ref().map(|a| a.allocation.ip))
107 .bind(server.egg.config_allocations.user_self_assign.start_port as i32)
108 .bind(server.egg.config_allocations.user_self_assign.end_port as i32)
109 .fetch_one(database.write())
110 .await?;
111
112 Ok(row.get("uuid"))
113 }
114
115 pub async fn by_uuid(
116 database: &crate::database::Database,
117 uuid: uuid::Uuid,
118 ) -> Result<Option<Self>, crate::database::DatabaseError> {
119 let row = sqlx::query(&format!(
120 r#"
121 SELECT {}
122 FROM server_allocations
123 JOIN node_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
124 WHERE server_allocations.uuid = $1
125 "#,
126 Self::columns_sql(None)
127 ))
128 .bind(uuid)
129 .fetch_optional(database.read())
130 .await?;
131
132 row.try_map(|row| Self::map(None, &row))
133 }
134
135 pub async fn by_server_uuid_uuid(
136 database: &crate::database::Database,
137 server_uuid: uuid::Uuid,
138 allocation_uuid: uuid::Uuid,
139 ) -> Result<Option<Self>, crate::database::DatabaseError> {
140 let row = sqlx::query(&format!(
141 r#"
142 SELECT {}
143 FROM server_allocations
144 JOIN node_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
145 WHERE server_allocations.server_uuid = $1 AND server_allocations.uuid = $2
146 "#,
147 Self::columns_sql(None)
148 ))
149 .bind(server_uuid)
150 .bind(allocation_uuid)
151 .fetch_optional(database.read())
152 .await?;
153
154 row.try_map(|row| Self::map(None, &row))
155 }
156
157 pub async fn by_server_uuid_with_pagination(
158 database: &crate::database::Database,
159 server_uuid: uuid::Uuid,
160 page: i64,
161 per_page: i64,
162 search: Option<&str>,
163 ) -> Result<super::Pagination<Self>, crate::database::DatabaseError> {
164 let offset = (page - 1) * per_page;
165
166 let rows = sqlx::query(&format!(
167 r#"
168 SELECT {}, COUNT(*) OVER() AS total_count
169 FROM server_allocations
170 JOIN node_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
171 WHERE server_allocations.server_uuid = $1
172 AND ($2 IS NULL OR host(node_allocations.ip) || ':' || node_allocations.port ILIKE '%' || $2 || '%' OR server_allocations.notes ILIKE '%' || $2 || '%')
173 ORDER BY server_allocations.created
174 LIMIT $3 OFFSET $4
175 "#,
176 Self::columns_sql(None)
177 ))
178 .bind(server_uuid)
179 .bind(search)
180 .bind(per_page)
181 .bind(offset)
182 .fetch_all(database.read())
183 .await?;
184
185 Ok(super::Pagination {
186 total: rows
187 .first()
188 .map_or(Ok(0), |row| row.try_get("total_count"))?,
189 per_page,
190 page,
191 data: rows
192 .into_iter()
193 .map(|row| Self::map(None, &row))
194 .try_collect_vec()?,
195 })
196 }
197
198 pub async fn count_by_server_uuid(
199 database: &crate::database::Database,
200 server_uuid: uuid::Uuid,
201 ) -> i64 {
202 sqlx::query_scalar(
203 r#"
204 SELECT COUNT(*)
205 FROM server_allocations
206 WHERE server_allocations.server_uuid = $1
207 "#,
208 )
209 .bind(server_uuid)
210 .fetch_one(database.read())
211 .await
212 .unwrap_or(0)
213 }
214
215 #[inline]
216 pub fn into_api_object(self, primary: Option<uuid::Uuid>) -> ApiServerAllocation {
217 ApiServerAllocation {
218 uuid: self.uuid,
219 ip: compact_str::format_compact!("{}", self.allocation.ip.ip()),
220 ip_alias: self.allocation.ip_alias,
221 port: self.allocation.port,
222 notes: self.notes,
223 is_primary: primary.is_some_and(|p| p == self.uuid),
224 created: self.created.and_utc(),
225 }
226 }
227}
228
229#[async_trait::async_trait]
230impl DeletableModel for ServerAllocation {
231 type DeleteOptions = ();
232
233 fn get_delete_handlers() -> &'static LazyLock<DeleteListenerList<Self>> {
234 static DELETE_LISTENERS: LazyLock<DeleteListenerList<ServerAllocation>> =
235 LazyLock::new(|| Arc::new(ModelHandlerList::default()));
236
237 &DELETE_LISTENERS
238 }
239
240 async fn delete(
241 &self,
242 state: &crate::State,
243 options: Self::DeleteOptions,
244 ) -> Result<(), anyhow::Error> {
245 let mut transaction = state.database.write().begin().await?;
246
247 self.run_delete_handlers(&options, state, &mut transaction)
248 .await?;
249
250 sqlx::query(
251 r#"
252 DELETE FROM server_allocations
253 WHERE server_allocations.uuid = $1
254 "#,
255 )
256 .bind(self.uuid)
257 .execute(&mut *transaction)
258 .await?;
259
260 transaction.commit().await?;
261
262 Ok(())
263 }
264}
265
266#[derive(ToSchema, Serialize)]
267#[schema(title = "ServerAllocation")]
268pub struct ApiServerAllocation {
269 pub uuid: uuid::Uuid,
270
271 pub ip: compact_str::CompactString,
272 pub ip_alias: Option<compact_str::CompactString>,
273 pub port: i32,
274
275 pub notes: Option<compact_str::CompactString>,
276 pub is_primary: bool,
277
278 pub created: chrono::DateTime<chrono::Utc>,
279}