1use std::collections::HashMap;
4use std::fmt::Write;
5use std::str::FromStr;
6use std::sync::Arc;
7use std::time::{Duration, SystemTime};
8
9use futures::future::{BoxFuture, FutureExt};
10use kube::runtime::events::Recorder;
11use serde::{Deserialize, Serialize};
12
13use sha2::{Digest, Sha256};
14use sqlx::postgres::{PgConnectOptions, PgPool, PgPoolOptions};
15use tokio::sync::{Mutex, RwLock};
16
17use crate::crd::{ConnectionAuth, ConnectionSpec, SecretKeySelector};
18use crate::observability::OperatorObservability;
19use crate::request_index::RequestIndex;
20
21const POOL_MAX_CONNECTIONS: u32 = 5;
26
27const POOL_ACQUIRE_TIMEOUT_SECS: u64 = 10;
30
31const _: () = assert!(POOL_MAX_CONNECTIONS >= 2);
32
33const POOL_IDLE_TIMEOUT_SECS: u64 = 60;
42
43const POOL_MAX_LIFETIME_SECS: u64 = 30 * 60;
48
49const POOL_MIN_CONNECTIONS: u32 = 0;
51
52const GCP_METADATA_TOKEN_ENDPOINT: &str =
53 "http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/token";
54const GCP_IAM_CREDENTIALS_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform";
55const GCP_TOKEN_CACHE_SKEW_SECS: u64 = 300;
56const GCP_IMPERSONATED_TOKEN_LIFETIME_SECS: u64 = 3600;
57const GCP_AUTH_HTTP_TIMEOUT_SECS: u64 = 10;
58
59#[derive(Clone)]
60struct CachedPool {
61 resource_version: Option<String>,
62 secret_fingerprint: Option<String>,
64 token_expires_at: Option<SystemTime>,
66 pool: PgPool,
67}
68
69struct ResolvedConnectionUrl {
70 database_url: String,
71 token_expires_at: Option<SystemTime>,
72 set_role: Option<String>,
76}
77
78pub fn build_set_role_stmt(role: &str) -> String {
85 format!("SET ROLE \"{}\"", role.replace('"', "\"\""))
86}
87
88const SET_ROLE_FAILURE_MARKER: &str = "pgroles:set-role-failed:";
91
92fn wrap_set_role_failure(role: &str, source: sqlx::Error) -> sqlx::Error {
97 sqlx::Error::Protocol(format!("{SET_ROLE_FAILURE_MARKER}{role}: {source}"))
98}
99
100fn classify_pool_connect_error(set_role: Option<&str>, err: sqlx::Error) -> ContextError {
103 if let (Some(role), sqlx::Error::Protocol(msg)) = (set_role, &err)
104 && msg.starts_with(SET_ROLE_FAILURE_MARKER)
105 {
106 return ContextError::SetRoleFailed {
107 role: role.to_string(),
108 source: err,
109 };
110 }
111 ContextError::DatabaseConnect { source: err }
112}
113
114#[derive(Clone)]
115struct GcpAccessToken {
116 token: String,
117 expires_at: SystemTime,
118}
119
120trait GcpAccessTokenProvider: Send + Sync {
121 fn fetch_token<'a>(
122 &'a self,
123 auth: &'a ConnectionAuth,
124 ) -> BoxFuture<'a, Result<GcpAccessToken, ContextError>>;
125}
126
127#[derive(Clone)]
128struct MetadataGcpAccessTokenProvider {
129 client: reqwest::Client,
130}
131
132impl Default for MetadataGcpAccessTokenProvider {
133 fn default() -> Self {
134 Self {
135 client: reqwest::Client::builder()
136 .no_proxy()
137 .timeout(Duration::from_secs(GCP_AUTH_HTTP_TIMEOUT_SECS))
138 .build()
139 .expect("GCP auth HTTP client should build"),
140 }
141 }
142}
143
144impl GcpAccessTokenProvider for MetadataGcpAccessTokenProvider {
145 fn fetch_token<'a>(
146 &'a self,
147 auth: &'a ConnectionAuth,
148 ) -> BoxFuture<'a, Result<GcpAccessToken, ContextError>> {
149 async move {
150 let scope = auth.gcp_scope();
151 if let Some(target) = auth.gcp_impersonate_service_account() {
152 self.fetch_impersonated_access_token(target, scope).await
153 } else {
154 self.fetch_metadata_access_token(scope).await
155 }
156 }
157 .boxed()
158 }
159}
160
161impl MetadataGcpAccessTokenProvider {
162 async fn fetch_metadata_access_token(
163 &self,
164 scope: &str,
165 ) -> Result<GcpAccessToken, ContextError> {
166 let response = self
167 .client
168 .get(GCP_METADATA_TOKEN_ENDPOINT)
169 .header("Metadata-Flavor", "Google")
170 .query(&[("scopes", scope)])
171 .send()
172 .await
173 .map_err(|source| ContextError::GcpAuthHttp {
174 endpoint: "metadata",
175 source,
176 })?;
177
178 let status = response.status();
179 if !status.is_success() {
180 let body = response_body_for_error(response).await;
181 return Err(ContextError::GcpAuthRejected {
182 endpoint: "metadata".to_string(),
183 status: status.as_u16(),
184 body,
185 });
186 }
187
188 let body: MetadataTokenResponse =
189 response
190 .json()
191 .await
192 .map_err(|source| ContextError::GcpAuthHttp {
193 endpoint: "metadata",
194 source,
195 })?;
196
197 if body.access_token.trim().is_empty() {
198 return Err(ContextError::GcpAuthInvalidResponse {
199 detail: "metadata token response omitted access_token".to_string(),
200 });
201 }
202 if body.expires_in == 0 {
203 return Err(ContextError::GcpAuthInvalidResponse {
204 detail: "metadata token response had zero expires_in".to_string(),
205 });
206 }
207
208 Ok(GcpAccessToken {
209 token: body.access_token,
210 expires_at: SystemTime::now() + Duration::from_secs(body.expires_in),
211 })
212 }
213
214 async fn fetch_impersonated_access_token(
215 &self,
216 target_service_account: &str,
217 scope: &str,
218 ) -> Result<GcpAccessToken, ContextError> {
219 let source = self
220 .fetch_metadata_access_token(GCP_IAM_CREDENTIALS_SCOPE)
221 .await?;
222 let encoded_target = percent_encoding::utf8_percent_encode(
223 target_service_account,
224 percent_encoding::NON_ALPHANUMERIC,
225 )
226 .to_string();
227 let endpoint = format!(
228 "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/{encoded_target}:generateAccessToken"
229 );
230 let request = GenerateAccessTokenRequest {
231 scope: vec![scope.to_string()],
232 lifetime: format!("{GCP_IMPERSONATED_TOKEN_LIFETIME_SECS}s"),
233 };
234
235 let response = self
236 .client
237 .post(&endpoint)
238 .bearer_auth(&source.token)
239 .json(&request)
240 .send()
241 .await
242 .map_err(|source| ContextError::GcpAuthHttp {
243 endpoint: "iamcredentials",
244 source,
245 })?;
246
247 let status = response.status();
248 if !status.is_success() {
249 let body = response_body_for_error(response).await;
250 return Err(ContextError::GcpAuthRejected {
251 endpoint: "iamcredentials".to_string(),
252 status: status.as_u16(),
253 body,
254 });
255 }
256
257 let body: GenerateAccessTokenResponse =
258 response
259 .json()
260 .await
261 .map_err(|source| ContextError::GcpAuthHttp {
262 endpoint: "iamcredentials",
263 source,
264 })?;
265
266 if body.access_token.trim().is_empty() {
267 return Err(ContextError::GcpAuthInvalidResponse {
268 detail: "IAMCredentials response omitted accessToken".to_string(),
269 });
270 }
271 let expires_at = parse_google_expire_time(&body.expire_time).ok_or_else(|| {
272 ContextError::GcpAuthInvalidResponse {
273 detail: format!(
274 "IAMCredentials response had invalid expireTime {:?}",
275 body.expire_time
276 ),
277 }
278 })?;
279
280 Ok(GcpAccessToken {
281 token: body.access_token,
282 expires_at,
283 })
284 }
285}
286
287#[derive(Deserialize)]
288struct MetadataTokenResponse {
289 access_token: String,
290 expires_in: u64,
291}
292
293#[derive(Serialize)]
294struct GenerateAccessTokenRequest {
295 scope: Vec<String>,
296 lifetime: String,
297}
298
299#[derive(Deserialize)]
300struct GenerateAccessTokenResponse {
301 #[serde(rename = "accessToken")]
302 access_token: String,
303 #[serde(rename = "expireTime")]
304 expire_time: String,
305}
306
307async fn response_body_for_error(response: reqwest::Response) -> String {
308 match response.text().await {
309 Ok(body) => truncate_for_error(body),
310 Err(error) => format!("failed to read error body: {error}"),
311 }
312}
313
314fn truncate_for_error(mut body: String) -> String {
315 const MAX_ERROR_BODY_BYTES: usize = 512;
316 if body.len() <= MAX_ERROR_BODY_BYTES {
317 return body;
318 }
319 let mut end = MAX_ERROR_BODY_BYTES;
320 while !body.is_char_boundary(end) {
321 end -= 1;
322 }
323 body.truncate(end);
324 body.push_str("...");
325 body
326}
327
328fn parse_google_expire_time(expire_time: &str) -> Option<SystemTime> {
329 expire_time
330 .parse::<jiff::Timestamp>()
331 .ok()
332 .map(SystemTime::from)
333}
334
335fn token_expires_after_skew(expires_at: Option<SystemTime>, now: SystemTime) -> bool {
336 let Some(expires_at) = expires_at else {
337 return true;
338 };
339 let Some(refresh_at) = now.checked_add(Duration::from_secs(GCP_TOKEN_CACHE_SKEW_SECS)) else {
340 return false;
341 };
342 expires_at > refresh_at
343}
344
345pub struct DatabaseLockGuard {
351 key: String,
352 locks: Arc<Mutex<HashMap<String, ()>>>,
353}
354
355impl Drop for DatabaseLockGuard {
356 fn drop(&mut self) {
357 if let Ok(mut map) = self.locks.try_lock() {
359 map.remove(&self.key);
360 tracing::debug!(database = %self.key, "released in-memory database lock");
361 } else {
362 let key = self.key.clone();
366 let locks = Arc::clone(&self.locks);
367 if let Ok(handle) = tokio::runtime::Handle::try_current() {
368 handle.spawn(async move {
369 locks.lock().await.remove(&key);
370 tracing::debug!(database = %key, "released in-memory database lock (deferred)");
371 });
372 tracing::debug!(
373 database = %self.key,
374 "deferred in-memory database lock release to background task"
375 );
376 } else {
377 let mut map = self.locks.blocking_lock();
380 map.remove(&key);
381 tracing::debug!(
382 database = %key,
383 "released in-memory database lock (fallback sync)"
384 );
385 }
386 }
387 }
388}
389
390#[derive(Clone)]
392pub struct OperatorContext {
393 pub kube_client: kube::Client,
395
396 pub event_recorder: Recorder,
398
399 pool_cache: Arc<RwLock<HashMap<String, CachedPool>>>,
401 database_locks: Arc<Mutex<HashMap<String, ()>>>,
407
408 pub observability: OperatorObservability,
410
411 pub request_index: RequestIndex,
413
414 pub watch_namespace: Option<String>,
416
417 gcp_token_provider: Arc<dyn GcpAccessTokenProvider>,
419}
420
421impl OperatorContext {
422 pub fn new_with_runtime_config(
425 kube_client: kube::Client,
426 observability: OperatorObservability,
427 event_recorder: Recorder,
428 request_index: RequestIndex,
429 watch_namespace: Option<String>,
430 ) -> Self {
431 Self {
432 kube_client,
433 event_recorder,
434 pool_cache: Arc::new(RwLock::new(HashMap::new())),
435 observability,
436 request_index,
437 watch_namespace,
438 database_locks: Arc::new(Mutex::new(HashMap::new())),
439 gcp_token_provider: Arc::new(MetadataGcpAccessTokenProvider::default()),
440 }
441 }
442
443 pub async fn try_lock_database(&self, database_identity: &str) -> Option<DatabaseLockGuard> {
449 let mut locks = self.database_locks.lock().await;
450 if locks.contains_key(database_identity) {
451 tracing::info!(
452 database = %database_identity,
453 "in-memory database lock contention — another reconcile is in progress"
454 );
455 return None;
456 }
457 locks.insert(database_identity.to_string(), ());
458 tracing::debug!(database = %database_identity, "acquired in-memory database lock");
459 Some(DatabaseLockGuard {
460 key: database_identity.to_string(),
461 locks: Arc::clone(&self.database_locks),
462 })
463 }
464
465 async fn resolve_param(
469 &self,
470 namespace: &str,
471 literal: &Option<String>,
472 secret: &Option<SecretKeySelector>,
473 ) -> Result<Option<String>, ContextError> {
474 if let Some(val) = literal {
475 return Ok(Some(val.clone()));
476 }
477 if let Some(sel) = secret {
478 return Ok(Some(
479 self.fetch_secret_value(namespace, &sel.name, &sel.key)
480 .await?,
481 ));
482 }
483 Ok(None)
484 }
485
486 pub async fn resolve_connection_url(
491 &self,
492 namespace: &str,
493 connection: &ConnectionSpec,
494 ) -> Result<String, ContextError> {
495 Ok(self
496 .resolve_connection_url_with_metadata(namespace, connection)
497 .await?
498 .database_url)
499 }
500
501 pub async fn resolve_database_target_fingerprint(
508 &self,
509 namespace: &str,
510 connection: &ConnectionSpec,
511 ) -> Result<String, ContextError> {
512 let (host, port, database) = if let Some(ref secret_ref) = connection.secret_ref {
513 let database_url = self
514 .fetch_secret_value(
515 namespace,
516 &secret_ref.name,
517 connection.effective_secret_key(),
518 )
519 .await?;
520 database_target_from_url(&database_url)
521 .map_err(|detail| ContextError::InvalidDatabaseUrl { detail })?
522 } else if let Some(ref params) = connection.params {
523 let host = self
524 .resolve_param(namespace, ¶ms.host, ¶ms.host_secret)
525 .await?
526 .ok_or_else(|| ContextError::EmptyResolvedValue {
527 field: "host".to_string(),
528 })?;
529 let port = match self
530 .resolve_param(
531 namespace,
532 ¶ms.port.map(|port| port.to_string()),
533 ¶ms.port_secret,
534 )
535 .await?
536 {
537 Some(value) => {
538 value
539 .parse::<u16>()
540 .map_err(|_| ContextError::InvalidResolvedPort {
541 value: value.clone(),
542 })?
543 }
544 None => 5432,
545 };
546 let database = self
547 .resolve_param(namespace, ¶ms.dbname, ¶ms.dbname_secret)
548 .await?
549 .ok_or_else(|| ContextError::EmptyResolvedValue {
550 field: "dbname".to_string(),
551 })?;
552 (host.to_ascii_lowercase(), port, database)
553 } else {
554 return Err(ContextError::SecretMissing {
555 name: "connection".to_string(),
556 key: "neither secretRef nor params is set".to_string(),
557 });
558 };
559
560 if host.trim().is_empty() {
561 return Err(ContextError::EmptyResolvedValue {
562 field: "host".to_string(),
563 });
564 }
565 if database.trim().is_empty() {
566 return Err(ContextError::EmptyResolvedValue {
567 field: "dbname".to_string(),
568 });
569 }
570 Ok(database_target_fingerprint(&host, port, &database))
571 }
572
573 async fn resolve_connection_url_with_metadata(
574 &self,
575 namespace: &str,
576 connection: &ConnectionSpec,
577 ) -> Result<ResolvedConnectionUrl, ContextError> {
578 if let Some(ref secret_ref) = connection.secret_ref {
579 let database_url = self
581 .fetch_secret_value(
582 namespace,
583 &secret_ref.name,
584 connection.effective_secret_key(),
585 )
586 .await?;
587 Ok(ResolvedConnectionUrl {
588 database_url,
589 token_expires_at: None,
590 set_role: None,
591 })
592 } else if let Some(ref params) = connection.params {
593 let host = self
595 .resolve_param(namespace, ¶ms.host, ¶ms.host_secret)
596 .await?
597 .ok_or_else(|| ContextError::EmptyResolvedValue {
598 field: "host".to_string(),
599 })?;
600 if host.trim().is_empty() {
601 return Err(ContextError::EmptyResolvedValue {
602 field: "host".to_string(),
603 });
604 }
605
606 let port_str = params.port.map(|p| p.to_string());
607 let port = self
608 .resolve_param(namespace, &port_str, ¶ms.port_secret)
609 .await?
610 .unwrap_or_else(|| "5432".to_string());
611 if port.trim().is_empty() {
612 return Err(ContextError::EmptyResolvedValue {
613 field: "port".to_string(),
614 });
615 }
616
617 let dbname = self
618 .resolve_param(namespace, ¶ms.dbname, ¶ms.dbname_secret)
619 .await?
620 .ok_or_else(|| ContextError::EmptyResolvedValue {
621 field: "dbname".to_string(),
622 })?;
623 if dbname.trim().is_empty() {
624 return Err(ContextError::EmptyResolvedValue {
625 field: "dbname".to_string(),
626 });
627 }
628
629 let username = self
630 .resolve_param(namespace, ¶ms.username, ¶ms.username_secret)
631 .await?
632 .ok_or_else(|| ContextError::EmptyResolvedValue {
633 field: "username".to_string(),
634 })?;
635 if username.trim().is_empty() {
636 return Err(ContextError::EmptyResolvedValue {
637 field: "username".to_string(),
638 });
639 }
640
641 let (password, token_expires_at) = if let Some(auth) = ¶ms.auth {
642 let token = self.gcp_token_provider.fetch_token(auth).await?;
643 (token.token, Some(token.expires_at))
644 } else {
645 let password = self
646 .resolve_param(namespace, ¶ms.password, ¶ms.password_secret)
647 .await?
648 .ok_or_else(|| ContextError::EmptyResolvedValue {
649 field: "password".to_string(),
650 })?;
651 (password, None)
652 };
653 if password.trim().is_empty() {
654 return Err(ContextError::EmptyResolvedValue {
655 field: "password".to_string(),
656 });
657 }
658
659 use percent_encoding::{NON_ALPHANUMERIC, utf8_percent_encode};
660 let encoded_username = utf8_percent_encode(&username, NON_ALPHANUMERIC).to_string();
661 let encoded_password = utf8_percent_encode(&password, NON_ALPHANUMERIC).to_string();
662
663 let mut url = format!(
664 "postgresql://{encoded_username}:{encoded_password}@{host}:{port}/{dbname}"
665 );
666
667 let ssl_mode = self
668 .resolve_param(namespace, ¶ms.ssl_mode, ¶ms.ssl_mode_secret)
669 .await?
670 .or_else(|| params.auth.as_ref().map(|_| "require".to_string()));
671 if let Some(ssl_mode) = ssl_mode {
672 if !crate::crd::VALID_SSL_MODES.contains(&ssl_mode.as_str()) {
675 return Err(ContextError::InvalidResolvedSslMode { value: ssl_mode });
676 }
677 url.push_str("?sslmode=");
678 url.push_str(&ssl_mode);
679 }
680
681 Ok(ResolvedConnectionUrl {
682 database_url: url,
683 token_expires_at,
684 set_role: params.set_role.clone(),
685 })
686 } else {
687 Err(ContextError::SecretMissing {
688 name: "connection".to_string(),
689 key: "neither secretRef nor params is set".to_string(),
690 })
691 }
692 }
693
694 pub async fn get_or_create_pool(
699 &self,
700 namespace: &str,
701 connection: &ConnectionSpec,
702 ) -> Result<PgPool, ContextError> {
703 let cache_key = connection.cache_key(namespace);
704
705 let (resource_version, secret_fingerprint) =
709 if let Some(ref secret_ref) = connection.secret_ref {
710 let secrets_api: kube::Api<k8s_openapi::api::core::v1::Secret> =
711 kube::Api::namespaced(self.kube_client.clone(), namespace);
712 let secret = secrets_api.get(&secret_ref.name).await.map_err(|err| {
713 ContextError::SecretFetch {
714 name: secret_ref.name.clone(),
715 namespace: namespace.to_string(),
716 source: err,
717 }
718 })?;
719 (secret.metadata.resource_version, None)
720 } else if connection.params.is_some() {
721 let mut secret_names = std::collections::BTreeSet::new();
724 connection.collect_secret_names(&mut secret_names);
725
726 if secret_names.is_empty() {
727 (None, Some(String::new()))
729 } else {
730 let secrets_api: kube::Api<k8s_openapi::api::core::v1::Secret> =
731 kube::Api::namespaced(self.kube_client.clone(), namespace);
732 let mut fingerprint_parts = Vec::new();
733 for name in &secret_names {
734 let secret = secrets_api.get(name).await.map_err(|err| {
735 ContextError::SecretFetch {
736 name: name.clone(),
737 namespace: namespace.to_string(),
738 source: err,
739 }
740 })?;
741 let rv = secret
742 .metadata
743 .resource_version
744 .unwrap_or_else(|| "unknown".to_string());
745 fingerprint_parts.push(format!("{name}={rv}"));
746 }
747 (None, Some(fingerprint_parts.join(",")))
748 }
749 } else {
750 (None, None)
751 };
752
753 {
755 let cache = self.pool_cache.read().await;
756 if let Some(cached) = cache.get(&cache_key) {
757 let version_matches = match (&resource_version, &cached.resource_version) {
760 (Some(current), Some(cached_rv)) => current == cached_rv,
761 _ => true,
762 };
763 let fingerprint_matches = match (&secret_fingerprint, &cached.secret_fingerprint) {
764 (Some(current), Some(cached_fp)) => current == cached_fp,
765 (None, None) => true,
766 _ => false,
767 };
768 let token_fresh =
769 token_expires_after_skew(cached.token_expires_at, SystemTime::now());
770 if version_matches && fingerprint_matches && token_fresh {
771 return Ok(cached.pool.clone());
772 }
773 }
774 }
775
776 let resolved = self
777 .resolve_connection_url_with_metadata(namespace, connection)
778 .await?;
779
780 let set_role = resolved.set_role.clone();
784 let pool = pool_options()
785 .after_connect(move |conn, _meta| {
786 let set_role = set_role.clone();
787 Box::pin(async move {
788 if let Some(role) = set_role {
789 let stmt = build_set_role_stmt(&role);
790 sqlx::Executor::execute(&mut *conn, stmt.as_str())
791 .await
792 .map_err(|err| wrap_set_role_failure(&role, err))?;
793 }
794 Ok(())
795 })
796 })
797 .connect(&resolved.database_url)
798 .await
799 .map_err(|err| classify_pool_connect_error(resolved.set_role.as_deref(), err))?;
800
801 let superseded = {
803 let mut cache = self.pool_cache.write().await;
804 cache.insert(
805 cache_key,
806 CachedPool {
807 resource_version,
808 secret_fingerprint,
809 token_expires_at: resolved.token_expires_at,
810 pool: pool.clone(),
811 },
812 )
813 };
814 close_unreachable_pool(superseded);
815
816 Ok(pool)
817 }
818
819 pub async fn fetch_secret_value(
823 &self,
824 namespace: &str,
825 secret_name: &str,
826 secret_key: &str,
827 ) -> Result<String, ContextError> {
828 let secrets_api: kube::Api<k8s_openapi::api::core::v1::Secret> =
829 kube::Api::namespaced(self.kube_client.clone(), namespace);
830
831 let secret =
832 secrets_api
833 .get(secret_name)
834 .await
835 .map_err(|err| ContextError::SecretFetch {
836 name: secret_name.to_string(),
837 namespace: namespace.to_string(),
838 source: err,
839 })?;
840
841 let data = secret.data.ok_or_else(|| ContextError::SecretMissing {
842 name: secret_name.to_string(),
843 key: secret_key.to_string(),
844 })?;
845
846 let value_bytes = data
847 .get(secret_key)
848 .ok_or_else(|| ContextError::SecretMissing {
849 name: secret_name.to_string(),
850 key: secret_key.to_string(),
851 })?;
852
853 String::from_utf8(value_bytes.0.clone()).map_err(|_| ContextError::SecretMissing {
854 name: secret_name.to_string(),
855 key: secret_key.to_string(),
856 })
857 }
858
859 pub async fn evict_pool(&self, namespace: &str, connection: &ConnectionSpec) {
861 let cache_key = connection.cache_key(namespace);
862 let evicted = {
863 let mut cache = self.pool_cache.write().await;
864 cache.remove(&cache_key)
865 };
866 close_unreachable_pool(evicted);
867 }
868}
869
870fn pool_options() -> PgPoolOptions {
872 PgPoolOptions::new()
873 .max_connections(POOL_MAX_CONNECTIONS)
874 .min_connections(POOL_MIN_CONNECTIONS)
875 .acquire_timeout(Duration::from_secs(POOL_ACQUIRE_TIMEOUT_SECS))
876 .idle_timeout(Duration::from_secs(POOL_IDLE_TIMEOUT_SECS))
877 .max_lifetime(Duration::from_secs(POOL_MAX_LIFETIME_SECS))
878}
879
880fn close_unreachable_pool(cached: Option<CachedPool>) {
888 let Some(cached) = cached else {
889 return;
890 };
891 tokio::spawn(async move {
892 cached.pool.close().await;
893 });
894}
895
896fn database_target_from_url(database_url: &str) -> Result<(String, u16, String), String> {
897 let options = PgConnectOptions::from_str(database_url).map_err(|source| source.to_string())?;
898 Ok((
899 options.get_host().to_ascii_lowercase(),
900 options.get_port(),
901 options
902 .get_database()
903 .unwrap_or_else(|| options.get_username())
904 .to_string(),
905 ))
906}
907
908fn database_target_fingerprint(host: &str, port: u16, database: &str) -> String {
909 let digest = Sha256::digest(format!("{}\0{port}\0{database}", host.to_ascii_lowercase()));
910 let mut fingerprint = String::with_capacity(7 + digest.len() * 2);
911 fingerprint.push_str("sha256:");
912 for byte in digest {
913 write!(&mut fingerprint, "{byte:02x}")
914 .expect("writing a database target fingerprint cannot fail");
915 }
916 fingerprint
917}
918
919#[derive(Debug, thiserror::Error)]
921pub enum ContextError {
922 #[error("failed to fetch Secret {namespace}/{name}: {source}")]
923 SecretFetch {
924 name: String,
925 namespace: String,
926 source: kube::Error,
927 },
928
929 #[error("Secret \"{name}\" does not contain key \"{key}\"")]
930 SecretMissing { name: String, key: String },
931
932 #[error("failed to connect to database: {source}")]
933 DatabaseConnect { source: sqlx::Error },
934
935 #[error("failed to apply SET ROLE \"{role}\" on pooled connection: {source}")]
936 SetRoleFailed { role: String, source: sqlx::Error },
937
938 #[error("connection param \"{field}\" resolved to an empty or whitespace-only value")]
939 EmptyResolvedValue { field: String },
940
941 #[error("connection URL is invalid: {detail}")]
942 InvalidDatabaseUrl { detail: String },
943
944 #[error("connection port resolved to invalid value \"{value}\"")]
945 InvalidResolvedPort { value: String },
946
947 #[error(
948 "connection param sslMode resolved to invalid value \"{value}\" (expected one of: disable, allow, prefer, require, verify-ca, verify-full)"
949 )]
950 InvalidResolvedSslMode { value: String },
951
952 #[error("failed to fetch GCP auth token from {endpoint}: {source}")]
953 GcpAuthHttp {
954 endpoint: &'static str,
955 source: reqwest::Error,
956 },
957
958 #[error("GCP auth token endpoint {endpoint} returned HTTP {status}: {body}")]
959 GcpAuthRejected {
960 endpoint: String,
961 status: u16,
962 body: String,
963 },
964
965 #[error("GCP auth token response was invalid: {detail}")]
966 GcpAuthInvalidResponse { detail: String },
967}
968
969impl ContextError {
970 pub fn is_secret_fetch_non_transient(&self) -> bool {
972 matches!(
973 self,
974 ContextError::SecretFetch {
975 source: kube::Error::Api(response),
976 ..
977 } if (400..500).contains(&response.code) && response.code != 429
978 )
979 }
980
981 pub fn is_gcp_auth_non_transient(&self) -> bool {
982 matches!(
983 self,
984 ContextError::GcpAuthRejected { status, .. }
985 if (400..500).contains(status) && *status != 429
986 ) || matches!(self, ContextError::GcpAuthInvalidResponse { .. })
987 }
988}
989
990#[cfg(test)]
991mod tests {
992 use super::*;
993
994 #[test]
1001 fn idle_connections_are_reaped_between_reconciles() {
1002 let options = pool_options();
1003
1004 assert_eq!(
1005 options.get_min_connections(),
1006 0,
1007 "a floor above zero would hold connections open on an idle database"
1008 );
1009
1010 let idle_timeout = options
1011 .get_idle_timeout()
1012 .expect("an unset idle timeout never reaps");
1013 assert!(
1014 idle_timeout < Duration::from_secs(crate::reconciler::DEFAULT_REQUEUE_SECS),
1015 "idle timeout {idle_timeout:?} must drain within the {}s requeue interval",
1016 crate::reconciler::DEFAULT_REQUEUE_SECS
1017 );
1018
1019 assert!(
1020 options.get_max_lifetime().is_some(),
1021 "connections need an age ceiling independent of activity"
1022 );
1023 }
1024
1025 #[test]
1026 fn database_target_fingerprint_excludes_credentials_and_options() {
1027 let first = database_target_from_url(
1028 "postgresql://alice:first@DB.EXAMPLE:6432/inventory?sslmode=require",
1029 )
1030 .expect("first URL");
1031 let second = database_target_from_url(
1032 "postgresql://bob:second@db.example:6432/inventory?application_name=test",
1033 )
1034 .expect("second URL");
1035
1036 assert_eq!(
1037 database_target_fingerprint(&first.0, first.1, &first.2),
1038 database_target_fingerprint(&second.0, second.1, &second.2)
1039 );
1040 }
1041
1042 #[test]
1043 fn database_target_fingerprint_changes_on_retarget() {
1044 let original = database_target_fingerprint("db.example", 5432, "inventory");
1045 assert_ne!(
1046 original,
1047 database_target_fingerprint("db.example", 5432, "billing")
1048 );
1049 assert_ne!(
1050 original,
1051 database_target_fingerprint("other.example", 5432, "inventory")
1052 );
1053 }
1054
1055 #[test]
1056 fn build_set_role_stmt_quotes_identifier() {
1057 assert_eq!(
1058 build_set_role_stmt("cloudsqlsuperuser"),
1059 "SET ROLE \"cloudsqlsuperuser\"",
1060 );
1061 }
1062
1063 #[test]
1064 fn classify_pool_connect_error_surfaces_set_role_failure_with_role() {
1065 let raw = wrap_set_role_failure(
1066 "cloudsqlsuperuser",
1067 sqlx::Error::Protocol("permission denied".to_string()),
1068 );
1069 let classified = classify_pool_connect_error(Some("cloudsqlsuperuser"), raw);
1070 assert!(matches!(
1071 classified,
1072 ContextError::SetRoleFailed { ref role, .. } if role == "cloudsqlsuperuser"
1073 ));
1074 }
1075
1076 #[test]
1077 fn classify_pool_connect_error_passes_through_unrelated_errors() {
1078 let err = sqlx::Error::PoolTimedOut;
1079 let classified = classify_pool_connect_error(Some("any_role"), err);
1080 assert!(matches!(classified, ContextError::DatabaseConnect { .. }));
1081 }
1082
1083 #[test]
1084 fn classify_pool_connect_error_without_set_role_is_database_connect() {
1085 let raw = wrap_set_role_failure("ghost", sqlx::Error::Protocol("oops".to_string()));
1088 let classified = classify_pool_connect_error(None, raw);
1089 assert!(matches!(classified, ContextError::DatabaseConnect { .. }));
1090 }
1091
1092 #[test]
1093 fn build_set_role_stmt_doubles_embedded_quote() {
1094 assert_eq!(build_set_role_stmt("a\"b"), "SET ROLE \"a\"\"b\"",);
1098 }
1099
1100 #[test]
1101 fn pool_cache_key_format() {
1102 let key = format!("{}/{}/{}", "prod", "pg-credentials", "DATABASE_URL");
1104 assert_eq!(key, "prod/pg-credentials/DATABASE_URL");
1105 }
1106
1107 #[test]
1108 fn secret_fetch_not_found_is_non_transient() {
1109 let error = ContextError::SecretFetch {
1110 name: "db-credentials".into(),
1111 namespace: "default".into(),
1112 source: kube::Error::Api(
1113 kube::core::Status::failure("secrets \"db-credentials\" not found", "NotFound")
1114 .with_code(404)
1115 .boxed(),
1116 ),
1117 };
1118
1119 assert!(error.is_secret_fetch_non_transient());
1120 }
1121
1122 #[test]
1123 fn secret_fetch_forbidden_is_non_transient() {
1124 let error = ContextError::SecretFetch {
1125 name: "db-credentials".into(),
1126 namespace: "default".into(),
1127 source: kube::Error::Api(
1128 kube::core::Status::failure("forbidden", "Forbidden")
1129 .with_code(403)
1130 .boxed(),
1131 ),
1132 };
1133
1134 assert!(error.is_secret_fetch_non_transient());
1135 }
1136
1137 #[test]
1138 fn secret_fetch_server_error_remains_transient() {
1139 let error = ContextError::SecretFetch {
1140 name: "db-credentials".into(),
1141 namespace: "default".into(),
1142 source: kube::Error::Api(
1143 kube::core::Status::failure("internal error", "InternalError")
1144 .with_code(500)
1145 .boxed(),
1146 ),
1147 };
1148
1149 assert!(!error.is_secret_fetch_non_transient());
1150 }
1151
1152 #[test]
1153 fn gcp_auth_client_error_is_non_transient() {
1154 let error = ContextError::GcpAuthRejected {
1155 endpoint: "metadata".into(),
1156 status: 403,
1157 body: "forbidden".into(),
1158 };
1159
1160 assert!(error.is_gcp_auth_non_transient());
1161 }
1162
1163 #[test]
1164 fn gcp_auth_rate_limit_remains_transient() {
1165 let error = ContextError::GcpAuthRejected {
1166 endpoint: "metadata".into(),
1167 status: 429,
1168 body: "rate limited".into(),
1169 };
1170
1171 assert!(!error.is_gcp_auth_non_transient());
1172 }
1173
1174 #[test]
1175 fn token_expiry_uses_five_minute_refresh_skew() {
1176 let now = SystemTime::UNIX_EPOCH + Duration::from_secs(1_000);
1177 assert!(token_expires_after_skew(
1178 Some(now + Duration::from_secs(GCP_TOKEN_CACHE_SKEW_SECS + 1)),
1179 now
1180 ));
1181 assert!(!token_expires_after_skew(
1182 Some(now + Duration::from_secs(GCP_TOKEN_CACHE_SKEW_SECS)),
1183 now
1184 ));
1185 }
1186
1187 #[test]
1188 fn parse_google_expire_time_accepts_rfc3339() {
1189 let parsed =
1190 parse_google_expire_time("2026-05-14T02:30:00Z").expect("expireTime should parse");
1191 assert_eq!(
1192 parsed
1193 .duration_since(SystemTime::UNIX_EPOCH)
1194 .unwrap()
1195 .as_secs(),
1196 1_778_725_800
1197 );
1198 }
1199
1200 #[test]
1201 fn truncate_for_error_keeps_utf8_boundary() {
1202 let body = "é".repeat(300);
1203 let truncated = truncate_for_error(body);
1204
1205 assert!(truncated.ends_with("..."));
1206 assert!(truncated.is_char_boundary(truncated.len() - 3));
1207 }
1208
1209 #[tokio::test]
1210 async fn try_lock_database_acquires_when_free() {
1211 let locks: Arc<Mutex<HashMap<String, ()>>> = Arc::new(Mutex::new(HashMap::new()));
1212 let ctx = OperatorContextLockHelper {
1213 database_locks: locks,
1214 };
1215 let guard = ctx.try_lock("db-a").await;
1216 assert!(guard.is_some(), "should acquire lock on free database");
1217 }
1218
1219 #[tokio::test]
1220 async fn try_lock_database_contention_returns_none() {
1221 let locks: Arc<Mutex<HashMap<String, ()>>> = Arc::new(Mutex::new(HashMap::new()));
1222 let ctx = OperatorContextLockHelper {
1223 database_locks: locks,
1224 };
1225
1226 let _guard1 = ctx
1227 .try_lock("db-a")
1228 .await
1229 .expect("first lock should succeed");
1230 let guard2 = ctx.try_lock("db-a").await;
1231 assert!(guard2.is_none(), "second lock on same database should fail");
1232 }
1233
1234 #[tokio::test]
1235 async fn try_lock_database_different_databases_independent() {
1236 let locks: Arc<Mutex<HashMap<String, ()>>> = Arc::new(Mutex::new(HashMap::new()));
1237 let ctx = OperatorContextLockHelper {
1238 database_locks: locks,
1239 };
1240
1241 let guard_a = ctx.try_lock("db-a").await;
1242 let guard_b = ctx.try_lock("db-b").await;
1243 assert!(guard_a.is_some(), "lock on db-a should succeed");
1244 assert!(
1245 guard_b.is_some(),
1246 "lock on db-b should succeed (different database)"
1247 );
1248 }
1249
1250 #[tokio::test]
1251 async fn try_lock_database_released_after_drop() {
1252 let locks: Arc<Mutex<HashMap<String, ()>>> = Arc::new(Mutex::new(HashMap::new()));
1253 let ctx = OperatorContextLockHelper {
1254 database_locks: Arc::clone(&locks),
1255 };
1256
1257 {
1258 let _guard = ctx.try_lock("db-a").await.expect("should acquire");
1259 }
1261
1262 let guard2 = ctx.try_lock("db-a").await;
1264 assert!(
1265 guard2.is_some(),
1266 "should re-acquire after previous guard dropped"
1267 );
1268 }
1269
1270 #[tokio::test]
1271 async fn try_lock_database_concurrent_contention() {
1272 let locks: Arc<Mutex<HashMap<String, ()>>> = Arc::new(Mutex::new(HashMap::new()));
1273
1274 let locks1 = Arc::clone(&locks);
1276 let locks2 = Arc::clone(&locks);
1277
1278 let handle1 = tokio::spawn(async move {
1279 let ctx = OperatorContextLockHelper {
1280 database_locks: locks1,
1281 };
1282 let guard = ctx.try_lock("shared-db").await;
1283 if guard.is_some() {
1284 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1286 }
1287 guard.is_some()
1288 });
1289
1290 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
1292
1293 let handle2 = tokio::spawn(async move {
1294 let ctx = OperatorContextLockHelper {
1295 database_locks: locks2,
1296 };
1297 let guard = ctx.try_lock("shared-db").await;
1298 guard.is_some()
1299 });
1300
1301 let (r1, r2) = tokio::join!(handle1, handle2);
1302 let acquired1 = r1.unwrap();
1303 let acquired2 = r2.unwrap();
1304
1305 assert!(
1307 acquired1 ^ acquired2,
1308 "exactly one of two concurrent locks should succeed: got ({acquired1}, {acquired2})"
1309 );
1310 }
1311
1312 struct OperatorContextLockHelper {
1314 database_locks: Arc<Mutex<HashMap<String, ()>>>,
1315 }
1316
1317 impl OperatorContextLockHelper {
1318 async fn try_lock(&self, database_identity: &str) -> Option<DatabaseLockGuard> {
1319 let mut locks = self.database_locks.lock().await;
1320 if locks.contains_key(database_identity) {
1321 return None;
1322 }
1323 locks.insert(database_identity.to_string(), ());
1324 Some(DatabaseLockGuard {
1325 key: database_identity.to_string(),
1326 locks: Arc::clone(&self.database_locks),
1327 })
1328 }
1329 }
1330
1331 #[tokio::test]
1332 async fn try_lock_database_high_concurrency_same_db() {
1333 let locks: Arc<Mutex<HashMap<String, ()>>> = Arc::new(Mutex::new(HashMap::new()));
1335 let concurrency = 50;
1336 let acquired_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
1337 let barrier = Arc::new(tokio::sync::Barrier::new(concurrency));
1338
1339 let mut handles = Vec::with_capacity(concurrency);
1340 for _ in 0..concurrency {
1341 let locks_clone = Arc::clone(&locks);
1342 let count = Arc::clone(&acquired_count);
1343 let bar = Arc::clone(&barrier);
1344 handles.push(tokio::spawn(async move {
1345 bar.wait().await;
1347 let ctx = OperatorContextLockHelper {
1348 database_locks: locks_clone,
1349 };
1350 let guard = ctx.try_lock("contested-db").await;
1351 if guard.is_some() {
1352 count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1353 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
1355 }
1356 }));
1357 }
1358
1359 for h in handles {
1360 h.await.unwrap();
1361 }
1362
1363 let total = acquired_count.load(std::sync::atomic::Ordering::SeqCst);
1365 assert_eq!(
1366 total, 1,
1367 "exactly one of {concurrency} concurrent tasks should acquire the lock, got {total}"
1368 );
1369 }
1370
1371 #[tokio::test]
1372 async fn try_lock_database_high_concurrency_different_dbs() {
1373 let locks: Arc<Mutex<HashMap<String, ()>>> = Arc::new(Mutex::new(HashMap::new()));
1375 let concurrency = 50;
1376 let acquired_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
1377 let barrier = Arc::new(tokio::sync::Barrier::new(concurrency));
1378
1379 let mut handles = Vec::with_capacity(concurrency);
1380 for i in 0..concurrency {
1381 let locks_clone = Arc::clone(&locks);
1382 let count = Arc::clone(&acquired_count);
1383 let bar = Arc::clone(&barrier);
1384 handles.push(tokio::spawn(async move {
1385 bar.wait().await;
1386 let ctx = OperatorContextLockHelper {
1387 database_locks: locks_clone,
1388 };
1389 let db_name = format!("db-{i}");
1390 let guard = ctx.try_lock(&db_name).await;
1391 if guard.is_some() {
1392 count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1393 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
1394 }
1395 }));
1396 }
1397
1398 for h in handles {
1399 h.await.unwrap();
1400 }
1401
1402 let total = acquired_count.load(std::sync::atomic::Ordering::SeqCst);
1403 assert_eq!(
1404 total, concurrency,
1405 "all {concurrency} tasks locking different dbs should succeed, got {total}"
1406 );
1407 }
1408
1409 #[tokio::test]
1410 async fn try_lock_database_acquire_release_cycle_under_contention() {
1411 let locks: Arc<Mutex<HashMap<String, ()>>> = Arc::new(Mutex::new(HashMap::new()));
1415 let concurrency = 20;
1416 let success_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
1417 let barrier = Arc::new(tokio::sync::Barrier::new(concurrency));
1418
1419 let mut handles = Vec::with_capacity(concurrency);
1420 for _ in 0..concurrency {
1421 let locks_clone = Arc::clone(&locks);
1422 let count = Arc::clone(&success_count);
1423 let bar = Arc::clone(&barrier);
1424 handles.push(tokio::spawn(async move {
1425 bar.wait().await;
1426 for _ in 0..100 {
1429 let ctx = OperatorContextLockHelper {
1430 database_locks: Arc::clone(&locks_clone),
1431 };
1432 if let Some(_guard) = ctx.try_lock("shared-db").await {
1433 count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1434 tokio::time::sleep(std::time::Duration::from_millis(1)).await;
1436 return;
1437 }
1438 tokio::time::sleep(std::time::Duration::from_millis(1)).await;
1439 }
1440 panic!("task failed to acquire lock after 100 retries");
1442 }));
1443 }
1444
1445 for h in handles {
1446 h.await.unwrap();
1447 }
1448
1449 let total = success_count.load(std::sync::atomic::Ordering::SeqCst);
1450 assert_eq!(
1451 total, concurrency,
1452 "all {concurrency} tasks should eventually acquire the lock"
1453 );
1454 }
1455}