fix(driver): do not kill query for query_row (#371)

databendlabs · Mar 21, 2024 · d7ec1e9 · d7ec1e9
1 parent bdcb149
commit d7ec1e9
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 49 deletions.
diff --git a/driver/src/conn.rs b/driver/src/conn.rs
@@ -105,13 +105,19 @@ pub trait Connection: DynClone + Send + Sync {
     }
 
     async fn exec(&self, sql: &str) -> Result<i64>;
-    async fn query_row(&self, sql: &str) -> Result<Option<Row>>;
+    async fn query_iter(&self, sql: &str) -> Result<RowIterator>;
+    async fn query_iter_ext(&self, sql: &str) -> Result<RowStatsIterator>;
+
+    async fn query_row(&self, sql: &str) -> Result<Option<Row>> {
+        let rows = self.query_all(sql).await?;
+        let row = rows.into_iter().next();
+        Ok(row)
+    }
+
     async fn query_all(&self, sql: &str) -> Result<Vec<Row>> {
         let rows = self.query_iter(sql).await?;
         rows.collect().await
     }
-    async fn query_iter(&self, sql: &str) -> Result<RowIterator>;
-    async fn query_iter_ext(&self, sql: &str) -> Result<RowStatsIterator>;
 
     /// Get presigned url for a given operation and stage location.
     /// The operation can be "UPLOAD" or "DOWNLOAD".

diff --git a/driver/src/flight_sql.rs b/driver/src/flight_sql.rs
@@ -65,12 +65,6 @@ impl Connection for FlightSQLConnection {
         Ok(affected_rows)
     }
 
-    async fn query_row(&self, sql: &str) -> Result<Option<Row>> {
-        let mut rows = self.query_iter(sql).await?;
-        let row = rows.try_next().await?;
-        Ok(row)
-    }
-
     async fn query_iter(&self, sql: &str) -> Result<RowIterator> {
         let rows_with_progress = self.query_iter_ext(sql).await?;
         let rows = rows_with_progress.filter_rows().await;

diff --git a/driver/src/rest_api.rs b/driver/src/rest_api.rs
@@ -76,27 +76,6 @@ impl Connection for RestAPIConnection {
         Ok(RowStatsIterator::new(Arc::new(schema), Box::pin(rows)))
     }
 
-    async fn query_row(&self, sql: &str) -> Result<Option<Row>> {
-        info!("query row: {}", sql);
-        let resp = self.client.start_query(sql).await?;
-        let resp = self.wait_for_data(resp).await?;
-        match resp.kill_uri {
-            Some(uri) => self
-                .client
-                .kill_query(&resp.id, &uri)
-                .await
-                .map_err(|e| e.into()),
-            None => Err(Error::InvalidResponse("kill_uri is empty".to_string())),
-        }?;
-        let schema = resp.schema.try_into()?;
-        if resp.data.is_empty() {
-            Ok(None)
-        } else {
-            let row = Row::try_from((Arc::new(schema), &resp.data[0]))?;
-            Ok(Some(row))
-        }
-    }
-
     async fn get_presigned_url(&self, operation: &str, stage: &str) -> Result<PresignedResponse> {
         info!("get presigned url: {} {}", operation, stage);
         let sql = format!("PRESIGN {} {}", operation, stage);
@@ -196,25 +175,6 @@ impl<'o> RestAPIConnection {
         Ok(Self { client })
     }
 
-    async fn wait_for_data(&self, pre: QueryResponse) -> Result<QueryResponse> {
-        if !pre.data.is_empty() {
-            return Ok(pre);
-        }
-        // preserve schema since it is not included in the final response in old servers
-        let pre_schema = pre.schema.clone();
-        let mut result = pre;
-        while let Some(next_uri) = result.next_uri {
-            result = self.client.query_page(&result.id, &next_uri).await?;
-            if !result.data.is_empty() {
-                break;
-            }
-        }
-        if result.schema.is_empty() {
-            result.schema = pre_schema;
-        }
-        Ok(result)
-    }
-
     async fn wait_for_schema(&self, pre: QueryResponse) -> Result<QueryResponse> {
         if !pre.data.is_empty() || !pre.schema.is_empty() {
             return Ok(pre);