cognite/
cognite_client.rs

1use reqwest::Client;
2use reqwest_middleware::{ClientBuilder, ClientWithMiddleware, Middleware};
3use std::env;
4use std::sync::Arc;
5
6use super::{ApiClient, Error, Result};
7use crate::api::core::sequences::SequencesResource;
8use crate::api::data_modeling::Models;
9use crate::api::iam::groups::GroupsResource;
10use crate::api::iam::sessions::SessionsResource;
11use crate::auth::AuthenticatorMiddleware;
12use crate::retry::CustomRetryMiddleware;
13use crate::AuthHeaderManager;
14use crate::{
15    assets::AssetsResource, datasets::DataSetsResource, events::EventsResource,
16    extpipes::ExtPipeRunsResource, extpipes::ExtPipesResource, files::Files,
17    labels::LabelsResource, raw::RawResource, relationships::RelationshipsResource,
18    time_series::TimeSeriesResource,
19};
20
21use crate::api::authenticator::{Authenticator, AuthenticatorConfig};
22
23macro_rules! env_or_error {
24    ($e: expr) => {
25        match env::var($e) {
26            Ok(el) => el,
27            Err(err) => {
28                let error_message =
29                    format!("{} is not defined in the environment. Error: {}", $e, err);
30                return Err(Error::EnvironmentVariableMissing(error_message));
31            }
32        }
33    };
34}
35
36macro_rules! env_or {
37    ($e: expr, $d: expr) => {
38        match env::var($e) {
39            Ok(el) => el,
40            Err(_) => $d,
41        }
42    };
43}
44
45macro_rules! env_or_none {
46    ($e: expr) => {
47        match env::var($e) {
48            Ok(el) => Some(el),
49            Err(_) => None,
50        }
51    };
52}
53
54#[derive(Default, Clone, Debug)]
55/// Configuration object for a cognite client.
56pub struct ClientConfig {
57    /// Maximum number of retries per request.
58    pub max_retries: u32,
59    /// Maximum delay between retries.
60    pub max_retry_delay_ms: Option<u64>,
61    /// Request timeout in milliseconds.
62    /// Note that this option does not work on wasm32 targets.
63    pub timeout_ms: Option<u64>,
64    /// Initial delay for exponential backoff, defaults to 125 milliseconds.
65    pub initial_delay_ms: Option<u64>,
66}
67
68#[derive(Clone)]
69/// Client object for the CDF API.
70pub struct CogniteClient {
71    /// Reference to an API client, which can let you make
72    /// your own requests to the CDF API.
73    pub api_client: Arc<ApiClient>,
74
75    /// CDF assets resource.
76    pub assets: AssetsResource,
77    /// CDF events resource.
78    pub events: EventsResource,
79    /// CDF files resource.
80    pub files: Files,
81    /// CDF time series resource.
82    pub time_series: TimeSeriesResource,
83    /// CDF groups resource.
84    pub groups: GroupsResource,
85    /// CDF raw resource.
86    pub raw: RawResource,
87    /// CDF data sets resource.
88    pub data_sets: DataSetsResource,
89    /// CDF labels resource.
90    pub labels: LabelsResource,
91    /// CDF relationships resource.
92    pub relationships: RelationshipsResource,
93    /// CDF extraction pipelines resource.
94    pub ext_pipes: ExtPipesResource,
95    /// CDF extraction pipeline runs resource.
96    pub ext_pipe_runs: ExtPipeRunsResource,
97    /// CDF sequences resource.
98    pub sequences: SequencesResource,
99    /// CDF sessions resource.
100    pub sessions: SessionsResource,
101    /// CDF data modeling resource.
102    pub models: Models,
103}
104
105static COGNITE_BASE_URL: &str = "COGNITE_BASE_URL";
106static COGNITE_PROJECT_NAME: &str = "COGNITE_PROJECT";
107static COGNITE_CLIENT_ID: &str = "COGNITE_CLIENT_ID";
108static COGNITE_CLIENT_SECRET: &str = "COGNITE_CLIENT_SECRET";
109static COGNITE_TOKEN_URL: &str = "COGNITE_TOKEN_URL";
110static COGNITE_RESOURCE: &str = "COGNITE_RESOURCE";
111static COGNITE_AUDIENCE: &str = "COGNITE_AUDIENCE";
112static COGNITE_SCOPES: &str = "COGNITE_SCOPES";
113
114impl CogniteClient {
115    /// Create a new cogntite client, taking OIDC credentials from the environment.
116    ///
117    /// # Arguments
118    ///
119    /// * `app_name` - The value used for the `x-cdp-app` header.
120    /// * `config` - Optional configuration for retries.
121    ///
122    /// This uses the environment variables
123    ///
124    /// * `COGNITE_BASE_URL`
125    /// * `COGNITE_PROJECT`
126    /// * `COGNITE_CLIENT_ID`
127    /// * `COGNITE_CLIENT_SECRET`
128    /// * `COGNITE_TOKEN_URL`
129    /// * `COGNITE_RESOURCE`
130    /// * `COGNITE_AUDIENCE`
131    /// * `COGNITE_SCOPES`
132    pub fn new_oidc(app_name: &str, config: Option<ClientConfig>) -> Result<Self> {
133        let api_base_url = env_or!(COGNITE_BASE_URL, "https://api.cognitedata.com/".to_string());
134        let project_name = env_or_error!(COGNITE_PROJECT_NAME);
135        let auth_config = AuthenticatorConfig {
136            client_id: env_or_error!(COGNITE_CLIENT_ID),
137            token_url: env_or_error!(COGNITE_TOKEN_URL),
138            secret: env_or_error!(COGNITE_CLIENT_SECRET),
139            resource: env_or_none!(COGNITE_RESOURCE),
140            audience: env_or_none!(COGNITE_AUDIENCE),
141            scopes: env_or_none!(COGNITE_SCOPES),
142            default_expires_in: None,
143        };
144
145        CogniteClient::new_from_oidc(&api_base_url, auth_config, &project_name, app_name, config)
146    }
147
148    /// Create a new cognite client, using a user-provided authentication manager.
149    ///
150    /// # Arguments
151    ///
152    /// * `api_base_url` - Base URL for the API. For example `https://api.cognitedata.com`
153    /// * `project_name` - Name of the CDF project to use.
154    /// * `auth` - Authentication provider.
155    /// * `app_name` - Value used for the `x-cdp-app` header.
156    /// * `config` - Optional configuration for retries.
157    pub fn new_custom_auth(
158        api_base_url: &str,
159        project_name: &str,
160        auth: AuthHeaderManager,
161        app_name: &str,
162        config: Option<ClientConfig>,
163    ) -> Result<Self> {
164        let api_base_path = format!("{}/api/{}/projects/{}", api_base_url, "v1", project_name);
165        let client = Self::get_client(config.unwrap_or_default(), auth, None, None)?;
166        let api_client = ApiClient::new(&api_base_path, app_name, client.clone());
167
168        Self::new_internal(api_client)
169    }
170
171    fn get_client(
172        config: ClientConfig,
173        authenticator: AuthHeaderManager,
174        client: Option<Client>,
175        middleware: Option<Vec<Arc<dyn Middleware>>>,
176    ) -> Result<ClientWithMiddleware> {
177        let client = if let Some(client) = client {
178            client
179        } else {
180            #[allow(unused_mut)]
181            let mut builder = Client::builder();
182            // We can add more here later
183            #[cfg(not(target_arch = "wasm32"))]
184            if let Some(timeout) = config.timeout_ms {
185                builder = builder.timeout(std::time::Duration::from_millis(timeout));
186            }
187
188            builder.build()?
189        };
190
191        let mut builder = ClientBuilder::new(client);
192        if config.max_retries > 0 {
193            builder = builder.with(CustomRetryMiddleware::new(
194                config.max_retries,
195                config.max_retry_delay_ms.unwrap_or(5 * 60 * 1000),
196                config.initial_delay_ms.unwrap_or(125),
197            ));
198        }
199        builder = builder.with(AuthenticatorMiddleware::new(authenticator)?);
200        if let Some(mw) = middleware {
201            for ware in mw {
202                builder = builder.with_arc(ware);
203            }
204        }
205        Ok(builder.build())
206    }
207
208    fn new_from_builder(
209        auth: AuthHeaderManager,
210        config: ClientConfig,
211        client: Option<Client>,
212        app_name: String,
213        project: String,
214        base_url: String,
215        middleware: Option<Vec<Arc<dyn Middleware>>>,
216    ) -> Result<Self> {
217        let api_base_path = format!("{}/api/{}/projects/{}", base_url, "v1", project);
218        let client = Self::get_client(config, auth, client, middleware)?;
219        let api_client = ApiClient::new(&api_base_path, &app_name, client.clone());
220        Self::new_internal(api_client)
221    }
222
223    fn new_internal(api_client: ApiClient) -> Result<Self> {
224        let ac = Arc::new(api_client);
225        Ok(CogniteClient {
226            api_client: ac.clone(),
227
228            assets: AssetsResource::new(ac.clone()),
229            events: EventsResource::new(ac.clone()),
230            files: Files::new(ac.clone()),
231            groups: GroupsResource::new(ac.clone()),
232            time_series: TimeSeriesResource::new(ac.clone()),
233            raw: RawResource::new(ac.clone()),
234            data_sets: DataSetsResource::new(ac.clone()),
235            labels: LabelsResource::new(ac.clone()),
236            relationships: RelationshipsResource::new(ac.clone()),
237            ext_pipes: ExtPipesResource::new(ac.clone()),
238            ext_pipe_runs: ExtPipeRunsResource::new(ac.clone()),
239            sequences: SequencesResource::new(ac.clone()),
240            sessions: SessionsResource::new(ac.clone()),
241            models: Models::new(ac),
242        })
243    }
244
245    /// Create a new cognite client using provided OIDC credentials.
246    ///
247    /// # Arguments
248    ///
249    /// * `api_base_url` - Base URL for the API. For example `https://api.cognitedata.com`
250    /// * `project_name` - Name of the CDF project to use.
251    /// * `auth_config` - Configuration for creating an OIDC authenticator.
252    /// * `app_name` - Value used for the `x-cdp-app` header.
253    /// * `config` - Optional configuration for retries.
254    pub fn new_from_oidc(
255        api_base_url: &str,
256        auth_config: AuthenticatorConfig,
257        project_name: &str,
258        app_name: &str,
259        config: Option<ClientConfig>,
260    ) -> Result<Self> {
261        let authenticator = Authenticator::new(auth_config);
262        let api_base_path = format!("{}/api/{}/projects/{}", api_base_url, "v1", project_name);
263        let auth = AuthHeaderManager::OIDCToken(Arc::new(authenticator));
264        let client = Self::get_client(config.unwrap_or_default(), auth, None, None)?;
265        let api_client = ApiClient::new(&api_base_path, app_name, client.clone());
266
267        Self::new_internal(api_client)
268    }
269
270    /// Create a builder with a fluent API for creating a cognite client.
271    pub fn builder() -> Builder {
272        Builder::default()
273    }
274}
275
276/// Fluent API for configuring a client.
277#[derive(Default)]
278pub struct Builder {
279    auth: Option<AuthHeaderManager>,
280    config: Option<ClientConfig>,
281    client: Option<Client>,
282    app_name: Option<String>,
283    project: Option<String>,
284    base_url: Option<String>,
285    custom_middleware: Option<Vec<Arc<dyn Middleware>>>,
286}
287
288impl Builder {
289    /// Set a custom authenticator.
290    ///
291    /// # Arguments
292    ///
293    /// * `auth` - Authenticator to use.
294    pub fn set_custom_auth(&mut self, auth: AuthHeaderManager) -> &mut Self {
295        self.auth = Some(auth);
296        self
297    }
298
299    /// Set an authenticator using OIDC client credentials.
300    ///
301    /// # Arguments
302    ///
303    /// * `auth` - Client credentials.
304    pub fn set_oidc_credentials(&mut self, auth: AuthenticatorConfig) -> &mut Self {
305        self.auth = Some(AuthHeaderManager::OIDCToken(Arc::new(Authenticator::new(
306            auth,
307        ))));
308        self
309    }
310
311    /// Set the CDF project to connect to.
312    ///
313    /// # Arguments
314    ///
315    /// * `project` - CDF project
316    pub fn set_project(&mut self, project: &str) -> &mut Self {
317        self.project = Some(project.to_owned());
318        self
319    }
320
321    /// Set the value of the `x-cdp-app` header.
322    pub fn set_app_name(&mut self, app_name: &str) -> &mut Self {
323        self.app_name = Some(app_name.to_owned());
324        self
325    }
326
327    /// Set the reqwest client used internally. If your application
328    /// connects to a large number of different CDF projects, or uses a large
329    /// number of different sets of credentials. It is recommended to share
330    /// a single reqwest client.
331    ///
332    /// # Arguments
333    ///
334    /// * `client` - reqwest client to use.
335    pub fn set_internal_client(&mut self, client: Client) -> &mut Self {
336        self.client = Some(client);
337        self
338    }
339
340    /// Set configuration for retries.
341    ///
342    /// # Arguments
343    ///
344    /// * `config` - Client configuration.
345    pub fn set_client_config(&mut self, config: ClientConfig) -> &mut Self {
346        self.config = Some(config);
347        self
348    }
349
350    /// Set the base URL used by the client.
351    ///
352    /// # Arguments
353    ///
354    /// * `base_url` - Cognite API base URL, for example `https://api.cognitedata.com`
355    pub fn set_base_url(&mut self, base_url: &str) -> &mut Self {
356        self.base_url = Some(base_url.to_owned());
357        self
358    }
359
360    /// Add some custom middleware.
361    ///
362    /// # Arguments
363    ///
364    /// * `middleware` - A reference to some reqwest middleware.
365    pub fn with_custom_middleware(&mut self, middleware: Arc<dyn Middleware>) -> &mut Self {
366        match &mut self.custom_middleware {
367            Some(x) => x.push(middleware),
368            None => self.custom_middleware = Some(vec![middleware]),
369        }
370        self
371    }
372
373    /// Create a cognite client. This may fail if not all required parameters are provided.
374    pub fn build(self) -> Result<CogniteClient> {
375        let auth = self
376            .auth
377            .ok_or_else(|| Error::Config("Some form of auth is required".to_string()))?;
378        let config = self.config.unwrap_or_default();
379        let client = self.client;
380        let app_name = self
381            .app_name
382            .ok_or_else(|| Error::Config("App name is required".to_string()))?;
383        let project = self
384            .project
385            .ok_or_else(|| Error::Config("Project is required".to_string()))?;
386        let base_url = self
387            .base_url
388            .unwrap_or_else(|| "https://api.cognitedata.com/".to_owned());
389
390        CogniteClient::new_from_builder(
391            auth,
392            config,
393            client,
394            app_name,
395            project,
396            base_url,
397            self.custom_middleware,
398        )
399    }
400}
cognite/cognite_client.rs

cognite/
cognite_client.rs
