google_cloud_bigquery/
query.rs

1pub use backon::*;
2
3use crate::{http, storage};
4
5#[derive(Debug, Clone)]
6pub struct QueryOption {
7    /// Exponential back off retry setting
8    pub(crate) retry: ExponentialBuilder,
9    /// true: use storage api is page token is empty
10    pub(crate) enable_storage_read: bool,
11}
12
13impl Default for QueryOption {
14    fn default() -> Self {
15        Self {
16            enable_storage_read: false,
17            retry: ExponentialBuilder::default().with_max_times(usize::MAX),
18        }
19    }
20}
21
22impl QueryOption {
23    pub fn with_retry(mut self, builder: ExponentialBuilder) -> Self {
24        self.retry = builder;
25        self
26    }
27    pub fn with_enable_storage_read(mut self, value: bool) -> Self {
28        self.enable_storage_read = value;
29        self
30    }
31}
32
33#[derive(thiserror::Error, Debug)]
34pub enum Error {
35    #[error(transparent)]
36    Http(#[from] http::query::Error),
37    #[error(transparent)]
38    Storage(#[from] storage::Error),
39}
40
41pub enum QueryResult<T: http::query::value::StructDecodable + storage::value::StructDecodable> {
42    Http(http::query::Iterator<T>),
43    Storage(storage::Iterator<T>),
44}
45
46pub struct Iterator<T: http::query::value::StructDecodable + storage::value::StructDecodable> {
47    pub(crate) inner: QueryResult<T>,
48    pub total_size: i64,
49}
50
51impl<T: http::query::value::StructDecodable + storage::value::StructDecodable> Iterator<T> {
52    pub async fn next(&mut self) -> Result<Option<T>, Error> {
53        Ok(match self.inner {
54            QueryResult::Storage(ref mut v) => v.next().await?,
55            QueryResult::Http(ref mut v) => v.next().await?,
56        })
57    }
58}
59
60pub mod row {
61    use crate::http::tabledata::list::Tuple;
62    use crate::{http, storage};
63    use arrow::array::ArrayRef;
64
65    #[derive(thiserror::Error, Debug)]
66    pub enum Error {
67        #[error(transparent)]
68        Http(#[from] http::query::row::Error),
69        #[error(transparent)]
70        Storage(#[from] storage::row::Error),
71    }
72
73    pub enum RowType {
74        Http(http::query::row::Row),
75        Storage(storage::row::Row),
76    }
77
78    pub struct Row {
79        inner: RowType,
80    }
81
82    impl Row {
83        pub fn column<T: http::query::value::Decodable + storage::value::Decodable>(
84            &self,
85            index: usize,
86        ) -> Result<T, Error> {
87            Ok(match &self.inner {
88                RowType::Http(row) => row.column(index)?,
89                RowType::Storage(row) => row.column(index)?,
90            })
91        }
92    }
93
94    impl http::query::value::StructDecodable for Row {
95        fn decode(value: Tuple) -> Result<Self, http::query::value::Error> {
96            Ok(Self {
97                inner: RowType::Http(http::query::row::Row::decode(value)?),
98            })
99        }
100    }
101
102    impl storage::value::StructDecodable for Row {
103        fn decode_arrow(fields: &[ArrayRef], row_no: usize) -> Result<Self, storage::value::Error> {
104            Ok(Self {
105                inner: RowType::Storage(storage::row::Row::decode_arrow(fields, row_no)?),
106            })
107        }
108    }
109}
110
111pub mod run {
112    #[derive(thiserror::Error, Debug)]
113    pub enum Error {
114        #[error(transparent)]
115        Http(#[from] crate::http::error::Error),
116        #[error("Retry exceeded with job incomplete")]
117        JobIncomplete,
118    }
119}