feat: add ttc member (tcp testing container) (#554)

* feat: add ttc member (tcp testing container)

* chore(cli): fix

* chore(cli): fix

* chore(cli): fix

* update

* update

* update

* update

* update

* update

* update

Author: sundy-li

.dockerignore

@@ -0,0 +1,3 @@
+target/
+frontend
+ttc/Dockerfile

.github/workflows/ttc.yml

@@ -0,0 +1,36 @@
+# .github/workflows/docker-publish.yml
+
+name: Docker
+
+on:
+  push:
+    branches:
+      - main
+    paths:
+      - 'core/**'
+      - 'driver/**'
+      - 'ttc/**'
+
+jobs:
+  docker:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v2
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v1
+
+      - name: Login to DockerHub
+        uses: docker/login-action@v3
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
+
+      - name: Build and push
+        uses: docker/build-push-action@v2
+        with:
+          push: true
+          file: ./ttc/Dockerfile
+          tags: datafuselabs/ttc-rust:latest

Cargo.toml

@@ -8,6 +8,7 @@ members = [
     "cli",
     "bindings/python",
     "bindings/nodejs",
+    "ttc",
 ]
 resolver = "2"
 

driver/src/conn.rs

@@ -17,6 +17,7 @@ use std::path::Path;
 use std::sync::Arc;
 
 use async_trait::async_trait;
+use databend_driver_core::raw_rows::{RawRow, RawRowIterator};
 use once_cell::sync::Lazy;
 use tokio::io::AsyncRead;
 use tokio_stream::StreamExt;
@@ -121,6 +122,19 @@ pub trait Connection: Send + Sync {
         rows.collect().await
     }
 
+    // raw data response query, only for test
+    async fn query_raw_iter(&self, _sql: &str) -> Result<RawRowIterator> {
+        Err(Error::BadArgument(
+            "Unsupported implement query_raw_iter".to_string(),
+        ))
+    }
+
+    // raw data response query, only for test
+    async fn query_raw_all(&self, sql: &str) -> Result<Vec<RawRow>> {
+        let rows = self.query_raw_iter(sql).await?;
+        rows.collect().await
+    }
+
     /// Get presigned url for a given operation and stage location.
     /// The operation can be "UPLOAD" or "DOWNLOAD".
     async fn get_presigned_url(&self, operation: &str, stage: &str) -> Result<PresignedResponse>;

driver/src/rest_api.rs

@@ -15,13 +15,15 @@
 use std::collections::{BTreeMap, VecDeque};
 use std::future::Future;
 use std::io::Cursor;
+use std::marker::PhantomData;
 use std::path::Path;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::task::{Context, Poll};
 
 use async_compression::tokio::write::ZstdEncoder;
 use async_trait::async_trait;
+use databend_driver_core::raw_rows::{RawRow, RawRowIterator, RawRowWithStats};
 use log::info;
 use tokio::fs::File;
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
@@ -83,10 +85,20 @@ impl Connection for RestAPIConnection {
         info!("query iter ext: {}", sql);
         let resp = self.client.start_query(sql).await?;
         let resp = self.wait_for_schema(resp).await?;
-        let (schema, rows) = RestAPIRows::from_response(self.client.clone(), resp)?;
+        let (schema, rows) = RestAPIRows::<RowWithStats>::from_response(self.client.clone(), resp)?;
         Ok(RowStatsIterator::new(Arc::new(schema), Box::pin(rows)))
     }
 
+    // raw data response query, only for test
+    async fn query_raw_iter(&self, sql: &str) -> Result<RawRowIterator> {
+        info!("query raw iter: {}", sql);
+        let resp = self.client.start_query(sql).await?;
+        let resp = self.wait_for_schema(resp).await?;
+        let (schema, rows) =
+            RestAPIRows::<RawRowWithStats>::from_response(self.client.clone(), resp)?;
+        Ok(RawRowIterator::new(Arc::new(schema), Box::pin(rows)))
+    }
+
     async fn get_presigned_url(&self, operation: &str, stage: &str) -> Result<PresignedResponse> {
         info!("get presigned url: {} {}", operation, stage);
         let sql = format!("PRESIGN {} {}", operation, stage);
@@ -254,7 +266,7 @@ impl<'o> RestAPIConnection {
 
 type PageFut = Pin<Box<dyn Future<Output = Result<QueryResponse>> + Send>>;
 
-pub struct RestAPIRows {
+pub struct RestAPIRows<T> {
     client: Arc<APIClient>,
     schema: SchemaRef,
     data: VecDeque<Vec<Option<String>>>,
@@ -263,9 +275,10 @@ pub struct RestAPIRows {
     node_id: Option<String>,
     next_uri: Option<String>,
     next_page: Option<PageFut>,
+    _phantom: std::marker::PhantomData<T>,
 }
 
-impl RestAPIRows {
+impl<T> RestAPIRows<T> {
     fn from_response(client: Arc<APIClient>, resp: QueryResponse) -> Result<(Schema, Self)> {
         let schema: Schema = resp.schema.try_into()?;
         let rows = Self {
@@ -277,24 +290,25 @@ impl RestAPIRows {
             data: resp.data.into(),
             stats: Some(ServerStats::from(resp.stats)),
             next_page: None,
+            _phantom: PhantomData,
         };
         Ok((schema, rows))
     }
 }
 
-impl Stream for RestAPIRows {
-    type Item = Result<RowWithStats>;
+impl<T: FromRowStats + std::marker::Unpin> Stream for RestAPIRows<T> {
+    type Item = Result<T>;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         if let Some(ss) = self.stats.take() {
-            return Poll::Ready(Some(Ok(RowWithStats::Stats(ss))));
+            return Poll::Ready(Some(Ok(T::from_stats(ss))));
         }
         // Skip to fetch next page if there is only one row left in buffer.
         // Therefore we could guarantee the `/final` called before the last row.
         if self.data.len() > 1 {
             if let Some(row) = self.data.pop_front() {
-                let row = Row::try_from((self.schema.clone(), row))?;
-                return Poll::Ready(Some(Ok(RowWithStats::Row(row))));
+                let row = T::try_from_row(row, self.schema.clone())?;
+                return Poll::Ready(Some(Ok(row)));
             }
         }
         match self.next_page {
@@ -307,8 +321,7 @@ impl Stream for RestAPIRows {
                     self.next_page = None;
                     let mut new_data = resp.data.into();
                     self.data.append(&mut new_data);
-                    let stats = ServerStats::from(resp.stats);
-                    Poll::Ready(Some(Ok(RowWithStats::Stats(stats))))
+                    Poll::Ready(Some(Ok(T::from_stats(resp.stats.into()))))
                 }
                 Poll::Ready(Err(e)) => {
                     self.next_page = None;
@@ -332,12 +345,38 @@ impl Stream for RestAPIRows {
                 }
                 None => match self.data.pop_front() {
                     Some(row) => {
-                        let row = Row::try_from((self.schema.clone(), row))?;
-                        Poll::Ready(Some(Ok(RowWithStats::Row(row))))
+                        let row = T::try_from_row(row, self.schema.clone())?;
+                        Poll::Ready(Some(Ok(row)))
                     }
                     None => Poll::Ready(None),
                 },
             },
         }
     }
 }
+
+trait FromRowStats: Send + Sync + Clone {
+    fn from_stats(stats: ServerStats) -> Self;
+    fn try_from_row(row: Vec<Option<String>>, schema: SchemaRef) -> Result<Self>;
+}
+
+impl FromRowStats for RowWithStats {
+    fn from_stats(stats: ServerStats) -> Self {
+        RowWithStats::Stats(stats)
+    }
+
+    fn try_from_row(row: Vec<Option<String>>, schema: SchemaRef) -> Result<Self> {
+        Ok(RowWithStats::Row(Row::try_from((schema, row))?))
+    }
+}
+
+impl FromRowStats for RawRowWithStats {
+    fn from_stats(stats: ServerStats) -> Self {
+        RawRowWithStats::Stats(stats)
+    }
+
+    fn try_from_row(row: Vec<Option<String>>, schema: SchemaRef) -> Result<Self> {
+        let rows = Row::try_from((schema, row.clone()))?;
+        Ok(RawRowWithStats::Row(RawRow::new(rows, row)))
+    }
+}

licenserc.toml

@@ -14,6 +14,7 @@ excludes = [
     ".prettierignore",
 
     "LICENSE",
+    "Dockerfile",
     "Makefile",
 
     # docs and generated files

sql/src/lib.rs

@@ -14,6 +14,7 @@
 
 mod cursor_ext;
 pub mod error;
+pub mod raw_rows;
 pub mod rows;
 pub mod schema;
 pub mod value;

sql/src/raw_rows.rs

@@ -0,0 +1,150 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::pin::Pin;
+use std::task::Context;
+use std::task::Poll;
+
+use tokio_stream::{Stream, StreamExt};
+
+use crate::error::Error;
+use crate::error::Result;
+use crate::rows::Row;
+use crate::rows::ServerStats;
+use crate::schema::SchemaRef;
+use crate::value::Value;
+
+#[derive(Clone, Debug)]
+pub enum RawRowWithStats {
+    Row(RawRow),
+    Stats(ServerStats),
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct RawRow {
+    pub row: Row,
+    pub raw_row: Vec<Option<String>>,
+}
+
+impl RawRow {
+    pub fn new(row: Row, raw_row: Vec<Option<String>>) -> Self {
+        Self { row, raw_row }
+    }
+
+    pub fn len(&self) -> usize {
+        self.raw_row.len()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.raw_row.is_empty()
+    }
+
+    pub fn values(&self) -> &[Option<String>] {
+        &self.raw_row
+    }
+
+    pub fn schema(&self) -> SchemaRef {
+        self.row.schema()
+    }
+}
+
+impl TryFrom<(SchemaRef, Vec<Option<String>>)> for RawRow {
+    type Error = Error;
+
+    fn try_from((schema, data): (SchemaRef, Vec<Option<String>>)) -> Result<Self> {
+        let mut values: Vec<Value> = Vec::new();
+        for (i, field) in schema.fields().iter().enumerate() {
+            let val: Option<&str> = data.get(i).and_then(|v| v.as_deref());
+            values.push(Value::try_from((&field.data_type, val))?);
+        }
+
+        let row = Row::new(schema, values);
+        Ok(RawRow::new(row, data))
+    }
+}
+
+impl IntoIterator for RawRow {
+    type Item = Option<String>;
+    type IntoIter = std::vec::IntoIter<Self::Item>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.raw_row.into_iter()
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct RawRows {
+    rows: Vec<RawRow>,
+}
+
+impl RawRows {
+    pub fn new(rows: Vec<RawRow>) -> Self {
+        Self { rows }
+    }
+
+    pub fn rows(&self) -> &[RawRow] {
+        &self.rows
+    }
+
+    pub fn len(&self) -> usize {
+        self.rows.len()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.rows.is_empty()
+    }
+}
+
+impl IntoIterator for RawRows {
+    type Item = RawRow;
+    type IntoIter = std::vec::IntoIter<Self::Item>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.rows.into_iter()
+    }
+}
+
+pub struct RawRowIterator {
+    schema: SchemaRef,
+    it: Pin<Box<dyn Stream<Item = Result<RawRow>> + Send>>,
+}
+
+impl RawRowIterator {
+    pub fn new(
+        schema: SchemaRef,
+        it: Pin<Box<dyn Stream<Item = Result<RawRowWithStats>> + Send>>,
+    ) -> Self {
+        let it = it.filter_map(|r| match r {
+            Ok(RawRowWithStats::Row(r)) => Some(Ok(r)),
+            Ok(_) => None,
+            Err(err) => Some(Err(err)),
+        });
+        Self {
+            schema,
+            it: Box::pin(it),
+        }
+    }
+
+    pub fn schema(&self) -> SchemaRef {
+        self.schema.clone()
+    }
+}
+
+impl Stream for RawRowIterator {
+    type Item = Result<RawRow>;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        Pin::new(&mut self.it).poll_next(cx)
+    }
+}