1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154

//! Manage BigQuery models.
use reqwest::Client;

use crate::auth::ServiceAccountAuthenticator;
use crate::error::BQError;
use crate::model::list_models_response::ListModelsResponse;
use crate::model::model::Model;
use crate::{process_response, urlencode};

/// A model API handler.
#[derive(Clone)]
pub struct ModelApi {
    client: Client,
    sa_auth: ServiceAccountAuthenticator,
}

impl ModelApi {
    pub(crate) fn new(client: Client, sa_auth: ServiceAccountAuthenticator) -> Self {
        Self { client, sa_auth }
    }

    /// Lists all models in the specified dataset. Requires the READER dataset role.
    /// # Arguments
    /// * `project_id` - Project ID of the models to list.
    /// * `dataset_id` - Dataset ID of the models to list.
    pub async fn list(
        &self,
        project_id: &str,
        dataset_id: &str,
        options: ListOptions,
    ) -> Result<ListModelsResponse, BQError> {
        let req_url = format!(
            "https://bigquery.googleapis.com/bigquery/v2/projects/{project_id}/datasets/{dataset_id}/models",
            project_id = urlencode(project_id),
            dataset_id = urlencode(dataset_id),
        );

        let access_token = self.sa_auth.access_token().await?;

        let request = self
            .client
            .get(req_url.as_str())
            .bearer_auth(access_token)
            .query(&options)
            .build()?;

        let resp = self.client.execute(request).await?;

        process_response(resp).await
    }

    /// Deletes the model specified by modelId from the dataset.
    /// # Arguments
    /// * `project_id`- Project ID of the model to delete
    /// * `dataset_id` - Dataset ID of the model to delete
    /// * `model_id` - Model ID of the model to delete
    pub async fn delete(&self, project_id: &str, dataset_id: &str, model_id: &str) -> Result<(), BQError> {
        let req_url = &format!(
            "https://bigquery.googleapis.com/bigquery/v2/projects/{project_id}/datasets/{dataset_id}/models/{model_id}",
            project_id = urlencode(project_id),
            dataset_id = urlencode(dataset_id),
            model_id = urlencode(model_id),
        );

        let access_token = self.sa_auth.access_token().await?;

        let request = self.client.delete(req_url).bearer_auth(access_token).build()?;
        let response = self.client.execute(request).await?;

        if response.status().is_success() {
            Ok(())
        } else {
            Err(BQError::ResponseError {
                error: response.json().await?,
            })
        }
    }

    /// Gets the specified model resource by model ID.
    /// # Arguments
    /// * `project_id`- Project ID of the requested model
    /// * `dataset_id` - Dataset ID of the requested model
    /// * `routine_id` - Routine ID of the requested model
    pub async fn get(&self, project_id: &str, dataset_id: &str, model_id: &str) -> Result<Model, BQError> {
        let req_url = &format!(
            "https://bigquery.googleapis.com/bigquery/v2/projects/{project_id}/datasets/{dataset_id}/models/{model_id}",
            project_id = urlencode(project_id),
            dataset_id = urlencode(dataset_id),
            model_id = urlencode(model_id),
        );

        let access_token = self.sa_auth.access_token().await?;

        let request = self.client.get(req_url).bearer_auth(access_token).build()?;
        let response = self.client.execute(request).await?;

        process_response(response).await
    }

    /// Patch specific fields in the specified model.
    /// # Arguments
    /// * `project_id`- Project ID of the model to patch
    /// * `dataset_id` - Dataset ID of the model to patch
    /// * `model_id` - Routine ID of the model to patch
    /// * `model` - Model to patch
    pub async fn update(
        &self,
        project_id: &str,
        dataset_id: &str,
        model_id: &str,
        model: Model,
    ) -> Result<Model, BQError> {
        let req_url = &format!(
            "https://bigquery.googleapis.com/bigquery/v2/projects/{project_id}/datasets/{dataset_id}/models/{model_id}",
            project_id = urlencode(project_id),
            dataset_id = urlencode(dataset_id),
            model_id = urlencode(model_id),
        );

        let access_token = self.sa_auth.access_token().await?;

        let request = self
            .client
            .put(req_url)
            .bearer_auth(access_token)
            .json(&model)
            .build()?;
        let response = self.client.execute(request).await?;

        process_response(response).await
    }
}

#[derive(Debug, Serialize, Deserialize, Default)]
#[serde(rename_all = "camelCase")]
pub struct ListOptions {
    max_results: Option<u64>,
    page_token: Option<String>,
}

impl ListOptions {
    /// The maximum number of results to return in a single response page. Leverage the page tokens
    /// to iterate through the entire collection.
    pub fn max_results(mut self, value: u64) -> Self {
        self.max_results = Some(value);
        self
    }

    /// Page token, returned by a previous call, to request the next page of results
    pub fn page_token(mut self, value: String) -> Self {
        self.page_token = Some(value);
        self
    }
}