From 6c9735b1327a28b57735f17824362ffa811a98f5 Mon Sep 17 00:00:00 2001
From: David Lutterkort <lutter@watzmann.net>
Date: Mon, 13 Jul 2020 13:45:13 -0700
Subject: [PATCH] store: Allow enforcing a timeout for SQL queries

---
 docs/environment-variables.md    |  3 +++
 store/postgres/src/relational.rs | 37 +++++++++++++++++++++++++-------
 2 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/docs/environment-variables.md b/docs/environment-variables.md
index 1cc22c0c107..49e188f163b 100644
--- a/docs/environment-variables.md
+++ b/docs/environment-variables.md
@@ -100,6 +100,9 @@ those.
 - `GRAPH_GRAPHQL_MAX_OPERATIONS_PER_CONNECTION`: maximum number of GraphQL
   operations per WebSocket connection. Any operation created after the limit
   will return an error to the client. Default: unlimited.
+- `GRAPH_SQL_STATEMENT_TIMEOUT`: the maximum number of seconds an
+  individual SQL query is allowed to take during GraphQL
+  execution. Default: unlimited
 
 ## Miscellaneous
 
diff --git a/store/postgres/src/relational.rs b/store/postgres/src/relational.rs
index fcf436a1976..20efb21cb4e 100644
--- a/store/postgres/src/relational.rs
+++ b/store/postgres/src/relational.rs
@@ -6,7 +6,7 @@
 //!
 //! The pivotal struct in this module is the `Layout` which handles all the
 //! information about mapping a GraphQL schema to database tables
-use diesel::connection::SimpleConnection;
+use diesel::{connection::SimpleConnection, Connection};
 use diesel::{debug_query, OptionalExtension, PgConnection, RunQueryDsl};
 use graph::prelude::{q, s};
 use inflector::Inflector;
@@ -64,6 +64,20 @@ lazy_static! {
             .map(|v| v.split(",").map(|s| s.to_owned()).collect())
             .unwrap_or(HashSet::new())
     };
+
+    /// `GRAPH_SQL_STATEMENT_TIMEOUT` is the timeout for queries in seconds.
+    /// If it is not set, no statement timeout will be enforced. The statement
+    /// timeout is local, i.e., can only be used within a transaction and
+    /// will be cleared at the end of the transaction
+    static ref STATEMENT_TIMEOUT: Option<String> = {
+        env::var("GRAPH_SQL_STATEMENT_TIMEOUT")
+        .ok()
+        .map(|s| {
+            u64::from_str(&s).unwrap_or_else(|_| {
+                panic!("GRAPH_SQL_STATEMENT_TIMEOUT must be a number, but is `{}`", s)
+            })
+        }).map(|timeout| format!("set local statement_timeout={}", timeout * 1000))
+    };
 }
 
 /// A string we use as a SQL name for a table or column. The important thing
@@ -622,13 +636,20 @@ impl Layout {
         let query_clone = query.clone();
 
         let start = Instant::now();
-        let values = query.load::<EntityData>(conn).map_err(|e| {
-            QueryExecutionError::ResolveEntitiesError(format!(
-                "{}, query = {:?}",
-                e,
-                debug_query(&query_clone).to_string()
-            ))
-        })?;
+        let values = conn
+            .transaction(|| {
+                if let Some(ref timeout_sql) = *STATEMENT_TIMEOUT {
+                    conn.batch_execute(timeout_sql)?;
+                }
+                query.load::<EntityData>(conn)
+            })
+            .map_err(|e| {
+                QueryExecutionError::ResolveEntitiesError(format!(
+                    "{}, query = {:?}",
+                    e,
+                    debug_query(&query_clone).to_string()
+                ))
+            })?;
         log_query_timing(logger, &query_clone, start.elapsed(), values.len());
         values
             .into_iter()