google_cloud_bigquery/http/
bigquery_dataset_client.rs

1use std::sync::Arc;
2
3use crate::http::bigquery_client::BigqueryClient;
4use crate::http::dataset;
5use crate::http::dataset::list::{DatasetOverview, ListDatasetsRequest, ListDatasetsResponse};
6use crate::http::dataset::Dataset;
7use crate::http::error::Error;
8
9#[derive(Debug, Clone)]
10pub struct BigqueryDatasetClient {
11    inner: Arc<BigqueryClient>,
12}
13
14impl BigqueryDatasetClient {
15    pub fn new(inner: Arc<BigqueryClient>) -> Self {
16        Self { inner }
17    }
18
19    /// https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets/create
20    #[cfg_attr(feature = "trace", tracing::instrument(skip_all))]
21    pub async fn create(&self, metadata: &Dataset) -> Result<Dataset, Error> {
22        let builder = dataset::insert::build(self.inner.endpoint(), self.inner.http(), metadata);
23        self.inner.send(builder).await
24    }
25
26    /// https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets/patch
27    #[cfg_attr(feature = "trace", tracing::instrument(skip_all))]
28    pub async fn patch(&self, metadata: &Dataset) -> Result<Dataset, Error> {
29        let builder = dataset::patch::build(self.inner.endpoint(), self.inner.http(), metadata);
30        self.inner.send(builder).await
31    }
32
33    /// https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets/delete
34    #[cfg_attr(feature = "trace", tracing::instrument(skip_all))]
35    pub async fn delete(&self, project_id: &str, dataset_id: &str) -> Result<(), Error> {
36        let builder = dataset::delete::build(self.inner.endpoint(), self.inner.http(), project_id, dataset_id);
37        self.inner.send_get_empty(builder).await
38    }
39
40    /// https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets/get
41    #[cfg_attr(feature = "trace", tracing::instrument(skip_all))]
42    pub async fn get(&self, project_id: &str, dataset_id: &str) -> Result<Dataset, Error> {
43        let builder = dataset::get::build(self.inner.endpoint(), self.inner.http(), project_id, dataset_id);
44        self.inner.send(builder).await
45    }
46
47    /// https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets/list
48    #[cfg_attr(feature = "trace", tracing::instrument(skip_all))]
49    pub async fn list(
50        &self,
51        project_id: &str,
52        req: Option<&ListDatasetsRequest>,
53    ) -> Result<Vec<DatasetOverview>, Error> {
54        let mut page_token: Option<String> = None;
55        let mut datasets = vec![];
56        loop {
57            let builder = dataset::list::build(self.inner.endpoint(), self.inner.http(), project_id, req, page_token);
58            let response: ListDatasetsResponse = self.inner.send(builder).await?;
59            datasets.extend(response.datasets);
60            if response.next_page_token.is_none() {
61                break;
62            }
63            page_token = response.next_page_token;
64        }
65        Ok(datasets)
66    }
67}
68
69#[cfg(test)]
70mod test {
71    use std::collections::HashMap;
72    use std::sync::Arc;
73
74    use serial_test::serial;
75
76    use crate::http::bigquery_client::test::{create_client, dataset_name};
77    use crate::http::bigquery_dataset_client::BigqueryDatasetClient;
78    use crate::http::dataset::list::ListDatasetsRequest;
79    use crate::http::dataset::{Access, Dataset, DatasetReference, SpecialGroup, StorageBillingModel};
80    use crate::http::types::{Collation, EncryptionConfiguration};
81
82    #[tokio::test]
83    #[serial]
84    pub async fn crud_dataset() {
85        let (client, project) = create_client().await;
86        let client = BigqueryDatasetClient::new(Arc::new(client));
87
88        // full prop dataset
89        let mut labels = HashMap::new();
90        labels.insert("key".to_string(), "value".to_string());
91        let ds2 = Dataset {
92            dataset_reference: DatasetReference {
93                dataset_id: dataset_name("crud_full"),
94                project_id: project.to_string(),
95            },
96            friendly_name: Some("gcr_test_friendly_name".to_string()),
97            description: Some("gcr_test_description".to_string()),
98            default_table_expiration_ms: Some(3600000),
99            default_partition_expiration_ms: Some(3600000),
100            labels: Some(labels),
101            access: vec![Access {
102                role: "READER".to_string(),
103                special_group: Some(SpecialGroup::AllAuthenticatedUsers),
104                ..Default::default()
105            }],
106            location: "asia-northeast1".to_string(),
107            default_encryption_configuration: Some(EncryptionConfiguration {
108                kms_key_name: Some(format!(
109                    "projects/{}/locations/asia-northeast1/keyRings/gcr_test/cryptoKeys/gcr_test",
110                    project.as_str()
111                )),
112            }),
113            is_case_insensitive: Some(true),
114            default_collation: Some(Collation::UndeterminedLocaleCaseInsensitive),
115            max_time_travel_hours: Some(48),
116            storage_billing_model: Some(StorageBillingModel::Physical), // 'Logical' doesn't return None in get request
117            ..Default::default()
118        };
119        let ds2 = client.create(&ds2).await.unwrap();
120
121        // minimum dataset
122        let mut ds1 = Dataset::default();
123        ds1.dataset_reference.dataset_id = dataset_name("crud_empty");
124        ds1.dataset_reference.project_id.clone_from(&project);
125        ds1 = client.create(&ds1).await.unwrap();
126
127        // test get
128        let mut res1 = client
129            .get(project.as_str(), &ds1.dataset_reference.dataset_id)
130            .await
131            .unwrap();
132        let res2 = client
133            .get(project.as_str(), &ds2.dataset_reference.dataset_id)
134            .await
135            .unwrap();
136
137        // default is 168
138        ds1.max_time_travel_hours = Some(168);
139        assert_eq!(ds1, res1);
140        assert_eq!(ds2, res2);
141
142        // test update
143        res1.description = Some("rust_test_empty_updated".to_string());
144        client.patch(&res1).await.unwrap();
145
146        // test list
147        let result = client.list(project.as_str(), None).await.unwrap();
148        assert!(result.len() >= 2);
149        let result = client
150            .list(
151                project.as_str(),
152                Some(&ListDatasetsRequest {
153                    max_results: Some(100),
154                    all: true,
155                    filter: "".to_string(),
156                }),
157            )
158            .await
159            .unwrap();
160        assert!(result.len() >= 2);
161
162        let result = client
163            .list(
164                project.as_str(),
165                Some(&ListDatasetsRequest {
166                    max_results: None,
167                    all: true,
168                    filter: "labels.key:value".to_string(),
169                }),
170            )
171            .await
172            .unwrap();
173        assert_eq!(1, result.len());
174        assert_eq!(res2.id, result[0].id);
175
176        // test delete
177        client
178            .delete(project.as_str(), ds1.dataset_reference.dataset_id.as_str())
179            .await
180            .unwrap();
181        client
182            .delete(project.as_str(), ds2.dataset_reference.dataset_id.as_str())
183            .await
184            .unwrap();
185    }
186}