1use crate::{models::UpdateQueryBuilder, prelude::*};
2use axum::http::StatusCode;
3use garde::Validate;
4use serde::{Deserialize, Serialize};
5use sqlx::{Row, postgres::PgRow};
6use std::{
7 collections::BTreeMap,
8 sync::{Arc, LazyLock},
9};
10use utoipa::ToSchema;
11#[derive(Serialize, Deserialize, Clone)]
12pub struct ServerAllocation {
13 pub uuid: uuid::Uuid,
14 pub allocation: super::node_allocation::NodeAllocation,
15
16 pub notes: Option<compact_str::CompactString>,
17
18 pub created: chrono::NaiveDateTime,
19
20 extension_data: super::ModelExtensionData,
21}
22
23impl BaseModel for ServerAllocation {
24 const NAME: &'static str = "server_allocation";
25
26 fn get_extension_list() -> &'static super::ModelExtensionList {
27 static EXTENSIONS: LazyLock<super::ModelExtensionList> =
28 LazyLock::new(|| parking_lot::RwLock::new(Vec::new()));
29
30 &EXTENSIONS
31 }
32
33 fn get_extension_data(&self) -> &super::ModelExtensionData {
34 &self.extension_data
35 }
36
37 #[inline]
38 fn base_columns(prefix: Option<&str>) -> BTreeMap<&'static str, compact_str::CompactString> {
39 let prefix = prefix.unwrap_or_default();
40
41 let mut columns = BTreeMap::from([
42 (
43 "server_allocations.uuid",
44 compact_str::format_compact!("{prefix}uuid"),
45 ),
46 (
47 "server_allocations.notes",
48 compact_str::format_compact!("{prefix}notes"),
49 ),
50 (
51 "server_allocations.created",
52 compact_str::format_compact!("{prefix}created"),
53 ),
54 ]);
55
56 columns.extend(super::node_allocation::NodeAllocation::base_columns(Some(
57 "allocation_",
58 )));
59
60 columns
61 }
62
63 #[inline]
64 fn map(prefix: Option<&str>, row: &PgRow) -> Result<Self, crate::database::DatabaseError> {
65 let prefix = prefix.unwrap_or_default();
66
67 Ok(Self {
68 uuid: row.try_get(compact_str::format_compact!("{prefix}uuid").as_str())?,
69 allocation: super::node_allocation::NodeAllocation::map(Some("allocation_"), row)?,
70 notes: row.try_get(compact_str::format_compact!("{prefix}notes").as_str())?,
71 created: row.try_get(compact_str::format_compact!("{prefix}created").as_str())?,
72 extension_data: Self::map_extensions(prefix, row)?,
73 })
74 }
75}
76
77impl ServerAllocation {
78 pub async fn create(
79 database: &crate::database::Database,
80 server_uuid: uuid::Uuid,
81 allocation_uuid: uuid::Uuid,
82 ) -> Result<uuid::Uuid, crate::database::DatabaseError> {
83 let row = sqlx::query(
84 r#"
85 INSERT INTO server_allocations (server_uuid, allocation_uuid)
86 VALUES ($1, $2)
87 RETURNING uuid
88 "#,
89 )
90 .bind(server_uuid)
91 .bind(allocation_uuid)
92 .fetch_one(database.write())
93 .await?;
94
95 Ok(row.try_get("uuid")?)
96 }
97
98 pub async fn create_random(
99 state: &crate::State,
100 server: &super::server::Server,
101 ) -> Result<uuid::Uuid, crate::database::DatabaseError> {
102 let egg_configuration = server.egg.configuration(&state.database).await?;
103
104 let Some(config_allocations) = egg_configuration.config_allocations else {
105 return Err(anyhow::Error::new(
106 crate::response::DisplayError::new(
107 "no egg allocation configuration found, cannot auto-assign allocation",
108 )
109 .with_status(StatusCode::EXPECTATION_FAILED),
110 )
111 .into());
112 };
113
114 let ip = server.allocation.as_ref().map(|a| a.allocation.ip);
115 let start_port = config_allocations.user_self_assign.start_port as i32;
116 let end_port = config_allocations.user_self_assign.end_port as i32;
117
118 let candidate_ips: Vec<sqlx::types::ipnetwork::IpNetwork> = sqlx::query_scalar(
119 r#"
120 SELECT DISTINCT node_allocations.ip
121 FROM node_allocations
122 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
123 WHERE
124 node_allocations.node_uuid = $1
125 AND ($2 IS NULL OR node_allocations.ip = $2)
126 AND node_allocations.port BETWEEN $3 AND $4
127 AND server_allocations.uuid IS NULL
128 "#,
129 )
130 .bind(server.node.uuid)
131 .bind(ip)
132 .bind(start_port)
133 .bind(end_port)
134 .fetch_all(state.database.read())
135 .await?;
136
137 if candidate_ips.is_empty() {
138 return Err(anyhow::Error::new(
139 crate::response::DisplayError::new("no node allocations are available")
140 .with_status(StatusCode::EXPECTATION_FAILED),
141 )
142 .into());
143 }
144
145 let exclude = match super::node_allocation::NodeAllocation::used_by_node(
146 state,
147 &server.node.fetch_cached(&state.database).await?,
148 &candidate_ips.iter().map(|ip| ip.ip()).collect::<Vec<_>>(),
149 )
150 .await
151 {
152 Ok(exclude) => exclude,
153 Err(err) => {
154 tracing::warn!(
155 node = %server.node.uuid,
156 "failed to resolve used ports, refusing to assign an allocation: {err:#}"
157 );
158
159 return Err(anyhow::Error::new(
160 crate::response::DisplayError::new(
161 "could not reach the node to check which ports are free",
162 )
163 .with_status(StatusCode::EXPECTATION_FAILED),
164 )
165 .into());
166 }
167 };
168
169 let row = sqlx::query(
170 r#"
171 INSERT INTO server_allocations (server_uuid, allocation_uuid)
172 VALUES ($1, (
173 SELECT node_allocations.uuid FROM node_allocations
174 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
175 WHERE
176 node_allocations.node_uuid = $2
177 AND ($3 IS NULL OR node_allocations.ip = $3)
178 AND node_allocations.port BETWEEN $4 AND $5
179 AND server_allocations.uuid IS NULL
180 AND NOT (node_allocations.uuid = ANY($6))
181 ORDER BY RANDOM()
182 LIMIT 1
183 ))
184 RETURNING uuid
185 "#,
186 )
187 .bind(server.uuid)
188 .bind(server.node.uuid)
189 .bind(ip)
190 .bind(start_port)
191 .bind(end_port)
192 .bind(&exclude)
193 .fetch_one(state.database.write())
194 .await;
195
196 let row = match row {
197 Ok(row) => row,
198 Err(err) if err.is_not_null_violation() => {
199 return Err(anyhow::Error::new(
200 crate::response::DisplayError::new(
201 "every free port on this node is currently in use",
202 )
203 .with_status(StatusCode::EXPECTATION_FAILED),
204 )
205 .into());
206 }
207 Err(err) => return Err(err.into()),
208 };
209
210 Ok(row.get("uuid"))
211 }
212
213 pub async fn by_uuid(
214 database: &crate::database::Database,
215 uuid: uuid::Uuid,
216 ) -> Result<Option<Self>, crate::database::DatabaseError> {
217 let row = sqlx::query(sqlx::AssertSqlSafe(format!(
218 r#"
219 SELECT {}
220 FROM server_allocations
221 JOIN node_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
222 WHERE server_allocations.uuid = $1
223 "#,
224 Self::columns_sql(None)
225 )))
226 .bind(uuid)
227 .fetch_optional(database.read())
228 .await?;
229
230 row.try_map(|row| Self::map(None, &row))
231 }
232
233 pub async fn by_server_uuid_uuid(
234 database: &crate::database::Database,
235 server_uuid: uuid::Uuid,
236 allocation_uuid: uuid::Uuid,
237 ) -> Result<Option<Self>, crate::database::DatabaseError> {
238 let row = sqlx::query(sqlx::AssertSqlSafe(format!(
239 r#"
240 SELECT {}
241 FROM server_allocations
242 JOIN node_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
243 WHERE server_allocations.server_uuid = $1 AND server_allocations.uuid = $2
244 "#,
245 Self::columns_sql(None)
246 )))
247 .bind(server_uuid)
248 .bind(allocation_uuid)
249 .fetch_optional(database.read())
250 .await?;
251
252 row.try_map(|row| Self::map(None, &row))
253 }
254
255 pub async fn by_server_uuid_with_pagination(
256 database: &crate::database::Database,
257 server_uuid: uuid::Uuid,
258 page: i64,
259 per_page: i64,
260 search: Option<&str>,
261 ) -> Result<super::Pagination<Self>, crate::database::DatabaseError> {
262 let offset = (page - 1) * per_page;
263
264 let rows = sqlx::query(sqlx::AssertSqlSafe(format!(
265 r#"
266 SELECT {}, COUNT(*) OVER() AS total_count
267 FROM server_allocations
268 JOIN node_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
269 WHERE server_allocations.server_uuid = $1
270 AND (
271 $2 IS NULL
272 OR host(node_allocations.ip) || ':' || node_allocations.port ILIKE '%' || $2 || '%'
273 OR (node_allocations.ip_alias IS NOT NULL AND node_allocations.ip_alias || ':' || node_allocations.port ILIKE '%' || $2 || '%')
274 OR server_allocations.notes ILIKE '%' || $2 || '%'
275 )
276 ORDER BY server_allocations.created
277 LIMIT $3 OFFSET $4
278 "#,
279 Self::columns_sql(None)
280 )))
281 .bind(server_uuid)
282 .bind(search)
283 .bind(per_page)
284 .bind(offset)
285 .fetch_all(database.read())
286 .await?;
287
288 Ok(super::Pagination {
289 total: rows
290 .first()
291 .map_or(Ok(0), |row| row.try_get("total_count"))?,
292 per_page,
293 page,
294 data: rows
295 .into_iter()
296 .map(|row| Self::map(None, &row))
297 .try_collect_vec()?,
298 })
299 }
300
301 pub async fn all_by_server_uuid(
302 database: &crate::database::Database,
303 server_uuid: uuid::Uuid,
304 ) -> Result<Vec<Self>, crate::database::DatabaseError> {
305 let rows = sqlx::query(sqlx::AssertSqlSafe(format!(
306 r#"
307 SELECT {}
308 FROM server_allocations
309 JOIN node_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
310 WHERE server_allocations.server_uuid = $1
311 ORDER BY server_allocations.created
312 "#,
313 Self::columns_sql(None)
314 )))
315 .bind(server_uuid)
316 .fetch_all(database.read())
317 .await?;
318
319 rows.into_iter()
320 .map(|row| Self::map(None, &row))
321 .try_collect_vec()
322 }
323
324 pub async fn count_by_server_uuid(
325 database: &crate::database::Database,
326 server_uuid: uuid::Uuid,
327 ) -> Result<i64, sqlx::Error> {
328 sqlx::query_scalar(
329 r#"
330 SELECT COUNT(*)
331 FROM server_allocations
332 WHERE server_allocations.server_uuid = $1
333 "#,
334 )
335 .bind(server_uuid)
336 .fetch_one(database.read())
337 .await
338 }
339}
340
341#[async_trait::async_trait]
342impl IntoApiObject for ServerAllocation {
343 type ApiObject = ApiServerAllocation;
344 type ExtraArgs<'a> = Option<uuid::Uuid>;
345
346 async fn into_api_object<'a>(
347 self,
348 state: &crate::State,
349 primary: Self::ExtraArgs<'a>,
350 ) -> Result<Self::ApiObject, crate::database::DatabaseError> {
351 let api_object = ApiServerAllocation::init_hooks(&self, state).await?;
352
353 let api_object = finish_extendible!(
354 ApiServerAllocation {
355 uuid: self.uuid,
356 ip: compact_str::format_compact!("{}", self.allocation.ip.ip()),
357 ip_alias: self.allocation.ip_alias,
358 port: self.allocation.port,
359 notes: self.notes,
360 is_primary: primary.is_some_and(|p| p == self.uuid),
361 created: self.created.and_utc(),
362 },
363 api_object,
364 state
365 )?;
366
367 Ok(api_object)
368 }
369}
370
371#[derive(ToSchema, Serialize, Deserialize, Validate, Default)]
372pub struct UpdateServerAllocationOptions {
373 #[garde(length(min = 1, max = 1024))]
374 #[schema(min_length = 1, max_length = 1024)]
375 #[serde(
376 default,
377 skip_serializing_if = "Option::is_none",
378 with = "::serde_with::rust::double_option"
379 )]
380 pub notes: Option<Option<compact_str::CompactString>>,
381}
382
383#[async_trait::async_trait]
384impl UpdatableModel for ServerAllocation {
385 type UpdateOptions = UpdateServerAllocationOptions;
386
387 fn get_update_handlers() -> &'static LazyLock<UpdateHandlerList<Self>> {
388 static UPDATE_LISTENERS: LazyLock<UpdateHandlerList<ServerAllocation>> =
389 LazyLock::new(|| Arc::new(ModelHandlerList::default()));
390
391 &UPDATE_LISTENERS
392 }
393
394 async fn update_with_transaction(
395 &mut self,
396 state: &crate::State,
397 mut options: Self::UpdateOptions,
398 transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
399 ) -> Result<(), crate::database::DatabaseError> {
400 options.validate()?;
401
402 let mut query_builder = UpdateQueryBuilder::new("server_allocations");
403
404 self.run_update_handlers(&mut options, &mut query_builder, state, transaction)
405 .await?;
406
407 query_builder
408 .set("notes", options.notes.as_ref().map(|n| n.as_ref()))
409 .where_eq("uuid", self.uuid);
410
411 query_builder.execute(&mut **transaction).await?;
412
413 if let Some(notes) = options.notes {
414 self.notes = notes;
415 }
416
417 self.run_after_update_handlers(state, transaction).await?;
418
419 Ok(())
420 }
421}
422
423#[async_trait::async_trait]
424impl DeletableModel for ServerAllocation {
425 type DeleteOptions = ();
426
427 fn get_delete_handlers() -> &'static LazyLock<DeleteHandlerList<Self>> {
428 static DELETE_LISTENERS: LazyLock<DeleteHandlerList<ServerAllocation>> =
429 LazyLock::new(|| Arc::new(ModelHandlerList::default()));
430
431 &DELETE_LISTENERS
432 }
433
434 async fn delete_with_transaction(
435 &self,
436 state: &crate::State,
437 options: Self::DeleteOptions,
438 transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
439 ) -> Result<(), anyhow::Error> {
440 self.run_delete_handlers(&options, state, transaction)
441 .await?;
442
443 sqlx::query(
444 r#"
445 DELETE FROM server_allocations
446 WHERE server_allocations.uuid = $1
447 "#,
448 )
449 .bind(self.uuid)
450 .execute(&mut **transaction)
451 .await?;
452
453 self.run_after_delete_handlers(&options, state, transaction)
454 .await?;
455
456 Ok(())
457 }
458}
459
460#[schema_extension_derive::extendible]
461#[init_args(ServerAllocation, crate::State)]
462#[hook_args(crate::State)]
463#[derive(ToSchema, Serialize)]
464#[schema(title = "ServerAllocation")]
465pub struct ApiServerAllocation {
466 pub uuid: uuid::Uuid,
467
468 pub ip: compact_str::CompactString,
469 pub ip_alias: Option<compact_str::CompactString>,
470 pub port: i32,
471
472 pub notes: Option<compact_str::CompactString>,
473 pub is_primary: bool,
474
475 pub created: chrono::DateTime<chrono::Utc>,
476}