refactor and simplify integration and sqllogic tests

Author: tonyalaribe

Cargo.lock

@@ -64,6 +64,7 @@ datafusion-common = "46.0.0"
 [dev-dependencies]
 serial_test = "3.2.0"
 tokio-postgres = { version = "0.7.10", features = ["with-chrono-0_4"] }
+scopeguard = "1.2.0"
 
 [features]
 default = []

Cargo.toml

@@ -1,16 +1,47 @@
 #[cfg(test)]
 mod integration {
+    use anyhow::Result;
     use dotenv::dotenv;
+    use scopeguard;
     use serial_test::serial;
-    use std::time::Duration;
+    use std::sync::Arc;
+    use std::time::{Duration, Instant};
     use timefusion::database::Database;
-    use tokio::sync::oneshot;
-    use tokio::time::sleep;
-    use tokio_postgres::NoTls;
+    use tokio::{sync::Notify, time::sleep};
+    use tokio_postgres::{Client, NoTls};
     use tokio_util::sync::CancellationToken;
     use uuid::Uuid;
 
-    async fn start_test_server() -> anyhow::Result<(oneshot::Sender<()>, String)> {
+    async fn connect_with_retry(timeout: Duration) -> Result<(Client, tokio::task::JoinHandle<()>), tokio_postgres::Error> {
+        let start = Instant::now();
+        let conn_string = "host=localhost port=5433 user=postgres password=postgres";
+        
+        while start.elapsed() < timeout {
+            match tokio_postgres::connect(conn_string, NoTls).await {
+                Ok((client, connection)) => {
+                    let handle = tokio::spawn(async move {
+                        if let Err(e) = connection.await {
+                            eprintln!("Connection error: {}", e);
+                        }
+                    });
+                    return Ok((client, handle));
+                },
+                Err(_) => sleep(Duration::from_millis(100)).await,
+            }
+        }
+        
+        // Final attempt
+        let (client, connection) = tokio_postgres::connect(conn_string, NoTls).await?;
+        let handle = tokio::spawn(async move {
+            if let Err(e) = connection.await {
+                eprintln!("Connection error: {}", e);
+            }
+        });
+        
+        Ok((client, handle))
+    }
+
+    async fn start_test_server() -> Result<(Arc<Notify>, String)> {
         let test_id = Uuid::new_v4().to_string();
         let _ = env_logger::builder().is_test(true).try_init();
         dotenv().ok();
@@ -20,186 +51,112 @@ mod integration {
             std::env::set_var("TIMEFUSION_TABLE_PREFIX", format!("test-{}", test_id));
         }
 
-        let (shutdown_tx, shutdown_rx) = oneshot::channel();
+        // Use a shareable notification
+        let shutdown_signal = Arc::new(Notify::new());
+        let shutdown_signal_clone = shutdown_signal.clone();
 
         tokio::spawn(async move {
-            let db = match Database::new().await {
-                Ok(db) => db,
-                Err(e) => {
-                    panic!("Failed to create database: {:?}", e);
-                }
-            };
-
+            let db = Database::new().await.expect("Failed to create database");
             let session_context = db.create_session_context();
             db.setup_session_context(&session_context).expect("Failed to setup session context");
 
             let shutdown_token = CancellationToken::new();
             let pg_server = db
-                .start_pgwire_server(
-                    session_context,
-                    5433,
-                    shutdown_token.clone(),
-                )
+                .start_pgwire_server(session_context, 5433, shutdown_token.clone())
                 .await
                 .expect("Failed to start PGWire server");
 
-            let _ = shutdown_rx.await;
+            // Wait for shutdown signal
+            shutdown_signal_clone.notified().await;
             shutdown_token.cancel();
             let _ = pg_server.await;
         });
 
-        // Give server time to start with retry approach
-        let start_time = std::time::Instant::now();
-        let timeout = Duration::from_secs(5);
+        // Wait for server to be ready
+        let _ = connect_with_retry(Duration::from_secs(5)).await?;
 
-        // Try a test connection to make sure server is up
-        while start_time.elapsed() < timeout {
-            match tokio_postgres::connect(
-                "host=localhost port=5433 user=postgres password=postgres", 
-                NoTls
-            ).await {
-                Ok((_client, connection)) => {
-                    // Connection succeeded, spawn a task to drive it and return
-                    tokio::spawn(async move {
-                        if let Err(e) = connection.await {
-                            eprintln!("Test connection error: {}", e);
-                        }
-                    });
-                    
-                    // Small delay before we release the client
-                    sleep(Duration::from_millis(100)).await;
-                    
-                    // Server is ready
-                    break;
-                },
-                Err(_) => {
-                    // Small backoff between retry attempts
-                    sleep(Duration::from_millis(100)).await;
-                }
-            }
-        }
-
-        Ok((shutdown_tx, test_id))
+        Ok((shutdown_signal, test_id))
     }
 
     #[tokio::test]
     #[serial]
-    async fn test_postgres_integration() -> anyhow::Result<()> {
-        let (shutdown_tx, test_id) = start_test_server().await?;
+    async fn test_postgres_integration() -> Result<()> {
+        let (shutdown_signal, test_id) = start_test_server().await?;
+        let shutdown = || { shutdown_signal.notify_one(); };
+        
+        // Use a guard to ensure we notify of shutdown even if the test panics
+        let shutdown_guard = scopeguard::guard((), |_| shutdown());
 
         // Connect to database
-        let connection_string = "host=localhost port=5433 user=postgres password=postgres";
-        let (client, connection) = match tokio_postgres::connect(connection_string, NoTls).await {
-            Ok((client, connection)) => (client, connection),
-            Err(e) => {
-                shutdown_tx.send(()).ok();
-                return Err(anyhow::anyhow!("Failed to connect to PostgreSQL: {:?}", e));
-            }
-        };
-
-        tokio::spawn(async move {
-            if let Err(e) = connection.await {
-                eprintln!("Connection error: {}", e);
-            }
-        });
-
-        // Allow connection to stabilize
-        sleep(Duration::from_millis(500)).await;
+        let (client, _) = connect_with_retry(Duration::from_secs(3)).await
+            .map_err(|e| anyhow::anyhow!("Failed to connect to PostgreSQL: {}", e))?;
 
         // Insert test data
         let timestamp_str = format!("'{}'", chrono::Utc::now().format("%Y-%m-%d %H:%M:%S"));
-
-        match client
-            .execute(
-                format!(
-                    "INSERT INTO otel_logs_and_spans (
-                    project_id, timestamp, id, 
-                    name, status_code, status_message, 
-                    level
-                ) VALUES ($1, {}, $2, $3, $4, $5, $6)",
-                    timestamp_str
-                )
-                .as_str(),
+        let insert_query = format!(
+            "INSERT INTO otel_logs_and_spans (project_id, timestamp, id, name, status_code, status_message, level) 
+             VALUES ($1, {}, $2, $3, $4, $5, $6)",
+            timestamp_str
+        );
+
+        // Run the test with proper error handling
+        let result = async {
+            // Insert initial record
+            client.execute(
+                &insert_query,
                 &[&"test_project", &test_id, &"test_span_name", &"OK", &"Test integration", &"INFO"],
-            )
-            .await
-        {
-            Ok(_) => (),
-            Err(e) => {
-                shutdown_tx.send(()).ok();
-                return Err(anyhow::anyhow!("Failed to insert test data: {:?}", e));
-            }
-        }
-
-        // Query back the data and check with the client directly
-        let rows = match client.query("SELECT COUNT(*) FROM otel_logs_and_spans WHERE id = $1", &[&test_id]).await {
-            Ok(rows) => rows,
-            Err(e) => {
-                shutdown_tx.send(()).ok();
-                return Err(anyhow::anyhow!("Failed to query data: {:?}", e));
+            ).await?;
+
+            // Verify record count
+            let rows = client.query(
+                "SELECT COUNT(*) FROM otel_logs_and_spans WHERE id = $1", 
+                &[&test_id]
+            ).await?;
+            
+            assert_eq!(rows[0].get::<_, i64>(0), 1, "Should have found exactly one row");
+
+            // Verify field values
+            let detail_rows = client.query(
+                "SELECT name, status_code FROM otel_logs_and_spans WHERE id = $1", 
+                &[&test_id]
+            ).await?;
+
+            assert_eq!(detail_rows.len(), 1, "Should have found exactly one detailed row");
+            assert_eq!(detail_rows[0].get::<_, String>(0), "test_span_name", "Name should match");
+            assert_eq!(detail_rows[0].get::<_, String>(1), "OK", "Status code should match");
+
+            // Insert multiple records in a batch
+            for i in 0..5 {
+                let span_id = Uuid::new_v4().to_string();
+                client.execute(
+                    &insert_query,
+                    &[
+                        &"test_project", 
+                        &span_id, 
+                        &format!("batch_span_{}", i), 
+                        &"OK", 
+                        &format!("Batch test {}", i), 
+                        &"INFO"
+                    ],
+                ).await?;
             }
-        };
 
-        assert_eq!(rows[0].get::<_, i64>(0), 1, "Should have found exactly one row");
+            // Query with filter to get total count
+            let count_rows = client.query(
+                "SELECT COUNT(*) FROM otel_logs_and_spans WHERE project_id = $1", 
+                &[&"test_project"]
+            ).await?;
 
-        // Verify field values
-        let detail_rows = match client.query("SELECT name, status_code FROM otel_logs_and_spans WHERE id = $1", &[&test_id]).await {
-            Ok(rows) => rows,
-            Err(e) => {
-                shutdown_tx.send(()).ok();
-                return Err(anyhow::anyhow!("Failed to query detailed data: {:?}", e));
-            }
-        };
-
-        assert_eq!(detail_rows.len(), 1, "Should have found exactly one detailed row");
-        assert_eq!(detail_rows[0].get::<_, String>(0), "test_span_name", "Name should match");
-        assert_eq!(detail_rows[0].get::<_, String>(1), "OK", "Status code should match");
-
-        // Insert multiple records
-        for i in 0..5 {
-            let span_id = Uuid::new_v4().to_string();
-            let batch_name = format!("batch_span_{}", i);
-            let batch_message = format!("Batch test {}", i);
-
-            match client
-                .execute(
-                    format!(
-                        "INSERT INTO otel_logs_and_spans (
-                        project_id, timestamp, id,
-                        name, status_code, status_message,
-                        level
-                    ) VALUES ($1, {}, $2, $3, $4, $5, $6)",
-                        timestamp_str
-                    )
-                    .as_str(),
-                    &[&"test_project", &span_id, &batch_name, &"OK", &batch_message, &"INFO"],
-                )
-                .await
-            {
-                Ok(_) => (),
-                Err(e) => {
-                    shutdown_tx.send(()).ok();
-                    return Err(anyhow::anyhow!("Failed to insert batch record: {:?}", e));
-                }
-            };
-        }
-
-        // Give a little time for all inserts to complete
-        sleep(Duration::from_millis(500)).await;
-
-        // Query with filter to get total count
-        let count_rows = match client.query("SELECT COUNT(*) FROM otel_logs_and_spans WHERE project_id = $1", &[&"test_project"]).await {
-            Ok(rows) => rows,
-            Err(e) => {
-                shutdown_tx.send(()).ok();
-                return Err(anyhow::anyhow!("Failed to query total count: {:?}", e));
-            }
-        };
-
-        assert_eq!(count_rows[0].get::<_, i64>(0), 6, "Should have a total of 6 records (1 initial + 5 batch)");
-
-        let _ = shutdown_tx.send(());
-
-        Ok(())
+            assert_eq!(count_rows[0].get::<_, i64>(0), 6, "Should have a total of 6 records (1 initial + 5 batch)");
+            
+            Ok::<_, tokio_postgres::Error>(())
+        }.await;
+        
+        // Drop the guard to ensure shutdown happens
+        std::mem::drop(shutdown_guard);
+        shutdown();
+        
+        // Map postgres errors to anyhow
+        result.map_err(|e| anyhow::anyhow!("Test failed: {}", e))
     }
 }