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 super::PaginatedQueryBuilder::new(
263 Self::columns_sql(None),
264 format!(
265 r#"
266 FROM server_allocations
267 JOIN node_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
268 WHERE server_allocations.server_uuid = $1
269 AND (
270 $2 IS NULL
271 OR host(node_allocations.ip) || ':' || node_allocations.port ILIKE '%' || $2 || '%'
272 OR (node_allocations.ip_alias IS NOT NULL AND node_allocations.ip_alias || ':' || node_allocations.port ILIKE '%' || $2 || '%')
273 OR server_allocations.notes ILIKE '%' || $2 || '%'
274 OR {search_uuid}
275 )
276 "#,
277 search_uuid = super::uuid_search_sql(2, "server_allocations.uuid"),
278 ),
279 "server_allocations.created, node_allocations.ip, node_allocations.port, server_allocations.uuid",
280 )
281 .bind(server_uuid)
282 .bind(search)
283 .fetch(database, page, per_page)
284 .await
285 }
286
287 pub async fn all_by_server_uuid(
288 database: &crate::database::Database,
289 server_uuid: uuid::Uuid,
290 ) -> Result<Vec<Self>, crate::database::DatabaseError> {
291 let rows = sqlx::query(sqlx::AssertSqlSafe(format!(
292 r#"
293 SELECT {}
294 FROM server_allocations
295 JOIN node_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
296 WHERE server_allocations.server_uuid = $1
297 ORDER BY server_allocations.created, node_allocations.ip, node_allocations.port, server_allocations.uuid
298 "#,
299 Self::columns_sql(None)
300 )))
301 .bind(server_uuid)
302 .fetch_all(database.read())
303 .await?;
304
305 rows.into_iter()
306 .map(|row| Self::map(None, &row))
307 .try_collect_vec()
308 }
309
310 pub async fn count_by_server_uuid(
311 database: &crate::database::Database,
312 server_uuid: uuid::Uuid,
313 ) -> Result<i64, sqlx::Error> {
314 sqlx::query_scalar(
315 r#"
316 SELECT COUNT(*)
317 FROM server_allocations
318 WHERE server_allocations.server_uuid = $1
319 "#,
320 )
321 .bind(server_uuid)
322 .fetch_one(database.read())
323 .await
324 }
325}
326
327#[async_trait::async_trait]
328impl IntoApiObject for ServerAllocation {
329 type ApiObject = ApiServerAllocation;
330 type ExtraArgs<'a> = Option<uuid::Uuid>;
331
332 async fn into_api_object<'a>(
333 self,
334 state: &crate::State,
335 primary: Self::ExtraArgs<'a>,
336 ) -> Result<Self::ApiObject, crate::database::DatabaseError> {
337 let api_object = ApiServerAllocation::init_hooks(&self, state).await?;
338
339 let api_object = finish_extendible!(
340 ApiServerAllocation {
341 uuid: self.uuid,
342 ip: compact_str::format_compact!("{}", self.allocation.ip.ip()),
343 ip_alias: self.allocation.ip_alias,
344 port: self.allocation.port,
345 notes: self.notes,
346 is_primary: primary.is_some_and(|p| p == self.uuid),
347 created: self.created.and_utc(),
348 },
349 api_object,
350 state
351 )?;
352
353 Ok(api_object)
354 }
355}
356
357#[derive(ToSchema, Serialize, Deserialize, Validate, Default)]
358pub struct UpdateServerAllocationOptions {
359 #[garde(length(min = 1, max = 1024))]
360 #[schema(min_length = 1, max_length = 1024)]
361 #[serde(
362 default,
363 skip_serializing_if = "Option::is_none",
364 with = "::serde_with::rust::double_option"
365 )]
366 pub notes: Option<Option<compact_str::CompactString>>,
367}
368
369#[async_trait::async_trait]
370impl UpdatableModel for ServerAllocation {
371 type UpdateOptions = UpdateServerAllocationOptions;
372
373 fn get_update_handlers() -> &'static LazyLock<UpdateHandlerList<Self>> {
374 static UPDATE_LISTENERS: LazyLock<UpdateHandlerList<ServerAllocation>> =
375 LazyLock::new(|| Arc::new(ModelHandlerList::default()));
376
377 &UPDATE_LISTENERS
378 }
379
380 async fn update_with_transaction(
381 &mut self,
382 state: &crate::State,
383 mut options: Self::UpdateOptions,
384 transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
385 ) -> Result<(), crate::database::DatabaseError> {
386 options.validate()?;
387
388 let mut query_builder = UpdateQueryBuilder::new("server_allocations");
389
390 self.run_update_handlers(&mut options, &mut query_builder, state, transaction)
391 .await?;
392
393 query_builder
394 .set("notes", options.notes.as_ref().map(|n| n.as_ref()))
395 .where_eq("uuid", self.uuid);
396
397 query_builder.execute(&mut **transaction).await?;
398
399 if let Some(notes) = options.notes {
400 self.notes = notes;
401 }
402
403 self.run_after_update_handlers(state, transaction).await?;
404
405 Ok(())
406 }
407}
408
409#[async_trait::async_trait]
410impl DeletableModel for ServerAllocation {
411 type DeleteOptions = ();
412
413 fn get_delete_handlers() -> &'static LazyLock<DeleteHandlerList<Self>> {
414 static DELETE_LISTENERS: LazyLock<DeleteHandlerList<ServerAllocation>> =
415 LazyLock::new(|| Arc::new(ModelHandlerList::default()));
416
417 &DELETE_LISTENERS
418 }
419
420 async fn delete_with_transaction(
421 &self,
422 state: &crate::State,
423 options: Self::DeleteOptions,
424 transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
425 ) -> Result<(), anyhow::Error> {
426 self.run_delete_handlers(&options, state, transaction)
427 .await?;
428
429 sqlx::query(
430 r#"
431 DELETE FROM server_allocations
432 WHERE server_allocations.uuid = $1
433 "#,
434 )
435 .bind(self.uuid)
436 .execute(&mut **transaction)
437 .await?;
438
439 self.run_after_delete_handlers(&options, state, transaction)
440 .await?;
441
442 Ok(())
443 }
444}
445
446#[schema_extension_derive::extendible]
447#[init_args(ServerAllocation, crate::State)]
448#[hook_args(crate::State)]
449#[derive(ToSchema, Serialize)]
450#[schema(title = "ServerAllocation")]
451pub struct ApiServerAllocation {
452 pub uuid: uuid::Uuid,
453
454 pub ip: compact_str::CompactString,
455 pub ip_alias: Option<compact_str::CompactString>,
456 pub port: i32,
457
458 pub notes: Option<compact_str::CompactString>,
459 pub is_primary: bool,
460
461 pub created: chrono::DateTime<chrono::Utc>,
462}