Move instantiate_tables to wrapper

maplibre · Oct 7, 2023 · c2591d9 · c2591d9
1 parent aa9ce8b
commit c2591d9
Showing 1 changed file with 18 additions and 11 deletions.
diff --git a/martin/src/pg/config.rs b/martin/src/pg/config.rs
@@ -4,7 +4,6 @@ use log::warn;
 use serde::{Deserialize, Serialize};
 use std::time::Duration;
 use tilejson::TileJSON;
-use tokio::time::timeout;
 
 use crate::config::{copy_unrecognized_config, UnrecognizedValues};
 use crate::pg::config_function::FuncInfoSources;
@@ -115,23 +114,31 @@ impl PgConfig {
     pub async fn resolve(&mut self, id_resolver: IdResolver) -> crate::Result<Sources> {
         let pg = PgBuilder::new(self, id_resolver).await?;
 
-        let instantiate_tables = pg.instantiate_tables();
-        pin_mut!(instantiate_tables);
-        if timeout(Duration::from_secs(5), &mut instantiate_tables)
-            .await
-            .is_err()
-        {
-            warn!("PostgreSQL table discovery is taking too long. Make sure your table geo columns have a GIS index, or use --disabling-bounds.");
-        }
-
         let ((mut tables, tbl_info), (funcs, func_info)) =
-            try_join(instantiate_tables, pg.instantiate_functions()).await?;
+            try_join(Self::instantiate_tables(&pg), pg.instantiate_functions()).await?;
 
         self.tables = Some(tbl_info);
         self.functions = Some(func_info);
         tables.extend(funcs);
         Ok(tables)
     }
+
+    async fn instantiate_tables(pg: &PgBuilder) -> Result<(Sources, TableInfoSources)> {
+        let instantiate_tables = pg.instantiate_tables();
+        pin_mut!(instantiate_tables);
+
+        let timeout_future = tokio::time::sleep(Duration::from_secs(5));
+
+        tokio::select! {
+            result = &mut instantiate_tables => {
+                result
+            }
+            () = timeout_future => {
+                warn!("PostgreSQL table discovery is taking too long. Make sure your table geo columns have a GIS index, or use --disabling-bounds.");
+                instantiate_tables.await
+            }
+        }
+    }
 }
 
 #[cfg(test)]