1use oauth2::{
2 basic::BasicTokenType, AccessToken, EmptyExtraTokenFields, RefreshToken, StandardTokenResponse,
3};
4use serde::{Deserialize, Serialize};
5use time::{Duration, OffsetDateTime};
6
7use crate::{Error, Result};
8
9pub type OAuthTokenResponse = StandardTokenResponse<EmptyExtraTokenFields, BasicTokenType>;
10
11#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
12pub struct Token {
13 pub access_token: String,
14 pub token_type: String,
15 #[serde(skip_serializing_if = "Option::is_none")]
16 pub refresh_token: Option<String>,
17 #[serde(
18 rename = "expiry",
19 alias = "expires_at",
20 with = "time::serde::rfc3339::option",
21 default,
22 skip_serializing_if = "Option::is_none"
23 )]
24 pub expires_at: Option<OffsetDateTime>,
25 #[serde(default, skip_serializing_if = "Vec::is_empty")]
26 pub scopes: Vec<String>,
27}
28
29impl Token {
30 pub fn needs_refresh(&self, now: OffsetDateTime, buffer: Duration) -> bool {
31 self.expires_at.is_some_and(|expiry| expiry - now <= buffer)
32 }
33
34 pub fn is_valid(&self, now: OffsetDateTime) -> bool {
35 !self.access_token.is_empty() && self.expires_at.is_none_or(|expiry| expiry > now)
36 }
37
38 pub fn access_token(&self) -> AccessToken {
39 AccessToken::new(self.access_token.clone())
40 }
41
42 pub fn refresh_token(&self) -> Option<RefreshToken> {
43 self.refresh_token.clone().map(RefreshToken::new)
44 }
45
46 pub fn from_response(
47 response: &OAuthTokenResponse,
48 now: OffsetDateTime,
49 previous: Option<&Token>,
50 ) -> Result<Self> {
51 use oauth2::TokenResponse;
52
53 let expires_at = response
54 .expires_in()
55 .and_then(|duration| Duration::try_from(duration).ok())
56 .and_then(|duration| now.checked_add(duration));
57 let token_type = match response.token_type() {
58 BasicTokenType::Bearer => "Bearer",
59 BasicTokenType::Mac => "MAC",
60 BasicTokenType::Extension(value) => value.as_str(),
61 };
62 let refresh_token = response
63 .refresh_token()
64 .map(|token| token.secret().to_owned())
65 .or_else(|| previous.and_then(|token| token.refresh_token.clone()));
66 let scopes = response
67 .scopes()
68 .map(|scopes| {
69 scopes
70 .iter()
71 .map(|scope| scope.as_ref().to_owned())
72 .collect()
73 })
74 .or_else(|| previous.map(|token| token.scopes.clone()))
75 .unwrap_or_default();
76
77 if response.access_token().secret().is_empty() {
78 return Err(Error::OAuth(
79 "token response did not contain an access token".into(),
80 ));
81 }
82 Ok(Self {
83 access_token: response.access_token().secret().to_owned(),
84 token_type: token_type.to_owned(),
85 refresh_token,
86 expires_at,
87 scopes,
88 })
89 }
90}
91
92#[cfg(test)]
93mod tests {
94 use super::Token;
95
96 #[test]
97 fn refresh_window_handles_full_binding_integer_range() {
98 let now = time::OffsetDateTime::now_utc();
99 let token = Token {
100 access_token: "access".into(),
101 token_type: "Bearer".into(),
102 refresh_token: None,
103 expires_at: Some(now),
104 scopes: vec![],
105 };
106 assert!(token.needs_refresh(now, time::Duration::seconds(i64::MAX)));
107 assert!(!token.needs_refresh(now, time::Duration::seconds(i64::MIN)));
108 }
109
110 #[test]
111 fn reads_databricks_cli_rfc3339_expiry() {
112 let token: Token = serde_json::from_str(
113 r#"{"access_token":"access","token_type":"Bearer","refresh_token":"refresh","expiry":"2026-07-30T18:35:12.447-04:00","expires_in":3600}"#,
114 )
115 .unwrap();
116
117 assert_eq!(
118 token.expires_at.unwrap().to_string(),
119 "2026-07-30 18:35:12.447 -04:00:00"
120 );
121 }
122}