Skip to main content

olai_stack_runtime/neon/
cli.rs

1//! [`NeonCli`]: the [`NeonControl`] that runs the Neon CLI.
2
3use std::collections::BTreeMap;
4use std::path::{Path, PathBuf};
5use std::process::Stdio;
6
7use async_trait::async_trait;
8use serde::Deserialize;
9use serde::de::DeserializeOwned;
10use tokio::process::Command;
11
12use super::{
13    Branch, Connection, CreateBranch, Function, NeonControl, NeonError, Storage, StorageCredential,
14};
15use crate::secrets::SecretValue;
16
17/// The environment variable naming the Neon CLI to run, ahead of every other place it is
18/// looked for.
19pub const NEON_CLI_VAR: &str = "NEON_CLI";
20
21/// Where the Neon CLI is, in order: `explicit`, `$NEON_CLI`, `neon` on `PATH`, the Homebrew
22/// prefixes and `~/.npm-global/bin`, a login shell's `command -v neon`, then the same for
23/// `neonctl`. An app started from the Finder has launchd's `PATH`, which holds none of them,
24/// hence the fixed places and the login shell.
25pub fn locate_neon(explicit: Option<&Path>) -> Option<PathBuf> {
26    if let Some(path) = explicit {
27        return path.is_file().then(|| path.to_owned());
28    }
29    if let Some(path) = std::env::var_os(NEON_CLI_VAR).map(PathBuf::from)
30        && path.is_file()
31    {
32        return Some(path);
33    }
34    ["neon", "neonctl"].into_iter().find_map(find_program)
35}
36
37fn find_program(name: &str) -> Option<PathBuf> {
38    let on_path = std::env::var_os("PATH")
39        .into_iter()
40        .flat_map(|p| std::env::split_paths(&p).collect::<Vec<_>>())
41        .map(|dir| dir.join(name));
42    let home = std::env::var_os("HOME").map(PathBuf::from);
43    let fixed = ["/opt/homebrew/bin", "/usr/local/bin"]
44        .into_iter()
45        .map(PathBuf::from)
46        .chain(home.map(|h| h.join(".npm-global/bin")))
47        .map(|dir| dir.join(name));
48    on_path
49        .chain(fixed)
50        .find(|p| p.is_file())
51        .or_else(|| login_shell_lookup(name))
52}
53
54/// `command -v <name>` in the user's login shell, which reads their profile's `PATH`.
55fn login_shell_lookup(name: &str) -> Option<PathBuf> {
56    let shell = std::env::var_os("SHELL")?;
57    let out = std::process::Command::new(shell)
58        .args(["-lc", &format!("command -v {name}")])
59        .stdin(Stdio::null())
60        .stderr(Stdio::null())
61        .output()
62        .ok()?;
63    let path = PathBuf::from(String::from_utf8(out.stdout).ok()?.trim());
64    (out.status.success() && path.is_file()).then_some(path)
65}
66
67/// The Neon CLI (`neon`, or `neonctl`), run with `--output json`. Signed in the way the CLI is
68/// (`~/.config/neon`, shared with the terminal) or with `NEON_API_KEY`.
69#[derive(Clone, Debug)]
70pub struct NeonCli {
71    program: PathBuf,
72}
73
74impl NeonCli {
75    /// The CLI [`locate_neon`] finds; [`NeonError::NotInstalled`] when there is none.
76    pub fn new() -> Result<Self, NeonError> {
77        locate_neon(None)
78            .map(Self::with_program)
79            .ok_or(NeonError::NotInstalled)
80    }
81
82    /// The CLI at `program`.
83    pub fn with_program(program: impl Into<PathBuf>) -> Self {
84        NeonCli {
85            program: program.into(),
86        }
87    }
88
89    /// The program it runs.
90    pub fn program(&self) -> &Path {
91        &self.program
92    }
93
94    /// Run `neon <args> --output json` and return its stdout.
95    async fn run(&self, args: &[&str]) -> Result<String, NeonError> {
96        let command = args.join(" ");
97        let out = Command::new(&self.program)
98            .args(args)
99            .args(["--output", "json"])
100            .stdin(Stdio::null())
101            .kill_on_drop(true)
102            .output()
103            .await
104            .map_err(|source| match source.kind() {
105                std::io::ErrorKind::NotFound => NeonError::NotInstalled,
106                _ => NeonError::Io {
107                    command: command.clone(),
108                    source,
109                },
110            })?;
111        let stderr = String::from_utf8_lossy(&out.stderr).trim().to_owned();
112        if !out.status.success() {
113            return Err(if is_auth_error(&stderr) {
114                NeonError::NotAuthenticated(stderr)
115            } else {
116                NeonError::Failed { command, stderr }
117            });
118        }
119        String::from_utf8(out.stdout).map_err(|e| NeonError::Parse {
120            command,
121            reason: e.to_string(),
122        })
123    }
124
125    /// Run `args` and parse the first JSON value it prints.
126    async fn json<T: DeserializeOwned>(&self, args: &[&str]) -> Result<T, NeonError> {
127        let stdout = self.run(args).await?;
128        first_json(&stdout).map_err(|reason| NeonError::Parse {
129            command: args.join(" "),
130            reason,
131        })
132    }
133}
134
135/// Whether `stderr` says the CLI isn't signed in.
136fn is_auth_error(stderr: &str) -> bool {
137    let s = stderr.to_ascii_lowercase();
138    [
139        "unauthorized",
140        "not authenticated",
141        "authentication",
142        "401",
143        "neon auth",
144    ]
145    .iter()
146    .any(|needle| s.contains(needle))
147}
148
149/// The first JSON value in `text`, which several verbs surround with lines of text (a
150/// `WARNING:` after `credentials create`, progress lines around `functions deploy`).
151pub(super) fn first_json<T: DeserializeOwned>(text: &str) -> Result<T, String> {
152    let mut last = String::from("no JSON in the output");
153    for (at, _) in text
154        .match_indices(['{', '['])
155        .filter(|(at, _)| *at == 0 || text[..*at].ends_with('\n'))
156    {
157        let mut values = serde_json::Deserializer::from_str(&text[at..]).into_iter::<T>();
158        match values.next() {
159            Some(Ok(value)) => return Ok(value),
160            Some(Err(e)) => last = e.to_string(),
161            None => {}
162        }
163    }
164    Err(last)
165}
166
167/// `GET /projects/<p>/branches`: the branches, and their annotations keyed by branch id.
168#[derive(Deserialize)]
169pub(super) struct ApiBranches {
170    pub branches: Vec<Branch>,
171    #[serde(default)]
172    pub annotations: BTreeMap<String, ApiAnnotation>,
173}
174
175#[derive(Deserialize)]
176pub(super) struct ApiAnnotation {
177    #[serde(default)]
178    pub value: BTreeMap<String, String>,
179}
180
181impl ApiBranches {
182    /// The branches with their annotations filled in.
183    pub fn into_branches(mut self) -> Vec<Branch> {
184        for branch in &mut self.branches {
185            if let Some(a) = self.annotations.remove(&branch.id) {
186                branch.annotations = a.value;
187            }
188        }
189        self.branches
190    }
191}
192
193/// `connection-string --extended`.
194#[derive(Deserialize)]
195pub(super) struct Extended {
196    pub connection_string: String,
197    pub host: String,
198    pub role: String,
199    pub password: String,
200    pub database: String,
201    #[serde(default)]
202    pub options: String,
203}
204
205impl From<Extended> for Connection {
206    fn from(e: Extended) -> Self {
207        Connection {
208            url: SecretValue::new(e.connection_string),
209            host: e.host,
210            role: e.role,
211            password: SecretValue::new(e.password),
212            database: e.database,
213            options: e.options,
214        }
215    }
216}
217
218/// `credentials create`.
219#[derive(Deserialize)]
220pub(super) struct CreatedCredential {
221    pub token_id: String,
222    pub s3_secret_access_key: String,
223}
224
225#[derive(Deserialize)]
226struct Named {
227    name: String,
228}
229
230/// A database of `databases list`.
231#[derive(Deserialize)]
232struct Database {
233    name: String,
234    owner_name: String,
235}
236
237/// The database Neon creates with a project. `connection-string` and `databases create` only
238/// pick a default while the branch has one database and one role; a branch that hosts more
239/// (another service's database, a read-only role) needs both named.
240const DEFAULT_DATABASE: &str = "neondb";
241
242impl NeonCli {
243    async fn database_list(&self, project: &str, branch: &str) -> Result<Vec<Database>, NeonError> {
244        self.json(&[
245            "databases",
246            "list",
247            "--project-id",
248            project,
249            "--branch",
250            branch,
251        ])
252        .await
253    }
254
255    /// The branch's main database: `neondb` when it has one, else the first.
256    async fn main_database(&self, project: &str, branch: &str) -> Result<Database, NeonError> {
257        let mut dbs = self.database_list(project, branch).await?;
258        if dbs.is_empty() {
259            return Err(NeonError::Failed {
260                command: "databases list".into(),
261                stderr: format!("branch {branch} has no database"),
262            });
263        }
264        let at = dbs
265            .iter()
266            .position(|d| d.name == DEFAULT_DATABASE)
267            .unwrap_or(0);
268        Ok(dbs.swap_remove(at))
269    }
270}
271
272#[async_trait]
273impl NeonControl for NeonCli {
274    async fn auth_status(&self) -> Result<String, NeonError> {
275        #[derive(Deserialize)]
276        struct Me {
277            #[serde(default)]
278            login: Option<String>,
279            #[serde(default)]
280            email: Option<String>,
281        }
282        let me: Me = self.json(&["me"]).await?;
283        Ok(me.login.or(me.email).unwrap_or_default())
284    }
285
286    async fn branches(&self, project: &str) -> Result<Vec<Branch>, NeonError> {
287        let path = format!("/projects/{project}/branches");
288        let api: ApiBranches = self.json(&["api", &path]).await?;
289        Ok(api.into_branches())
290    }
291
292    async fn create_branch(&self, project: &str, spec: &CreateBranch) -> Result<Branch, NeonError> {
293        #[derive(Deserialize)]
294        struct Created {
295            branch: Branch,
296        }
297        let annotations = serde_json::to_string(&spec.annotations).unwrap_or_default();
298        let mut args = vec![
299            "branches",
300            "create",
301            "--project-id",
302            project,
303            "--name",
304            &spec.name,
305            "--no-secrets",
306        ];
307        if let Some(parent) = &spec.parent {
308            args.extend(["--parent", parent]);
309        }
310        if let Some(at) = &spec.expires_at {
311            args.extend(["--expires-at", at]);
312        }
313        if !spec.annotations.is_empty() {
314            args.extend(["--annotation", &annotations]);
315        }
316        let mut created: Created = self.json(&args).await?;
317        // The CLI stores annotations but doesn't print them back.
318        created.branch.annotations = spec.annotations.clone();
319        Ok(created.branch)
320    }
321
322    async fn delete_branch(&self, project: &str, id: &str) -> Result<(), NeonError> {
323        self.run(&["branches", "delete", id, "--project-id", project])
324            .await
325            .map(drop)
326    }
327
328    async fn reset_branch(&self, project: &str, id: &str) -> Result<(), NeonError> {
329        self.run(&["branches", "reset", id, "--parent", "--project-id", project])
330            .await
331            .map(drop)
332    }
333
334    async fn connection(
335        &self,
336        project: &str,
337        branch: &str,
338        database: Option<&str>,
339        pooled: bool,
340    ) -> Result<Connection, NeonError> {
341        let mut args = vec![
342            "connection-string",
343            branch,
344            "--project-id",
345            project,
346            "--extended",
347        ];
348        if pooled {
349            args.push("--pooled");
350        }
351        let main;
352        let db = match database {
353            Some(db) => db,
354            None => {
355                main = self.main_database(project, branch).await?;
356                main.name.as_str()
357            }
358        };
359        args.extend(["--database-name", db]);
360        let extended: Extended = self.json(&args).await?;
361        Ok(extended.into())
362    }
363
364    async fn databases(&self, project: &str, branch: &str) -> Result<Vec<String>, NeonError> {
365        let dbs = self.database_list(project, branch).await?;
366        Ok(dbs.into_iter().map(|d| d.name).collect())
367    }
368
369    async fn create_database(
370        &self,
371        project: &str,
372        branch: &str,
373        name: &str,
374    ) -> Result<(), NeonError> {
375        // Owned like the main database, by the role `connection-string` connects as.
376        let owner = self.main_database(project, branch).await?.owner_name;
377        self.run(&[
378            "databases",
379            "create",
380            "--name",
381            name,
382            "--owner-name",
383            &owner,
384            "--project-id",
385            project,
386            "--branch",
387            branch,
388        ])
389        .await
390        .map(drop)
391    }
392
393    async fn functions(&self, project: &str, branch: &str) -> Result<Vec<Function>, NeonError> {
394        self.json(&[
395            "functions",
396            "list",
397            "--project-id",
398            project,
399            "--branch",
400            branch,
401        ])
402        .await
403    }
404
405    async fn storage(&self, project: &str, branch: &str) -> Result<Storage, NeonError> {
406        let path = format!("/projects/{project}/branches/{branch}/storage");
407        self.json(&["api", &path]).await
408    }
409
410    async fn buckets(&self, project: &str, branch: &str) -> Result<Vec<String>, NeonError> {
411        let buckets: Vec<Named> = self
412            .json(&[
413                "buckets",
414                "list",
415                "--project-id",
416                project,
417                "--branch",
418                branch,
419            ])
420            .await?;
421        Ok(buckets.into_iter().map(|b| b.name).collect())
422    }
423
424    async fn create_bucket(
425        &self,
426        project: &str,
427        branch: &str,
428        name: &str,
429    ) -> Result<(), NeonError> {
430        // Prints a sentence, not JSON.
431        self.run(&[
432            "buckets",
433            "create",
434            name,
435            "--project-id",
436            project,
437            "--branch",
438            branch,
439        ])
440        .await
441        .map(drop)
442    }
443
444    async fn create_credential(
445        &self,
446        project: &str,
447        branch: &str,
448        name: &str,
449    ) -> Result<StorageCredential, NeonError> {
450        let created: CreatedCredential = self
451            .json(&[
452                "credentials",
453                "create",
454                "--project-id",
455                project,
456                "--branch",
457                branch,
458                "--name",
459                name,
460                "--scope",
461                "storage:read",
462                "--scope",
463                "storage:write",
464            ])
465            .await?;
466        Ok(StorageCredential {
467            token_id: created.token_id,
468            secret_access_key: SecretValue::new(created.s3_secret_access_key),
469        })
470    }
471
472    async fn revoke_credential(
473        &self,
474        project: &str,
475        branch: &str,
476        token_id: &str,
477    ) -> Result<(), NeonError> {
478        self.run(&[
479            "credentials",
480            "revoke",
481            token_id,
482            "--project-id",
483            project,
484            "--branch",
485            branch,
486        ])
487        .await
488        .map(drop)
489    }
490}
491
492#[cfg(test)]
493mod tests {
494    use super::*;
495
496    #[test]
497    fn the_first_json_value_wins_around_text() {
498        let v: serde_json::Value =
499            first_json("Function deployment triggered.\n{\"id\": 1}\nDone.\n").unwrap();
500        assert_eq!(v["id"], 1);
501        let v: serde_json::Value = first_json("{\"a\": [1]}\nWARNING: store it\n").unwrap();
502        assert_eq!(v["a"][0], 1);
503        // A `{` inside a line of text isn't the start of a value.
504        let v: serde_json::Value = first_json("note {curly}\n[2]").unwrap();
505        assert_eq!(v[0], 2);
506        assert!(first_json::<serde_json::Value>("Bucket created").is_err());
507    }
508
509    #[test]
510    fn auth_errors_are_told_apart() {
511        assert!(is_auth_error("ERROR: Unauthorized. Run `neon auth`"));
512        assert!(!is_auth_error("ERROR: branch not found"));
513    }
514}