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(|| std::sync::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(&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(&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(&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 ($2 IS NULL OR host(node_allocations.ip) || ':' || node_allocations.port ILIKE '%' || $2 || '%' OR server_allocations.notes ILIKE '%' || $2 || '%')
198 ORDER BY server_allocations.created
199 LIMIT $3 OFFSET $4
200 "#,
201 Self::columns_sql(None)
202 ))
203 .bind(server_uuid)
204 .bind(search)
205 .bind(per_page)
206 .bind(offset)
207 .fetch_all(database.read())
208 .await?;
209
210 Ok(super::Pagination {
211 total: rows
212 .first()
213 .map_or(Ok(0), |row| row.try_get("total_count"))?,
214 per_page,
215 page,
216 data: rows
217 .into_iter()
218 .map(|row| Self::map(None, &row))
219 .try_collect_vec()?,
220 })
221 }
222
223 pub async fn all_by_server_uuid(
224 database: &crate::database::Database,
225 server_uuid: uuid::Uuid,
226 ) -> Result<Vec<Self>, crate::database::DatabaseError> {
227 let rows = sqlx::query(&format!(
228 r#"
229 SELECT {}
230 FROM server_allocations
231 JOIN node_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
232 WHERE server_allocations.server_uuid = $1
233 ORDER BY server_allocations.created
234 "#,
235 Self::columns_sql(None)
236 ))
237 .bind(server_uuid)
238 .fetch_all(database.read())
239 .await?;
240
241 rows.into_iter()
242 .map(|row| Self::map(None, &row))
243 .try_collect_vec()
244 }
245
246 pub async fn count_by_server_uuid(
247 database: &crate::database::Database,
248 server_uuid: uuid::Uuid,
249 ) -> Result<i64, sqlx::Error> {
250 sqlx::query_scalar(
251 r#"
252 SELECT COUNT(*)
253 FROM server_allocations
254 WHERE server_allocations.server_uuid = $1
255 "#,
256 )
257 .bind(server_uuid)
258 .fetch_one(database.read())
259 .await
260 }
261}
262
263#[async_trait::async_trait]
264impl IntoApiObject for ServerAllocation {
265 type ApiObject = ApiServerAllocation;
266 type ExtraArgs<'a> = Option<uuid::Uuid>;
267
268 async fn into_api_object<'a>(
269 self,
270 state: &crate::State,
271 primary: Self::ExtraArgs<'a>,
272 ) -> Result<Self::ApiObject, crate::database::DatabaseError> {
273 let api_object = ApiServerAllocation::init_hooks(&self, state).await?;
274
275 let api_object = finish_extendible!(
276 ApiServerAllocation {
277 uuid: self.uuid,
278 ip: compact_str::format_compact!("{}", self.allocation.ip.ip()),
279 ip_alias: self.allocation.ip_alias,
280 port: self.allocation.port,
281 notes: self.notes,
282 is_primary: primary.is_some_and(|p| p == self.uuid),
283 created: self.created.and_utc(),
284 },
285 api_object,
286 state
287 )?;
288
289 Ok(api_object)
290 }
291}
292
293#[async_trait::async_trait]
294impl DeletableModel for ServerAllocation {
295 type DeleteOptions = ();
296
297 fn get_delete_handlers() -> &'static LazyLock<DeleteHandlerList<Self>> {
298 static DELETE_LISTENERS: LazyLock<DeleteHandlerList<ServerAllocation>> =
299 LazyLock::new(|| Arc::new(ModelHandlerList::default()));
300
301 &DELETE_LISTENERS
302 }
303
304 async fn delete_with_transaction(
305 &self,
306 state: &crate::State,
307 options: Self::DeleteOptions,
308 transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
309 ) -> Result<(), anyhow::Error> {
310 self.run_delete_handlers(&options, state, transaction)
311 .await?;
312
313 sqlx::query(
314 r#"
315 DELETE FROM server_allocations
316 WHERE server_allocations.uuid = $1
317 "#,
318 )
319 .bind(self.uuid)
320 .execute(&mut **transaction)
321 .await?;
322
323 self.run_after_delete_handlers(&options, state, transaction)
324 .await?;
325
326 Ok(())
327 }
328}
329
330#[schema_extension_derive::extendible]
331#[init_args(ServerAllocation, crate::State)]
332#[hook_args(crate::State)]
333#[derive(ToSchema, Serialize)]
334#[schema(title = "ServerAllocation")]
335pub struct ApiServerAllocation {
336 pub uuid: uuid::Uuid,
337
338 pub ip: compact_str::CompactString,
339 pub ip_alias: Option<compact_str::CompactString>,
340 pub port: i32,
341
342 pub notes: Option<compact_str::CompactString>,
343 pub is_primary: bool,
344
345 pub created: chrono::DateTime<chrono::Utc>,
346}