graph, store: Add GRAPH_STORE_SETUP_TIMEOUT for database setup

Add a separate timeout for database setup operations (migrations, schema
creation, FDW configuration) which can legitimately take longer than
normal runtime operations. This fixes timeout issues during runner tests
when using the default 5s GRAPH_STORE_CONNECTION_TIMEOUT.

The new GRAPH_STORE_SETUP_TIMEOUT defaults to 30s and is used only during
database initialization via the new get_for_setup() method, preserving
the fast-fail behavior (5s) for production runtime operations.

Author: lutter

docs/environment-variables.md

@@ -216,6 +216,9 @@ those.
   decisions. Set to `true` to turn simulation on, defaults to `false`
 - `GRAPH_STORE_CONNECTION_TIMEOUT`: How long to wait to connect to a
   database before assuming the database is down in ms. Defaults to 5000ms.
+- `GRAPH_STORE_SETUP_TIMEOUT`: Timeout for database setup operations
+  (migrations, schema creation) in milliseconds. Defaults to 30000ms (30s).
+  Setup operations can legitimately take longer than normal runtime operations.
 - `GRAPH_STORE_CONNECTION_UNAVAILABLE_RETRY`: When a database shard is marked
   unavailable due to connection timeouts, this controls how often to allow a
   single probe request through to check if the database has recovered. Only one

graph/src/env/store.rs

@@ -63,6 +63,10 @@ pub struct EnvVarsStore {
     /// Set by the environment variable `GRAPH_STORE_CONNECTION_TIMEOUT` (expressed
     /// in milliseconds). The default value is 5000ms.
     pub connection_timeout: Duration,
+    /// Set by `GRAPH_STORE_SETUP_TIMEOUT` (in milliseconds). Default: 30000ms.
+    /// Used during database setup (migrations, schema creation) which can
+    /// legitimately take longer than normal operations.
+    pub setup_timeout: Duration,
     /// Set by the environment variable `GRAPH_STORE_CONNECTION_MIN_IDLE`. No
     /// default value is provided.
     pub connection_min_idle: Option<u32>,
@@ -214,6 +218,7 @@ impl TryFrom<InnerStore> for EnvVarsStore {
             ),
             recent_blocks_cache_capacity: x.recent_blocks_cache_capacity,
             connection_timeout: Duration::from_millis(x.connection_timeout_in_millis),
+            setup_timeout: Duration::from_millis(x.setup_timeout_in_millis),
             connection_min_idle: x.connection_min_idle,
             connection_idle_timeout: Duration::from_secs(x.connection_idle_timeout_in_secs),
             write_queue_size: x.write_queue_size,
@@ -299,6 +304,8 @@ pub struct InnerStore {
     // configured differently for each pool.
     #[envconfig(from = "GRAPH_STORE_CONNECTION_TIMEOUT", default = "5000")]
     connection_timeout_in_millis: u64,
+    #[envconfig(from = "GRAPH_STORE_SETUP_TIMEOUT", default = "30000")]
+    setup_timeout_in_millis: u64,
     #[envconfig(from = "GRAPH_STORE_CONNECTION_MIN_IDLE")]
     connection_min_idle: Option<u32>,
     #[envconfig(from = "GRAPH_STORE_CONNECTION_IDLE_TIMEOUT", default = "600")]

store/postgres/src/pool/coordinator.rs

@@ -265,7 +265,7 @@ impl PoolCoordinator {
 
         let primary = self.primary()?;
 
-        let mut pconn = primary.get().await?;
+        let mut pconn = primary.get_for_setup().await?;
 
         let states: Vec<_> = states
             .into_iter()

store/postgres/src/pool/mod.rs

@@ -620,6 +620,18 @@ impl PoolInner {
         self.get_from_pool(&self.pool, None, Duration::ZERO).await
     }
 
+    /// Get a connection using the setup timeout. Use only during database
+    /// initialization where operations can legitimately take longer.
+    async fn get_for_setup(&self) -> Result<AsyncPgConnection, StoreError> {
+        let setup_timeouts = Timeouts {
+            wait: Some(ENV_VARS.store.setup_timeout),
+            create: Some(ENV_VARS.store.setup_timeout),
+            recycle: Some(ENV_VARS.store.setup_timeout),
+        };
+        self.get_from_pool(&self.pool, Some(setup_timeouts), Duration::ZERO)
+            .await
+    }
+
     /// Get the pool for fdw connections. It is an error if none is configured
     fn fdw_pool(&self, logger: &Logger) -> Result<&AsyncPool, StoreError> {
         let pool = match &self.fdw_pool {
@@ -701,7 +713,7 @@ impl PoolInner {
     }
 
     async fn locale_check(&self, logger: &Logger) -> Result<(), StoreError> {
-        let mut conn = self.get().await?;
+        let mut conn = self.get_for_setup().await?;
         let _: () = if let Err(msg) = catalog::Locale::load(&mut conn).await?.suitable() {
             if self.shard == *PRIMARY_SHARD && primary::is_empty(&mut conn).await? {
                 const MSG: &str = "Database does not use C locale. \
@@ -751,7 +763,7 @@ impl PoolInner {
 
     async fn configure_fdw(&self, servers: &[ForeignServer]) -> Result<(), StoreError> {
         info!(&self.logger, "Setting up fdw");
-        let mut conn = self.get().await?;
+        let mut conn = self.get_for_setup().await?;
         conn.batch_execute("create extension if not exists postgres_fdw")
             .await?;
         conn.transaction(|conn| {
@@ -790,7 +802,10 @@ impl PoolInner {
         // careful that block_on only gets called on a blocking thread to
         // avoid errors from the tokio runtime
         let logger = self.logger.cheap_clone();
-        let mut conn = self.get().await.map(AsyncConnectionWrapper::from)?;
+        let mut conn = self
+            .get_for_setup()
+            .await
+            .map(AsyncConnectionWrapper::from)?;
 
         tokio::task::spawn_blocking(move || {
             diesel::Connection::transaction::<_, StoreError, _>(&mut conn, |conn| {
@@ -808,7 +823,7 @@ impl PoolInner {
         }
 
         info!(&self.logger, "Dropping cross-shard views");
-        let mut conn = self.get().await?;
+        let mut conn = self.get_for_setup().await?;
         conn.transaction(|conn| {
             async {
                 let query = format!("drop schema if exists {} cascade", CROSS_SHARD_NSP);
@@ -845,7 +860,7 @@ impl PoolInner {
             return Ok(());
         }
 
-        let mut conn = self.get().await?;
+        let mut conn = self.get_for_setup().await?;
         let sharded = Namespace::special(CROSS_SHARD_NSP);
         if catalog::has_namespace(&mut conn, &sharded).await? {
             // We dropped the namespace before, but another node must have
@@ -897,7 +912,7 @@ impl PoolInner {
     pub async fn remap(&self, server: &ForeignServer) -> Result<(), StoreError> {
         if server.shard == *PRIMARY_SHARD {
             info!(&self.logger, "Mapping primary");
-            let mut conn = self.get().await?;
+            let mut conn = self.get_for_setup().await?;
             conn.transaction(|conn| ForeignServer::map_primary(conn, &self.shard).scope_boxed())
                 .await?;
         }
@@ -907,7 +922,7 @@ impl PoolInner {
                 "Mapping metadata from {}",
                 server.shard.as_str()
             );
-            let mut conn = self.get().await?;
+            let mut conn = self.get_for_setup().await?;
             conn.transaction(|conn| server.map_metadata(conn).scope_boxed())
                 .await?;
         }
@@ -919,7 +934,7 @@ impl PoolInner {
             return Ok(false);
         }
 
-        let mut conn = self.get().await?;
+        let mut conn = self.get_for_setup().await?;
         server.needs_remap(&mut conn).await
     }
 }