cognite/
cognite_client.rs

1use reqwest::Client;
2use reqwest_middleware::{ClientBuilder, ClientWithMiddleware, Middleware};
3use std::env;
4use std::sync::Arc;
5use std::time::Duration;
6
7use super::{ApiClient, Error, Result};
8use crate::api::core::sequences::SequencesResource;
9use crate::api::data_modeling::Models;
10use crate::api::iam::groups::GroupsResource;
11use crate::api::iam::sessions::SessionsResource;
12use crate::auth::AuthenticatorMiddleware;
13use crate::retry::CustomRetryMiddleware;
14use crate::AuthHeaderManager;
15use crate::{
16    assets::AssetsResource, datasets::DataSetsResource, events::EventsResource,
17    extpipes::ExtPipeRunsResource, extpipes::ExtPipesResource, files::Files,
18    labels::LabelsResource, raw::RawResource, relationships::RelationshipsResource,
19    time_series::TimeSeriesResource,
20};
21
22use crate::api::authenticator::{Authenticator, AuthenticatorConfig};
23
24macro_rules! env_or_error {
25    ($e: expr) => {
26        match env::var($e) {
27            Ok(el) => el,
28            Err(err) => {
29                let error_message =
30                    format!("{} is not defined in the environment. Error: {}", $e, err);
31                return Err(Error::EnvironmentVariableMissing(error_message));
32            }
33        }
34    };
35}
36
37macro_rules! env_or {
38    ($e: expr, $d: expr) => {
39        match env::var($e) {
40            Ok(el) => el,
41            Err(_) => $d,
42        }
43    };
44}
45
46macro_rules! env_or_none {
47    ($e: expr) => {
48        match env::var($e) {
49            Ok(el) => Some(el),
50            Err(_) => None,
51        }
52    };
53}
54
55#[derive(Default, Clone, Debug)]
56/// Configuration object for a cognite client.
57pub struct ClientConfig {
58    /// Maximum number of retries per request.
59    pub max_retries: u32,
60    /// Maximum delay between retries.
61    pub max_retry_delay_ms: Option<u64>,
62    /// Request timeout in milliseconds.
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            let mut builder = Client::builder();
181            // We can add more here later
182            if let Some(timeout) = config.timeout_ms {
183                builder = builder.timeout(Duration::from_millis(timeout));
184            }
185
186            builder.build()?
187        };
188
189        let mut builder = ClientBuilder::new(client);
190        if config.max_retries > 0 {
191            builder = builder.with(CustomRetryMiddleware::new(
192                config.max_retries,
193                config.max_retry_delay_ms.unwrap_or(5 * 60 * 1000),
194                config.initial_delay_ms.unwrap_or(125),
195            ));
196        }
197        builder = builder.with(AuthenticatorMiddleware::new(authenticator)?);
198        if let Some(mw) = middleware {
199            for ware in mw {
200                builder = builder.with_arc(ware);
201            }
202        }
203        Ok(builder.build())
204    }
205
206    fn new_from_builder(
207        auth: AuthHeaderManager,
208        config: ClientConfig,
209        client: Option<Client>,
210        app_name: String,
211        project: String,
212        base_url: String,
213        middleware: Option<Vec<Arc<dyn Middleware>>>,
214    ) -> Result<Self> {
215        let api_base_path = format!("{}/api/{}/projects/{}", base_url, "v1", project);
216        let client = Self::get_client(config, auth, client, middleware)?;
217        let api_client = ApiClient::new(&api_base_path, &app_name, client.clone());
218        Self::new_internal(api_client)
219    }
220
221    fn new_internal(api_client: ApiClient) -> Result<Self> {
222        let ac = Arc::new(api_client);
223        Ok(CogniteClient {
224            api_client: ac.clone(),
225
226            assets: AssetsResource::new(ac.clone()),
227            events: EventsResource::new(ac.clone()),
228            files: Files::new(ac.clone()),
229            groups: GroupsResource::new(ac.clone()),
230            time_series: TimeSeriesResource::new(ac.clone()),
231            raw: RawResource::new(ac.clone()),
232            data_sets: DataSetsResource::new(ac.clone()),
233            labels: LabelsResource::new(ac.clone()),
234            relationships: RelationshipsResource::new(ac.clone()),
235            ext_pipes: ExtPipesResource::new(ac.clone()),
236            ext_pipe_runs: ExtPipeRunsResource::new(ac.clone()),
237            sequences: SequencesResource::new(ac.clone()),
238            sessions: SessionsResource::new(ac.clone()),
239            models: Models::new(ac),
240        })
241    }
242
243    /// Create a new cognite client using provided OIDC credentials.
244    ///
245    /// # Arguments
246    ///
247    /// * `api_base_url` - Base URL for the API. For example `https://api.cognitedata.com`
248    /// * `project_name` - Name of the CDF project to use.
249    /// * `auth_config` - Configuration for creating an OIDC authenticator.
250    /// * `app_name` - Value used for the `x-cdp-app` header.
251    /// * `config` - Optional configuration for retries.
252    pub fn new_from_oidc(
253        api_base_url: &str,
254        auth_config: AuthenticatorConfig,
255        project_name: &str,
256        app_name: &str,
257        config: Option<ClientConfig>,
258    ) -> Result<Self> {
259        let authenticator = Authenticator::new(auth_config);
260        let api_base_path = format!("{}/api/{}/projects/{}", api_base_url, "v1", project_name);
261        let auth = AuthHeaderManager::OIDCToken(Arc::new(authenticator));
262        let client = Self::get_client(config.unwrap_or_default(), auth, None, None)?;
263        let api_client = ApiClient::new(&api_base_path, app_name, client.clone());
264
265        Self::new_internal(api_client)
266    }
267
268    /// Create a builder with a fluent API for creating a cognite client.
269    pub fn builder() -> Builder {
270        Builder::default()
271    }
272}
273
274/// Fluent API for configuring a client.
275#[derive(Default)]
276pub struct Builder {
277    auth: Option<AuthHeaderManager>,
278    config: Option<ClientConfig>,
279    client: Option<Client>,
280    app_name: Option<String>,
281    project: Option<String>,
282    base_url: Option<String>,
283    custom_middleware: Option<Vec<Arc<dyn Middleware>>>,
284}
285
286impl Builder {
287    /// Set a custom authenticator.
288    ///
289    /// # Arguments
290    ///
291    /// * `auth` - Authenticator to use.
292    pub fn set_custom_auth(&mut self, auth: AuthHeaderManager) -> &mut Self {
293        self.auth = Some(auth);
294        self
295    }
296
297    /// Set an authenticator using OIDC client credentials.
298    ///
299    /// # Arguments
300    ///
301    /// * `auth` - Client credentials.
302    pub fn set_oidc_credentials(&mut self, auth: AuthenticatorConfig) -> &mut Self {
303        self.auth = Some(AuthHeaderManager::OIDCToken(Arc::new(Authenticator::new(
304            auth,
305        ))));
306        self
307    }
308
309    /// Set the CDF project to connect to.
310    ///
311    /// # Arguments
312    ///
313    /// * `project` - CDF project
314    pub fn set_project(&mut self, project: &str) -> &mut Self {
315        self.project = Some(project.to_owned());
316        self
317    }
318
319    /// Set the value of the `x-cdp-app` header.
320    pub fn set_app_name(&mut self, app_name: &str) -> &mut Self {
321        self.app_name = Some(app_name.to_owned());
322        self
323    }
324
325    /// Set the reqwest client used internally. If your application
326    /// connects to a large number of different CDF projects, or uses a large
327    /// number of different sets of credentials. It is recommended to share
328    /// a single reqwest client.
329    ///
330    /// # Arguments
331    ///
332    /// * `client` - reqwest client to use.
333    pub fn set_internal_client(&mut self, client: Client) -> &mut Self {
334        self.client = Some(client);
335        self
336    }
337
338    /// Set configuration for retries.
339    ///
340    /// # Arguments
341    ///
342    /// * `config` - Client configuration.
343    pub fn set_client_config(&mut self, config: ClientConfig) -> &mut Self {
344        self.config = Some(config);
345        self
346    }
347
348    /// Set the base URL used by the client.
349    ///
350    /// # Arguments
351    ///
352    /// * `base_url` - Cognite API base URL, for example `https://api.cognitedata.com`
353    pub fn set_base_url(&mut self, base_url: &str) -> &mut Self {
354        self.base_url = Some(base_url.to_owned());
355        self
356    }
357
358    /// Add some custom middleware.
359    ///
360    /// # Arguments
361    ///
362    /// * `middleware` - A reference to some reqwest middleware.
363    pub fn with_custom_middleware(&mut self, middleware: Arc<dyn Middleware>) -> &mut Self {
364        match &mut self.custom_middleware {
365            Some(x) => x.push(middleware),
366            None => self.custom_middleware = Some(vec![middleware]),
367        }
368        self
369    }
370
371    /// Create a cognite client. This may fail if not all required parameters are provided.
372    pub fn build(self) -> Result<CogniteClient> {
373        let auth = self
374            .auth
375            .ok_or_else(|| Error::Config("Some form of auth is required".to_string()))?;
376        let config = self.config.unwrap_or_default();
377        let client = self.client;
378        let app_name = self
379            .app_name
380            .ok_or_else(|| Error::Config("App name is required".to_string()))?;
381        let project = self
382            .project
383            .ok_or_else(|| Error::Config("Project is required".to_string()))?;
384        let base_url = self
385            .base_url
386            .unwrap_or_else(|| "https://api.cognitedata.com/".to_owned());
387
388        CogniteClient::new_from_builder(
389            auth,
390            config,
391            client,
392            app_name,
393            project,
394            base_url,
395            self.custom_middleware,
396        )
397    }
398}
cognite/cognite_client.rs

cognite/
cognite_client.rs
