1use crate::prelude::*;
2use axum::http::StatusCode;
3use serde::{Deserialize, Serialize};
4use sqlx::{Row, postgres::PgRow};
5use std::{
6 collections::BTreeMap,
7 sync::{Arc, LazyLock},
8};
9#[derive(Serialize, Deserialize, Clone)]
10pub struct ServerAllocation {
11 pub uuid: uuid::Uuid,
12 pub allocation: super::node_allocation::NodeAllocation,
13
14 pub notes: Option<compact_str::CompactString>,
15
16 pub created: chrono::NaiveDateTime,
17
18 extension_data: super::ModelExtensionData,
19}
20
21impl BaseModel for ServerAllocation {
22 const NAME: &'static str = "server_allocation";
23
24 fn get_extension_list() -> &'static super::ModelExtensionList {
25 static EXTENSIONS: LazyLock<super::ModelExtensionList> =
26 LazyLock::new(|| parking_lot::RwLock::new(Vec::new()));
27
28 &EXTENSIONS
29 }
30
31 fn get_extension_data(&self) -> &super::ModelExtensionData {
32 &self.extension_data
33 }
34
35 #[inline]
36 fn base_columns(prefix: Option<&str>) -> BTreeMap<&'static str, compact_str::CompactString> {
37 let prefix = prefix.unwrap_or_default();
38
39 let mut columns = BTreeMap::from([
40 (
41 "server_allocations.uuid",
42 compact_str::format_compact!("{prefix}uuid"),
43 ),
44 (
45 "server_allocations.notes",
46 compact_str::format_compact!("{prefix}notes"),
47 ),
48 (
49 "server_allocations.created",
50 compact_str::format_compact!("{prefix}created"),
51 ),
52 ]);
53
54 columns.extend(super::node_allocation::NodeAllocation::base_columns(Some(
55 "allocation_",
56 )));
57
58 columns
59 }
60
61 #[inline]
62 fn map(prefix: Option<&str>, row: &PgRow) -> Result<Self, crate::database::DatabaseError> {
63 let prefix = prefix.unwrap_or_default();
64
65 Ok(Self {
66 uuid: row.try_get(compact_str::format_compact!("{prefix}uuid").as_str())?,
67 allocation: super::node_allocation::NodeAllocation::map(Some("allocation_"), row)?,
68 notes: row.try_get(compact_str::format_compact!("{prefix}notes").as_str())?,
69 created: row.try_get(compact_str::format_compact!("{prefix}created").as_str())?,
70 extension_data: Self::map_extensions(prefix, row)?,
71 })
72 }
73}
74
75impl ServerAllocation {
76 pub async fn create(
77 database: &crate::database::Database,
78 server_uuid: uuid::Uuid,
79 allocation_uuid: uuid::Uuid,
80 ) -> Result<uuid::Uuid, crate::database::DatabaseError> {
81 let row = sqlx::query(
82 r#"
83 INSERT INTO server_allocations (server_uuid, allocation_uuid)
84 VALUES ($1, $2)
85 RETURNING uuid
86 "#,
87 )
88 .bind(server_uuid)
89 .bind(allocation_uuid)
90 .fetch_one(database.write())
91 .await?;
92
93 Ok(row.try_get("uuid")?)
94 }
95
96 pub async fn create_random(
97 database: &crate::database::Database,
98 server: &super::server::Server,
99 ) -> Result<uuid::Uuid, crate::database::DatabaseError> {
100 let egg_configuration = server.egg.configuration(database).await?;
101
102 let Some(config_allocations) = egg_configuration.config_allocations else {
103 return Err(anyhow::Error::new(
104 crate::response::DisplayError::new(
105 "no egg allocation configuration found, cannot auto-assign allocation",
106 )
107 .with_status(StatusCode::EXPECTATION_FAILED),
108 )
109 .into());
110 };
111
112 let row = sqlx::query(
113 r#"
114 INSERT INTO server_allocations (server_uuid, allocation_uuid)
115 VALUES ($1, (
116 SELECT node_allocations.uuid FROM node_allocations
117 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
118 WHERE
119 node_allocations.node_uuid = $2
120 AND ($3 IS NULL OR node_allocations.ip = $3)
121 AND node_allocations.port BETWEEN $4 AND $5
122 AND server_allocations.uuid IS NULL
123 ORDER BY RANDOM()
124 LIMIT 1
125 ))
126 RETURNING uuid
127 "#,
128 )
129 .bind(server.uuid)
130 .bind(server.node.uuid)
131 .bind(server.allocation.as_ref().map(|a| a.allocation.ip))
132 .bind(config_allocations.user_self_assign.start_port as i32)
133 .bind(config_allocations.user_self_assign.end_port as i32)
134 .fetch_one(database.write())
135 .await?;
136
137 Ok(row.get("uuid"))
138 }
139
140 pub async fn by_uuid(
141 database: &crate::database::Database,
142 uuid: uuid::Uuid,
143 ) -> Result<Option<Self>, crate::database::DatabaseError> {
144 let row = sqlx::query(sqlx::AssertSqlSafe(format!(
145 r#"
146 SELECT {}
147 FROM server_allocations
148 JOIN node_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
149 WHERE server_allocations.uuid = $1
150 "#,
151 Self::columns_sql(None)
152 )))
153 .bind(uuid)
154 .fetch_optional(database.read())
155 .await?;
156
157 row.try_map(|row| Self::map(None, &row))
158 }
159
160 pub async fn by_server_uuid_uuid(
161 database: &crate::database::Database,
162 server_uuid: uuid::Uuid,
163 allocation_uuid: uuid::Uuid,
164 ) -> Result<Option<Self>, crate::database::DatabaseError> {
165 let row = sqlx::query(sqlx::AssertSqlSafe(format!(
166 r#"
167 SELECT {}
168 FROM server_allocations
169 JOIN node_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
170 WHERE server_allocations.server_uuid = $1 AND server_allocations.uuid = $2
171 "#,
172 Self::columns_sql(None)
173 )))
174 .bind(server_uuid)
175 .bind(allocation_uuid)
176 .fetch_optional(database.read())
177 .await?;
178
179 row.try_map(|row| Self::map(None, &row))
180 }
181
182 pub async fn by_server_uuid_with_pagination(
183 database: &crate::database::Database,
184 server_uuid: uuid::Uuid,
185 page: i64,
186 per_page: i64,
187 search: Option<&str>,
188 ) -> Result<super::Pagination<Self>, crate::database::DatabaseError> {
189 let offset = (page - 1) * per_page;
190
191 let rows = sqlx::query(sqlx::AssertSqlSafe(format!(
192 r#"
193 SELECT {}, COUNT(*) OVER() AS total_count
194 FROM server_allocations
195 JOIN node_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
196 WHERE server_allocations.server_uuid = $1
197 AND (
198 $2 IS NULL
199 OR host(node_allocations.ip) || ':' || node_allocations.port ILIKE '%' || $2 || '%'
200 OR (node_allocations.ip_alias IS NOT NULL AND node_allocations.ip_alias || ':' || node_allocations.port ILIKE '%' || $2 || '%')
201 OR server_allocations.notes ILIKE '%' || $2 || '%'
202 )
203 ORDER BY server_allocations.created
204 LIMIT $3 OFFSET $4
205 "#,
206 Self::columns_sql(None)
207 )))
208 .bind(server_uuid)
209 .bind(search)
210 .bind(per_page)
211 .bind(offset)
212 .fetch_all(database.read())
213 .await?;
214
215 Ok(super::Pagination {
216 total: rows
217 .first()
218 .map_or(Ok(0), |row| row.try_get("total_count"))?,
219 per_page,
220 page,
221 data: rows
222 .into_iter()
223 .map(|row| Self::map(None, &row))
224 .try_collect_vec()?,
225 })
226 }
227
228 pub async fn all_by_server_uuid(
229 database: &crate::database::Database,
230 server_uuid: uuid::Uuid,
231 ) -> Result<Vec<Self>, crate::database::DatabaseError> {
232 let rows = sqlx::query(sqlx::AssertSqlSafe(format!(
233 r#"
234 SELECT {}
235 FROM server_allocations
236 JOIN node_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
237 WHERE server_allocations.server_uuid = $1
238 ORDER BY server_allocations.created
239 "#,
240 Self::columns_sql(None)
241 )))
242 .bind(server_uuid)
243 .fetch_all(database.read())
244 .await?;
245
246 rows.into_iter()
247 .map(|row| Self::map(None, &row))
248 .try_collect_vec()
249 }
250
251 pub async fn count_by_server_uuid(
252 database: &crate::database::Database,
253 server_uuid: uuid::Uuid,
254 ) -> Result<i64, sqlx::Error> {
255 sqlx::query_scalar(
256 r#"
257 SELECT COUNT(*)
258 FROM server_allocations
259 WHERE server_allocations.server_uuid = $1
260 "#,
261 )
262 .bind(server_uuid)
263 .fetch_one(database.read())
264 .await
265 }
266}
267
268#[async_trait::async_trait]
269impl IntoApiObject for ServerAllocation {
270 type ApiObject = ApiServerAllocation;
271 type ExtraArgs<'a> = Option<uuid::Uuid>;
272
273 async fn into_api_object<'a>(
274 self,
275 state: &crate::State,
276 primary: Self::ExtraArgs<'a>,
277 ) -> Result<Self::ApiObject, crate::database::DatabaseError> {
278 let api_object = ApiServerAllocation::init_hooks(&self, state).await?;
279
280 let api_object = finish_extendible!(
281 ApiServerAllocation {
282 uuid: self.uuid,
283 ip: compact_str::format_compact!("{}", self.allocation.ip.ip()),
284 ip_alias: self.allocation.ip_alias,
285 port: self.allocation.port,
286 notes: self.notes,
287 is_primary: primary.is_some_and(|p| p == self.uuid),
288 created: self.created.and_utc(),
289 },
290 api_object,
291 state
292 )?;
293
294 Ok(api_object)
295 }
296}
297
298#[async_trait::async_trait]
299impl DeletableModel for ServerAllocation {
300 type DeleteOptions = ();
301
302 fn get_delete_handlers() -> &'static LazyLock<DeleteHandlerList<Self>> {
303 static DELETE_LISTENERS: LazyLock<DeleteHandlerList<ServerAllocation>> =
304 LazyLock::new(|| Arc::new(ModelHandlerList::default()));
305
306 &DELETE_LISTENERS
307 }
308
309 async fn delete_with_transaction(
310 &self,
311 state: &crate::State,
312 options: Self::DeleteOptions,
313 transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
314 ) -> Result<(), anyhow::Error> {
315 self.run_delete_handlers(&options, state, transaction)
316 .await?;
317
318 sqlx::query(
319 r#"
320 DELETE FROM server_allocations
321 WHERE server_allocations.uuid = $1
322 "#,
323 )
324 .bind(self.uuid)
325 .execute(&mut **transaction)
326 .await?;
327
328 self.run_after_delete_handlers(&options, state, transaction)
329 .await?;
330
331 Ok(())
332 }
333}
334
335#[schema_extension_derive::extendible]
336#[init_args(ServerAllocation, crate::State)]
337#[hook_args(crate::State)]
338#[derive(ToSchema, Serialize)]
339#[schema(title = "ServerAllocation")]
340pub struct ApiServerAllocation {
341 pub uuid: uuid::Uuid,
342
343 pub ip: compact_str::CompactString,
344 pub ip_alias: Option<compact_str::CompactString>,
345 pub port: i32,
346
347 pub notes: Option<compact_str::CompactString>,
348 pub is_primary: bool,
349
350 pub created: chrono::DateTime<chrono::Utc>,
351}