[DTS-6979] Add more unit tests for async query (#922)

[DTS-6979] Add more unit tests for async query

Author: andyl-db

server/src/main/scala/io/delta/sharing/server/DeltaSharingService.scala

@@ -192,6 +192,10 @@ class DeltaSharingService(serverConfig: ServerConfig) {
 
   private val rand = new scala.util.Random()
 
+  // Map to track poll count per queryId for testing async query behavior
+  private val queryPollCounters =
+    new java.util.concurrent.ConcurrentHashMap[String, java.util.concurrent.atomic.AtomicInteger]()
+
   private val sharedTableManager = new SharedTableManager(serverConfig)
 
   private val deltaSharedTableLoader = new DeltaSharedTableLoader(serverConfig)
@@ -375,22 +379,42 @@ class DeltaSharingService(serverConfig: ServerConfig) {
       throw new DeltaSharingIllegalArgumentException("expected error")
     }
 
-    // simulate async query with 50% chance of return
-    // asynchronously for a specific table
-    // client should be able to handle both cases and for server
-    // test please use other table names.
-    if(rand.nextInt(100) > 50 && table == "table2" && !request.pageToken.isDefined) {
+    // Track poll count for this queryId and return results after 5 polls
+    val pollCounter = queryPollCounters.computeIfAbsent(
+      queryId, _ => new java.util.concurrent.atomic.AtomicInteger(0))
+    val pollCount = pollCounter.incrementAndGet()
+
+    // Special handling for query ID change mid-query test
+    val returnedQueryId = if (table.endsWith("_change_query_id_to_be_null")) {
+      // Return null query ID to test client error handling
+      null
+    } else if (table.endsWith("_change_query_id")) {
+      // Change query ID on the 2nd poll to simulate server-side query ID change
+      s"${queryId}_modified"
+    } else {
+      queryId
+    }
+
+    // Keep returning pending status until we've been polled more than 5 times
+    if(pollCount <= 3 && !request.pageToken.isDefined) {
         streamingOutput(
           Some(0),
           "parquet",
           Seq(
-            SingleAction(queryStatus = QueryStatus(queryId))
+            SingleAction(queryStatus = QueryStatus(returnedQueryId))
           )
         )
       } else {
 
+      // Test case: Use a bad table name to trigger error during loadTable (on 2nd poll)
+      val tableToLoad = if (table.endsWith("_bad_table") && pollCount == 2) {
+        "nonexistent_bad_table"
+      } else {
+        table
+      }
+
       // we are reusing the table here to simulate a view query result
-      val tableConfig = sharedTableManager.getTable(share, schema, table)
+      val tableConfig = sharedTableManager.getTable(share, schema, tableToLoad)
       val capabilitiesMap = getDeltaSharingCapabilitiesMap(
         req.headers().get(DELTA_SHARING_CAPABILITIES_HEADER))
       val responseFormatSet = getResponseFormatSet(capabilitiesMap)
@@ -482,7 +506,10 @@ class DeltaSharingService(serverConfig: ServerConfig) {
 
     val requestFileIdHash = getRequestFileIdHash(req)
     if(getAsyncQuery(capabilitiesMap)) {
-      val queryId = s"${share}_${schema}_${table}"
+      // Generate unique queryId and initialize poll counter
+      val queryId = s"${share}_${schema}_${table}_${System.currentTimeMillis()}_" +
+        s"${java.util.UUID.randomUUID().toString}"
+      queryPollCounters.put(queryId, new java.util.concurrent.atomic.AtomicInteger(0))
 
       streamingOutput(
         Some(0),

spark/src/test/scala/io/delta/sharing/spark/DeltaSharingSuite.scala

@@ -603,6 +603,112 @@ class DeltaSharingSuite extends QueryTest with SharedSparkSession with DeltaShar
       .collect()
     assert(TestDeltaSharingClient.limits === Seq(1L))
   }
+
+  integrationTest("async query timeout") {
+    // Test that async query times out when it exceeds the specified timeout
+    withSQLConf(
+      "spark.delta.sharing.network.useAsyncQuery" -> "true",
+      "spark.delta.sharing.network.asyncQueryTimeout" -> "1", // 1ms timeout
+      "spark.delta.sharing.network.asyncQueryPollIntervalMillis" -> "100"
+    ) {
+      val tablePath = testProfileFile.getCanonicalPath + "#share_azure.default.table_wasb"
+
+      // Note: This test requires server-side support for async queries that take longer
+      // than timeout In a real scenario, this would timeout if the server keeps returning
+      // queryStatus
+      val ex = intercept[IllegalStateException] {
+        spark.read.format("deltaSharing").load(tablePath).collect()
+      }
+      assert(ex.getMessage.contains("Query is timed out after") ||
+             ex.getMessage.contains("1000 ms"))
+    }
+  }
+
+  integrationTest("async query initial query table errors") {
+    // Test error conditions in the initial query table request (DeltaSharingClient.scala:L650)
+    withSQLConf("spark.delta.sharing.network.useAsyncQuery" -> "true") {
+      val tablePath = testProfileFile.getCanonicalPath + "#share1.default.nonexistent_table"
+
+      // Test 1: Invalid table should throw UnexpectedHttpStatus
+      val ex1 = intercept[io.delta.sharing.client.util.UnexpectedHttpStatus] {
+        spark.read.format("deltaSharing").load(tablePath).collect()
+      }
+      assert(ex1.getMessage.contains("404") || ex1.getMessage.contains("Not Found"))
+    }
+  }
+
+  integrationTest("async query polling errors") {
+    // Test error conditions during polling (DeltaSharingClient.scala:L1099)
+    withSQLConf(
+      "spark.delta.sharing.network.useAsyncQuery" -> "true",
+      "spark.delta.sharing.network.asyncQueryPollIntervalMillis" -> "100"
+    ) {
+      // Test 1: Inconsistent queryId in polling response
+      // Server will change the queryId mid-query for tables ending with "_change_query_id"
+      val tablePath1 = testProfileFile.getCanonicalPath +
+        "#share_azure.default.table_wasb_change_query_id"
+      val ex1 = intercept[IllegalStateException] {
+        spark.read.format("deltaSharing").load(tablePath1).collect()
+      }
+      assert(ex1.getMessage.contains("QueryId is not consistent") ||
+             ex1.getMessage.contains("query"))
+
+      // Test 2: Missing queryId when query is still pending
+      // Server will return null queryId for tables ending with "_change_query_id_to_be_null"
+      val tablePath2 = testProfileFile.getCanonicalPath +
+        "#share_azure.default.table_wasb_change_query_id_to_be_null"
+      val ex2 = intercept[IllegalStateException] {
+        spark.read.format("deltaSharing").load(tablePath2).collect()
+      }
+      assert(ex2.getMessage.contains("QueryId is not returned") ||
+             ex2.getMessage.contains("null"))
+    }
+  }
+
+  integrationTest("async query handles load table failures") {
+    withSQLConf(
+      "spark.delta.sharing.network.useAsyncQuery" -> "true",
+      "spark.delta.sharing.network.asyncQueryPollIntervalMillis" -> "100"
+    ) {
+      // Test: Server encounters error loading table after polling
+      // Server will try to load a non-existent table on the 2nd poll for tables ending with
+      // "_bad_table", which will cause loadTable to fail
+      val tablePath3 = testProfileFile.getCanonicalPath +
+        "#share_azure.default.table_wasb_bad_table"
+      val ex3 = intercept[io.delta.sharing.client.util.UnexpectedHttpStatus] {
+        spark.read.format("deltaSharing").load(tablePath3).collect()
+      }
+      assert(ex3.getMessage.contains("404") || ex3.getMessage.contains("400") ||
+             ex3.getMessage.contains("does not exist"))
+    }
+  }
+
+  integrationTest("async query error handling in initial response") {
+    // Test various error conditions that can occur in getNDJsonWithAsync
+    withSQLConf("spark.delta.sharing.network.useAsyncQuery" -> "true") {
+      // Test 1: Invalid version parameter
+      val tablePath = testProfileFile.getCanonicalPath + "#share8.default.cdf_table_cdf_enabled"
+      val ex1 = intercept[io.delta.sharing.client.util.UnexpectedHttpStatus] {
+        spark.read
+          .format("deltaSharing")
+          .option("versionAsOf", "999999") // Non-existent version
+          .load(tablePath)
+          .collect()
+      }
+      assert(ex1.getMessage.contains("400") || ex1.getMessage.contains("404"))
+
+      // Test 2: Invalid timestamp parameter
+      val ex2 = intercept[io.delta.sharing.client.util.UnexpectedHttpStatus] {
+        spark.read
+          .format("deltaSharing")
+          .option("timestampAsOf", "1970-01-01 00:00:00")
+          .load(tablePath)
+          .collect()
+      }
+      assert(ex2.getMessage.contains("400") ||
+             ex2.getMessage.contains("The provided timestamp"))
+    }
+  }
 }
 
 class DeltaSharingWithParquetIOCacheEnabledSuite extends DeltaSharingSuite {