1use std::{path::PathBuf, str::FromStr, sync::Arc, time::Duration};
2
3use futures_util::TryStreamExt;
4use saddle_core::{ComponentLifecycle, LifecycleFuture, OperationId, Result};
5use saddle_observability::{CallKind, Observer};
6use sqlx::{
7 Connection, MySqlPool,
8 mysql::{MySqlConnectOptions, MySqlConnection, MySqlPoolOptions},
9};
10
11use crate::{
12 CallContext, DbRow, MAX_RESULT_BYTES, Statement, Transaction, TransactionFuture, WriteResult,
13 cleanup::CleanupCoordinator,
14 error::{
15 invalid_config, map_operation_error, name_mapping_bypass, result_limit_exceeded,
16 transaction_begin_failed,
17 },
18 name_mapping::{
19 MappingStartupConfig, OperationSql, PhysicalOperationPlans, StaticLogicalTable, freeze,
20 },
21 row::row_payload_bytes,
22};
23
24pub const MAX_QUERY_ROWS: usize = 10_000;
25pub const MAX_INBOUND_PACKET_BYTES: u64 = 8_388_608;
32
33pub struct DatabaseStartupInjection {
39 url: String,
40 mapping_directory: PathBuf,
41}
42
43#[derive(Clone, Copy, Debug, Eq, PartialEq)]
44pub enum DatabaseStartupInjectionError {
45 InvalidConnectionEnvironment,
46 MissingConnectionSecret,
47 InvalidConnectionSecret,
48 InvalidMappingDirectory,
49 MappingDirectoryUnavailable,
50}
51
52impl DatabaseStartupInjectionError {
53 pub const fn code(self) -> &'static str {
54 match self {
55 Self::InvalidConnectionEnvironment => "db.startup.connection_env_invalid",
56 Self::MissingConnectionSecret => "db.startup.connection_secret_missing",
57 Self::InvalidConnectionSecret => "db.startup.connection_secret_invalid",
58 Self::InvalidMappingDirectory => "db.startup.mapping_dir_invalid",
59 Self::MappingDirectoryUnavailable => "db.startup.mapping_dir_unavailable",
60 }
61 }
62}
63
64impl DatabaseStartupInjection {
65 #[doc(hidden)]
68 pub fn load(
69 config_directory: &std::path::Path,
70 connection_environment: &str,
71 mapping_directory: &std::path::Path,
72 ) -> std::result::Result<Self, DatabaseStartupInjectionError> {
73 if !valid_environment_name(connection_environment) {
74 return Err(DatabaseStartupInjectionError::InvalidConnectionEnvironment);
75 }
76 let url = std::env::var(connection_environment).map_err(|error| match error {
77 std::env::VarError::NotPresent => {
78 DatabaseStartupInjectionError::MissingConnectionSecret
79 }
80 std::env::VarError::NotUnicode(_) => {
81 DatabaseStartupInjectionError::InvalidConnectionSecret
82 }
83 })?;
84 if url.trim().is_empty() {
85 return Err(DatabaseStartupInjectionError::InvalidConnectionSecret);
86 }
87 if mapping_directory.as_os_str().is_empty()
88 || mapping_directory.is_absolute()
89 || mapping_directory
90 .components()
91 .any(|component| !matches!(component, std::path::Component::Normal(_)))
92 {
93 return Err(DatabaseStartupInjectionError::InvalidMappingDirectory);
94 }
95 let mapping_directory = config_directory.join(mapping_directory);
96 if !mapping_directory.is_dir() {
97 return Err(DatabaseStartupInjectionError::MappingDirectoryUnavailable);
98 }
99 Ok(Self {
100 url,
101 mapping_directory,
102 })
103 }
104
105 #[doc(hidden)]
106 pub fn into_database_config(self) -> DatabaseConfig {
107 DatabaseConfig::new(self.url).name_mapping_directory(self.mapping_directory)
108 }
109}
110
111fn valid_environment_name(value: &str) -> bool {
112 let mut bytes = value.bytes();
113 let Some(first) = bytes.next() else {
114 return false;
115 };
116 (first.is_ascii_alphabetic() || first == b'_')
117 && bytes.all(|byte| byte.is_ascii_alphanumeric() || byte == b'_')
118}
119
120#[derive(Clone)]
122pub struct DatabaseConfig {
123 url: String,
124 max_connections: u32,
125 acquire_timeout: Duration,
126 name_mappings: Option<MappingStartupConfig>,
127}
128
129impl DatabaseConfig {
130 pub fn new(url: impl Into<String>) -> Self {
131 Self {
132 url: url.into(),
133 max_connections: 16,
134 acquire_timeout: Duration::from_secs(5),
135 name_mappings: None,
136 }
137 }
138 pub fn max_connections(mut self, value: u32) -> Self {
139 self.max_connections = value;
140 self
141 }
142 pub fn acquire_timeout(mut self, value: Duration) -> Self {
143 self.acquire_timeout = value;
144 self
145 }
146
147 pub fn name_mapping_directory(mut self, directory: impl Into<PathBuf>) -> Self {
151 self.name_mappings
152 .get_or_insert_with(|| MappingStartupConfig::new(PathBuf::new()))
153 .set_directory(directory.into());
154 self
155 }
156
157 pub fn register_logical_table<T: StaticLogicalTable>(mut self) -> Self {
159 self.name_mappings
160 .get_or_insert_with(|| MappingStartupConfig::new(PathBuf::new()))
161 .register_table::<T>();
162 self
163 }
164
165 #[doc(hidden)]
166 pub fn register_query_operation<O: crate::internal::StaticQueryOptionalOperation>(
167 mut self,
168 ) -> Self {
169 self.name_mappings
170 .get_or_insert_with(|| MappingStartupConfig::new(PathBuf::new()))
171 .register_query::<O>(O::OPERATION, O::LOGICAL_TABLE, O::LOGICAL_COLUMNS, O::SQL);
172 self
173 }
174
175 #[doc(hidden)]
176 pub fn register_write_operation<O: crate::internal::StaticWriteOperation>(mut self) -> Self {
177 self.name_mappings
178 .get_or_insert_with(|| MappingStartupConfig::new(PathBuf::new()))
179 .register_write::<O>(O::OPERATION, O::LOGICAL_TABLE, O::LOGICAL_COLUMNS, O::SQL);
180 self
181 }
182
183 pub fn validate_name_mappings(
186 &self,
187 ) -> std::result::Result<(), crate::DatabaseNameMappingError> {
188 freeze(self.name_mappings.clone()).map(|_| ())
189 }
190
191 pub(crate) fn verified_connections(mut self, value: u32) -> Self {
192 self.max_connections = value;
193 self
194 }
195
196 pub(crate) fn deployment_url(&self) -> &str {
197 &self.url
198 }
199
200 pub(crate) fn options(&self) -> Result<MySqlConnectOptions> {
201 if self.max_connections == 0 {
202 return Err(invalid_config("max connections must be greater than zero"));
203 }
204 if self.acquire_timeout.is_zero() {
205 return Err(invalid_config("acquire timeout must be greater than zero"));
206 }
207 MySqlConnectOptions::from_str(&self.url)
208 .map_err(|_| invalid_config("database URL is not a valid MySQL/MariaDB URL"))
209 }
210}
211
212#[derive(Clone)]
214pub struct Database {
215 pub(crate) pool: MySqlPool,
216 observer: Observer,
217 cleanup: Arc<CleanupCoordinator>,
218 physical_plans: Arc<PhysicalOperationPlans>,
219}
220
221impl Database {
222 pub async fn connect(config: DatabaseConfig, observer: Observer) -> Result<Self> {
224 let physical_plans = freeze(config.name_mappings.clone())
225 .map_err(|_| invalid_config("database name mapping is invalid"))?;
226 let options = config.options()?;
227 let mut preflight = MySqlConnection::connect_with(&options)
228 .await
229 .map_err(map_operation_error)?;
230 let server_packet_limit = sqlx::query_scalar::<_, u64>("SELECT @@max_allowed_packet")
231 .fetch_one(&mut preflight)
232 .await
233 .map_err(map_operation_error)?;
234 preflight.close().await.map_err(map_operation_error)?;
235 if server_packet_limit > MAX_INBOUND_PACKET_BYTES {
236 return Err(invalid_config(
237 "server max_allowed_packet exceeds the V1 inbound allocation limit",
238 ));
239 }
240 let pool = MySqlPoolOptions::new()
241 .max_connections(config.max_connections)
242 .acquire_timeout(config.acquire_timeout)
243 .idle_timeout(None)
244 .max_lifetime(None)
245 .after_connect(|connection, _metadata| {
246 Box::pin(async move {
247 let packet_limit = sqlx::query_scalar::<_, u64>("SELECT @@max_allowed_packet")
248 .fetch_one(connection)
249 .await?;
250 if packet_limit > MAX_INBOUND_PACKET_BYTES {
251 return Err(sqlx::Error::Protocol(
252 "server packet limit exceeds Saddle V1 allocation boundary".to_owned(),
253 ));
254 }
255 Ok(())
256 })
257 })
258 .connect_with(options)
259 .await
260 .map_err(map_operation_error)?;
261 let mut preopened = Vec::new();
262 preopened
263 .try_reserve_exact(config.max_connections as usize)
264 .map_err(|_| invalid_config("database connection profile is too large"))?;
265 while preopened.len() < config.max_connections as usize {
266 preopened.push(pool.acquire().await.map_err(map_operation_error)?);
267 }
268 for connection in &mut preopened {
269 connection.return_to_pool().await;
270 }
271 Ok(Self {
272 pool,
273 observer,
274 cleanup: CleanupCoordinator::start(),
275 physical_plans: Arc::new(physical_plans),
276 })
277 }
278
279 pub(crate) fn query_sql<O: crate::internal::StaticQueryOptionalOperation>(
280 &self,
281 ) -> OperationSql {
282 self.physical_plans.query::<O>()
283 }
284
285 pub(crate) fn write_sql<O: crate::internal::StaticWriteOperation>(&self) -> OperationSql {
286 self.physical_plans.write::<O>()
287 }
288
289 pub(crate) fn name_mappings_frozen(&self) -> bool {
290 self.physical_plans.enabled()
291 }
292
293 pub async fn query_all(
294 &self,
295 parent: &CallContext,
296 statement: Statement,
297 ) -> Result<Vec<DbRow>> {
298 if self.physical_plans.enabled() {
299 return Err(name_mapping_bypass());
300 }
301 statement.validate()?;
302 let operation = statement.operation().to_owned();
303 let call = self.observer.start_child_call(
304 parent,
305 CallKind::Database,
306 "database",
307 "database",
308 OperationId::from(operation),
309 );
310 let result = async {
311 let mut stream = statement.query().fetch(&self.pool);
312 let mut rows = Vec::new();
313 let mut result_bytes = 0_usize;
314 while let Some(row) = stream.try_next().await.map_err(map_operation_error)? {
315 if rows.len() == MAX_QUERY_ROWS {
316 return Err(result_limit_exceeded());
317 }
318 result_bytes = result_bytes.saturating_add(row_payload_bytes(&row)?);
319 if result_bytes > MAX_RESULT_BYTES {
320 return Err(result_limit_exceeded());
321 }
322 rows.push(DbRow(row));
323 }
324 Ok(rows)
325 }
326 .await;
327 finish_call(call, &result);
328 result
329 }
330
331 pub async fn query_optional(
332 &self,
333 parent: &CallContext,
334 statement: Statement,
335 ) -> Result<Option<DbRow>> {
336 if self.physical_plans.enabled() {
337 return Err(name_mapping_bypass());
338 }
339 statement.validate()?;
340 let operation = statement.operation().to_owned();
341 let call = self.observer.start_child_call(
342 parent,
343 CallKind::Database,
344 "database",
345 "database",
346 OperationId::from(operation),
347 );
348 let result = statement
349 .query()
350 .fetch_optional(&self.pool)
351 .await
352 .map_err(map_operation_error)
353 .and_then(|row| {
354 row.map(|row| {
355 row_payload_bytes(&row)?;
356 Ok(DbRow(row))
357 })
358 .transpose()
359 });
360 finish_call(call, &result);
361 result
362 }
363
364 pub async fn write(&self, parent: &CallContext, statement: Statement) -> Result<WriteResult> {
365 if self.physical_plans.enabled() {
366 return Err(name_mapping_bypass());
367 }
368 statement.validate()?;
369 let operation = statement.operation().to_owned();
370 let call = self.observer.start_child_call(
371 parent,
372 CallKind::Database,
373 "database",
374 "database",
375 OperationId::from(operation),
376 );
377 let result = statement
378 .query()
379 .execute(&self.pool)
380 .await
381 .map(|result| WriteResult::new(result.rows_affected(), result.last_insert_id()))
382 .map_err(map_operation_error);
383 finish_call(call, &result);
384 result
385 }
386
387 pub async fn transaction<T, F>(
392 &self,
393 parent: &CallContext,
394 operation: impl Into<OperationId>,
395 work: F,
396 ) -> Result<T>
397 where
398 T: Send,
399 F: for<'a> FnOnce(&'a mut Transaction) -> TransactionFuture<'a, T> + Send,
400 {
401 if self.physical_plans.enabled() {
402 return Err(name_mapping_bypass());
403 }
404 let operation = operation.into();
405 crate::statement::validate_operation(operation.as_str())?;
406 let call = self.observer.start_child_call(
407 parent,
408 CallKind::Transaction,
409 "database",
410 "database",
411 operation,
412 );
413 let cleanup = match self.cleanup.transaction_sender() {
414 Some(cleanup) => cleanup,
415 None => {
416 let error = transaction_begin_failed();
417 call.fail(&error);
418 return Err(error);
419 }
420 };
421 let begin = self.observer.start_child_call(
422 call.context(),
423 CallKind::Transaction,
424 "database",
425 "database",
426 "begin",
427 );
428 let raw = match self.pool.begin().await {
429 Ok(transaction) => {
430 begin.succeed();
431 transaction
432 }
433 Err(_) => {
434 let error = transaction_begin_failed();
435 begin.fail(&error);
436 call.fail(&error);
437 return Err(error);
438 }
439 };
440 let mut transaction = Transaction::new(raw, self.observer.clone(), call, cleanup);
441 match work(&mut transaction).await {
442 Ok(value) => transaction.commit().await.map(|()| value),
443 Err(work_error) => Err(transaction.rollback(work_error).await),
444 }
445 }
446
447 pub(crate) async fn close(&self) -> Result<()> {
448 let cleanup = self.cleanup.shutdown().await;
449 self.pool.close().await;
450 cleanup
451 }
452
453 #[cfg(test)]
454 pub(crate) fn connect_lazy(config: DatabaseConfig, observer: Observer) -> Result<Self> {
455 let physical_plans = freeze(config.name_mappings.clone())
456 .map_err(|_| invalid_config("database name mapping is invalid"))?;
457 let options = config.options()?;
458 let pool = MySqlPoolOptions::new()
459 .max_connections(config.max_connections)
460 .acquire_timeout(config.acquire_timeout)
461 .connect_lazy_with(options);
462 Ok(Self {
463 pool,
464 observer,
465 cleanup: CleanupCoordinator::start(),
466 physical_plans: Arc::new(physical_plans),
467 })
468 }
469}
470
471impl ComponentLifecycle for Database {
472 fn name(&self) -> &'static str {
473 "database"
474 }
475 fn start(&self) -> LifecycleFuture<'_> {
476 Box::pin(async { Ok(()) })
477 }
478 fn shutdown(&self) -> LifecycleFuture<'_> {
479 Box::pin(async move { self.close().await })
480 }
481}
482
483fn finish_call<T>(call: saddle_observability::ActiveCall, result: &Result<T>) {
484 match result {
485 Ok(_) => call.succeed(),
486 Err(error) => call.fail(error),
487 }
488}
489
490#[cfg(test)]
491mod tests {
492 use saddle_core::{ApplicationId, ErrorKind, ModuleId, ServiceId, SpanId, TraceId};
493 use saddle_observability::ObserverConfig;
494 use serde_json::Value;
495 use std::{
496 io,
497 sync::{Arc, Mutex},
498 };
499
500 use super::*;
501 use crate::SaddleError;
502
503 struct MappedTable;
504
505 impl crate::StaticLogicalTable for MappedTable {
506 const TABLE: &'static str = "逻辑表";
507 const COLUMNS: &'static [&'static str] = &["逻辑列"];
508 }
509
510 #[derive(Clone, Default)]
511 struct Capture(Arc<Mutex<Vec<u8>>>);
512 impl io::Write for Capture {
513 fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
514 self.0.lock().unwrap().extend_from_slice(bytes);
515 Ok(bytes.len())
516 }
517 fn flush(&mut self) -> io::Result<()> {
518 Ok(())
519 }
520 }
521 fn context() -> CallContext {
522 CallContext::new(
523 ApplicationId::from("shop"),
524 ModuleId::from("orders"),
525 ServiceId::from("orders"),
526 OperationId::from("create"),
527 TraceId::from_u128(1),
528 SpanId::from_u64(2),
529 )
530 }
531
532 #[test]
533 fn configuration_rejects_invalid_bounds_without_exposing_url() {
534 let observer = Observer::with_writer(ObserverConfig::default(), io::sink()).unwrap();
535 let error = Database::connect_lazy(
536 DatabaseConfig::new("mysql://secret@localhost/db").max_connections(0),
537 observer,
538 )
539 .err()
540 .unwrap();
541 assert_eq!(error.code(), "db.invalid_config");
542 assert!(!error.to_string().contains("secret"));
543 }
544
545 #[test]
546 fn unified_startup_injection_resolves_secret_and_relative_mapping_directory() {
547 let directory =
548 std::env::temp_dir().join(format!("saddle-db-alpha10-config-{}", std::process::id()));
549 let mappings = directory.join("mappings");
550 let _ = std::fs::remove_dir_all(&directory);
551 std::fs::create_dir_all(&mappings).unwrap();
552 let executable = std::env::current_exe().unwrap();
553 for mode in ["happy", "missing", "bad-path"] {
554 let mut child = std::process::Command::new(&executable);
555 child
556 .args([
557 "--ignored",
558 "--exact",
559 "database::tests::unified_startup_injection_child",
560 ])
561 .env("SADDLE_ALPHA10_INJECTION_MODE", mode)
562 .env("SADDLE_ALPHA10_CONFIG_ROOT", &directory)
563 .env_remove("SADDLE_ALPHA10_DATABASE_TEST");
564 if mode != "missing" {
565 child.env(
566 "SADDLE_ALPHA10_DATABASE_TEST",
567 "mysql://deployment-secret@localhost/database",
568 );
569 }
570 assert!(
571 child.status().unwrap().success(),
572 "child mode {mode} failed"
573 );
574 }
575 std::fs::remove_dir_all(directory).unwrap();
576 }
577
578 #[test]
579 #[ignore = "executed in isolated child processes by the parent test"]
580 fn unified_startup_injection_child() {
581 let root = PathBuf::from(std::env::var_os("SADDLE_ALPHA10_CONFIG_ROOT").unwrap());
582 let mode = std::env::var("SADDLE_ALPHA10_INJECTION_MODE").unwrap();
583 let mapping = if mode == "bad-path" {
584 std::path::Path::new("../mappings")
585 } else {
586 std::path::Path::new("mappings")
587 };
588 let result = DatabaseStartupInjection::load(&root, "SADDLE_ALPHA10_DATABASE_TEST", mapping);
589 match mode.as_str() {
590 "happy" => {
591 let injection = result.unwrap();
592 assert_eq!(injection.mapping_directory, root.join("mappings"));
593 assert!(injection.url.contains("deployment-secret"));
594 }
595 "missing" => assert_eq!(
596 result.err(),
597 Some(DatabaseStartupInjectionError::MissingConnectionSecret)
598 ),
599 "bad-path" => assert_eq!(
600 result.err(),
601 Some(DatabaseStartupInjectionError::InvalidMappingDirectory)
602 ),
603 _ => panic!("unknown child mode"),
604 }
605 }
606
607 #[tokio::test]
608 async fn mapping_enabled_rejects_legacy_raw_statement_before_database_io() {
609 let directory =
610 std::env::temp_dir().join(format!("saddle-db-alpha7-bypass-{}", std::process::id()));
611 let _ = std::fs::remove_dir_all(&directory);
612 std::fs::create_dir(&directory).unwrap();
613 std::fs::write(
614 directory.join("table.json"),
615 r#"{"table":{"from":"逻辑表","to":"physical_table"},"columns":[{"from":"逻辑列","to":"physical_column"}]}"#,
616 )
617 .unwrap();
618 let observer = Observer::with_writer(ObserverConfig::default(), io::sink()).unwrap();
619 let database = Database::connect_lazy(
620 DatabaseConfig::new("mysql://localhost/unused")
621 .name_mapping_directory(&directory)
622 .register_logical_table::<MappedTable>(),
623 observer,
624 )
625 .unwrap();
626 let error = database
627 .write(
628 &context(),
629 Statement::new("legacy.write", "SELECT 1").unwrap(),
630 )
631 .await
632 .unwrap_err();
633 assert_eq!(error.code(), "db.name_mapping_required");
634 database.shutdown().await.unwrap();
635 std::fs::remove_dir_all(directory).unwrap();
636 }
637
638 #[tokio::test]
639 async fn closed_pool_query_has_stable_error_and_trace_record() {
640 let capture = Capture::default();
641 let observer = Observer::with_writer(ObserverConfig::default(), capture.clone()).unwrap();
642 let database = Database::connect_lazy(
643 DatabaseConfig::new("mysql://localhost/db"),
644 observer.clone(),
645 )
646 .unwrap();
647 database.close().await.unwrap();
648 let error = match database
649 .query_all(
650 &context(),
651 Statement::new("orders.list", "SELECT 1").unwrap(),
652 )
653 .await
654 {
655 Ok(_) => panic!("closed pool query unexpectedly succeeded"),
656 Err(error) => error,
657 };
658 assert_eq!(error.kind(), ErrorKind::Unavailable);
659 assert_eq!(error.code(), "db.connection_unavailable");
660 observer.flush().await.unwrap();
661 let output = String::from_utf8(capture.0.lock().unwrap().clone()).unwrap();
662 let records: Vec<Value> = output
663 .lines()
664 .map(|line| serde_json::from_str(line).unwrap())
665 .collect();
666 assert_eq!(records[0]["call_kind"], "database");
667 assert_eq!(records[1]["error_code"], "db.connection_unavailable");
668 assert_eq!(records[0]["trace_id"], context().trace_id().to_string());
669 assert!(!output.contains("SELECT 1"));
670 }
671
672 #[tokio::test]
673 async fn transaction_begin_failure_records_phase_and_stable_error() {
674 let capture = Capture::default();
675 let observer = Observer::with_writer(ObserverConfig::default(), capture.clone()).unwrap();
676 let database = Database::connect_lazy(
677 DatabaseConfig::new("mysql://localhost/db"),
678 observer.clone(),
679 )
680 .unwrap();
681 database.pool.close().await;
682 let error = database
683 .transaction(&context(), "orders.create", |_transaction| {
684 Box::pin(async { Ok::<_, SaddleError>(()) })
685 })
686 .await
687 .unwrap_err();
688 assert_eq!(error.code(), "db.transaction_begin_failed");
689 observer.flush().await.unwrap();
690 let output = String::from_utf8(capture.0.lock().unwrap().clone()).unwrap();
691 let records: Vec<Value> = output
692 .lines()
693 .map(|line| serde_json::from_str(line).unwrap())
694 .collect();
695 assert_eq!(records.len(), 4);
696 assert_eq!(records[0]["operation"], "orders.create");
697 assert_eq!(records[1]["operation"], "begin");
698 assert_eq!(records[2]["error_code"], "db.transaction_begin_failed");
699 assert_eq!(records[3]["error_code"], "db.transaction_begin_failed");
700 assert!(
701 records
702 .iter()
703 .all(|record| record["trace_id"] == context().trace_id().to_string())
704 );
705 }
706
707 #[allow(dead_code)]
708 async fn transaction_usage_compiles(database: &Database, context: &CallContext) -> Result<()> {
709 database
710 .transaction(context, "orders.create", |transaction| {
711 Box::pin(async move {
712 transaction
713 .write(
714 Statement::new("orders.insert", "INSERT INTO orders(id) VALUES (?)")?
715 .bind(1_u64)?,
716 )
717 .await?;
718 transaction
719 .query_optional(
720 Statement::new("orders.find", "SELECT id FROM orders WHERE id = ?")?
721 .bind(1_u64)?,
722 )
723 .await?;
724 Ok(())
725 })
726 })
727 .await
728 }
729}