gcp_bigquery_client/
model_api.rs

1//! Manage BigQuery models.
2use std::sync::Arc;
3
4use reqwest::Client;
5
6use crate::auth::Authenticator;
7use crate::error::BQError;
8use crate::model::list_models_response::ListModelsResponse;
9use crate::model::model::Model;
10use crate::{process_response, urlencode, BIG_QUERY_V2_URL};
11
12/// A model API handler.
13#[derive(Clone)]
14pub struct ModelApi {
15    client: Client,
16    auth: Arc<dyn Authenticator>,
17    base_url: String,
18}
19
20impl ModelApi {
21    pub(crate) fn new(client: Client, auth: Arc<dyn Authenticator>) -> Self {
22        Self {
23            client,
24            auth,
25            base_url: BIG_QUERY_V2_URL.to_string(),
26        }
27    }
28
29    pub(crate) fn with_base_url(&mut self, base_url: String) -> &mut Self {
30        self.base_url = base_url;
31        self
32    }
33
34    /// Lists all models in the specified dataset. Requires the READER dataset role.
35    /// # Arguments
36    /// * `project_id` - Project ID of the models to list.
37    /// * `dataset_id` - Dataset ID of the models to list.
38    pub async fn list(
39        &self,
40        project_id: &str,
41        dataset_id: &str,
42        options: ListOptions,
43    ) -> Result<ListModelsResponse, BQError> {
44        let req_url = format!(
45            "{base_url}/projects/{project_id}/datasets/{dataset_id}/models",
46            base_url = self.base_url,
47            project_id = urlencode(project_id),
48            dataset_id = urlencode(dataset_id),
49        );
50
51        let access_token = self.auth.access_token().await?;
52
53        let request = self
54            .client
55            .get(req_url.as_str())
56            .bearer_auth(access_token)
57            .query(&options)
58            .build()?;
59
60        let resp = self.client.execute(request).await?;
61
62        process_response(resp).await
63    }
64
65    /// Deletes the model specified by modelId from the dataset.
66    /// # Arguments
67    /// * `project_id`- Project ID of the model to delete
68    /// * `dataset_id` - Dataset ID of the model to delete
69    /// * `model_id` - Model ID of the model to delete
70    pub async fn delete(&self, project_id: &str, dataset_id: &str, model_id: &str) -> Result<(), BQError> {
71        let req_url = &format!(
72            "{base_url}/projects/{project_id}/datasets/{dataset_id}/models/{model_id}",
73            base_url = self.base_url,
74            project_id = urlencode(project_id),
75            dataset_id = urlencode(dataset_id),
76            model_id = urlencode(model_id),
77        );
78
79        let access_token = self.auth.access_token().await?;
80
81        let request = self.client.delete(req_url).bearer_auth(access_token).build()?;
82        let response = self.client.execute(request).await?;
83
84        if response.status().is_success() {
85            Ok(())
86        } else {
87            Err(BQError::ResponseError {
88                error: response.json().await?,
89            })
90        }
91    }
92
93    /// Gets the specified model resource by model ID.
94    /// # Arguments
95    /// * `project_id`- Project ID of the requested model
96    /// * `dataset_id` - Dataset ID of the requested model
97    /// * `routine_id` - Routine ID of the requested model
98    pub async fn get(&self, project_id: &str, dataset_id: &str, model_id: &str) -> Result<Model, BQError> {
99        let req_url = &format!(
100            "{base_url}/projects/{project_id}/datasets/{dataset_id}/models/{model_id}",
101            base_url = self.base_url,
102            project_id = urlencode(project_id),
103            dataset_id = urlencode(dataset_id),
104            model_id = urlencode(model_id),
105        );
106
107        let access_token = self.auth.access_token().await?;
108
109        let request = self.client.get(req_url).bearer_auth(access_token).build()?;
110        let response = self.client.execute(request).await?;
111
112        process_response(response).await
113    }
114
115    /// Patch specific fields in the specified model.
116    /// # Arguments
117    /// * `project_id`- Project ID of the model to patch
118    /// * `dataset_id` - Dataset ID of the model to patch
119    /// * `model_id` - Routine ID of the model to patch
120    /// * `model` - Model to patch
121    pub async fn update(
122        &self,
123        project_id: &str,
124        dataset_id: &str,
125        model_id: &str,
126        model: Model,
127    ) -> Result<Model, BQError> {
128        let req_url = &format!(
129            "{base_url}/projects/{project_id}/datasets/{dataset_id}/models/{model_id}",
130            base_url = self.base_url,
131            project_id = urlencode(project_id),
132            dataset_id = urlencode(dataset_id),
133            model_id = urlencode(model_id),
134        );
135
136        let access_token = self.auth.access_token().await?;
137
138        let request = self
139            .client
140            .put(req_url)
141            .bearer_auth(access_token)
142            .json(&model)
143            .build()?;
144        let response = self.client.execute(request).await?;
145
146        process_response(response).await
147    }
148}
149
150#[derive(Debug, Serialize, Deserialize, Default)]
151#[serde(rename_all = "camelCase")]
152pub struct ListOptions {
153    max_results: Option<u64>,
154    page_token: Option<String>,
155}
156
157impl ListOptions {
158    /// The maximum number of results to return in a single response page. Leverage the page tokens
159    /// to iterate through the entire collection.
160    pub fn max_results(mut self, value: u64) -> Self {
161        self.max_results = Some(value);
162        self
163    }
164
165    /// Page token, returned by a previous call, to request the next page of results
166    pub fn page_token(mut self, value: String) -> Self {
167        self.page_token = Some(value);
168        self
169    }
170}
gcp_bigquery_client/model_api.rs

gcp_bigquery_client/
model_api.rs
