Pipe through static connection

Author: danieljharvey

Cargo.lock

@@ -31,9 +31,9 @@ similar_names = "allow"
 too_many_lines = "allow"
 
 [workspace.dependencies]
-ndc-models = { git = "https://github.com/hasura/ndc-spec.git", tag = "v0.2.3" }
-ndc-sdk = { git = "https://github.com/hasura/ndc-sdk-rs.git", rev = "e2a1aeb" }
-ndc-test = { git = "https://github.com/hasura/ndc-spec.git", tag = "v0.2.3" }
+ndc-models = { git = "https://github.com/hasura/ndc-spec.git", rev = "3fdff9e" } 
+ndc-sdk = { git = "https://github.com/hasura/ndc-sdk-rs.git", rev = "2584f5d55049df60951bcdf2ba9c506a2ea8ff1d" }
+ndc-test = { git = "https://github.com/hasura/ndc-spec.git", rev = "3fdff9e" } 
 
 anyhow = "1"
 async-trait = "0.1"

Cargo.toml

@@ -12,7 +12,7 @@ use crate::error::{
 use crate::values::{IsolationLevel, PoolSettings};
 use crate::version3;
 use crate::version4;
-use crate::version5;
+use crate::version5::{self, DynamicConnectionSettings};
 use crate::VersionTag;
 use schemars::{gen::SchemaSettings, schema::RootSchema};
 
@@ -56,6 +56,16 @@ impl ParsedConfiguration {
     }
 }
 
+#[derive(Debug)]
+pub enum ConnectionSettings {
+    Static {
+        connection_uri: String,
+    },
+    Dynamic {
+        dynamic_connection_settings: DynamicConnectionSettings,
+    },
+}
+
 /// The 'Configuration' type collects all the information necessary to serve queries at runtime.
 ///
 /// 'ParsedConfiguration' deals with a multitude of different concrete version formats, and each
@@ -71,8 +81,8 @@ impl ParsedConfiguration {
 pub struct Configuration {
     pub metadata: metadata::Metadata,
     pub configuration_version_tag: VersionTag,
+    pub connection: ConnectionSettings,
     pub pool_settings: PoolSettings,
-    pub connection_uri: String,
     pub isolation_level: IsolationLevel,
     pub mutations_version: Option<metadata::mutations::MutationsVersion>,
     pub mutations_prefix: Option<String>,

crates/configuration/src/configuration.rs

@@ -12,8 +12,8 @@ pub mod version5;
 
 pub use configuration::{
     generate_latest_schema, introspect, make_runtime_configuration, parse_configuration,
-    upgrade_to_latest_version, write_parsed_configuration, Configuration, ParsedConfiguration,
-    DEFAULT_CONNECTION_URI_VARIABLE,
+    upgrade_to_latest_version, write_parsed_configuration, Configuration, ConnectionSettings,
+    ParsedConfiguration, DEFAULT_CONNECTION_URI_VARIABLE,
 };
 pub use values::{ConnectionUri, IsolationLevel, PoolSettings, Secret};
 

crates/configuration/src/lib.rs

@@ -597,7 +597,7 @@ pub fn make_runtime_configuration(
     Ok(crate::Configuration {
         metadata: convert_metadata(configuration.metadata),
         pool_settings: configuration.connection_settings.pool_settings,
-        connection_uri,
+        connection: crate::ConnectionSettings::Static { connection_uri },
         isolation_level: configuration.connection_settings.isolation_level,
         mutations_version: convert_mutations_version(configuration.mutations_version),
         configuration_version_tag: VersionTag::Version3,

crates/configuration/src/version3/mod.rs

@@ -30,7 +30,7 @@ pub fn make_runtime_configuration(
     Ok(crate::Configuration {
         metadata: convert_metadata(parsed_config.metadata),
         pool_settings: parsed_config.connection_settings.pool_settings,
-        connection_uri,
+        connection: crate::ConnectionSettings::Static { connection_uri },
         isolation_level: parsed_config.connection_settings.isolation_level,
         mutations_version: convert_mutations_version(parsed_config.mutations_version),
         configuration_version_tag: VersionTag::Version4,

crates/configuration/src/version4/to_runtime_configuration.rs

@@ -27,10 +27,18 @@ pub fn make_runtime_configuration(
             })
         }
     }?;
+
+    let connection = match parsed_config.dynamic_connections {
+        Some(dynamic_connection_settings) => crate::ConnectionSettings::Dynamic {
+            dynamic_connection_settings,
+        },
+        None => crate::ConnectionSettings::Static { connection_uri },
+    };
+
     Ok(crate::Configuration {
         metadata: convert_metadata(parsed_config.metadata),
+        connection,
         pool_settings: parsed_config.connection_settings.pool_settings,
-        connection_uri,
         isolation_level: parsed_config.connection_settings.isolation_level,
         mutations_version: convert_mutations_version(parsed_config.mutations_version),
         configuration_version_tag: VersionTag::Version4,

crates/configuration/src/version5/to_runtime_configuration.rs

@@ -39,7 +39,10 @@ impl Connector for Postgres {
     /// the number of idle connections in a connection pool
     /// can be polled but not updated directly.
     fn fetch_metrics(_configuration: &Self::Configuration, state: &Self::State) -> Result<()> {
-        state.query_metrics.update_pool_metrics(&state.pool);
+        match &state.pool {
+            state::Pool::Static { pool, .. } => state.query_metrics.update_pool_metrics(&pool),
+            state::Pool::Dynamic(_) => {}
+        };
         Ok(())
     }
 
@@ -279,9 +282,9 @@ impl<Env: Environment + Send + Sync + 'static> ConnectorSetup for PostgresSetup<
     ) -> Result<<Self::Connector as Connector>::State> {
         // create the state
         state::create_state(
-            &configuration.connection_uri,
-            &self.environment,
+            &configuration.connection,
             &configuration.pool_settings,
+            &self.environment,
             metrics,
             configuration.configuration_version_tag,
         )

crates/connectors/ndc-postgres/src/connector.rs

@@ -99,12 +99,15 @@ async fn execute_mutation(
     state: &state::State,
     plan: sql::execution_plan::ExecutionPlan<sql::execution_plan::Mutations>,
 ) -> Result<JsonResponse<models::MutationResponse>, query_engine_execution::error::Error> {
-    query_engine_execution::mutation::execute(
-        &state.pool,
-        &state.database_info,
-        &state.query_metrics,
-        plan,
-    )
-    .await
-    .map(JsonResponse::Serialized)
+    let (pool, database_info) = match &state.pool {
+        state::Pool::Static {
+            pool,
+            database_info,
+        } => (pool, database_info),
+        _ => todo!("Dynamic connect for execute_mutation"),
+    };
+
+    query_engine_execution::mutation::execute(&pool, &database_info, &state.query_metrics, plan)
+        .await
+        .map(JsonResponse::Serialized)
 }

crates/connectors/ndc-postgres/src/mutation.rs

@@ -5,6 +5,7 @@
 
 use std::collections::BTreeMap;
 
+use query_engine_execution::database_info;
 use tracing::{info_span, Instrument};
 
 use ndc_postgres_configuration as configuration;
@@ -43,9 +44,17 @@ pub async fn explain(
 
         // Execute an explain query.
         let results = async {
+            let (pool, database_info) = match &state.pool {
+                state::Pool::Static {
+                    pool,
+                    database_info,
+                } => (pool, database_info),
+                _ => todo!("Dynamic connect for explain_mutation"),
+            };
+
             query_engine_execution::mutation::explain(
-                &state.pool,
-                &state.database_info,
+                &pool,
+                &database_info,
                 &state.query_metrics,
                 plan,
             )

crates/connectors/ndc-postgres/src/mutation/explain.rs

@@ -80,12 +80,15 @@ async fn execute_query(
     state: &state::State,
     plan: sql::execution_plan::ExecutionPlan<sql::execution_plan::Query>,
 ) -> Result<JsonResponse<models::QueryResponse>, query_engine_execution::error::Error> {
-    query_engine_execution::query::execute(
-        &state.pool,
-        &state.database_info,
-        &state.query_metrics,
-        plan,
-    )
-    .await
-    .map(JsonResponse::Serialized)
+    let (pool, database_info) = match &state.pool {
+        state::Pool::Static {
+            pool,
+            database_info,
+        } => (pool, database_info),
+        _ => todo!("Dynamic connect for execute_query"),
+    };
+
+    query_engine_execution::query::execute(&pool, &database_info, &state.query_metrics, plan)
+        .await
+        .map(JsonResponse::Serialized)
 }

crates/connectors/ndc-postgres/src/query.rs

@@ -5,6 +5,7 @@
 
 use std::collections::BTreeMap;
 
+use query_engine_execution::database_info;
 use tracing::{info_span, Instrument};
 
 use ndc_postgres_configuration as configuration;
@@ -42,9 +43,16 @@ pub async fn explain(
 
         // Execute an explain query.
         let (query, plan) = async {
+            let (pool, database_info) = match &state.pool {
+                state::Pool::Static {
+                    pool,
+                    database_info,
+                } => (pool, database_info),
+                _ => todo!("Dynamic connect for explain"),
+            };
             query_engine_execution::query::explain(
-                &state.pool,
-                &state.database_info,
+                &pool,
+                &database_info,
                 &state.query_metrics,
                 plan,
             )

crates/connectors/ndc-postgres/src/query/explain.rs

@@ -8,11 +8,14 @@ mod mutation;
 
 use std::collections::BTreeMap;
 
+use ndc_postgres_configuration::version5::DynamicConnectionSettings;
+use ndc_postgres_configuration::ConnectionSettings;
 use ndc_sdk::connector;
 use ndc_sdk::models;
 
 use helpers::*;
 use ndc_postgres_configuration as configuration;
+use ndc_sdk::models::RequestLevelArguments;
 use ndc_sdk::models::ScalarTypeName;
 use query_engine_metadata::metadata;
 use query_engine_metadata::metadata::TypeRepresentation;
@@ -418,6 +421,37 @@ pub fn get_schema(
             extraction_functions: BTreeMap::new(),
         });
 
+    // if we are using dynamic connection strings then we're going to need some
+    // request_level_arguments to tell us which ones to use
+    let request_arguments = match config.connection {
+        ConnectionSettings::Dynamic {
+            dynamic_connection_settings: DynamicConnectionSettings::NamedFromList,
+        } => {
+            let request_level_arguments = BTreeMap::from_iter(vec![(
+                "connection_identifier".into(),
+                models::ArgumentInfo {
+                    argument_type: models::Type::Named {
+                        name: "String".into(),
+                    },
+                    description: Some(
+                        "Identifier to choose which connection string to use for this request."
+                            .into(),
+                    ),
+                },
+            )]);
+            RequestLevelArguments {
+                query_arguments: request_level_arguments.clone(),
+                mutation_arguments: request_level_arguments,
+                relational_query_arguments: BTreeMap::new(),
+            }
+        }
+        ConnectionSettings::Static { .. } => RequestLevelArguments {
+            query_arguments: BTreeMap::new(),
+            mutation_arguments: BTreeMap::new(),
+            relational_query_arguments: BTreeMap::new(),
+        },
+    };
+
     Ok(models::SchemaResponse {
         collections,
         procedures,
@@ -431,6 +465,7 @@ pub fn get_schema(
                 }),
             }),
         }),
+        request_arguments,
     })
 }