Skip to content

Commit

Permalink
Turn SQL into optional cargo feature
Browse files Browse the repository at this point in the history
  • Loading branch information
lablans committed Sep 4, 2024
1 parent 2e240fb commit 3fe7ddb
Show file tree
Hide file tree
Showing 4 changed files with 46 additions and 13 deletions.
8 changes: 5 additions & 3 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,7 @@ laplace_rs = {git = "https://github.com/samply/laplace-rs.git", tag = "v0.3.0" }
uuid = "1.8.0"
rand = { default-features = false, version = "0.8.5" }
futures-util = { version = "0.3", default-features = false, features = ["std"] }
sqlx = { version = "0.8.2", features = [ "runtime-tokio", "postgres", "macros", "chrono"] }
tryhard = "0.5"
kurtbuilds_sqlx_serde = "0.3.1"

# Logging
tracing = { version = "0.1.37", default_features = false }
Expand All @@ -34,12 +32,16 @@ once_cell = "1.18"
# Command Line Interface
clap = { version = "4", default_features = false, features = ["std", "env", "derive", "help", "color"] }

# Query via SQL
sqlx = { version = "0.8.2", features = [ "runtime-tokio", "postgres", "macros", "chrono"], optional = true }
kurtbuilds_sqlx_serde = { version = "0.3.1", optional = true }


[features]
default = []
bbmri = []
dktk = []
dktk = ["query-sql"]
query-sql = ["dep:sqlx", "dep:kurtbuilds_sqlx_serde"]

[dev-dependencies]
pretty_assertions = "1.4.0"
Expand Down
4 changes: 4 additions & 0 deletions src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,9 @@ pub enum Obfuscate {
pub enum EndpointType {
Blaze,
Omop,
#[cfg(feature = "query-sql")]
BlazeAndSql,
#[cfg(feature = "query-sql")]
Sql,
}

Expand All @@ -29,7 +31,9 @@ impl fmt::Display for EndpointType {
match self {
EndpointType::Blaze => write!(f, "blaze"),
EndpointType::Omop => write!(f, "omop"),
#[cfg(feature = "query-sql")]
EndpointType::BlazeAndSql => write!(f, "blaze_and_sql"),
#[cfg(feature = "query-sql")]
EndpointType::Sql => write!(f, "sql"),
}
}
Expand Down
5 changes: 3 additions & 2 deletions src/errors.rs
Original file line number Diff line number Diff line change
Expand Up @@ -63,12 +63,13 @@ pub enum FocusError {
MissingExporterTaskType,
#[error("Cannot connect to database: {0}")]
CannotConnectToDatabase(String),
#[error("Error executing query: {0}")]
ErrorExecutingQuery(sqlx::Error),
#[error("QueryResultBad: {0}")]
QueryResultBad(String),
#[error("Query not allowed: {0}")]
QueryNotAllowed(String),
#[cfg(feature = "query-sql")]
#[error("Error executing query: {0}")]
ErrorExecutingQuery(sqlx::Error),
}

impl FocusError {
Expand Down
42 changes: 34 additions & 8 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,20 +8,21 @@ mod errors;
mod graceful_shutdown;
mod logger;

mod db;
mod exporter;
mod intermediate_rep;
mod projects;
mod task_processing;
mod util;

#[cfg(feature = "query-sql")]
mod db;

use base64::engine::general_purpose;
use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
use beam_lib::{TaskRequest, TaskResult};
use futures_util::future::BoxFuture;
use futures_util::FutureExt;
use laplace_rs::ObfCache;
use sqlx::PgPool;
use tokio::sync::Mutex;

use crate::blaze::{parse_blaze_query_payload_ast, AstQuery};
Expand Down Expand Up @@ -118,22 +119,45 @@ pub async fn main() -> ExitCode {
}
}

async fn main_loop() -> ExitCode {
let db_pool = if let Some(connection_string) = CONFIG.postgres_connection_string.clone() {
#[cfg(not(feature = "query-sql"))]
type DbPool = bool;

#[cfg(feature = "query-sql")]
type DbPool = sqlx::PgPool;

#[cfg(not(feature = "query-sql"))]
async fn get_db_pool() -> Result<Option<DbPool>,ExitCode> {
Ok(None)
}

#[cfg(feature = "query-sql")]
async fn get_db_pool() -> Result<Option<DbPool>,ExitCode> {
if let Some(connection_string) = CONFIG.postgres_connection_string.clone() {
match db::get_pg_connection_pool(&connection_string, 8).await {
Err(e) => {
error!("Error connecting to database: {}", e);
return ExitCode::from(8);
Err(ExitCode::from(8))
}
Ok(pool) => Some(pool),
Ok(pool) => Ok(Some(pool)),
}
} else {
None
Ok(None)
}
}

async fn main_loop() -> ExitCode {
let db_pool = match get_db_pool().await {
Ok(pool) => pool,
Err(code) => {
return code;
},
};
let endpoint_service_available: fn() -> BoxFuture<'static, bool> = match CONFIG.endpoint_type {
EndpointType::Blaze => || blaze::check_availability().boxed(),
EndpointType::Omop => || async { true }.boxed(), // TODO health check
#[cfg(feature = "query-sql")]
EndpointType::BlazeAndSql => || blaze::check_availability().boxed(),
#[cfg(feature = "query-sql")]
EndpointType::Sql => || async { true }.boxed(),
};
let mut failures = 0;
Expand Down Expand Up @@ -169,7 +193,7 @@ async fn process_task(
task: &BeamTask,
obf_cache: Arc<Mutex<ObfCache>>,
report_cache: Arc<Mutex<ReportCache>>,
db_pool: Option<PgPool>,
db_pool: Option<DbPool>,
) -> Result<BeamResult, FocusError> {
debug!("Processing task {}", task.id);

Expand Down Expand Up @@ -217,6 +241,7 @@ async fn process_task(
)
.await
},
#[cfg(feature = "query-sql")]
EndpointType::BlazeAndSql => {
let mut generated_from_ast: bool = false;
let data = base64_decode(&task.body)?;
Expand Down Expand Up @@ -260,6 +285,7 @@ async fn process_task(
.await
}
},
#[cfg(feature="query-sql")]
EndpointType::Sql => {
let data = base64_decode(&task.body)?;
let query_maybe: Result<db::SqlQuery, serde_json::Error> = serde_json::from_slice(&(data));
Expand Down

0 comments on commit 3fe7ddb

Please sign in to comment.