1use crate::prelude::*;
2use compact_str::ToCompactString;
3use serde::{Deserialize, Serialize};
4use sqlx::{Row, postgres::PgRow};
5use std::{
6 collections::{BTreeMap, HashMap},
7 sync::LazyLock,
8};
9#[derive(Serialize, Deserialize, Clone)]
10pub struct NodeAllocation {
11 pub uuid: uuid::Uuid,
12 pub server: Option<Fetchable<super::server::Server>>,
13
14 pub ip: sqlx::types::ipnetwork::IpNetwork,
15 pub ip_alias: Option<compact_str::CompactString>,
16 pub port: i32,
17
18 pub created: chrono::NaiveDateTime,
19
20 extension_data: super::ModelExtensionData,
21}
22
23impl BaseModel for NodeAllocation {
24 const NAME: &'static str = "node_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 BTreeMap::from([
42 (
43 "node_allocations.uuid",
44 compact_str::format_compact!("{prefix}uuid"),
45 ),
46 (
47 "node_allocations.ip",
48 compact_str::format_compact!("{prefix}ip"),
49 ),
50 (
51 "node_allocations.ip_alias",
52 compact_str::format_compact!("{prefix}ip_alias"),
53 ),
54 (
55 "node_allocations.port",
56 compact_str::format_compact!("{prefix}port"),
57 ),
58 (
59 "node_allocations.created",
60 compact_str::format_compact!("{prefix}created"),
61 ),
62 ])
63 }
64
65 #[inline]
66 fn map(prefix: Option<&str>, row: &PgRow) -> Result<Self, crate::database::DatabaseError> {
67 let prefix = prefix.unwrap_or_default();
68
69 Ok(Self {
70 uuid: row.try_get(compact_str::format_compact!("{prefix}uuid").as_str())?,
71 server: if let Ok(server_uuid) = row.try_get::<uuid::Uuid, _>("server_uuid") {
72 Some(super::server::Server::get_fetchable(server_uuid))
73 } else {
74 None
75 },
76 ip: row.try_get(compact_str::format_compact!("{prefix}ip").as_str())?,
77 ip_alias: row.try_get(compact_str::format_compact!("{prefix}ip_alias").as_str())?,
78 port: row.try_get(compact_str::format_compact!("{prefix}port").as_str())?,
79 created: row.try_get(compact_str::format_compact!("{prefix}created").as_str())?,
80 extension_data: Self::map_extensions(prefix, row)?,
81 })
82 }
83}
84
85impl NodeAllocation {
86 pub async fn create(
87 database: &crate::database::Database,
88 node_uuid: uuid::Uuid,
89 ip: &sqlx::types::ipnetwork::IpNetwork,
90 ip_alias: Option<&str>,
91 port: i32,
92 ) -> Result<(), crate::database::DatabaseError> {
93 sqlx::query(
94 r#"
95 INSERT INTO node_allocations (node_uuid, ip, ip_alias, port)
96 VALUES ($1, $2, $3, $4)
97 "#,
98 )
99 .bind(node_uuid)
100 .bind(ip)
101 .bind(ip_alias)
102 .bind(port)
103 .execute(database.write())
104 .await?;
105
106 Ok(())
107 }
108
109 pub async fn used_by_node(
110 state: &crate::State,
111 node: &super::node::Node,
112 ips: &[std::net::IpAddr],
113 ) -> Result<Vec<uuid::Uuid>, anyhow::Error> {
114 if ips.is_empty() {
115 return Ok(Vec::new());
116 }
117
118 let (used_ips, used_ports): (Vec<_>, Vec<_>) = node
119 .used_ports(state, ips)
120 .await?
121 .into_iter()
122 .flat_map(|(ip, ports)| {
123 ports
124 .into_iter()
125 .map(move |port| (sqlx::types::ipnetwork::IpNetwork::from(ip), port as i32))
126 })
127 .unzip();
128
129 if used_ips.is_empty() {
130 return Ok(Vec::new());
131 }
132
133 Ok(sqlx::query_scalar(
134 r#"
135 SELECT node_allocations.uuid
136 FROM node_allocations
137 JOIN UNNEST($2::inet[], $3::int[]) AS used(ip, port)
138 ON used.ip = node_allocations.ip AND used.port = node_allocations.port
139 WHERE node_allocations.node_uuid = $1
140 "#,
141 )
142 .bind(node.uuid)
143 .bind(&used_ips)
144 .bind(&used_ports)
145 .fetch_all(state.database.read())
146 .await?)
147 }
148
149 pub async fn used_by_node_any_ip(
150 state: &crate::State,
151 node: &super::node::Node,
152 ) -> Result<Vec<uuid::Uuid>, anyhow::Error> {
153 let ips: Vec<sqlx::types::ipnetwork::IpNetwork> = sqlx::query_scalar(
154 r#"
155 SELECT DISTINCT node_allocations.ip
156 FROM node_allocations
157 WHERE node_allocations.node_uuid = $1
158 "#,
159 )
160 .bind(node.uuid)
161 .fetch_all(state.database.read())
162 .await?;
163
164 Self::used_by_node(
165 state,
166 node,
167 &ips.iter().map(|ip| ip.ip()).collect::<Vec<_>>(),
168 )
169 .await
170 }
171
172 pub async fn get_random(
173 database: &crate::database::Database,
174 node_uuid: uuid::Uuid,
175 start_port: u16,
176 end_port: u16,
177 amount: i64,
178 exclude: &[uuid::Uuid],
179 ) -> Result<Vec<uuid::Uuid>, crate::database::DatabaseError> {
180 let rows = sqlx::query(
181 r#"
182 WITH eligible_ips AS (
183 SELECT node_allocations.ip
184 FROM node_allocations
185 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
186 WHERE
187 node_allocations.node_uuid = $1
188 AND node_allocations.port BETWEEN $2 AND $3
189 AND server_allocations.uuid IS NULL
190 AND NOT (node_allocations.uuid = ANY($5))
191 GROUP BY node_allocations.ip
192 HAVING COUNT(*) >= $4
193 ),
194 random_ip AS (
195 SELECT ip FROM eligible_ips ORDER BY RANDOM() LIMIT 1
196 )
197 SELECT node_allocations.uuid
198 FROM node_allocations
199 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
200 WHERE
201 node_allocations.node_uuid = $1
202 AND node_allocations.port BETWEEN $2 AND $3
203 AND server_allocations.uuid IS NULL
204 AND NOT (node_allocations.uuid = ANY($5))
205 AND node_allocations.ip = (SELECT ip FROM random_ip)
206 ORDER BY RANDOM()
207 LIMIT $4
208 "#,
209 )
210 .bind(node_uuid)
211 .bind(start_port as i32)
212 .bind(end_port as i32)
213 .bind(amount)
214 .bind(exclude)
215 .fetch_all(database.write())
216 .await?;
217
218 if rows.len() != amount as usize {
219 return Err(anyhow::anyhow!("only found {} available allocations", rows.len()).into());
220 }
221
222 Ok(rows
223 .into_iter()
224 .map(|row| row.get::<uuid::Uuid, _>("uuid"))
225 .collect())
226 }
227
228 pub async fn get_random_ip(
229 database: &crate::database::Database,
230 node_uuid: uuid::Uuid,
231 ip: &sqlx::types::ipnetwork::IpNetwork,
232 start_port: u16,
233 end_port: u16,
234 amount: i64,
235 exclude: &[uuid::Uuid],
236 ) -> Result<Vec<uuid::Uuid>, crate::database::DatabaseError> {
237 let rows = sqlx::query(
238 r#"
239 SELECT node_allocations.uuid
240 FROM node_allocations
241 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
242 WHERE
243 node_allocations.node_uuid = $1
244 AND node_allocations.ip = $2
245 AND node_allocations.port BETWEEN $3 AND $4
246 AND server_allocations.uuid IS NULL
247 AND NOT (node_allocations.uuid = ANY($6))
248 ORDER BY RANDOM()
249 LIMIT $5
250 "#,
251 )
252 .bind(node_uuid)
253 .bind(ip)
254 .bind(start_port as i32)
255 .bind(end_port as i32)
256 .bind(amount)
257 .bind(exclude)
258 .fetch_all(database.write())
259 .await?;
260
261 if rows.len() != amount as usize {
262 return Err(anyhow::anyhow!(
263 "only found {} available allocations on this IP",
264 rows.len()
265 )
266 .into());
267 }
268
269 Ok(rows
270 .into_iter()
271 .map(|row| row.get::<uuid::Uuid, _>("uuid"))
272 .collect())
273 }
274
275 pub async fn get_random_dedicated(
276 database: &crate::database::Database,
277 node_uuid: uuid::Uuid,
278 start_port: u16,
279 end_port: u16,
280 amount: i64,
281 exclude: &[uuid::Uuid],
282 ) -> Result<Vec<uuid::Uuid>, crate::database::DatabaseError> {
283 let rows = sqlx::query(
284 r#"
285 WITH eligible_ips AS (
286 SELECT node_allocations.ip
287 FROM node_allocations
288 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
289 WHERE node_allocations.node_uuid = $1
290 GROUP BY node_allocations.ip
291 HAVING
292 COUNT(server_allocations.uuid) = 0
293 AND COUNT(*) FILTER (WHERE node_allocations.uuid = ANY($5)) = 0
294 AND SUM(CASE WHEN node_allocations.port BETWEEN $2 AND $3 THEN 1 ELSE 0 END) >= $4
295 ),
296 random_ip AS (
297 SELECT ip FROM eligible_ips ORDER BY RANDOM() LIMIT 1
298 )
299 SELECT node_allocations.uuid
300 FROM node_allocations
301 WHERE
302 node_allocations.node_uuid = $1
303 AND node_allocations.port BETWEEN $2 AND $3
304 AND node_allocations.ip = (SELECT ip FROM random_ip)
305 ORDER BY RANDOM()
306 LIMIT $4
307 "#,
308 )
309 .bind(node_uuid)
310 .bind(start_port as i32)
311 .bind(end_port as i32)
312 .bind(amount)
313 .bind(exclude)
314 .fetch_all(database.write())
315 .await?;
316
317 if rows.len() != amount as usize {
318 return Err(anyhow::anyhow!(
319 "only found {} available dedicated allocations",
320 rows.len()
321 )
322 .into());
323 }
324
325 Ok(rows
326 .into_iter()
327 .map(|row| row.get::<uuid::Uuid, _>("uuid"))
328 .collect())
329 }
330
331 pub async fn get_preserved(
332 database: &crate::database::Database,
333 node_uuid: uuid::Uuid,
334 ports: &[i32],
335 exclude: &[uuid::Uuid],
336 ) -> Result<
337 Option<(sqlx::types::ipnetwork::IpNetwork, Vec<(uuid::Uuid, i32)>)>,
338 crate::database::DatabaseError,
339 > {
340 let rows = sqlx::query(
341 r#"
342 WITH best_ip AS (
343 SELECT node_allocations.ip
344 FROM node_allocations
345 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
346 WHERE
347 node_allocations.node_uuid = $1
348 AND node_allocations.port = ANY($2)
349 AND server_allocations.uuid IS NULL
350 AND NOT (node_allocations.uuid = ANY($3))
351 GROUP BY node_allocations.ip
352 ORDER BY COUNT(DISTINCT node_allocations.port) DESC, RANDOM()
353 LIMIT 1
354 )
355 SELECT node_allocations.uuid, node_allocations.ip, node_allocations.port
356 FROM node_allocations
357 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
358 WHERE
359 node_allocations.node_uuid = $1
360 AND node_allocations.port = ANY($2)
361 AND server_allocations.uuid IS NULL
362 AND NOT (node_allocations.uuid = ANY($3))
363 AND node_allocations.ip = (SELECT ip FROM best_ip)
364 "#,
365 )
366 .bind(node_uuid)
367 .bind(ports)
368 .bind(exclude)
369 .fetch_all(database.write())
370 .await?;
371
372 let Some(first) = rows.first() else {
373 return Ok(None);
374 };
375
376 let ip = first.get::<sqlx::types::ipnetwork::IpNetwork, _>("ip");
377
378 Ok(Some((
379 ip,
380 rows.into_iter()
381 .map(|row| (row.get::<uuid::Uuid, _>("uuid"), row.get::<i32, _>("port")))
382 .collect(),
383 )))
384 }
385
386 pub async fn by_node_uuid_ip_port_unused(
387 database: &crate::database::Database,
388 node_uuid: uuid::Uuid,
389 ip: &sqlx::types::ipnetwork::IpNetwork,
390 port: i32,
391 exclude: &[uuid::Uuid],
392 ) -> Result<Option<Self>, crate::database::DatabaseError> {
393 let row = sqlx::query(sqlx::AssertSqlSafe(format!(
394 r#"
395 SELECT {}
396 FROM node_allocations
397 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
398 WHERE node_allocations.node_uuid = $1 AND node_allocations.ip = $2 AND node_allocations.port = $3 AND server_allocations.uuid IS NULL
399 AND NOT (node_allocations.uuid = ANY($4))
400 "#,
401 Self::columns_sql(None)
402 )))
403 .bind(node_uuid)
404 .bind(ip)
405 .bind(port)
406 .bind(exclude)
407 .fetch_optional(database.read())
408 .await?;
409
410 row.try_map(|row| Self::map(None, &row))
411 }
412
413 pub async fn get_from_deployment<'a>(
414 state: &crate::State,
415 deployment: &'a super::egg_configuration::EggConfigAllocationsDeployment,
416 node: &super::node::Node,
417 variables: &mut HashMap<&'a str, compact_str::CompactString>,
418 ) -> Result<(Option<uuid::Uuid>, Vec<uuid::Uuid>), crate::database::DatabaseError> {
419 let database = &state.database;
420 let node_uuid = node.uuid;
421
422 let used = Self::used_by_node_any_ip(state, node).await?;
423
424 let mut primary = None;
425 let mut additional = Vec::new();
426
427 const MAX_ITER: usize = 100;
428
429 macro_rules! is_unused {
430 ($uuid:expr) => {
431 primary.as_ref().map_or(true, |p| p.uuid != $uuid) && !additional.contains(&$uuid)
432 };
433 }
434
435 macro_rules! get_random {
436 ($start_port:expr, $end_port:expr) => {{
437 let mut exclude = used.clone();
438 exclude.extend_from_slice(&additional);
439 if let Some(primary) = &primary {
440 exclude.push(primary.uuid);
441 Self::get_random_ip(
442 database,
443 node_uuid,
444 &primary.ip,
445 $start_port,
446 $end_port,
447 1,
448 &exclude,
449 )
450 .await?
451 } else {
452 Self::get_random(database, node_uuid, $start_port, $end_port, 1, &exclude)
453 .await?
454 }
455 }};
456 }
457
458 let mut success = false;
459
460 'primary: for i in 0..MAX_ITER {
461 if i != 0 && deployment.primary.is_none() {
462 break;
463 }
464
465 additional.clear();
466 variables.clear();
467
468 if let Some(primary_allocation) = &deployment.primary {
469 let random = if deployment.dedicated {
470 Self::get_random_dedicated(
471 database,
472 node_uuid,
473 primary_allocation.start_port,
474 primary_allocation.end_port,
475 1,
476 &used,
477 )
478 .await?
479 } else {
480 Self::get_random(
481 database,
482 node_uuid,
483 primary_allocation.start_port,
484 primary_allocation.end_port,
485 1,
486 &used,
487 )
488 .await?
489 };
490
491 let Some(allocation) = random.into_iter().next() else {
492 return Err(anyhow::anyhow!("no available primary allocation found").into());
493 };
494 let allocation =
495 match Self::by_node_uuid_uuid(database, node_uuid, allocation).await? {
496 Some(allocation) => allocation,
497 None => {
498 return Err(
499 anyhow::anyhow!("allocated primary allocation not found").into()
500 );
501 }
502 };
503
504 if let Some(variable_name) = &primary_allocation.assign_to_variable {
505 variables.insert(variable_name, allocation.port.to_compact_string());
506 }
507
508 primary = Some(allocation);
509 }
510
511 for additional_allocation in &deployment.additional {
512 match additional_allocation.mode {
513 super::egg_configuration::EggConfigAllocationDeploymentAdditionalAllocationMode::Random => {
514 let random = get_random!(1, u16::MAX);
515
516 let Some(allocation) = random.into_iter().next() else {
517 return Err(anyhow::anyhow!("no available additional allocation found").into());
518 };
519 additional.push(allocation);
520
521 if let Some(variable_name) = &additional_allocation.assign_to_variable {
522 let allocation = match Self::by_node_uuid_uuid(database, node_uuid, allocation).await? {
523 Some(allocation) => allocation,
524 None => {
525 return Err(
526 anyhow::anyhow!("allocated additional allocation not found").into()
527 );
528 }
529 };
530
531 variables.insert(variable_name, allocation.port.to_compact_string());
532 }
533 },
534 super::egg_configuration::EggConfigAllocationDeploymentAdditionalAllocationMode::Range {
535 start_port,
536 end_port,
537 } => {
538 let random = get_random!(start_port, end_port);
539
540 let Some(allocation) = random.into_iter().next() else {
541 return Err(anyhow::anyhow!("no available additional allocation found").into());
542 };
543 additional.push(allocation);
544
545 if let Some(variable_name) = &additional_allocation.assign_to_variable {
546 let allocation = match Self::by_node_uuid_uuid(database, node_uuid, allocation).await? {
547 Some(allocation) => allocation,
548 None => {
549 return Err(
550 anyhow::anyhow!("allocated additional allocation not found").into()
551 );
552 }
553 };
554
555 variables.insert(variable_name, allocation.port.to_compact_string());
556 }
557 }
558 super::egg_configuration::EggConfigAllocationDeploymentAdditionalAllocationMode::AddPrimary { value } => {
559 let primary = match &primary {
560 Some(primary) => primary,
561 None => {
562 return Err(anyhow::anyhow!("primary allocation is required for `add_primary` mode").into());
563 }
564 };
565
566 let allocation_port = primary.port + value as i32;
567
568 let allocation = match Self::by_node_uuid_ip_port_unused(database, node_uuid, &primary.ip, allocation_port, &used).await? {
569 Some(allocation) => allocation,
570 None => continue 'primary,
571 };
572 if !is_unused!(allocation.uuid) {
573 return Err(anyhow::anyhow!("allocated additional allocation is already in use").into());
574 }
575 additional.push(allocation.uuid);
576
577 if let Some(variable_name) = &additional_allocation.assign_to_variable {
578 variables.insert(variable_name, allocation.port.to_compact_string());
579 }
580 }
581 super::egg_configuration::EggConfigAllocationDeploymentAdditionalAllocationMode::SubtractPrimary { value } => {
582 let primary = match &primary {
583 Some(primary) => primary,
584 None => {
585 return Err(anyhow::anyhow!("primary allocation is required for `subtract_primary` mode").into());
586 }
587 };
588
589 let allocation_port = primary.port - value as i32;
590
591 let allocation = match Self::by_node_uuid_ip_port_unused(database, node_uuid, &primary.ip, allocation_port, &used).await? {
592 Some(allocation) => allocation,
593 None => continue 'primary,
594 };
595 if !is_unused!(allocation.uuid) {
596 return Err(anyhow::anyhow!("allocated additional allocation is already in use").into());
597 }
598 additional.push(allocation.uuid);
599
600 if let Some(variable_name) = &additional_allocation.assign_to_variable {
601 variables.insert(variable_name, allocation.port.to_compact_string());
602 }
603 }
604 super::egg_configuration::EggConfigAllocationDeploymentAdditionalAllocationMode::MultiplyPrimary { value } => {
605 let primary = match &primary {
606 Some(primary) => primary,
607 None => {
608 return Err(anyhow::anyhow!("primary allocation is required for `multiply_primary` mode").into());
609 }
610 };
611
612 let allocation_port = (primary.port as f64 * value) as i32;
613
614 let allocation = match Self::by_node_uuid_ip_port_unused(database, node_uuid, &primary.ip, allocation_port, &used).await? {
615 Some(allocation) => allocation,
616 None => continue 'primary,
617 };
618 if !is_unused!(allocation.uuid) {
619 return Err(anyhow::anyhow!("allocated additional allocation is already in use").into());
620 }
621 additional.push(allocation.uuid);
622
623 if let Some(variable_name) = &additional_allocation.assign_to_variable {
624 variables.insert(variable_name, allocation.port.to_compact_string());
625 }
626 }
627 super::egg_configuration::EggConfigAllocationDeploymentAdditionalAllocationMode::DividePrimary { value } => {
628 let primary = match &primary {
629 Some(primary) => primary,
630 None => {
631 return Err(anyhow::anyhow!("primary allocation is required for `divide_primary` mode").into());
632 }
633 };
634
635 let allocation_port = (primary.port as f64 / value) as i32;
636
637 let allocation = match Self::by_node_uuid_ip_port_unused(database, node_uuid, &primary.ip, allocation_port, &used).await? {
638 Some(allocation) => allocation,
639 None => continue 'primary,
640 };
641 if !is_unused!(allocation.uuid) {
642 return Err(anyhow::anyhow!("allocated additional allocation is already in use").into());
643 }
644 additional.push(allocation.uuid);
645
646 if let Some(variable_name) = &additional_allocation.assign_to_variable {
647 variables.insert(variable_name, allocation.port.to_compact_string());
648 }
649 }
650 }
651 }
652
653 success = true;
654 break;
655 }
656
657 if !success {
658 return Err(anyhow::anyhow!(
659 "could not satisfy all additional allocation rules after {MAX_ITER} attempts"
660 )
661 .into());
662 }
663
664 Ok((primary.map(|p| p.uuid), additional))
665 }
666
667 pub async fn by_node_uuid_uuid(
668 database: &crate::database::Database,
669 node_uuid: uuid::Uuid,
670 uuid: uuid::Uuid,
671 ) -> Result<Option<Self>, crate::database::DatabaseError> {
672 let row = sqlx::query(sqlx::AssertSqlSafe(format!(
673 r#"
674 SELECT {}
675 FROM node_allocations
676 WHERE node_allocations.node_uuid = $1 AND node_allocations.uuid = $2
677 "#,
678 Self::columns_sql(None)
679 )))
680 .bind(node_uuid)
681 .bind(uuid)
682 .fetch_optional(database.read())
683 .await?;
684
685 row.try_map(|row| Self::map(None, &row))
686 }
687
688 pub async fn available_by_node_uuid_with_pagination(
689 database: &crate::database::Database,
690 node_uuid: uuid::Uuid,
691 page: i64,
692 per_page: i64,
693 search: Option<&str>,
694 ) -> Result<super::Pagination<Self>, crate::database::DatabaseError> {
695 let offset = (page - 1) * per_page;
696
697 let rows = sqlx::query(sqlx::AssertSqlSafe(format!(
698 r#"
699 SELECT {}, COUNT(*) OVER() AS total_count
700 FROM node_allocations
701 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
702 WHERE node_allocations.node_uuid = $1 AND server_allocations.uuid IS NULL
703 AND (
704 $2 IS NULL
705 OR host(node_allocations.ip) || ':' || node_allocations.port ILIKE '%' || $2 || '%'
706 OR (node_allocations.ip_alias IS NOT NULL AND node_allocations.ip_alias || ':' || node_allocations.port ILIKE '%' || $2 || '%')
707 )
708 ORDER BY node_allocations.ip, node_allocations.port
709 LIMIT $3 OFFSET $4
710 "#,
711 Self::columns_sql(None)
712 )))
713 .bind(node_uuid)
714 .bind(search)
715 .bind(per_page)
716 .bind(offset)
717 .fetch_all(database.read())
718 .await?;
719
720 Ok(super::Pagination {
721 total: rows
722 .first()
723 .map_or(Ok(0), |row| row.try_get("total_count"))?,
724 per_page,
725 page,
726 data: rows
727 .into_iter()
728 .map(|row| Self::map(None, &row))
729 .try_collect_vec()?,
730 })
731 }
732
733 pub async fn by_node_uuid_with_pagination(
734 database: &crate::database::Database,
735 node_uuid: uuid::Uuid,
736 page: i64,
737 per_page: i64,
738 search: Option<&str>,
739 ) -> Result<super::Pagination<Self>, crate::database::DatabaseError> {
740 let offset = (page - 1) * per_page;
741
742 let rows = sqlx::query(sqlx::AssertSqlSafe(format!(
743 r#"
744 SELECT {}, server_allocations.server_uuid, COUNT(*) OVER() AS total_count
745 FROM node_allocations
746 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
747 WHERE node_allocations.node_uuid = $1
748 AND (
749 $2 IS NULL OR host(node_allocations.ip) || ':' || node_allocations.port ILIKE '%' || $2 || '%'
750 OR (node_allocations.ip_alias IS NOT NULL AND node_allocations.ip_alias || ':' || node_allocations.port ILIKE '%' || $2 || '%')
751 OR server_allocations.notes ILIKE '%' || $2 || '%'
752 )
753 ORDER BY node_allocations.ip, node_allocations.port
754 LIMIT $3 OFFSET $4
755 "#,
756 Self::columns_sql(None)
757 )))
758 .bind(node_uuid)
759 .bind(search)
760 .bind(per_page)
761 .bind(offset)
762 .fetch_all(database.read())
763 .await?;
764
765 Ok(super::Pagination {
766 total: rows
767 .first()
768 .map_or(Ok(0), |row| row.try_get("total_count"))?,
769 per_page,
770 page,
771 data: rows
772 .into_iter()
773 .map(|row| Self::map(None, &row))
774 .try_collect_vec()?,
775 })
776 }
777
778 pub async fn delete_by_uuids(
779 database: &crate::database::Database,
780 node_uuid: uuid::Uuid,
781 uuids: &[uuid::Uuid],
782 ) -> Result<u64, crate::database::DatabaseError> {
783 let deleted = sqlx::query(
784 r#"
785 DELETE FROM node_allocations
786 WHERE node_allocations.node_uuid = $1 AND node_allocations.uuid = ANY($2)
787 "#,
788 )
789 .bind(node_uuid)
790 .bind(uuids)
791 .execute(database.write())
792 .await?
793 .rows_affected();
794
795 Ok(deleted)
796 }
797}
798
799#[async_trait::async_trait]
800impl IntoAdminApiObject for NodeAllocation {
801 type AdminApiObject = AdminApiNodeAllocation;
802 type ExtraArgs<'a> = &'a crate::storage::StorageUrlRetriever<'a>;
803
804 async fn into_admin_api_object<'a>(
805 self,
806 state: &crate::State,
807 storage_url_retriever: Self::ExtraArgs<'a>,
808 ) -> Result<Self::AdminApiObject, crate::database::DatabaseError> {
809 let api_object = AdminApiNodeAllocation::init_hooks(&self, state).await?;
810
811 let server = match self.server {
812 Some(fetchable) => Some(
813 fetchable
814 .fetch_cached(&state.database)
815 .await?
816 .into_admin_api_object(state, storage_url_retriever)
817 .await?,
818 ),
819 None => None,
820 };
821
822 let api_object = finish_extendible!(
823 AdminApiNodeAllocation {
824 uuid: self.uuid,
825 server,
826 ip: self.ip.ip().to_compact_string(),
827 ip_alias: self.ip_alias,
828 port: self.port,
829 created: self.created.and_utc(),
830 },
831 api_object,
832 state
833 )?;
834
835 Ok(api_object)
836 }
837}
838
839#[schema_extension_derive::extendible]
840#[init_args(NodeAllocation, crate::State)]
841#[hook_args(crate::State)]
842#[derive(ToSchema, Serialize)]
843#[schema(title = "NodeAllocation")]
844pub struct AdminApiNodeAllocation {
845 pub uuid: uuid::Uuid,
846 pub server: Option<super::server::AdminApiServer>,
847
848 pub ip: compact_str::CompactString,
849 pub ip_alias: Option<compact_str::CompactString>,
850 pub port: i32,
851
852 pub created: chrono::DateTime<chrono::Utc>,
853}