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 state: &crate::State,
98 server: &super::server::Server,
99 ) -> Result<uuid::Uuid, crate::database::DatabaseError> {
100 let egg_configuration = server.egg.configuration(&state.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 ip = server.allocation.as_ref().map(|a| a.allocation.ip);
113 let start_port = config_allocations.user_self_assign.start_port as i32;
114 let end_port = config_allocations.user_self_assign.end_port as i32;
115
116 let candidate_ips: Vec<sqlx::types::ipnetwork::IpNetwork> = sqlx::query_scalar(
117 r#"
118 SELECT DISTINCT node_allocations.ip
119 FROM node_allocations
120 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
121 WHERE
122 node_allocations.node_uuid = $1
123 AND ($2 IS NULL OR node_allocations.ip = $2)
124 AND node_allocations.port BETWEEN $3 AND $4
125 AND server_allocations.uuid IS NULL
126 "#,
127 )
128 .bind(server.node.uuid)
129 .bind(ip)
130 .bind(start_port)
131 .bind(end_port)
132 .fetch_all(state.database.read())
133 .await?;
134
135 if candidate_ips.is_empty() {
136 return Err(anyhow::Error::new(
137 crate::response::DisplayError::new("no node allocations are available")
138 .with_status(StatusCode::EXPECTATION_FAILED),
139 )
140 .into());
141 }
142
143 let exclude = match super::node_allocation::NodeAllocation::used_by_node(
144 state,
145 &server.node.fetch_cached(&state.database).await?,
146 &candidate_ips.iter().map(|ip| ip.ip()).collect::<Vec<_>>(),
147 )
148 .await
149 {
150 Ok(exclude) => exclude,
151 Err(err) => {
152 tracing::warn!(
153 node = %server.node.uuid,
154 "failed to resolve used ports, refusing to assign an allocation: {err:#}"
155 );
156
157 return Err(anyhow::Error::new(
158 crate::response::DisplayError::new(
159 "could not reach the node to check which ports are free",
160 )
161 .with_status(StatusCode::EXPECTATION_FAILED),
162 )
163 .into());
164 }
165 };
166
167 let row = sqlx::query(
168 r#"
169 INSERT INTO server_allocations (server_uuid, allocation_uuid)
170 VALUES ($1, (
171 SELECT node_allocations.uuid FROM node_allocations
172 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
173 WHERE
174 node_allocations.node_uuid = $2
175 AND ($3 IS NULL OR node_allocations.ip = $3)
176 AND node_allocations.port BETWEEN $4 AND $5
177 AND server_allocations.uuid IS NULL
178 AND NOT (node_allocations.uuid = ANY($6))
179 ORDER BY RANDOM()
180 LIMIT 1
181 ))
182 RETURNING uuid
183 "#,
184 )
185 .bind(server.uuid)
186 .bind(server.node.uuid)
187 .bind(ip)
188 .bind(start_port)
189 .bind(end_port)
190 .bind(&exclude)
191 .fetch_one(state.database.write())
192 .await;
193
194 let row = match row {
195 Ok(row) => row,
196 Err(err) if err.is_not_null_violation() => {
197 return Err(anyhow::Error::new(
198 crate::response::DisplayError::new(
199 "every free port on this node is currently in use",
200 )
201 .with_status(StatusCode::EXPECTATION_FAILED),
202 )
203 .into());
204 }
205 Err(err) => return Err(err.into()),
206 };
207
208 Ok(row.get("uuid"))
209 }
210
211 pub async fn by_uuid(
212 database: &crate::database::Database,
213 uuid: uuid::Uuid,
214 ) -> Result<Option<Self>, crate::database::DatabaseError> {
215 let row = sqlx::query(sqlx::AssertSqlSafe(format!(
216 r#"
217 SELECT {}
218 FROM server_allocations
219 JOIN node_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
220 WHERE server_allocations.uuid = $1
221 "#,
222 Self::columns_sql(None)
223 )))
224 .bind(uuid)
225 .fetch_optional(database.read())
226 .await?;
227
228 row.try_map(|row| Self::map(None, &row))
229 }
230
231 pub async fn by_server_uuid_uuid(
232 database: &crate::database::Database,
233 server_uuid: uuid::Uuid,
234 allocation_uuid: uuid::Uuid,
235 ) -> Result<Option<Self>, crate::database::DatabaseError> {
236 let row = sqlx::query(sqlx::AssertSqlSafe(format!(
237 r#"
238 SELECT {}
239 FROM server_allocations
240 JOIN node_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
241 WHERE server_allocations.server_uuid = $1 AND server_allocations.uuid = $2
242 "#,
243 Self::columns_sql(None)
244 )))
245 .bind(server_uuid)
246 .bind(allocation_uuid)
247 .fetch_optional(database.read())
248 .await?;
249
250 row.try_map(|row| Self::map(None, &row))
251 }
252
253 pub async fn by_server_uuid_with_pagination(
254 database: &crate::database::Database,
255 server_uuid: uuid::Uuid,
256 page: i64,
257 per_page: i64,
258 search: Option<&str>,
259 ) -> Result<super::Pagination<Self>, crate::database::DatabaseError> {
260 let offset = (page - 1) * per_page;
261
262 let rows = sqlx::query(sqlx::AssertSqlSafe(format!(
263 r#"
264 SELECT {}, COUNT(*) OVER() AS total_count
265 FROM server_allocations
266 JOIN node_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
267 WHERE server_allocations.server_uuid = $1
268 AND (
269 $2 IS NULL
270 OR host(node_allocations.ip) || ':' || node_allocations.port ILIKE '%' || $2 || '%'
271 OR (node_allocations.ip_alias IS NOT NULL AND node_allocations.ip_alias || ':' || node_allocations.port ILIKE '%' || $2 || '%')
272 OR server_allocations.notes ILIKE '%' || $2 || '%'
273 )
274 ORDER BY server_allocations.created
275 LIMIT $3 OFFSET $4
276 "#,
277 Self::columns_sql(None)
278 )))
279 .bind(server_uuid)
280 .bind(search)
281 .bind(per_page)
282 .bind(offset)
283 .fetch_all(database.read())
284 .await?;
285
286 Ok(super::Pagination {
287 total: rows
288 .first()
289 .map_or(Ok(0), |row| row.try_get("total_count"))?,
290 per_page,
291 page,
292 data: rows
293 .into_iter()
294 .map(|row| Self::map(None, &row))
295 .try_collect_vec()?,
296 })
297 }
298
299 pub async fn all_by_server_uuid(
300 database: &crate::database::Database,
301 server_uuid: uuid::Uuid,
302 ) -> Result<Vec<Self>, crate::database::DatabaseError> {
303 let rows = sqlx::query(sqlx::AssertSqlSafe(format!(
304 r#"
305 SELECT {}
306 FROM server_allocations
307 JOIN node_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
308 WHERE server_allocations.server_uuid = $1
309 ORDER BY server_allocations.created
310 "#,
311 Self::columns_sql(None)
312 )))
313 .bind(server_uuid)
314 .fetch_all(database.read())
315 .await?;
316
317 rows.into_iter()
318 .map(|row| Self::map(None, &row))
319 .try_collect_vec()
320 }
321
322 pub async fn count_by_server_uuid(
323 database: &crate::database::Database,
324 server_uuid: uuid::Uuid,
325 ) -> Result<i64, sqlx::Error> {
326 sqlx::query_scalar(
327 r#"
328 SELECT COUNT(*)
329 FROM server_allocations
330 WHERE server_allocations.server_uuid = $1
331 "#,
332 )
333 .bind(server_uuid)
334 .fetch_one(database.read())
335 .await
336 }
337}
338
339#[async_trait::async_trait]
340impl IntoApiObject for ServerAllocation {
341 type ApiObject = ApiServerAllocation;
342 type ExtraArgs<'a> = Option<uuid::Uuid>;
343
344 async fn into_api_object<'a>(
345 self,
346 state: &crate::State,
347 primary: Self::ExtraArgs<'a>,
348 ) -> Result<Self::ApiObject, crate::database::DatabaseError> {
349 let api_object = ApiServerAllocation::init_hooks(&self, state).await?;
350
351 let api_object = finish_extendible!(
352 ApiServerAllocation {
353 uuid: self.uuid,
354 ip: compact_str::format_compact!("{}", self.allocation.ip.ip()),
355 ip_alias: self.allocation.ip_alias,
356 port: self.allocation.port,
357 notes: self.notes,
358 is_primary: primary.is_some_and(|p| p == self.uuid),
359 created: self.created.and_utc(),
360 },
361 api_object,
362 state
363 )?;
364
365 Ok(api_object)
366 }
367}
368
369#[async_trait::async_trait]
370impl DeletableModel for ServerAllocation {
371 type DeleteOptions = ();
372
373 fn get_delete_handlers() -> &'static LazyLock<DeleteHandlerList<Self>> {
374 static DELETE_LISTENERS: LazyLock<DeleteHandlerList<ServerAllocation>> =
375 LazyLock::new(|| Arc::new(ModelHandlerList::default()));
376
377 &DELETE_LISTENERS
378 }
379
380 async fn delete_with_transaction(
381 &self,
382 state: &crate::State,
383 options: Self::DeleteOptions,
384 transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
385 ) -> Result<(), anyhow::Error> {
386 self.run_delete_handlers(&options, state, transaction)
387 .await?;
388
389 sqlx::query(
390 r#"
391 DELETE FROM server_allocations
392 WHERE server_allocations.uuid = $1
393 "#,
394 )
395 .bind(self.uuid)
396 .execute(&mut **transaction)
397 .await?;
398
399 self.run_after_delete_handlers(&options, state, transaction)
400 .await?;
401
402 Ok(())
403 }
404}
405
406#[schema_extension_derive::extendible]
407#[init_args(ServerAllocation, crate::State)]
408#[hook_args(crate::State)]
409#[derive(ToSchema, Serialize)]
410#[schema(title = "ServerAllocation")]
411pub struct ApiServerAllocation {
412 pub uuid: uuid::Uuid,
413
414 pub ip: compact_str::CompactString,
415 pub ip_alias: Option<compact_str::CompactString>,
416 pub port: i32,
417
418 pub notes: Option<compact_str::CompactString>,
419 pub is_primary: bool,
420
421 pub created: chrono::DateTime<chrono::Utc>,
422}