Support railway connect for service databases (#458)

* support railway connect for service databases

* lint

* remove unused types

* indicate legacy plugin

.cargo/config

@@ -1,2 +1,3 @@
 [alias]
-lint = "clippy --all-features --all-targets"
+lint = "clippy --all-features --all-targets"
+lint-fix = "clippy --fix --allow-dirty --allow-staged --all-targets --all-features -- -D warnings"

src/commands/connect.rs

@@ -1,25 +1,39 @@
 use anyhow::bail;
+use std::{collections::BTreeMap, fmt::Display};
 use tokio::process::Command;
 use which::which;
 
-use crate::controllers::variables::get_plugin_variables;
-use crate::controllers::{environment::get_matched_environment, project::get_project};
+use crate::controllers::project::get_plugin_or_service;
+use crate::controllers::{
+    environment::get_matched_environment,
+    project::{get_project, PluginOrService},
+    variables::get_plugin_or_service_variables,
+};
 use crate::errors::RailwayError;
-use crate::util::prompt::{prompt_select, PromptPlugin};
+use crate::util::prompt::prompt_select;
 
 use super::{queries::project::PluginType, *};
 
 /// Connect to a plugin's shell (psql for Postgres, mongosh for MongoDB, etc.)
 #[derive(Parser)]
 pub struct Args {
     /// The name of the plugin to connect to
-    plugin_name: Option<String>,
+    service_name: Option<String>,
 
     /// Environment to pull variables from (defaults to linked environment)
     #[clap(short, long)]
     environment: Option<String>,
 }
 
+impl Display for PluginOrService {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            PluginOrService::Plugin(plugin) => write!(f, "{} (legacy)", plugin.friendly_name),
+            PluginOrService::Service(service) => write!(f, "{}", service.name),
+        }
+    }
+}
+
 pub async fn command(args: Args, _json: bool) -> Result<()> {
     let configs = Configs::new()?;
     let client = GQLClient::new_authorized(&configs)?;
@@ -31,46 +45,104 @@ pub async fn command(args: Args, _json: bool) -> Result<()> {
 
     let project = get_project(&client, &configs, linked_project.project.clone()).await?;
 
-    let plugin = match args.plugin_name {
-        Some(name) => {
-            &project
-                .plugins
-                .edges
-                .iter()
-                .find(|edge| edge.node.friendly_name == name)
-                .ok_or_else(|| RailwayError::PluginNotFound(name))?
-                .node
-        }
-        None => {
-            let plugins: Vec<_> = project
-                .plugins
-                .edges
-                .iter()
-                .map(|p| PromptPlugin(&p.node))
-                .collect();
-            if plugins.is_empty() {
-                return Err(RailwayError::ProjectHasNoPlugins.into());
+    let plugin_or_service = args
+        .service_name
+        .clone()
+        .map(|name| get_plugin_or_service(&project, name))
+        .unwrap_or_else(|| {
+            let mut nodes_to_prompt: Vec<PluginOrService> = Vec::new();
+            for plugin in &project.plugins.edges {
+                nodes_to_prompt.push(PluginOrService::Plugin(plugin.node.clone()));
+            }
+            for service in &project.services.edges {
+                nodes_to_prompt.push(PluginOrService::Service(service.node.clone()));
             }
-            prompt_select("Select a plugin", plugins)
-                .context("No plugin selected")?
-                .0
-        }
-    };
+
+            if nodes_to_prompt.is_empty() {
+                return Err(RailwayError::ProjectHasNoServicesOrPlugins.into());
+            }
+
+            prompt_select("Select service", nodes_to_prompt).context("No service selected")
+        })?;
 
     let environment_id = get_matched_environment(&project, environment)?.id;
 
-    let variables = get_plugin_variables(
+    let variables = get_plugin_or_service_variables(
         &client,
         &configs,
         linked_project.project,
-        environment_id,
-        plugin.id.clone(),
+        environment_id.clone(),
+        &plugin_or_service,
     )
     .await?;
 
+    let plugin_type = plugin_or_service
+        .get_plugin_type(environment_id)
+        .ok_or_else(|| RailwayError::UnknownDatabaseType(plugin_or_service.get_name()))?;
+
+    let (cmd_name, args) = get_connect_command(plugin_type, variables)?;
+
+    if which(cmd_name.clone()).is_err() {
+        bail!("{} must be installed to continue", cmd_name);
+    }
+
+    Command::new(cmd_name.as_str())
+        .args(args)
+        .spawn()?
+        .wait()
+        .await?;
+
+    Ok(())
+}
+
+impl PluginOrService {
+    pub fn get_name(&self) -> String {
+        match self {
+            PluginOrService::Plugin(plugin) => plugin.friendly_name.clone(),
+            PluginOrService::Service(service) => service.name.clone(),
+        }
+    }
+
+    pub fn get_plugin_type(&self, environment_id: String) -> Option<PluginType> {
+        match self {
+            PluginOrService::Plugin(plugin) => Some(plugin.name.clone()),
+            PluginOrService::Service(service) => {
+                let service_instance = service
+                    .service_instances
+                    .edges
+                    .iter()
+                    .find(|si| si.node.environment_id == environment_id);
+
+                service_instance
+                    .and_then(|si| si.node.source.clone())
+                    .and_then(|source| source.image)
+                    .map(|image: String| image.to_lowercase())
+                    .and_then(|image: String| {
+                        if image.contains("postgres") {
+                            Some(PluginType::postgresql)
+                        } else if image.contains("redis") {
+                            Some(PluginType::redis)
+                        } else if image.contains("mongo") {
+                            Some(PluginType::mongodb)
+                        } else if image.contains("mysql") {
+                            Some(PluginType::mysql)
+                        } else {
+                            None
+                        }
+                    })
+            }
+        }
+    }
+}
+
+fn get_connect_command(
+    plugin_type: PluginType,
+    variables: BTreeMap<String, String>,
+) -> Result<(String, Vec<String>)> {
     let pass_arg; // Hack to get ownership of formatted string outside match
     let default = &"".to_string();
-    let (cmd_name, args): (&str, Vec<&str>) = match &plugin.name {
+
+    let (cmd_name, args): (&str, Vec<&str>) = match &plugin_type {
         PluginType::postgresql => (
             "psql",
             vec![variables.get("DATABASE_URL").unwrap_or(default)],
@@ -104,11 +176,8 @@ pub async fn command(args: Args, _json: bool) -> Result<()> {
         PluginType::Other(o) => bail!("Unsupported plugin type {}", o),
     };
 
-    if which(cmd_name).is_err() {
-        bail!("{} must be installed to continue", cmd_name);
-    }
-
-    Command::new(cmd_name).args(args).spawn()?.wait().await?;
-
-    Ok(())
+    Ok((
+        cmd_name.to_string(),
+        args.iter().map(|s| s.to_string()).collect(),
+    ))
 }

src/controllers/project.rs

@@ -3,12 +3,22 @@ use reqwest::Client;
 use crate::{
     client::post_graphql,
     commands::{
-        queries::{self},
+        queries::{
+            self,
+            project::{
+                ProjectProject, ProjectProjectPluginsEdgesNode, ProjectProjectServicesEdgesNode,
+            },
+        },
         Configs,
     },
     errors::RailwayError,
 };
-use anyhow::Result;
+use anyhow::{bail, Result};
+
+pub enum PluginOrService {
+    Plugin(ProjectProjectPluginsEdgesNode),
+    Service(ProjectProjectServicesEdgesNode),
+}
 
 pub async fn get_project(
     client: &Client,
@@ -32,3 +42,30 @@ pub async fn get_project(
 
     Ok(project)
 }
+
+pub fn get_plugin_or_service(
+    project: &ProjectProject,
+    service_or_plugin_name: String,
+) -> Result<PluginOrService> {
+    let service = project
+        .services
+        .edges
+        .iter()
+        .find(|edge| edge.node.name.to_lowercase() == service_or_plugin_name);
+
+    let plugin = project
+        .plugins
+        .edges
+        .iter()
+        .find(|edge| edge.node.friendly_name.to_lowercase() == service_or_plugin_name);
+
+    if let Some(service) = service {
+        return Ok(PluginOrService::Service(service.node.clone()));
+    } else if let Some(plugin) = plugin {
+        return Ok(PluginOrService::Plugin(plugin.node.clone()));
+    }
+
+    bail!(RailwayError::ServiceOrPluginNotFound(
+        service_or_plugin_name
+    ))
+}

src/controllers/variables.rs

@@ -6,6 +6,8 @@ use anyhow::Result;
 use reqwest::Client;
 use std::collections::BTreeMap;
 
+use super::project::PluginOrService;
+
 pub async fn get_service_variables(
     client: &Client,
     configs: &Configs,
@@ -71,3 +73,42 @@ pub async fn get_plugin_variables(
 
     Ok(variables)
 }
+
+pub async fn get_plugin_or_service_variables(
+    client: &Client,
+    configs: &Configs,
+    project_id: String,
+    environment_id: String,
+    plugin_or_service: &PluginOrService,
+) -> Result<BTreeMap<String, String>> {
+    let variables = match plugin_or_service {
+        PluginOrService::Plugin(plugin) => {
+            let query = queries::variables_for_plugin::Variables {
+                project_id: project_id.clone(),
+                environment_id: environment_id.clone(),
+                plugin_id: plugin.id.clone(),
+            };
+
+            post_graphql::<queries::VariablesForPlugin, _>(client, configs.get_backboard(), query)
+                .await?
+                .variables
+        }
+        PluginOrService::Service(service) => {
+            let query = queries::variables_for_service_deployment::Variables {
+                project_id: project_id.clone(),
+                environment_id: environment_id.clone(),
+                service_id: service.id.clone(),
+            };
+
+            post_graphql::<queries::VariablesForServiceDeployment, _>(
+                client,
+                configs.get_backboard(),
+                query,
+            )
+            .await?
+            .variables_for_service_deployment
+        }
+    };
+
+    Ok(variables)
+}

src/errors.rs

@@ -41,8 +41,11 @@ pub enum RailwayError {
     #[error("Service \"{0}\" not found.\nRun `railway service` to connect to a service.")]
     ServiceNotFound(String),
 
-    #[error("Project has no plugins.\nRun `railway add` to add a plugin.")]
-    ProjectHasNoPlugins,
+    #[error("Service or plugin \"{0}\" not found.")]
+    ServiceOrPluginNotFound(String),
+
+    #[error("Project has no services or plugins.")]
+    ProjectHasNoServicesOrPlugins,
 
     #[error("No service linked and no plugins found\nRun `railway service` to link a service")]
     NoServiceLinked,
@@ -55,4 +58,7 @@ pub enum RailwayError {
 
     #[error("{0}")]
     FailedToUpload(String),
+
+    #[error("Could not determine database type for service {0}")]
+    UnknownDatabaseType(String),
 }

src/gql/queries/strings/Project.graphql

@@ -24,6 +24,20 @@ query Project($id: String!) {
         node {
           id
           name
+
+          serviceInstances {
+            edges {
+              node {
+                id
+                serviceId
+                environmentId
+                source {
+                  repo
+                  image
+                }
+              }
+            }
+          }
         }
       }
     }