1 use native_tls::TlsConnector;
2 use postgres_native_tls::MakeTlsConnector;
3 use serde::{Deserialize, Serialize};
4 use std::{env, fmt, num::ParseIntError};
5 use thiserror::Error;
6 use tokio::task::JoinHandle;
7 use tokio_postgres::{
8 Client, Config as PostgresConfig, Error as PostgresError, NoTls, Row,
9 config::SslMode as PostgresSslMode,
10 };
11
12 #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
13 pub enum GresiqEnvironment {
14 Development,
15 Preview,
16 Production,
17 }
18
19 impl GresiqEnvironment {
20 pub fn as_str(self) -> &'static str {
21 match self {
22 Self::Development => "development",
23 Self::Preview => "preview",
24 Self::Production => "production",
25 }
26 }
27 }
28
29 impl fmt::Display for GresiqEnvironment {
30 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
31 formatter.write_str(self.as_str())
32 }
33 }
34
35 impl std::str::FromStr for GresiqEnvironment {
36 type Err = GresiqConfigError;
37
38 fn from_str(value: &str) -> Result<Self, Self::Err> {
39 match value.trim().to_ascii_lowercase().as_str() {
40 "development" | "dev" => Ok(Self::Development),
41 "preview" | "staging" => Ok(Self::Preview),
42 "production" | "prod" => Ok(Self::Production),
43 _ => Err(GresiqConfigError::InvalidEnvironment(value.to_owned())),
44 }
45 }
46 }
47
48 #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
49 pub enum ComputeTier {
50 Burst,
51 Autoscale,
52 Dedicated,
53 }
54
55 impl ComputeTier {
56 pub fn as_str(self) -> &'static str {
57 match self {
58 Self::Burst => "burst",
59 Self::Autoscale => "autoscale",
60 Self::Dedicated => "dedicated",
61 }
62 }
63 }
64
65 impl fmt::Display for ComputeTier {
66 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
67 formatter.write_str(self.as_str())
68 }
69 }
70
71 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
72 pub struct BranchTarget {
73 pub project_slug: String,
74 pub branch_name: String,
75 pub region: String,
76 pub environment: GresiqEnvironment,
77 }
78
79 impl BranchTarget {
80 pub fn summary(&self) -> String {
81 format!(
82 "{}/{}/{} ({})",
83 self.project_slug, self.branch_name, self.region, self.environment
84 )
85 }
86 }
87
88 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
89 pub struct Credentials {
90 pub username: String,
91 pub password: String,
92 pub database_name: String,
93 }
94
95 impl Credentials {
96 pub fn redacted_password(&self) -> &'static str {
97 "********"
98 }
99 }
100
101 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
102 pub struct Endpoint {
103 pub host: String,
104 pub port: u16,
105 pub ssl_mode: SslMode,
106 }
107
108 #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
109 pub enum SslMode {
110 Disable,
111 Prefer,
112 Require,
113 }
114
115 impl SslMode {
116 pub fn as_str(self) -> &'static str {
117 match self {
118 Self::Disable => "disable",
119 Self::Prefer => "prefer",
120 Self::Require => "require",
121 }
122 }
123 }
124
125 impl fmt::Display for SslMode {
126 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
127 formatter.write_str(self.as_str())
128 }
129 }
130
131 impl std::str::FromStr for SslMode {
132 type Err = GresiqConfigError;
133
134 fn from_str(value: &str) -> Result<Self, Self::Err> {
135 match value.trim().to_ascii_lowercase().as_str() {
136 "disable" => Ok(Self::Disable),
137 "prefer" => Ok(Self::Prefer),
138 "require" => Ok(Self::Require),
139 _ => Err(GresiqConfigError::InvalidSslMode(value.to_owned())),
140 }
141 }
142 }
143
144 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
145 pub struct ConnectionPolicy {
146 pub compute_tier: ComputeTier,
147 pub application_name: String,
148 pub connect_timeout_seconds: u64,
149 }
150
151 impl Default for ConnectionPolicy {
152 fn default() -> Self {
153 Self {
154 compute_tier: ComputeTier::Autoscale,
155 application_name: "gresiq".to_owned(),
156 connect_timeout_seconds: 10,
157 }
158 }
159 }
160
161 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
162 pub struct GresiqConfig {
163 pub branch_target: BranchTarget,
164 pub endpoint: Endpoint,
165 pub credentials: Credentials,
166 pub connection_policy: ConnectionPolicy,
167 }
168
169 impl GresiqConfig {
170 pub fn from_env() -> Result<Self, GresiqConfigError> {
171 let environment = env::var("SMBCLOUD_ENVIRONMENT")
172 .unwrap_or_else(|_| "development".to_owned())
173 .parse()?;
174
175 let ssl_mode = env::var("GRESIQ_SSL_MODE")
176 .unwrap_or_else(|_| "require".to_owned())
177 .parse()?;
178
179 let port = read_optional_port("GRESIQ_PORT")?.unwrap_or(5432);
180
181 let connect_timeout_seconds =
182 read_optional_u64("GRESIQ_CONNECT_TIMEOUT_SECONDS")?.unwrap_or(10);
183
184 let compute_tier = match env::var("GRESIQ_COMPUTE_TIER")
185 .unwrap_or_else(|_| "autoscale".to_owned())
186 .trim()
187 .to_ascii_lowercase()
188 .as_str()
189 {
190 "burst" => ComputeTier::Burst,
191 "autoscale" => ComputeTier::Autoscale,
192 "dedicated" => ComputeTier::Dedicated,
193 other => {
194 return Err(GresiqConfigError::InvalidComputeTier(other.to_owned()));
195 }
196 };
197
198 Ok(Self {
199 branch_target: BranchTarget {
200 project_slug: read_required_env("GRESIQ_PROJECT_SLUG")?,
201 branch_name: env::var("GRESIQ_BRANCH").unwrap_or_else(|_| "main".to_owned()),
202 region: env::var("GRESIQ_REGION").unwrap_or_else(|_| "eu-north-1".to_owned()),
203 environment,
204 },
205 endpoint: Endpoint {
206 host: read_required_env("GRESIQ_HOST")?,
207 port,
208 ssl_mode,
209 },
210 credentials: Credentials {
211 username: read_required_env("GRESIQ_USERNAME")?,
212 password: read_required_env("GRESIQ_PASSWORD")?,
213 database_name: env::var("GRESIQ_DATABASE")
214 .unwrap_or_else(|_| "postgres".to_owned()),
215 },
216 connection_policy: ConnectionPolicy {
217 compute_tier,
218 application_name: env::var("GRESIQ_APPLICATION_NAME")
219 .unwrap_or_else(|_| "gresiq".to_owned()),
220 connect_timeout_seconds,
221 },
222 })
223 }
224
225 pub fn connection_string(&self) -> String {
226 format!(
227 "postgresql://{}:{}@{}:{}/{}?sslmode={}&application_name={}",
228 self.credentials.username,
229 self.credentials.password,
230 self.endpoint.host,
231 self.endpoint.port,
232 self.credentials.database_name,
233 self.endpoint.ssl_mode,
234 self.connection_policy.application_name
235 )
236 }
237
238 pub fn redacted_connection_string(&self) -> String {
239 format!(
240 "postgresql://{}:{}@{}:{}/{}?sslmode={}&application_name={}",
241 self.credentials.username,
242 self.credentials.redacted_password(),
243 self.endpoint.host,
244 self.endpoint.port,
245 self.credentials.database_name,
246 self.endpoint.ssl_mode,
247 self.connection_policy.application_name
248 )
249 }
250
251 pub fn postgres_config(&self) -> PostgresConfig {
252 let mut config = PostgresConfig::new();
253 config.host(&self.endpoint.host);
254 config.port(self.endpoint.port);
255 config.user(&self.credentials.username);
256 config.password(&self.credentials.password);
257 config.dbname(&self.credentials.database_name);
258 config.application_name(&self.connection_policy.application_name);
259 config.connect_timeout(std::time::Duration::from_secs(
260 self.connection_policy.connect_timeout_seconds,
261 ));
262 config.ssl_mode(match self.endpoint.ssl_mode {
263 SslMode::Disable => PostgresSslMode::Disable,
264 SslMode::Prefer => PostgresSslMode::Prefer,
265 SslMode::Require => PostgresSslMode::Require,
266 });
267 config
268 }
269
270 pub async fn connect(&self) -> Result<GresiqConnection, GresiqConnectError> {
271 let postgres_config = self.postgres_config();
272
273 match self.endpoint.ssl_mode {
274 SslMode::Disable => {
275 let (client, connection) = postgres_config.connect(NoTls).await?;
276 let connection_task = tokio::spawn(async move { connection.await });
277
278 Ok(GresiqConnection {
279 client,
280 connection_task: Some(connection_task),
281 })
282 }
283 SslMode::Prefer | SslMode::Require => {
284 let tls_connector = TlsConnector::builder().build()?;
285 let tls_connector = MakeTlsConnector::new(tls_connector);
286 let (client, connection) = postgres_config.connect(tls_connector).await?;
287 let connection_task = tokio::spawn(async move { connection.await });
288
289 Ok(GresiqConnection {
290 client,
291 connection_task: Some(connection_task),
292 })
293 }
294 }
295 }
296 }
297
298 #[derive(Debug)]
299 pub struct GresiqConnection {
300 client: Client,
301 connection_task: Option<JoinHandle<Result<(), PostgresError>>>,
302 }
303
304 impl GresiqConnection {
305 pub fn client(&self) -> &Client {
306 &self.client
307 }
308
309 pub async fn health_check(&self) -> Result<GresiqHealthCheck, GresiqConnectError> {
310 let row = self
311 .client
312 .query_one("select current_database(), current_user, version()", &[])
313 .await?;
314
315 Ok(GresiqHealthCheck::from_row(&row))
316 }
317
318 pub async fn close(mut self) -> Result<(), GresiqConnectError> {
319 if let Some(connection_task) = self.connection_task.take() {
320 connection_task.await??;
321 }
322 Ok(())
323 }
324 }
325
326 impl Drop for GresiqConnection {
327 fn drop(&mut self) {
328 if let Some(connection_task) = self.connection_task.take() {
329 connection_task.abort();
330 }
331 }
332 }
333
334 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
335 pub struct GresiqHealthCheck {
336 pub database_name: String,
337 pub username: String,
338 pub server_version: String,
339 }
340
341 impl GresiqHealthCheck {
342 fn from_row(row: &Row) -> Self {
343 Self {
344 database_name: row.get(0),
345 username: row.get(1),
346 server_version: row.get(2),
347 }
348 }
349 }
350
351 #[derive(Debug, Error)]
352 pub enum GresiqConfigError {
353 #[error("missing required environment variable `{0}`")]
354 MissingEnvVar(&'static str),
355 #[error("environment variable `{0}` contains invalid unicode")]
356 InvalidEnvUnicode(&'static str),
357 #[error("failed to parse environment variable `{name}` as an integer")]
358 InvalidNumber {
359 name: &'static str,
360 #[source]
361 source: ParseIntError,
362 },
363 #[error("invalid smbCloud environment `{0}`")]
364 InvalidEnvironment(String),
365 #[error("invalid GresIQ ssl mode `{0}`")]
366 InvalidSslMode(String),
367 #[error("invalid GresIQ compute tier `{0}`")]
368 InvalidComputeTier(String),
369 }
370
371 #[derive(Debug, Error)]
372 pub enum GresiqConnectError {
373 #[error(transparent)]
374 Tls(#[from] native_tls::Error),
375 #[error(transparent)]
376 Postgres(#[from] PostgresError),
377 #[error("postgres connection task failed")]
378 Join(#[from] tokio::task::JoinError),
379 }
380
381 fn read_required_env(name: &'static str) -> Result<String, GresiqConfigError> {
382 env::var(name).map_err(|_| GresiqConfigError::MissingEnvVar(name))
383 }
384
385 fn read_optional_port(name: &'static str) -> Result<Option<u16>, GresiqConfigError> {
386 read_optional_env(name)?
387 .map(|value| {
388 value
389 .parse()
390 .map_err(|source| GresiqConfigError::InvalidNumber { name, source })
391 })
392 .transpose()
393 }
394
395 fn read_optional_u64(name: &'static str) -> Result<Option<u64>, GresiqConfigError> {
396 read_optional_env(name)?
397 .map(|value| {
398 value
399 .parse()
400 .map_err(|source| GresiqConfigError::InvalidNumber { name, source })
401 })
402 .transpose()
403 }
404
405 fn read_optional_env(name: &'static str) -> Result<Option<String>, GresiqConfigError> {
406 match env::var(name) {
407 Ok(value) => Ok(Some(value)),
408 Err(env::VarError::NotPresent) => Ok(None),
409 Err(env::VarError::NotUnicode(_)) => Err(GresiqConfigError::InvalidEnvUnicode(name)),
410 }
411 }
412
413 #[cfg(test)]
414 mod tests {
415 use super::*;
416
417 #[test]
418 fn redacted_connection_string_hides_the_password() {
419 let config = GresiqConfig {
420 branch_target: BranchTarget {
421 project_slug: "acme-store".to_owned(),
422 branch_name: "main".to_owned(),
423 region: "eu-north-1".to_owned(),
424 environment: GresiqEnvironment::Production,
425 },
426 endpoint: Endpoint {
427 host: "main.acme-store.gresiq.smbcloud.xyz".to_owned(),
428 port: 5432,
429 ssl_mode: SslMode::Require,
430 },
431 credentials: Credentials {
432 username: "postgres".to_owned(),
433 password: "super-secret".to_owned(),
434 database_name: "app".to_owned(),
435 },
436 connection_policy: ConnectionPolicy::default(),
437 };
438
439 assert!(!config.redacted_connection_string().contains("super-secret"));
440 assert!(config.redacted_connection_string().contains("********"));
441 }
442
443 #[test]
444 fn branch_summary_matches_expected_shape() {
445 let branch_target = BranchTarget {
446 project_slug: "acme-store".to_owned(),
447 branch_name: "preview".to_owned(),
448 region: "eu-west-1".to_owned(),
449 environment: GresiqEnvironment::Preview,
450 };
451
452 assert_eq!(
453 branch_target.summary(),
454 "acme-store/preview/eu-west-1 (preview)"
455 );
456 }
457 }