1use crate::{prelude::*, storage::StorageUrlRetriever};
2use serde::{Deserialize, Serialize};
3use sqlx::{Row, postgres::PgRow};
4use std::collections::BTreeMap;
5use utoipa::ToSchema;
6
7#[derive(Serialize, Deserialize, Clone)]
8pub struct NodeAllocation {
9 pub uuid: uuid::Uuid,
10 pub server: Option<Fetchable<super::server::Server>>,
11
12 pub ip: sqlx::types::ipnetwork::IpNetwork,
13 pub ip_alias: Option<compact_str::CompactString>,
14 pub port: i32,
15
16 pub created: chrono::NaiveDateTime,
17}
18
19impl BaseModel for NodeAllocation {
20 const NAME: &'static str = "node_allocation";
21
22 #[inline]
23 fn columns(prefix: Option<&str>) -> BTreeMap<&'static str, compact_str::CompactString> {
24 let prefix = prefix.unwrap_or_default();
25
26 BTreeMap::from([
27 (
28 "node_allocations.uuid",
29 compact_str::format_compact!("{prefix}uuid"),
30 ),
31 (
32 "node_allocations.ip",
33 compact_str::format_compact!("{prefix}ip"),
34 ),
35 (
36 "node_allocations.ip_alias",
37 compact_str::format_compact!("{prefix}ip_alias"),
38 ),
39 (
40 "node_allocations.port",
41 compact_str::format_compact!("{prefix}port"),
42 ),
43 (
44 "node_allocations.created",
45 compact_str::format_compact!("{prefix}created"),
46 ),
47 ])
48 }
49
50 #[inline]
51 fn map(prefix: Option<&str>, row: &PgRow) -> Result<Self, crate::database::DatabaseError> {
52 let prefix = prefix.unwrap_or_default();
53
54 Ok(Self {
55 uuid: row.try_get(compact_str::format_compact!("{prefix}uuid").as_str())?,
56 server: if let Ok(server_uuid) = row.try_get::<uuid::Uuid, _>("server_uuid") {
57 Some(super::server::Server::get_fetchable(server_uuid))
58 } else {
59 None
60 },
61 ip: row.try_get(compact_str::format_compact!("{prefix}ip").as_str())?,
62 ip_alias: row.try_get(compact_str::format_compact!("{prefix}ip_alias").as_str())?,
63 port: row.try_get(compact_str::format_compact!("{prefix}port").as_str())?,
64 created: row.try_get(compact_str::format_compact!("{prefix}created").as_str())?,
65 })
66 }
67}
68
69impl NodeAllocation {
70 pub async fn create(
71 database: &crate::database::Database,
72 node_uuid: uuid::Uuid,
73 ip: &sqlx::types::ipnetwork::IpNetwork,
74 ip_alias: Option<&str>,
75 port: i32,
76 ) -> Result<(), crate::database::DatabaseError> {
77 sqlx::query(
78 r#"
79 INSERT INTO node_allocations (node_uuid, ip, ip_alias, port)
80 VALUES ($1, $2, $3, $4)
81 "#,
82 )
83 .bind(node_uuid)
84 .bind(ip)
85 .bind(ip_alias)
86 .bind(port)
87 .execute(database.write())
88 .await?;
89
90 Ok(())
91 }
92
93 pub async fn get_random(
94 database: &crate::database::Database,
95 node_uuid: uuid::Uuid,
96 start_port: u16,
97 end_port: u16,
98 amount: i64,
99 ) -> Result<Vec<uuid::Uuid>, crate::database::DatabaseError> {
100 let rows = sqlx::query(
101 r#"
102 WITH eligible_ips AS (
103 SELECT node_allocations.ip
104 FROM node_allocations
105 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
106 WHERE
107 node_allocations.node_uuid = $1
108 AND node_allocations.port BETWEEN $2 AND $3
109 AND server_allocations.uuid IS NULL
110 GROUP BY node_allocations.ip
111 HAVING COUNT(*) >= $4
112 ),
113 random_ip AS (
114 SELECT ip FROM eligible_ips ORDER BY RANDOM() LIMIT 1
115 )
116 SELECT node_allocations.uuid
117 FROM node_allocations
118 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
119 WHERE
120 node_allocations.node_uuid = $1
121 AND node_allocations.port BETWEEN $2 AND $3
122 AND server_allocations.uuid IS NULL
123 AND node_allocations.ip = (SELECT ip FROM random_ip)
124 ORDER BY RANDOM()
125 LIMIT $4
126 "#,
127 )
128 .bind(node_uuid)
129 .bind(start_port as i32)
130 .bind(end_port as i32)
131 .bind(amount)
132 .fetch_all(database.write())
133 .await?;
134
135 if rows.len() != amount as usize {
136 return Err(anyhow::anyhow!("only found {} available allocations", rows.len()).into());
137 }
138
139 Ok(rows
140 .into_iter()
141 .map(|row| row.get::<uuid::Uuid, _>("uuid"))
142 .collect())
143 }
144
145 pub async fn by_node_uuid_uuid(
146 database: &crate::database::Database,
147 node_uuid: uuid::Uuid,
148 uuid: uuid::Uuid,
149 ) -> Result<Option<Self>, crate::database::DatabaseError> {
150 let row = sqlx::query(&format!(
151 r#"
152 SELECT {}
153 FROM node_allocations
154 WHERE node_allocations.node_uuid = $1 AND node_allocations.uuid = $2
155 "#,
156 Self::columns_sql(None)
157 ))
158 .bind(node_uuid)
159 .bind(uuid)
160 .fetch_optional(database.read())
161 .await?;
162
163 row.try_map(|row| Self::map(None, &row))
164 }
165
166 pub async fn available_by_node_uuid_with_pagination(
167 database: &crate::database::Database,
168 node_uuid: uuid::Uuid,
169 page: i64,
170 per_page: i64,
171 search: Option<&str>,
172 ) -> Result<super::Pagination<Self>, crate::database::DatabaseError> {
173 let offset = (page - 1) * per_page;
174
175 let rows = sqlx::query(&format!(
176 r#"
177 SELECT {}, COUNT(*) OVER() AS total_count
178 FROM node_allocations
179 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
180 WHERE
181 ($2 IS NULL OR host(node_allocations.ip) || ':' || node_allocations.port ILIKE '%' || $2 || '%')
182 AND (node_allocations.node_uuid = $1 AND server_allocations.uuid IS NULL)
183 ORDER BY node_allocations.ip, node_allocations.port
184 LIMIT $3 OFFSET $4
185 "#,
186 Self::columns_sql(None)
187 ))
188 .bind(node_uuid)
189 .bind(search)
190 .bind(per_page)
191 .bind(offset)
192 .fetch_all(database.read())
193 .await?;
194
195 Ok(super::Pagination {
196 total: rows
197 .first()
198 .map_or(Ok(0), |row| row.try_get("total_count"))?,
199 per_page,
200 page,
201 data: rows
202 .into_iter()
203 .map(|row| Self::map(None, &row))
204 .try_collect_vec()?,
205 })
206 }
207
208 pub async fn by_node_uuid_with_pagination(
209 database: &crate::database::Database,
210 node_uuid: uuid::Uuid,
211 page: i64,
212 per_page: i64,
213 search: Option<&str>,
214 ) -> Result<super::Pagination<Self>, crate::database::DatabaseError> {
215 let offset = (page - 1) * per_page;
216
217 let rows = sqlx::query(&format!(
218 r#"
219 SELECT {}, server_allocations.server_uuid, COUNT(*) OVER() AS total_count
220 FROM node_allocations
221 LEFT JOIN server_allocations ON server_allocations.allocation_uuid = node_allocations.uuid
222 WHERE node_allocations.node_uuid = $1 AND ($2 IS NULL OR host(node_allocations.ip) || ':' || node_allocations.port ILIKE '%' || $2 || '%')
223 ORDER BY node_allocations.ip, node_allocations.port
224 LIMIT $3 OFFSET $4
225 "#,
226 Self::columns_sql(None)
227 ))
228 .bind(node_uuid)
229 .bind(search)
230 .bind(per_page)
231 .bind(offset)
232 .fetch_all(database.read())
233 .await?;
234
235 Ok(super::Pagination {
236 total: rows
237 .first()
238 .map_or(Ok(0), |row| row.try_get("total_count"))?,
239 per_page,
240 page,
241 data: rows
242 .into_iter()
243 .map(|row| Self::map(None, &row))
244 .try_collect_vec()?,
245 })
246 }
247
248 pub async fn delete_by_uuids(
249 database: &crate::database::Database,
250 node_uuid: uuid::Uuid,
251 uuids: &[uuid::Uuid],
252 ) -> Result<u64, crate::database::DatabaseError> {
253 let deleted = sqlx::query(
254 r#"
255 DELETE FROM node_allocations
256 WHERE node_allocations.node_uuid = $1 AND node_allocations.uuid = ANY($2)
257 "#,
258 )
259 .bind(node_uuid)
260 .bind(uuids)
261 .execute(database.write())
262 .await?
263 .rows_affected();
264
265 Ok(deleted)
266 }
267
268 #[inline]
269 pub async fn into_admin_api_object(
270 self,
271 database: &crate::database::Database,
272 storage_url_retriever: &StorageUrlRetriever<'_>,
273 ) -> Result<AdminApiNodeAllocation, crate::database::DatabaseError> {
274 let server = match self.server {
275 Some(fetchable) => Some(
276 fetchable
277 .fetch_cached(database)
278 .await?
279 .into_admin_api_object(database, storage_url_retriever)
280 .await?,
281 ),
282 None => None,
283 };
284
285 Ok(AdminApiNodeAllocation {
286 uuid: self.uuid,
287 server,
288 ip: compact_str::format_compact!("{}", self.ip.ip()),
289 ip_alias: self.ip_alias,
290 port: self.port,
291 created: self.created.and_utc(),
292 })
293 }
294}
295
296#[derive(ToSchema, Serialize)]
297#[schema(title = "NodeAllocation")]
298pub struct AdminApiNodeAllocation {
299 pub uuid: uuid::Uuid,
300 pub server: Option<super::server::AdminApiServer>,
301
302 pub ip: compact_str::CompactString,
303 pub ip_alias: Option<compact_str::CompactString>,
304 pub port: i32,
305
306 pub created: chrono::DateTime<chrono::Utc>,
307}