Skip to content

Commit 0fca92b

Browse files
committed
store: Limit maximum execution time for SQL queries
1 parent 09823bc commit 0fca92b

File tree

2 files changed

+11
-6
lines changed

2 files changed

+11
-6
lines changed

store/postgres/src/deployment_store.rs

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ use crate::detail::ErrorDetail;
5454
use crate::dynds::DataSourcesTable;
5555
use crate::primary::{DeploymentId, Primary};
5656
use crate::relational::index::{CreateIndex, IndexList, Method};
57-
use crate::relational::{self, Layout, LayoutCache, SqlName, Table};
57+
use crate::relational::{self, Layout, LayoutCache, SqlName, Table, STATEMENT_TIMEOUT};
5858
use crate::relational_queries::{FromEntityData, JSONData};
5959
use crate::{advisory_lock, catalog, retry};
6060
use crate::{detail, ConnectionPool};
@@ -297,12 +297,17 @@ impl DeploymentStore {
297297
query: &str,
298298
) -> Result<Vec<SqlQueryObject>, QueryExecutionError> {
299299
let query = format!("select to_jsonb(sub.*) as data from ({}) as sub", query);
300-
301300
let query = diesel::sql_query(query);
302301

303-
// Execute the provided SQL query
304-
let results = query
305-
.load::<JSONData>(conn)
302+
let results = conn
303+
.transaction(|conn| {
304+
if let Some(ref timeout_sql) = *STATEMENT_TIMEOUT {
305+
conn.batch_execute(timeout_sql)?;
306+
}
307+
308+
// Execute the provided SQL query
309+
query.load::<JSONData>(conn)
310+
})
306311
.map_err(|e| QueryExecutionError::SqlError(e.to_string()))?;
307312

308313
Ok(results

store/postgres/src/relational.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,7 @@ pub const STRING_PREFIX_SIZE: usize = 256;
9494
pub const BYTE_ARRAY_PREFIX_SIZE: usize = 64;
9595

9696
lazy_static! {
97-
static ref STATEMENT_TIMEOUT: Option<String> = ENV_VARS
97+
pub(crate) static ref STATEMENT_TIMEOUT: Option<String> = ENV_VARS
9898
.graphql
9999
.sql_statement_timeout
100100
.map(|duration| format!("set local statement_timeout={}", duration.as_millis()));

0 commit comments

Comments
 (0)