|
| 1 | +use anyhow::Context as _; |
| 2 | +use configuration::{serialized::Schema, Configuration, ConfigurationOptions}; |
| 3 | +use deadpool_postgres::{Manager, ManagerConfig, Pool, RecyclingMethod}; |
| 4 | +use postgres_native_tls::MakeTlsConnector; |
| 5 | + |
| 6 | +const SOURCE: &str = "MONGODB"; |
| 7 | + |
| 8 | +/// Reads connector configuration from a PostgreSQL config store, |
| 9 | +/// using the shared config_tables schema with a `raw_schema` column |
| 10 | +/// that stores the connector's native schema JSON per collection. |
| 11 | +#[derive(Clone)] |
| 12 | +pub struct PostgresConfigurationStore { |
| 13 | + pool: Pool, |
| 14 | + connector_id: String, |
| 15 | + schema: String, |
| 16 | +} |
| 17 | + |
| 18 | +impl std::fmt::Debug for PostgresConfigurationStore { |
| 19 | + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { |
| 20 | + f.debug_struct("PostgresConfigurationStore") |
| 21 | + .field("connector_id", &self.connector_id) |
| 22 | + .field("schema", &self.schema) |
| 23 | + .finish() |
| 24 | + } |
| 25 | +} |
| 26 | + |
| 27 | +impl PostgresConfigurationStore { |
| 28 | + pub fn new(url: String, connector_id: String, schema: String) -> anyhow::Result<Self> { |
| 29 | + let tls_connector = native_tls::TlsConnector::builder() |
| 30 | + .build() |
| 31 | + .context("failed to build TLS connector")?; |
| 32 | + let tls = MakeTlsConnector::new(tls_connector); |
| 33 | + |
| 34 | + let pg_config: tokio_postgres::Config = url |
| 35 | + .parse() |
| 36 | + .context("failed to parse postgres connection URL")?; |
| 37 | + |
| 38 | + let manager_config = ManagerConfig { |
| 39 | + recycling_method: RecyclingMethod::Fast, |
| 40 | + }; |
| 41 | + let manager = Manager::from_config(pg_config, tls, manager_config); |
| 42 | + let pool = Pool::builder(manager) |
| 43 | + .max_size(4) |
| 44 | + .build() |
| 45 | + .context("failed to build postgres connection pool")?; |
| 46 | + |
| 47 | + Ok(Self { |
| 48 | + pool, |
| 49 | + connector_id, |
| 50 | + schema, |
| 51 | + }) |
| 52 | + } |
| 53 | + |
| 54 | + async fn get_client(&self) -> anyhow::Result<deadpool_postgres::Client> { |
| 55 | + self.pool |
| 56 | + .get() |
| 57 | + .await |
| 58 | + .map_err(|e| anyhow::anyhow!("failed to get postgres connection from pool: {e}")) |
| 59 | + } |
| 60 | + |
| 61 | + /// Read the schema for a single collection by name. |
| 62 | + /// Returns a Configuration containing only that collection and its associated object types. |
| 63 | + pub async fn read_collection_configuration( |
| 64 | + &self, |
| 65 | + collection_name: &str, |
| 66 | + ) -> anyhow::Result<Configuration> { |
| 67 | + self.read_collections_configuration(&[collection_name]) |
| 68 | + .await |
| 69 | + } |
| 70 | + |
| 71 | + /// Read schemas for multiple collections by name and merge them into a single Configuration. |
| 72 | + /// This is needed when a query involves relationships to other collections. |
| 73 | + pub async fn read_collections_configuration( |
| 74 | + &self, |
| 75 | + collection_names: &[&str], |
| 76 | + ) -> anyhow::Result<Configuration> { |
| 77 | + let client = self.get_client().await?; |
| 78 | + |
| 79 | + let mut merged_schema = Schema::default(); |
| 80 | + |
| 81 | + for &collection_name in collection_names { |
| 82 | + let query = format!( |
| 83 | + r#"SELECT name, raw_schema |
| 84 | + FROM "{}".config_tables |
| 85 | + WHERE UPPER(source) = UPPER($1) |
| 86 | + AND connector_id = $2 |
| 87 | + AND name = $3 |
| 88 | + AND is_deleted = false |
| 89 | + ORDER BY updated_at DESC |
| 90 | + LIMIT 1"#, |
| 91 | + self.schema |
| 92 | + ); |
| 93 | + |
| 94 | + let row = client |
| 95 | + .query_opt(&query, &[&SOURCE, &self.connector_id, &collection_name]) |
| 96 | + .await |
| 97 | + .with_context(|| { |
| 98 | + format!("failed to query config_tables for collection {collection_name}") |
| 99 | + })? |
| 100 | + .ok_or_else(|| { |
| 101 | + anyhow::anyhow!("collection {collection_name} not found in config store") |
| 102 | + })?; |
| 103 | + |
| 104 | + let name: String = row.get(0); |
| 105 | + let raw_schema_json: serde_json::Value = row.get(1); |
| 106 | + let schema: Schema = serde_json::from_value(raw_schema_json) |
| 107 | + .with_context(|| format!("failed to parse raw_schema for collection {name}"))?; |
| 108 | + |
| 109 | + merged_schema.collections.extend(schema.collections); |
| 110 | + merged_schema.object_types.extend(schema.object_types); |
| 111 | + } |
| 112 | + |
| 113 | + Configuration::validate( |
| 114 | + merged_schema, |
| 115 | + Default::default(), |
| 116 | + Default::default(), |
| 117 | + ConfigurationOptions::default(), |
| 118 | + ) |
| 119 | + } |
| 120 | + |
| 121 | + /// Read the connection URI from config_metadata. |
| 122 | + /// Returns the value if stored as {"value": "..."} or the env var name if {"variable": "..."}. |
| 123 | + /// Falls back to MONGODB_DATABASE_URI env var if not found. |
| 124 | + pub async fn read_connection_uri(&self) -> anyhow::Result<ConnectionUri> { |
| 125 | + let client = self.get_client().await?; |
| 126 | + |
| 127 | + let query = format!( |
| 128 | + r#"SELECT value |
| 129 | + FROM "{}".config_metadata |
| 130 | + WHERE UPPER(source) = UPPER($1) |
| 131 | + AND connector_id = $2 |
| 132 | + AND key = $3 |
| 133 | + ORDER BY updated_at DESC |
| 134 | + LIMIT 1"#, |
| 135 | + self.schema |
| 136 | + ); |
| 137 | + |
| 138 | + let row = client |
| 139 | + .query_opt(&query, &[&SOURCE, &self.connector_id, &"connection_uri"]) |
| 140 | + .await |
| 141 | + .context("failed to query config_metadata for connection_uri")?; |
| 142 | + |
| 143 | + match row { |
| 144 | + Some(row) => { |
| 145 | + let value: serde_json::Value = row.get(0); |
| 146 | + let uri: ConnectionUri = |
| 147 | + serde_json::from_value(value).context("failed to parse connection_uri")?; |
| 148 | + Ok(uri) |
| 149 | + } |
| 150 | + None => Ok(ConnectionUri::Variable { |
| 151 | + variable: DEFAULT_DATABASE_URI_ENV_VAR.to_string(), |
| 152 | + }), |
| 153 | + } |
| 154 | + } |
| 155 | +} |
| 156 | + |
| 157 | +/// Connection URI as stored in config_metadata. |
| 158 | +#[derive(Debug, serde::Deserialize)] |
| 159 | +#[serde(untagged)] |
| 160 | +pub enum ConnectionUri { |
| 161 | + /// Direct value: {"value": "mongodb://..."} |
| 162 | + Value { value: String }, |
| 163 | + /// Environment variable reference: {"variable": "MONGODB_DATABASE_URI"} |
| 164 | + Variable { variable: String }, |
| 165 | +} |
| 166 | + |
| 167 | +impl ConnectionUri { |
| 168 | + /// Resolve to the actual URI string. |
| 169 | + pub fn resolve(&self) -> anyhow::Result<String> { |
| 170 | + match self { |
| 171 | + ConnectionUri::Value { value } => Ok(value.clone()), |
| 172 | + ConnectionUri::Variable { variable } => std::env::var(variable) |
| 173 | + .map_err(|_| anyhow::anyhow!("environment variable {variable} is not set")), |
| 174 | + } |
| 175 | + } |
| 176 | +} |
| 177 | + |
| 178 | +pub const DEFAULT_DATABASE_URI_ENV_VAR: &str = "MONGODB_DATABASE_URI"; |
0 commit comments