olai_stack_runtime/neon/
source.rs

1//! [`NeonBranchSource`]: a Neon branch's overlay on an instance's environment.
2
3use std::collections::BTreeMap;
4use std::sync::Arc;
5
6use async_trait::async_trait;
7use olai_stack_state::SecretRef;
8use serde::{Deserialize, Serialize};
9
10use super::{NeonControl, NeonError};
11use crate::secrets::{SecretProviders, SecretValue};
12use crate::source::{Overlay, Source, SourceCtx, SourceError};
13
14/// The name instances bind the Neon source by.
15pub const NEON_SOURCE: &str = "neon";
16
17/// What the storage variables the planner reads start with (`NEON_S3_ENDPOINT`, …).
18pub(super) const S3_VAR_PREFIX: &str = "NEON_S3_";
19
20/// An instance's binding to a Neon branch: its [`SourceBinding`] config for
21/// [`NEON_SOURCE`], written by [`bind`](super::bind). It holds no secret, only where the
22/// storage credential's secret is kept.
23///
24/// [`SourceBinding`]: olai_stack_state::SourceBinding
25#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
26pub struct NeonBinding {
27    /// The Neon project.
28    pub project_id: String,
29    /// The branch's id.
30    pub branch: String,
31    /// The branch's name, for showing.
32    #[serde(default)]
33    pub branch_name: String,
34    /// The functions the plan runs on the branch: slug → the variable its URL goes in.
35    #[serde(default)]
36    pub functions: BTreeMap<String, String>,
37    /// The storage credential minted for the instance, when the plan uses object storage.
38    #[serde(default, skip_serializing_if = "Option::is_none")]
39    pub storage_credential: Option<StorageCredentialRef>,
40}
41
42/// A storage credential as a binding refers to it: its id (the S3 access key id) and where its
43/// secret lives.
44#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
45pub struct StorageCredentialRef {
46    /// The token id.
47    pub token_id: String,
48    /// The secret access key's reference.
49    pub secret: SecretRef,
50}
51
52/// The overlay of the Neon branch an instance is bound to (its source config is a
53/// [`NeonBinding`]):
54///
55/// - Postgres: `PGHOST` (the direct host), `PGUSER`, `PGDATABASE` and `PGPASSWORD`, which the
56///   plan's `hosted-postgres` reads, and `DATABASE_URL` (through the pooler) and
57///   `DATABASE_URL_UNPOOLED` for whatever reads those;
58/// - object storage, when the binding has a credential: `NEON_S3_ENDPOINT`,
59///   `NEON_S3_REGION`, `NEON_S3_ACCESS_KEY_ID` and `NEON_S3_SECRET_ACCESS_KEY`;
60/// - every bound function's URL, in its variable (`UNITYCATALOG_FUNCTION_URL`). A function
61///   gone from the branch since binding is an error.
62///
63/// Neon's credentials don't expire, so the overlay never needs a refresh.
64pub struct NeonBranchSource {
65    neon: Arc<dyn NeonControl>,
66    secrets: SecretProviders,
67}
68
69impl NeonBranchSource {
70    /// The source, asking `neon`, with the storage credential's secret resolved by `secrets`.
71    pub fn new(neon: Arc<dyn NeonControl>, secrets: SecretProviders) -> Self {
72        NeonBranchSource { neon, secrets }
73    }
74
75    async fn overlay_for(&self, binding: &NeonBinding) -> Result<Overlay, NeonError> {
76        let (project, branch) = (&binding.project_id, &binding.branch);
77        let mut overlay = Overlay::default();
78
79        let direct = self.neon.connection(project, branch, None, false).await?;
80        let pooled_url = pooled(direct.url.expose(), &direct.host);
81        overlay.vars.extend([
82            ("PGHOST".to_owned(), direct.host.clone()),
83            ("PGUSER".to_owned(), direct.role.clone()),
84            ("PGDATABASE".to_owned(), direct.database.clone()),
85        ]);
86        overlay.credentials.extend([
87            ("PGPASSWORD".to_owned(), direct.password.clone()),
88            ("DATABASE_URL".to_owned(), SecretValue::new(pooled_url)),
89            ("DATABASE_URL_UNPOOLED".to_owned(), direct.url.clone()),
90        ]);
91
92        if let Some(credential) = &binding.storage_credential {
93            let storage = self.neon.storage(project, branch).await?;
94            if !storage.enabled {
95                return Err(NeonError::NoStorage(binding.branch_name.clone()));
96            }
97            overlay.endpoints.insert(
98                format!("{S3_VAR_PREFIX}ENDPOINT"),
99                storage.s3_endpoint.trim_end_matches('/').to_owned(),
100            );
101            overlay
102                .vars
103                .insert(format!("{S3_VAR_PREFIX}REGION"), storage.region);
104            let secret =
105                self.secrets
106                    .resolve(&credential.secret)
107                    .map_err(|e| NeonError::Failed {
108                        command: "credentials (stored secret)".into(),
109                        stderr: e.to_string(),
110                    })?;
111            overlay.credentials.extend([
112                (
113                    format!("{S3_VAR_PREFIX}ACCESS_KEY_ID"),
114                    SecretValue::new(credential.token_id.clone()),
115                ),
116                (format!("{S3_VAR_PREFIX}SECRET_ACCESS_KEY"), secret),
117            ]);
118        }
119
120        if !binding.functions.is_empty() {
121            let deployed = self.neon.functions(project, branch).await?;
122            let mut missing = Vec::new();
123            for (slug, var) in &binding.functions {
124                match deployed.iter().find(|f| &f.slug == slug) {
125                    Some(f) => {
126                        overlay.endpoints.insert(
127                            var.clone(),
128                            f.invocation_url.trim_end_matches('/').to_owned(),
129                        );
130                    }
131                    None => missing.push(slug.clone()),
132                }
133            }
134            if !missing.is_empty() {
135                return Err(NeonError::MissingFunctions {
136                    branch: binding.branch_name.clone(),
137                    slugs: missing,
138                });
139            }
140        }
141        Ok(overlay)
142    }
143}
144
145/// `url` through the connection pooler: Neon's pooled host is the direct one with `-pooler`
146/// after the endpoint id.
147fn pooled(url: &str, host: &str) -> String {
148    match host.split_once('.') {
149        Some((endpoint, rest)) if !endpoint.ends_with("-pooler") => {
150            url.replacen(host, &format!("{endpoint}-pooler.{rest}"), 1)
151        }
152        _ => url.to_owned(),
153    }
154}
155
156#[async_trait]
157impl Source for NeonBranchSource {
158    fn name(&self) -> &str {
159        NEON_SOURCE
160    }
161
162    async fn overlay(&self, ctx: &SourceCtx) -> Result<Overlay, SourceError> {
163        let fail = |reason: String| SourceError {
164            source_name: NEON_SOURCE.into(),
165            reason,
166        };
167        let binding: NeonBinding = serde_json::from_value(ctx.config.clone())
168            .map_err(|e| fail(format!("the instance's binding: {e}")))?;
169        if binding.project_id.is_empty() || binding.branch.is_empty() {
170            return Err(fail("the instance is bound to no Neon branch".into()));
171        }
172        self.overlay_for(&binding)
173            .await
174            .map_err(|e| fail(e.to_string()))
175    }
176}
177
178#[cfg(test)]
179mod tests {
180    use super::*;
181
182    #[test]
183    fn the_pooled_host_adds_pooler_to_the_endpoint() {
184        let host = "ep-snowy-mud-bar7hnwg.c-8.eu-central-1.aws.neon.tech";
185        let url = format!("postgresql://owner@{host}/neondb?sslmode=require");
186        assert_eq!(
187            pooled(&url, host),
188            "postgresql://owner@ep-snowy-mud-bar7hnwg-pooler.c-8.eu-central-1.aws.neon.tech/neondb?sslmode=require"
189        );
190        let already = "ep-x-pooler.example.neon.tech";
191        assert_eq!(pooled(already, already), already);
192    }
193}
olai_stack_runtime/neon/source.rs

olai_stack_runtime/neon/
source.rs
