databend_driver/
rest_api.rs

1// Copyright 2021 Datafuse Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use async_trait::async_trait;
16use log::info;
17use std::collections::{BTreeMap, VecDeque};
18use std::marker::PhantomData;
19use std::path::Path;
20use std::pin::Pin;
21use std::sync::Arc;
22use std::task::{Context, Poll};
23use std::time::Instant;
24use tokio::fs::File;
25use tokio::io::BufReader;
26use tokio_stream::Stream;
27
28use crate::client::LoadMethod;
29use crate::conn::{ConnectionInfo, IConnection, Reader};
30use databend_client::APIClient;
31use databend_client::Pages;
32use databend_driver_core::error::{Error, Result};
33use databend_driver_core::raw_rows::{RawRow, RawRowIterator, RawRowWithStats};
34use databend_driver_core::rows::{Row, RowIterator, RowStatsIterator, RowWithStats, ServerStats};
35use databend_driver_core::schema::{Schema, SchemaRef};
36
37const LOAD_PLACEHOLDER: &str = "@_databend_load";
38
39#[derive(Clone)]
40pub struct RestAPIConnection {
41    client: Arc<APIClient>,
42}
43
44impl RestAPIConnection {
45    fn gen_temp_stage_location(&self) -> Result<String> {
46        let now = chrono::Utc::now()
47            .timestamp_nanos_opt()
48            .ok_or_else(|| Error::IO("Failed to get current timestamp".to_string()))?;
49        Ok(format!("@~/client/load/{now}"))
50    }
51
52    async fn load_data_with_stage(
53        &self,
54        sql: &str,
55        data: Reader,
56        size: u64,
57    ) -> Result<ServerStats> {
58        let location = self.gen_temp_stage_location()?;
59        self.upload_to_stage(&location, data, size).await?;
60        if self.client.capability().streaming_load {
61            let sql = sql.replace(LOAD_PLACEHOLDER, &location);
62            let page = self.client.query_all(&sql).await?;
63            Ok(ServerStats::from(page.stats))
64        } else {
65            let file_format_options = Self::default_file_format_options();
66            let copy_options = Self::default_copy_options();
67            let stats = self
68                .client
69                .insert_with_stage(sql, &location, file_format_options, copy_options)
70                .await?;
71            Ok(ServerStats::from(stats))
72        }
73    }
74
75    async fn load_data_with_streaming(
76        &self,
77        sql: &str,
78        data: Reader,
79        size: u64,
80    ) -> Result<ServerStats> {
81        let start = Instant::now();
82        let response = self
83            .client
84            .streaming_load(sql, data, "<no_filename>")
85            .await?;
86        Ok(ServerStats {
87            total_rows: 0,
88            total_bytes: 0,
89            read_rows: response.stats.rows,
90            read_bytes: size as usize,
91            write_rows: response.stats.rows,
92            write_bytes: response.stats.bytes,
93            running_time_ms: start.elapsed().as_millis() as f64,
94            spill_file_nums: 0,
95            spill_bytes: 0,
96        })
97    }
98    async fn load_data_with_options(
99        &self,
100        sql: &str,
101        data: Reader,
102        size: u64,
103        file_format_options: Option<BTreeMap<&str, &str>>,
104        copy_options: Option<BTreeMap<&str, &str>>,
105    ) -> Result<ServerStats> {
106        let location = self.gen_temp_stage_location()?;
107        let file_format_options =
108            file_format_options.unwrap_or_else(Self::default_file_format_options);
109        let copy_options = copy_options.unwrap_or_else(Self::default_copy_options);
110        self.upload_to_stage(&location, Box::new(data), size)
111            .await?;
112        let stats = self
113            .client
114            .insert_with_stage(sql, &location, file_format_options, copy_options)
115            .await?;
116        Ok(ServerStats::from(stats))
117    }
118}
119
120#[async_trait]
121impl IConnection for RestAPIConnection {
122    async fn info(&self) -> ConnectionInfo {
123        ConnectionInfo {
124            handler: "RestAPI".to_string(),
125            host: self.client.host().to_string(),
126            port: self.client.port(),
127            user: self.client.username(),
128            catalog: self.client.current_catalog(),
129            database: self.client.current_database(),
130            warehouse: self.client.current_warehouse(),
131        }
132    }
133
134    fn last_query_id(&self) -> Option<String> {
135        self.client.last_query_id()
136    }
137
138    async fn close(&self) -> Result<()> {
139        self.client.close().await;
140        Ok(())
141    }
142
143    async fn exec(&self, sql: &str) -> Result<i64> {
144        info!("exec: {}", sql);
145        let page = self.client.query_all(sql).await?;
146        Ok(page.stats.progresses.write_progress.rows as i64)
147    }
148
149    async fn kill_query(&self, query_id: &str) -> Result<()> {
150        Ok(self.client.kill_query(query_id).await?)
151    }
152
153    async fn query_iter(&self, sql: &str) -> Result<RowIterator> {
154        info!("query iter: {}", sql);
155        let rows_with_progress = self.query_iter_ext(sql).await?;
156        let rows = rows_with_progress.filter_rows().await;
157        Ok(rows)
158    }
159
160    async fn query_iter_ext(&self, sql: &str) -> Result<RowStatsIterator> {
161        info!("query iter ext: {}", sql);
162        let pages = self.client.start_query(sql, true).await?;
163        let (schema, rows) = RestAPIRows::<RowWithStats>::from_pages(pages).await?;
164        Ok(RowStatsIterator::new(Arc::new(schema), Box::pin(rows)))
165    }
166
167    // raw data response query, only for test
168    async fn query_raw_iter(&self, sql: &str) -> Result<RawRowIterator> {
169        info!("query raw iter: {}", sql);
170        let pages = self.client.start_query(sql, true).await?;
171        let (schema, rows) = RestAPIRows::<RawRowWithStats>::from_pages(pages).await?;
172        Ok(RawRowIterator::new(Arc::new(schema), Box::pin(rows)))
173    }
174
175    async fn upload_to_stage(&self, stage: &str, data: Reader, size: u64) -> Result<()> {
176        self.client.upload_to_stage(stage, data, size).await?;
177        Ok(())
178    }
179
180    async fn load_data(
181        &self,
182        sql: &str,
183        data: Reader,
184        size: u64,
185        method: LoadMethod,
186    ) -> Result<ServerStats> {
187        let sql = sql.trim_end();
188        let sql = sql.trim_end_matches(';');
189        info!("load data: {}, size: {}, method: {method:?}", sql, size);
190        let sql_low = sql.to_lowercase();
191        let has_place_holder = sql_low.contains(LOAD_PLACEHOLDER);
192        let sql = match (self.client.capability().streaming_load, has_place_holder) {
193            (false, false) => {
194                // todo: deprecate this later
195                return self
196                    .load_data_with_options(sql, data, size, None, None)
197                    .await;
198            }
199            (false, true) => return Err(Error::BadArgument(
200                "Please upgrade your server to >= 1.2.781 to support insert from @_databend_load"
201                    .to_string(),
202            )),
203            (true, false) => {
204                format!("{sql} from @_databend_load file_format=(type=csv)")
205            }
206            (true, true) => sql.to_string(),
207        };
208
209        match method {
210            LoadMethod::Streaming => self.load_data_with_streaming(&sql, data, size).await,
211            LoadMethod::Stage => self.load_data_with_stage(&sql, data, size).await,
212        }
213    }
214
215    async fn load_file(&self, sql: &str, fp: &Path, method: LoadMethod) -> Result<ServerStats> {
216        info!("load file: {}, file: {:?}", sql, fp,);
217        let file = File::open(fp).await?;
218        let metadata = file.metadata().await?;
219        let size = metadata.len();
220        let data = BufReader::new(file);
221        self.load_data(sql, Box::new(data), size, method).await
222    }
223
224    async fn load_file_with_options(
225        &self,
226        sql: &str,
227        fp: &Path,
228        file_format_options: Option<BTreeMap<&str, &str>>,
229        copy_options: Option<BTreeMap<&str, &str>>,
230    ) -> Result<ServerStats> {
231        let file = File::open(fp).await?;
232        let metadata = file.metadata().await?;
233        let size = metadata.len();
234        let data = BufReader::new(file);
235        self.load_data_with_options(sql, Box::new(data), size, file_format_options, copy_options)
236            .await
237    }
238
239    async fn stream_load(
240        &self,
241        sql: &str,
242        data: Vec<Vec<&str>>,
243        method: LoadMethod,
244    ) -> Result<ServerStats> {
245        info!("stream load: {}; rows: {:?}", sql, data.len());
246        let mut wtr = csv::WriterBuilder::new().from_writer(vec![]);
247        for row in data {
248            wtr.write_record(row)
249                .map_err(|e| Error::BadArgument(e.to_string()))?;
250        }
251        let bytes = wtr.into_inner().map_err(|e| Error::IO(e.to_string()))?;
252        let size = bytes.len() as u64;
253        let reader = Box::new(std::io::Cursor::new(bytes));
254        let stats = if self.client.capability().streaming_load {
255            let sql = format!("{sql} from @_databend_load file_format = (type = csv)");
256            self.load_data(&sql, reader, size, method).await?
257        } else {
258            self.load_data_with_options(sql, reader, size, None, None)
259                .await?
260        };
261        Ok(stats)
262    }
263}
264
265impl<'o> RestAPIConnection {
266    pub async fn try_create(dsn: &str, name: String) -> Result<Self> {
267        let client = APIClient::new(dsn, Some(name)).await?;
268        Ok(Self { client })
269    }
270
271    fn default_file_format_options() -> BTreeMap<&'o str, &'o str> {
272        vec![
273            ("type", "CSV"),
274            ("field_delimiter", ","),
275            ("record_delimiter", "\n"),
276            ("skip_header", "0"),
277        ]
278        .into_iter()
279        .collect()
280    }
281
282    fn default_copy_options() -> BTreeMap<&'o str, &'o str> {
283        vec![("purge", "true")].into_iter().collect()
284    }
285}
286
287pub struct RestAPIRows<T> {
288    pages: Pages,
289
290    schema: SchemaRef,
291    data: VecDeque<Vec<Option<String>>>,
292    stats: Option<ServerStats>,
293
294    _phantom: std::marker::PhantomData<T>,
295}
296
297impl<T> RestAPIRows<T> {
298    async fn from_pages(pages: Pages) -> Result<(Schema, Self)> {
299        let (pages, schema) = pages.wait_for_schema(true).await?;
300        let schema: Schema = schema.try_into()?;
301        let rows = Self {
302            pages,
303            schema: Arc::new(schema.clone()),
304            data: Default::default(),
305            stats: None,
306            _phantom: PhantomData,
307        };
308        Ok((schema, rows))
309    }
310}
311
312impl<T: FromRowStats + std::marker::Unpin> Stream for RestAPIRows<T> {
313    type Item = Result<T>;
314
315    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
316        if let Some(ss) = self.stats.take() {
317            return Poll::Ready(Some(Ok(T::from_stats(ss))));
318        }
319        // Skip to fetch next page if there is only one row left in buffer.
320        // Therefore, we could guarantee the `/final` called before the last row.
321        if self.data.len() > 1 {
322            if let Some(row) = self.data.pop_front() {
323                let row = T::try_from_row(row, self.schema.clone())?;
324                return Poll::Ready(Some(Ok(row)));
325            }
326        }
327
328        match Pin::new(&mut self.pages).poll_next(cx) {
329            Poll::Ready(Some(Ok(page))) => {
330                if self.schema.fields().is_empty() {
331                    self.schema = Arc::new(page.schema.try_into()?);
332                }
333                let mut new_data = page.data.into();
334                self.data.append(&mut new_data);
335                Poll::Ready(Some(Ok(T::from_stats(page.stats.into()))))
336            }
337            Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e.into()))),
338            Poll::Ready(None) => match self.data.pop_front() {
339                Some(row) => {
340                    let row = T::try_from_row(row, self.schema.clone())?;
341                    Poll::Ready(Some(Ok(row)))
342                }
343                None => Poll::Ready(None),
344            },
345            Poll::Pending => Poll::Pending,
346        }
347    }
348}
349
350trait FromRowStats: Send + Sync + Clone {
351    fn from_stats(stats: ServerStats) -> Self;
352    fn try_from_row(row: Vec<Option<String>>, schema: SchemaRef) -> Result<Self>;
353}
354
355impl FromRowStats for RowWithStats {
356    fn from_stats(stats: ServerStats) -> Self {
357        RowWithStats::Stats(stats)
358    }
359
360    fn try_from_row(row: Vec<Option<String>>, schema: SchemaRef) -> Result<Self> {
361        Ok(RowWithStats::Row(Row::try_from((schema, row))?))
362    }
363}
364
365impl FromRowStats for RawRowWithStats {
366    fn from_stats(stats: ServerStats) -> Self {
367        RawRowWithStats::Stats(stats)
368    }
369
370    fn try_from_row(row: Vec<Option<String>>, schema: SchemaRef) -> Result<Self> {
371        let rows = Row::try_from((schema, row.clone()))?;
372        Ok(RawRowWithStats::Row(RawRow::new(rows, row)))
373    }
374}
databend_driver/rest_api.rs

databend_driver/
rest_api.rs
