1use crate::{
2 crypt::EncryptedString,
3 models::{InsertQueryBuilder, UpdateQueryBuilder},
4 prelude::*,
5};
6use garde::Validate;
7use rand::distr::SampleString;
8use serde::{Deserialize, Serialize};
9use sqlx::{Row, postgres::PgRow};
10use std::{
11 collections::BTreeMap,
12 sync::{Arc, LazyLock},
13};
14use utoipa::ToSchema;
15
16#[inline]
17fn default_true() -> bool {
18 true
19}
20
21#[derive(ToSchema, Serialize, Deserialize, Validate, Clone)]
22pub struct DatabaseAgentHostTypeSettings {
23 #[garde(skip)]
24 #[serde(default = "default_true")]
25 pub enabled: bool,
26
27 #[garde(
28 length(chars, min = 3, max = 255),
29 inner(custom(crate::utils::validate_host))
30 )]
31 #[schema(min_length = 3, max_length = 255)]
32 #[serde(default)]
33 pub public_host: Option<compact_str::CompactString>,
34 #[garde(range(min = 1))]
35 #[schema(minimum = 1)]
36 #[serde(default)]
37 pub public_port: Option<u16>,
38}
39
40impl Default for DatabaseAgentHostTypeSettings {
41 fn default() -> Self {
42 Self {
43 enabled: true,
44 public_host: None,
45 public_port: None,
46 }
47 }
48}
49
50#[derive(ToSchema, Serialize, Deserialize, Validate, Clone, Default)]
51pub struct DatabaseAgentHostTypes {
52 #[garde(dive)]
53 #[serde(default)]
54 pub postgres: DatabaseAgentHostTypeSettings,
55 #[garde(dive)]
56 #[serde(default)]
57 pub mariadb: DatabaseAgentHostTypeSettings,
58 #[garde(dive)]
59 #[serde(default)]
60 pub mongodb: DatabaseAgentHostTypeSettings,
61 #[garde(dive)]
62 #[serde(default)]
63 pub redis: DatabaseAgentHostTypeSettings,
64}
65
66impl DatabaseAgentHostTypes {
67 #[inline]
68 pub fn get(&self, r#type: db_agent_api::DatabaseAgentType) -> &DatabaseAgentHostTypeSettings {
69 match r#type {
70 db_agent_api::DatabaseAgentType::Postgres => &self.postgres,
71 db_agent_api::DatabaseAgentType::Mariadb => &self.mariadb,
72 db_agent_api::DatabaseAgentType::Mongodb => &self.mongodb,
73 db_agent_api::DatabaseAgentType::Redis => &self.redis,
74 }
75 }
76}
77
78#[derive(Serialize, Deserialize, Clone)]
79pub struct DatabaseAgentHost {
80 pub uuid: uuid::Uuid,
81
82 pub name: compact_str::CompactString,
83 pub description: Option<compact_str::CompactString>,
84
85 pub deployment_enabled: bool,
86 pub maintenance_enabled: bool,
87
88 pub url: reqwest::Url,
89
90 pub memory: i64,
91 pub disk: i64,
92
93 pub types: DatabaseAgentHostTypes,
94
95 pub token: EncryptedString,
96
97 pub created: chrono::NaiveDateTime,
98
99 extension_data: super::ModelExtensionData,
100}
101
102impl BaseModel for DatabaseAgentHost {
103 const NAME: &'static str = "database_agent_host";
104
105 fn get_extension_list() -> &'static super::ModelExtensionList {
106 static EXTENSIONS: LazyLock<super::ModelExtensionList> =
107 LazyLock::new(|| parking_lot::RwLock::new(Vec::new()));
108
109 &EXTENSIONS
110 }
111
112 fn get_extension_data(&self) -> &super::ModelExtensionData {
113 &self.extension_data
114 }
115
116 #[inline]
117 fn base_columns(prefix: Option<&str>) -> BTreeMap<&'static str, compact_str::CompactString> {
118 let prefix = prefix.unwrap_or_default();
119
120 BTreeMap::from([
121 (
122 "database_agent_hosts.uuid",
123 compact_str::format_compact!("{prefix}uuid"),
124 ),
125 (
126 "database_agent_hosts.name",
127 compact_str::format_compact!("{prefix}name"),
128 ),
129 (
130 "database_agent_hosts.description",
131 compact_str::format_compact!("{prefix}description"),
132 ),
133 (
134 "database_agent_hosts.deployment_enabled",
135 compact_str::format_compact!("{prefix}deployment_enabled"),
136 ),
137 (
138 "database_agent_hosts.maintenance_enabled",
139 compact_str::format_compact!("{prefix}maintenance_enabled"),
140 ),
141 (
142 "database_agent_hosts.url",
143 compact_str::format_compact!("{prefix}url"),
144 ),
145 (
146 "database_agent_hosts.memory",
147 compact_str::format_compact!("{prefix}memory"),
148 ),
149 (
150 "database_agent_hosts.disk",
151 compact_str::format_compact!("{prefix}disk"),
152 ),
153 (
154 "database_agent_hosts.types",
155 compact_str::format_compact!("{prefix}types"),
156 ),
157 (
158 "database_agent_hosts.token",
159 compact_str::format_compact!("{prefix}token"),
160 ),
161 (
162 "database_agent_hosts.created",
163 compact_str::format_compact!("{prefix}created"),
164 ),
165 ])
166 }
167
168 #[inline]
169 fn map(prefix: Option<&str>, row: &PgRow) -> Result<Self, crate::database::DatabaseError> {
170 let prefix = prefix.unwrap_or_default();
171
172 Ok(Self {
173 uuid: row.try_get(compact_str::format_compact!("{prefix}uuid").as_str())?,
174 name: row.try_get(compact_str::format_compact!("{prefix}name").as_str())?,
175 description: row
176 .try_get(compact_str::format_compact!("{prefix}description").as_str())?,
177 deployment_enabled: row
178 .try_get(compact_str::format_compact!("{prefix}deployment_enabled").as_str())?,
179 maintenance_enabled: row
180 .try_get(compact_str::format_compact!("{prefix}maintenance_enabled").as_str())?,
181 url: row
182 .try_get::<String, _>(compact_str::format_compact!("{prefix}url").as_str())?
183 .parse()
184 .map_err(anyhow::Error::new)?,
185 memory: row.try_get(compact_str::format_compact!("{prefix}memory").as_str())?,
186 disk: row.try_get(compact_str::format_compact!("{prefix}disk").as_str())?,
187 types: serde_json::from_value(
188 row.try_get(compact_str::format_compact!("{prefix}types").as_str())?,
189 )?,
190 token: row.try_get(compact_str::format_compact!("{prefix}token").as_str())?,
191 created: row.try_get(compact_str::format_compact!("{prefix}created").as_str())?,
192 extension_data: Self::map_extensions(prefix, row)?,
193 })
194 }
195
196 fn cache_invalidation_keys(&self) -> Vec<compact_str::CompactString> {
197 vec![compact_str::format_compact!(
198 "{}::{}",
199 Self::NAME,
200 self.uuid
201 )]
202 }
203}
204
205impl DatabaseAgentHost {
206 pub async fn all_with_pagination(
207 database: &crate::database::Database,
208 page: i64,
209 per_page: i64,
210 search: Option<&str>,
211 ) -> Result<super::Pagination<Self>, crate::database::DatabaseError> {
212 super::PaginatedQueryBuilder::new(
213 Self::columns_sql(None),
214 format!(
215 r#"
216 FROM database_agent_hosts
217 WHERE {search}
218 "#,
219 search = super::search_sql(
220 1,
221 &["database_agent_hosts.name"],
222 &["database_agent_hosts.uuid"]
223 ),
224 ),
225 "database_agent_hosts.created, database_agent_hosts.uuid",
226 )
227 .bind(search)
228 .fetch(database, page, per_page)
229 .await
230 }
231
232 pub async fn by_node_most_eligible(
233 database: &crate::database::Database,
234 node: &super::node::Node,
235 r#type: db_agent_api::DatabaseAgentType,
236 memory: i64,
237 disk: i64,
238 ) -> Result<Vec<Self>, crate::database::DatabaseError> {
239 let rows = sqlx::query(sqlx::AssertSqlSafe(format!(
240 r#"
241 WITH database_usage AS (
242 SELECT
243 server_database_instances.database_agent_host_uuid,
244 COALESCE(SUM(COALESCE(server_database_instances.memory, t.memory)), 0)::BIGINT AS used_memory,
245 COALESCE(SUM(COALESCE(server_database_instances.disk, t.disk)), 0)::BIGINT AS used_disk
246 FROM server_database_instances
247 LEFT JOIN database_agent_templates t ON t.uuid = server_database_instances.database_agent_template_uuid
248 GROUP BY server_database_instances.database_agent_host_uuid
249 )
250 SELECT {}
251 FROM database_agent_hosts
252 LEFT JOIN database_usage u ON database_agent_hosts.uuid = u.database_agent_host_uuid
253 WHERE (
254 EXISTS (
255 SELECT 1 FROM node_database_agent_hosts
256 WHERE node_database_agent_hosts.database_agent_host_uuid = database_agent_hosts.uuid AND node_database_agent_hosts.node_uuid = $1
257 )
258 OR EXISTS (
259 SELECT 1 FROM location_database_agent_hosts
260 WHERE location_database_agent_hosts.database_agent_host_uuid = database_agent_hosts.uuid AND location_database_agent_hosts.location_uuid = $2
261 )
262 )
263 AND database_agent_hosts.deployment_enabled
264 AND NOT database_agent_hosts.maintenance_enabled
265 AND COALESCE((database_agent_hosts.types -> $3 ->> 'enabled')::BOOL, TRUE)
266 AND (database_agent_hosts.memory = 0 OR COALESCE(u.used_memory, 0) + $4 <= database_agent_hosts.memory)
267 AND (database_agent_hosts.disk = 0 OR COALESCE(u.used_disk, 0) + $5 <= database_agent_hosts.disk)
268 ORDER BY
269 GREATEST(
270 CASE WHEN database_agent_hosts.memory = 0 THEN 0 ELSE (COALESCE(u.used_memory, 0) + $4)::FLOAT / database_agent_hosts.memory END,
271 CASE WHEN database_agent_hosts.disk = 0 THEN 0 ELSE (COALESCE(u.used_disk, 0) + $5)::FLOAT / database_agent_hosts.disk END
272 )
273 "#,
274 Self::columns_sql(None)
275 )))
276 .bind(node.uuid)
277 .bind(node.location.uuid)
278 .bind(r#type.as_str())
279 .bind(memory)
280 .bind(disk)
281 .fetch_all(database.read())
282 .await?;
283
284 rows.into_iter()
285 .map(|row| Self::map(None, &row))
286 .try_collect_vec()
287 }
288
289 #[inline]
290 pub fn generate_token() -> String {
291 rand::distr::Alphanumeric.sample_string(&mut rand::rng(), 64)
292 }
293
294 pub async fn reset_token(&self, state: &crate::State) -> Result<String, anyhow::Error> {
295 let (token, encrypted_token) =
296 EncryptedString::from_plaintext_with_input(Self::generate_token(), &state.database)
297 .await?;
298
299 sqlx::query(
300 r#"
301 UPDATE database_agent_hosts
302 SET token = $2
303 WHERE database_agent_hosts.uuid = $1
304 "#,
305 )
306 .bind(self.uuid)
307 .bind(encrypted_token)
308 .execute(state.database.write())
309 .await?;
310
311 Self::invalidate_cached(&state.database, self.uuid).await;
312
313 Ok(token)
314 }
315
316 #[inline]
317 pub async fn api_client(
318 &self,
319 database: &crate::database::Database,
320 ) -> Result<db_agent_api::client::DbAgentClient, anyhow::Error> {
321 Ok(db_agent_api::client::DbAgentClient::new(
322 self.url.to_string(),
323 self.token.decrypt(database).await?.into(),
324 ))
325 }
326
327 pub async fn fetch_configuration(
331 &self,
332 database: &crate::database::Database,
333 ) -> Result<db_agent_api::Config, anyhow::Error> {
334 database
335 .cache
336 .cached(
337 &format!("database_agent_host::{}::configuration::raw", self.uuid),
338 120,
339 || async {
340 Ok::<_, anyhow::Error>(db_agent_api::Config(
341 self.api_client(database).await?.get_system_config().await?,
342 ))
343 },
344 )
345 .await
346 }
347
348 pub async fn fetch_database_resources(
352 &self,
353 database: &crate::database::Database,
354 ) -> Result<std::collections::HashMap<uuid::Uuid, db_agent_api::ResourceUsage>, anyhow::Error>
355 {
356 database
357 .cache
358 .cached(
359 &format!("database_agent_host::{}::database_resources", self.uuid),
360 15,
361 || async {
362 let resources = self
363 .api_client(database)
364 .await?
365 .get_instances_utilization()
366 .await?;
367
368 Ok::<_, anyhow::Error>(resources.into_iter().collect())
369 },
370 )
371 .await
372 }
373
374 pub async fn update_configuration(
378 &self,
379 database: &crate::database::Database,
380 config_patch: &serde_json::Value,
381 ) -> Result<bool, anyhow::Error> {
382 let response = self
383 .api_client(database)
384 .await?
385 .patch_system_config(config_patch)
386 .await?;
387 if !response.applied {
388 return Ok(false);
389 }
390
391 database
392 .cache
393 .invalidate(&format!(
394 "database_agent_host::{}::configuration::raw",
395 self.uuid
396 ))
397 .await?;
398
399 Ok(true)
400 }
401}
402
403#[async_trait::async_trait]
404impl IntoAdminApiObject for DatabaseAgentHost {
405 type AdminApiObject = AdminApiDatabaseAgentHost;
406 type ExtraArgs<'a> = ();
407
408 async fn into_admin_api_object<'a>(
409 self,
410 state: &crate::State,
411 _args: Self::ExtraArgs<'a>,
412 ) -> Result<Self::AdminApiObject, crate::database::DatabaseError> {
413 let api_object = AdminApiDatabaseAgentHost::init_hooks(&self, state).await?;
414
415 let api_object = finish_extendible!(
416 AdminApiDatabaseAgentHost {
417 uuid: self.uuid,
418 name: self.name,
419 description: self.description,
420 deployment_enabled: self.deployment_enabled,
421 maintenance_enabled: self.maintenance_enabled,
422 url: self.url.to_string(),
423 memory: self.memory,
424 disk: self.disk,
425 types: self.types,
426 created: self.created.and_utc(),
427 },
428 api_object,
429 state
430 )?;
431
432 Ok(api_object)
433 }
434}
435
436#[async_trait::async_trait]
437impl ByUuid for DatabaseAgentHost {
438 async fn by_uuid(
439 database: &crate::database::Database,
440 uuid: uuid::Uuid,
441 ) -> Result<Self, crate::database::DatabaseError> {
442 let row = sqlx::query(sqlx::AssertSqlSafe(format!(
443 r#"
444 SELECT {}
445 FROM database_agent_hosts
446 WHERE database_agent_hosts.uuid = $1
447 "#,
448 Self::columns_sql(None)
449 )))
450 .bind(uuid)
451 .fetch_one(database.read())
452 .await?;
453
454 Self::map(None, &row)
455 }
456
457 async fn by_uuid_with_transaction(
458 transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
459 uuid: uuid::Uuid,
460 ) -> Result<Self, crate::database::DatabaseError> {
461 let row = sqlx::query(sqlx::AssertSqlSafe(format!(
462 r#"
463 SELECT {}
464 FROM database_agent_hosts
465 WHERE database_agent_hosts.uuid = $1
466 "#,
467 Self::columns_sql(None)
468 )))
469 .bind(uuid)
470 .fetch_one(&mut **transaction)
471 .await?;
472
473 Self::map(None, &row)
474 }
475}
476
477#[derive(ToSchema, Deserialize, Validate)]
478pub struct CreateDatabaseAgentHostOptions {
479 #[garde(length(chars, min = 1, max = 255))]
480 #[schema(min_length = 1, max_length = 255)]
481 pub name: compact_str::CompactString,
482 #[garde(length(chars, min = 1, max = 1024))]
483 #[schema(min_length = 1, max_length = 1024)]
484 pub description: Option<compact_str::CompactString>,
485
486 #[garde(skip)]
487 pub deployment_enabled: bool,
488 #[garde(skip)]
489 pub maintenance_enabled: bool,
490
491 #[garde(length(chars, min = 3, max = 255), url)]
492 #[schema(min_length = 3, max_length = 255, format = "uri")]
493 pub url: compact_str::CompactString,
494
495 #[garde(range(min = 0))]
496 #[schema(minimum = 0)]
497 pub memory: i64,
498 #[garde(range(min = 0))]
499 #[schema(minimum = 0)]
500 pub disk: i64,
501
502 #[garde(dive)]
503 #[serde(default)]
504 pub types: DatabaseAgentHostTypes,
505}
506
507#[async_trait::async_trait]
508impl CreatableModel for DatabaseAgentHost {
509 type CreateOptions<'a> = CreateDatabaseAgentHostOptions;
510 type CreateResult = Self;
511
512 fn get_create_handlers() -> &'static LazyLock<CreateListenerList<Self>> {
513 static CREATE_LISTENERS: LazyLock<CreateListenerList<DatabaseAgentHost>> =
514 LazyLock::new(|| Arc::new(ModelHandlerList::default()));
515
516 &CREATE_LISTENERS
517 }
518
519 async fn create_with_transaction(
520 state: &crate::State,
521 mut options: Self::CreateOptions<'_>,
522 transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
523 ) -> Result<Self, crate::database::DatabaseError> {
524 options.validate()?;
525
526 let mut query_builder = InsertQueryBuilder::new("database_agent_hosts");
527
528 Self::run_create_handlers(&mut options, &mut query_builder, state, transaction).await?;
529
530 let token = Self::generate_token();
531
532 query_builder
533 .set("name", &options.name)
534 .set("description", &options.description)
535 .set("deployment_enabled", options.deployment_enabled)
536 .set("maintenance_enabled", options.maintenance_enabled)
537 .set("url", &options.url)
538 .set("memory", options.memory)
539 .set("disk", options.disk)
540 .set("types", serde_json::to_value(&options.types)?)
541 .set(
542 "token",
543 EncryptedString::from_plaintext(token, &state.database).await?,
544 );
545
546 let row = query_builder
547 .returning(&Self::columns_sql(None))
548 .fetch_one(&mut **transaction)
549 .await?;
550 let mut database_agent_host = Self::map(None, &row)?;
551
552 Self::run_after_create_handlers(&mut database_agent_host, &options, state, transaction)
553 .await?;
554
555 Ok(database_agent_host)
556 }
557}
558
559#[derive(ToSchema, Serialize, Deserialize, Validate, Clone, Default)]
560pub struct UpdateDatabaseAgentHostOptions {
561 #[garde(length(chars, min = 1, max = 255))]
562 #[schema(min_length = 1, max_length = 255)]
563 pub name: Option<compact_str::CompactString>,
564 #[garde(length(chars, min = 1, max = 1024))]
565 #[schema(min_length = 1, max_length = 1024)]
566 #[serde(
567 default,
568 skip_serializing_if = "Option::is_none",
569 with = "::serde_with::rust::double_option"
570 )]
571 pub description: Option<Option<compact_str::CompactString>>,
572
573 #[garde(skip)]
574 pub deployment_enabled: Option<bool>,
575 #[garde(skip)]
576 pub maintenance_enabled: Option<bool>,
577
578 #[garde(length(chars, min = 3, max = 255), url)]
579 #[schema(min_length = 3, max_length = 255, format = "uri")]
580 pub url: Option<compact_str::CompactString>,
581
582 #[garde(range(min = 0))]
583 #[schema(minimum = 0)]
584 pub memory: Option<i64>,
585 #[garde(range(min = 0))]
586 #[schema(minimum = 0)]
587 pub disk: Option<i64>,
588
589 #[garde(dive)]
590 pub types: Option<DatabaseAgentHostTypes>,
591}
592
593#[async_trait::async_trait]
594impl UpdatableModel for DatabaseAgentHost {
595 type UpdateOptions = UpdateDatabaseAgentHostOptions;
596
597 fn get_update_handlers() -> &'static LazyLock<UpdateHandlerList<Self>> {
598 static UPDATE_LISTENERS: LazyLock<UpdateHandlerList<DatabaseAgentHost>> =
599 LazyLock::new(|| Arc::new(ModelHandlerList::default()));
600
601 &UPDATE_LISTENERS
602 }
603
604 async fn update_with_transaction(
605 &mut self,
606 state: &crate::State,
607 mut options: Self::UpdateOptions,
608 transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
609 ) -> Result<(), crate::database::DatabaseError> {
610 options.validate()?;
611
612 let mut query_builder = UpdateQueryBuilder::new("database_agent_hosts");
613
614 self.run_update_handlers(&mut options, &mut query_builder, state, transaction)
615 .await?;
616
617 query_builder
618 .set("name", options.name.as_ref())
619 .set(
620 "description",
621 options.description.as_ref().map(|d| d.as_ref()),
622 )
623 .set("deployment_enabled", options.deployment_enabled)
624 .set("maintenance_enabled", options.maintenance_enabled)
625 .set("url", options.url.as_ref())
626 .set("memory", options.memory)
627 .set("disk", options.disk)
628 .set(
629 "types",
630 options
631 .types
632 .as_ref()
633 .map(serde_json::to_value)
634 .transpose()?,
635 )
636 .where_eq("uuid", self.uuid);
637
638 query_builder.execute(&mut **transaction).await?;
639
640 if let Some(name) = options.name {
641 self.name = name;
642 }
643 if let Some(description) = options.description {
644 self.description = description;
645 }
646 if let Some(deployment_enabled) = options.deployment_enabled {
647 self.deployment_enabled = deployment_enabled;
648 }
649 if let Some(maintenance_enabled) = options.maintenance_enabled {
650 self.maintenance_enabled = maintenance_enabled;
651 }
652 if let Some(url) = options.url {
653 self.url = url.parse().map_err(anyhow::Error::new)?;
654 }
655 if let Some(memory) = options.memory {
656 self.memory = memory;
657 }
658 if let Some(disk) = options.disk {
659 self.disk = disk;
660 }
661 if let Some(types) = options.types {
662 self.types = types;
663 }
664
665 self.run_after_update_handlers(state, transaction).await?;
666
667 Ok(())
668 }
669}
670
671#[async_trait::async_trait]
672impl DeletableModel for DatabaseAgentHost {
673 type DeleteOptions = ();
674
675 fn get_delete_handlers() -> &'static LazyLock<DeleteHandlerList<Self>> {
676 static DELETE_LISTENERS: LazyLock<DeleteHandlerList<DatabaseAgentHost>> =
677 LazyLock::new(|| Arc::new(ModelHandlerList::default()));
678
679 &DELETE_LISTENERS
680 }
681
682 async fn delete_with_transaction(
683 &self,
684 state: &crate::State,
685 options: Self::DeleteOptions,
686 transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
687 ) -> Result<(), anyhow::Error> {
688 self.run_delete_handlers(&options, state, transaction)
689 .await?;
690
691 sqlx::query(
692 r#"
693 DELETE FROM database_agent_hosts
694 WHERE database_agent_hosts.uuid = $1
695 "#,
696 )
697 .bind(self.uuid)
698 .execute(&mut **transaction)
699 .await?;
700
701 self.run_after_delete_handlers(&options, state, transaction)
702 .await?;
703
704 Ok(())
705 }
706}
707
708#[schema_extension_derive::extendible]
709#[init_args(DatabaseAgentHost, crate::State)]
710#[hook_args(crate::State)]
711#[derive(ToSchema, Serialize)]
712#[schema(title = "DatabaseAgentHost")]
713pub struct AdminApiDatabaseAgentHost {
714 pub uuid: uuid::Uuid,
715
716 pub name: compact_str::CompactString,
717 pub description: Option<compact_str::CompactString>,
718
719 pub deployment_enabled: bool,
720 pub maintenance_enabled: bool,
721
722 #[schema(format = "uri")]
723 pub url: String,
724
725 pub memory: i64,
726 pub disk: i64,
727
728 pub types: DatabaseAgentHostTypes,
729
730 pub created: chrono::DateTime<chrono::Utc>,
731}