Convert All Usage of LocalDateTime to Instant

databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java

@@ -8,7 +8,9 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import java.io.IOException;
 import java.io.InputStream;
+import java.time.Instant;
 import java.time.LocalDateTime;
+import java.time.ZoneOffset;
 import java.time.format.DateTimeFormatter;
 import java.time.format.DateTimeParseException;
 import java.util.Arrays;
@@ -36,21 +38,40 @@ public CliTokenSource(
     this.env = env;
   }
 
-  static LocalDateTime parseExpiry(String expiry) {
+  /**
+   * Parses an expiry time string and returns the corresponding {@link Instant}.
+   *
+   * <p>The expiry time string is always in UTC. Any time zone or offset information present in the
+   * input is ignored.
+   *
+   * <p>The method attempts to parse the input using several common date-time formats, including
+   * ISO-8601 and patterns with varying sub-second precision.
+   *
+   * @param expiry the expiry time string to parse, which must represent a UTC time
+   * @return the parsed {@link Instant}
+   * @throws DateTimeParseException if the input string cannot be parsed
+   */
+  static Instant parseExpiry(String expiry) {
+    DateTimeParseException lastException = null;
+    // Try to parse the expiry as an ISO-8601 string in UTC first
+    try {
+      return Instant.parse(expiry);
+    } catch (DateTimeParseException e) {
+      lastException = e;
+    }
+
     String multiplePrecisionPattern =
         "[SSSSSSSSS][SSSSSSSS][SSSSSSS][SSSSSS][SSSSS][SSSS][SSS][SS][S]";
     List<String> datePatterns =
         Arrays.asList(
             "yyyy-MM-dd HH:mm:ss",
             "yyyy-MM-dd HH:mm:ss." + multiplePrecisionPattern,
-            "yyyy-MM-dd'T'HH:mm:ss." + multiplePrecisionPattern + "XXX",
-            "yyyy-MM-dd'T'HH:mm:ss." + multiplePrecisionPattern + "'Z'");
-    DateTimeParseException lastException = null;
+            "yyyy-MM-dd'T'HH:mm:ss." + multiplePrecisionPattern + "XXX");
     for (String pattern : datePatterns) {
       try {
         DateTimeFormatter formatter = DateTimeFormatter.ofPattern(pattern);
         LocalDateTime dateTime = LocalDateTime.parse(expiry, formatter);
-        return dateTime;
+        return dateTime.atZone(ZoneOffset.UTC).toInstant();
       } catch (DateTimeParseException e) {
         lastException = e;
       }
@@ -83,7 +104,7 @@ protected Token refresh() {
       String tokenType = jsonNode.get(tokenTypeField).asText();
       String accessToken = jsonNode.get(accessTokenField).asText();
       String expiry = jsonNode.get(expiryField).asText();
-      LocalDateTime expiresOn = parseExpiry(expiry);
+      Instant expiresOn = parseExpiry(expiry);
       return new Token(accessToken, tokenType, expiresOn);
     } catch (DatabricksException e) {
       throw e;

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java

@@ -3,7 +3,7 @@
 import com.databricks.sdk.core.DatabricksException;
 import com.databricks.sdk.core.http.HttpClient;
 import com.google.common.base.Strings;
-import java.time.LocalDateTime;
+import java.time.Instant;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
@@ -166,7 +166,7 @@ public Token refresh() {
       throw e;
     }
 
-    LocalDateTime expiry = LocalDateTime.now().plusSeconds(response.getExpiresIn());
+    Instant expiry = Instant.now().plusSeconds(response.getExpiresIn());
     return new Token(
         response.getAccessToken(), response.getTokenType(), response.getRefreshToken(), expiry);
   }

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/EndpointTokenSource.java

@@ -2,7 +2,7 @@
 
 import com.databricks.sdk.core.DatabricksException;
 import com.databricks.sdk.core.http.HttpClient;
-import java.time.LocalDateTime;
+import java.time.Instant;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
@@ -87,7 +87,7 @@ protected Token refresh() {
       throw e;
     }
 
-    LocalDateTime expiry = LocalDateTime.now().plusSeconds(oauthResponse.getExpiresIn());
+    Instant expiry = Instant.now().plusSeconds(oauthResponse.getExpiresIn());
     return new Token(
         oauthResponse.getAccessToken(),
         oauthResponse.getTokenType(),

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OidcTokenSource.java

@@ -8,7 +8,7 @@
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableMap;
 import java.io.IOException;
-import java.time.LocalDateTime;
+import java.time.Instant;
 
 /**
  * {@code OidcTokenSource} is responsible for obtaining OAuth tokens using the OpenID Connect (OIDC)
@@ -77,7 +77,7 @@ protected Token refresh() {
     if (resp.getErrorCode() != null) {
       throw new IllegalArgumentException(resp.getErrorCode() + ": " + resp.getErrorSummary());
     }
-    LocalDateTime expiry = LocalDateTime.now().plusSeconds(resp.getExpiresIn());
+    Instant expiry = Instant.now().plusSeconds(resp.getExpiresIn());
     return new Token(resp.getAccessToken(), resp.getTokenType(), resp.getRefreshToken(), expiry);
   }
 }

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java

@@ -5,8 +5,7 @@
 import com.databricks.sdk.core.http.FormRequest;
 import com.databricks.sdk.core.http.HttpClient;
 import com.databricks.sdk.core.http.Request;
-import java.time.LocalDateTime;
-import java.time.temporal.ChronoUnit;
+import java.time.Instant;
 import java.util.Base64;
 import java.util.Map;
 import org.apache.http.HttpHeaders;
@@ -70,7 +69,7 @@ protected static Token retrieveToken(
       if (resp.getErrorCode() != null) {
         throw new IllegalArgumentException(resp.getErrorCode() + ": " + resp.getErrorSummary());
       }
-      LocalDateTime expiry = LocalDateTime.now().plus(resp.getExpiresIn(), ChronoUnit.SECONDS);
+      Instant expiry = Instant.now().plusSeconds(resp.getExpiresIn());
       return new Token(resp.getAccessToken(), resp.getTokenType(), resp.getRefreshToken(), expiry);
     } catch (Exception e) {
       throw new DatabricksException("Failed to refresh credentials: " + e.getMessage(), e);

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java

@@ -4,8 +4,7 @@
 import com.databricks.sdk.core.utils.SystemClockSupplier;
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
-import java.time.LocalDateTime;
-import java.time.temporal.ChronoUnit;
+import java.time.Instant;
 import java.util.Objects;
 
 public class Token {
@@ -19,21 +18,20 @@ public class Token {
    * The expiry time of the token.
    *
    * <p>OAuth token responses include the duration of the lifetime of the access token. When the
-   * token is retrieved, this is converted to a LocalDateTime tracking the expiry time of the token
-   * with respect to the current clock.
+   * token is retrieved, this is converted to an Instant tracking the expiry time of the token with
+   * respect to the current clock.
    */
-  @JsonProperty private LocalDateTime expiry;
+  @JsonProperty private Instant expiry;
 
   private final ClockSupplier clockSupplier;
 
   /** Constructor for non-refreshable tokens (e.g. M2M). */
-  public Token(String accessToken, String tokenType, LocalDateTime expiry) {
+  public Token(String accessToken, String tokenType, Instant expiry) {
     this(accessToken, tokenType, null, expiry, new SystemClockSupplier());
   }
 
   /** Constructor for non-refreshable tokens (e.g. M2M) with ClockSupplier */
-  public Token(
-      String accessToken, String tokenType, LocalDateTime expiry, ClockSupplier clockSupplier) {
+  public Token(String accessToken, String tokenType, Instant expiry, ClockSupplier clockSupplier) {
     this(accessToken, tokenType, null, expiry, clockSupplier);
   }
 
@@ -43,7 +41,7 @@ public Token(
       @JsonProperty("accessToken") String accessToken,
       @JsonProperty("tokenType") String tokenType,
       @JsonProperty("refreshToken") String refreshToken,
-      @JsonProperty("expiry") LocalDateTime expiry) {
+      @JsonProperty("expiry") Instant expiry) {
     this(accessToken, tokenType, refreshToken, expiry, new SystemClockSupplier());
   }
 
@@ -52,7 +50,7 @@ public Token(
       String accessToken,
       String tokenType,
       String refreshToken,
-      LocalDateTime expiry,
+      Instant expiry,
       ClockSupplier clockSupplier) {
     Objects.requireNonNull(accessToken, "accessToken must be defined");
     Objects.requireNonNull(tokenType, "tokenType must be defined");
@@ -71,8 +69,8 @@ public boolean isExpired() {
     }
     // Azure Databricks rejects tokens that expire in 30 seconds or less,
     // so we refresh the token 40 seconds before it expires.
-    LocalDateTime potentiallyExpired = expiry.minus(40, ChronoUnit.SECONDS);
-    LocalDateTime now = LocalDateTime.now(clockSupplier.getClock());
+    Instant potentiallyExpired = expiry.minusSeconds(40);
+    Instant now = Instant.now(clockSupplier.getClock());
     return potentiallyExpired.isBefore(now);
   }
 

databricks-sdk-java/src/test/java/com/databricks/sdk/core/AzureCliCredentialsProviderTest.java

@@ -7,7 +7,6 @@
 
 import com.databricks.sdk.core.oauth.Token;
 import com.databricks.sdk.core.oauth.TokenSource;
-import java.time.LocalDateTime;
 import java.util.Arrays;
 import java.util.List;
 import org.junit.jupiter.api.Test;
@@ -25,7 +24,7 @@ class AzureCliCredentialsProviderTest {
   private static CliTokenSource mockTokenSource() {
     CliTokenSource tokenSource = Mockito.mock(CliTokenSource.class);
     Mockito.when(tokenSource.getToken())
-        .thenReturn(new Token(TOKEN, TOKEN_TYPE, LocalDateTime.now()));
+        .thenReturn(new Token(TOKEN, TOKEN_TYPE, java.time.Instant.now()));
     return tokenSource;
   }
 

databricks-sdk-java/src/test/java/com/databricks/sdk/core/CliTokenSourceTest.java

@@ -2,28 +2,28 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
-import java.time.LocalDateTime;
+import java.time.Instant;
 import java.time.format.DateTimeParseException;
 import org.junit.jupiter.api.Test;
 
 public class CliTokenSourceTest {
 
   @Test
   public void testParseExpiryWithoutTruncate() {
-    LocalDateTime parsedDateTime = CliTokenSource.parseExpiry("2023-07-17T09:02:22.330612218Z");
-    assertEquals(LocalDateTime.of(2023, 7, 17, 9, 2, 22, 330612218), parsedDateTime);
+    Instant parsedInstant = CliTokenSource.parseExpiry("2023-07-17T09:02:22.330612218Z");
+    assertEquals(Instant.parse("2023-07-17T09:02:22.330612218Z"), parsedInstant);
   }
 
   @Test
   public void testParseExpiryWithTruncate() {
-    LocalDateTime parsedDateTime = CliTokenSource.parseExpiry("2023-07-17T09:02:22.33061221Z");
-    assertEquals(LocalDateTime.of(2023, 7, 17, 9, 2, 22, 330612210), parsedDateTime);
+    Instant parsedInstant = CliTokenSource.parseExpiry("2023-07-17T09:02:22.33061221Z");
+    assertEquals(Instant.parse("2023-07-17T09:02:22.330612210Z"), parsedInstant);
   }
 
   @Test
   public void testParseExpiryWithTruncateAndLessNanoSecondDigits() {
-    LocalDateTime parsedDateTime = CliTokenSource.parseExpiry("2023-07-17T09:02:22.330612Z");
-    assertEquals(LocalDateTime.of(2023, 7, 17, 9, 2, 22, 330612000), parsedDateTime);
+    Instant parsedInstant = CliTokenSource.parseExpiry("2023-07-17T09:02:22.330612Z");
+    assertEquals(Instant.parse("2023-07-17T09:02:22.330612000Z"), parsedInstant);
   }
 
   @Test
@@ -34,4 +34,25 @@ public void testParseExpiryWithMoreThanNineNanoSecondDigits() {
       assert (e.getMessage().contains("could not be parsed"));
     }
   }
+
+  @Test
+  public void testParseExpiryWithSpaceFormat() {
+    Instant parsedInstant = CliTokenSource.parseExpiry("2023-07-17 09:02:22");
+    assertEquals(Instant.parse("2023-07-17T09:02:22Z"), parsedInstant);
+  }
+
+  @Test
+  public void testParseExpiryWithSpaceFormatAndMillis() {
+    Instant parsedInstant = CliTokenSource.parseExpiry("2023-07-17 09:02:22.123");
+    assertEquals(Instant.parse("2023-07-17T09:02:22.123Z"), parsedInstant);
+  }
+
+  @Test
+  public void testParseExpiryWithInvalidFormat() {
+    try {
+      CliTokenSource.parseExpiry("17-07-2023 09:02:22");
+    } catch (DateTimeParseException e) {
+      assert (e.getMessage().contains("could not be parsed"));
+    }
+  }
 }

databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java

@@ -12,7 +12,6 @@
 import com.databricks.sdk.core.oauth.TokenSource;
 import com.databricks.sdk.core.utils.Environment;
 import java.io.IOException;
-import java.time.LocalDateTime;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -232,7 +231,7 @@ public void testGetTokenSourceWithOAuth() {
     HttpClient httpClient = mock(HttpClient.class);
     TokenSource mockTokenSource = mock(TokenSource.class);
     when(mockTokenSource.getToken())
-        .thenReturn(new Token("test-token", "Bearer", LocalDateTime.now().plusHours(1)));
+        .thenReturn(new Token("test-token", "Bearer", java.time.Instant.now().plusSeconds(3600)));
     OAuthHeaderFactory mockHeaderFactory = OAuthHeaderFactory.fromTokenSource(mockTokenSource);
     CredentialsProvider mockProvider = mock(CredentialsProvider.class);
     when(mockProvider.authType()).thenReturn("test");

databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java

@@ -8,7 +8,7 @@
 import com.databricks.sdk.core.http.Response;
 import java.io.IOException;
 import java.net.URL;
-import java.time.LocalDateTime;
+import java.time.Instant;
 import java.util.stream.Stream;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -29,8 +29,7 @@ public class DataPlaneTokenSourceTest {
 
   private static Stream<Arguments> provideDataPlaneTokenScenarios() throws Exception {
     // Mock DatabricksOAuthTokenSource for control plane token
-    Token cpToken =
-        new Token(TEST_CP_TOKEN, TEST_TOKEN_TYPE, null, LocalDateTime.now().plusSeconds(600));
+    Token cpToken = new Token(TEST_CP_TOKEN, TEST_TOKEN_TYPE, null, Instant.now().plusSeconds(600));
     DatabricksOAuthTokenSource mockCpTokenSource = mock(DatabricksOAuthTokenSource.class);
     when(mockCpTokenSource.getToken()).thenReturn(cpToken);
 
@@ -81,9 +80,8 @@ private static Stream<Arguments> provideDataPlaneTokenScenarios() throws Excepti
                 "dp-access-token1",
                 TEST_TOKEN_TYPE,
                 TEST_REFRESH_TOKEN,
-                LocalDateTime.now().plusSeconds(TEST_EXPIRES_IN)),
-            null // No exception
-            ),
+                Instant.now().plusSeconds(TEST_EXPIRES_IN)),
+            null),
         Arguments.of(
             "Success: endpoint2/auth2 (different cache key)",
             TEST_ENDPOINT_2,
@@ -95,7 +93,7 @@ private static Stream<Arguments> provideDataPlaneTokenScenarios() throws Excepti
                 "dp-access-token2",
                 TEST_TOKEN_TYPE,
                 TEST_REFRESH_TOKEN,
-                LocalDateTime.now().plusSeconds(TEST_EXPIRES_IN)),
+                Instant.now().plusSeconds(TEST_EXPIRES_IN)),
             null),
         Arguments.of(
             "Error response from endpoint",
@@ -203,7 +201,7 @@ void testDataPlaneTokenSource(
   @Test
   void testEndpointTokenSourceCaching() throws Exception {
     Token cpToken =
-        new Token(TEST_CP_TOKEN, TEST_TOKEN_TYPE, null, LocalDateTime.now().plusSeconds(3600));
+        new Token(TEST_CP_TOKEN, TEST_TOKEN_TYPE, null, Instant.now().plusSeconds(3600));
     DatabricksOAuthTokenSource mockCpTokenSource = mock(DatabricksOAuthTokenSource.class);
     when(mockCpTokenSource.getToken()).thenReturn(cpToken);
 

databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/EndpointTokenSourceTest.java

@@ -9,7 +9,7 @@
 import com.databricks.sdk.core.http.Response;
 import java.io.IOException;
 import java.net.URL;
-import java.time.LocalDateTime;
+import java.time.Instant;
 import java.util.stream.Stream;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
@@ -48,7 +48,7 @@ private static Stream<Arguments> provideEndpointTokenScenarios() throws Exceptio
     String malformedJson = "{not valid json}";
 
     // Mock DatabricksOAuthTokenSource for control plane token
-    Token cpToken = new Token(TEST_CP_TOKEN, TEST_TOKEN_TYPE, LocalDateTime.now().plusMinutes(10));
+    Token cpToken = new Token(TEST_CP_TOKEN, TEST_TOKEN_TYPE, Instant.now().plusSeconds(600));
     DatabricksOAuthTokenSource mockCpTokenSource = mock(DatabricksOAuthTokenSource.class);
     when(mockCpTokenSource.getToken()).thenReturn(cpToken);