Convert All Usage of LocalDateTime to Instant

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

@@ -8,7 +8,10 @@
 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.OffsetDateTime;
+import java.time.ZoneId;
 import java.time.format.DateTimeFormatter;
 import java.time.format.DateTimeParseException;
 import java.util.Arrays;
@@ -36,26 +39,45 @@ public CliTokenSource(
     this.env = env;
   }
 
-  static LocalDateTime parseExpiry(String expiry) {
+  /**
+   * Parses an expiry time string and returns the corresponding {@link Instant}. The method attempts
+   * to parse the input in the following order: 1. RFC 3339/ISO 8601 format with offset (e.g.
+   * "2024-03-20T10:30:00Z") - If the timestamp is compliant with ISO 8601 and RFC 3339, the
+   * timezone indicator (Z or +/-HH:mm) is used as part of the conversion process to UTC 2. Local
+   * date-time format "yyyy-MM-dd HH:mm:ss" (e.g. "2024-03-20 10:30:00") 3. Local date-time format
+   * with optional fractional seconds of varying precision (e.g. "2024-03-20 10:30:00.123")
+   *
+   * <p>For timestamps without a timezone indicator (formats 2 and 3), the system's default time
+   * zone is assumed and used for the conversion to UTC.
+   *
+   * @param expiry expiry time string in one of the supported formats
+   * @return the parsed {@link Instant}
+   * @throws DateTimeParseException if the input string cannot be parsed in any of the supported
+   *     formats
+   */
+  static Instant parseExpiry(String expiry) {
+    DateTimeParseException parseException;
+    try {
+      return OffsetDateTime.parse(expiry).toInstant();
+    } catch (DateTimeParseException e) {
+      parseException = 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;
+        Arrays.asList("yyyy-MM-dd HH:mm:ss", "yyyy-MM-dd HH:mm:ss." + multiplePrecisionPattern);
     for (String pattern : datePatterns) {
       try {
         DateTimeFormatter formatter = DateTimeFormatter.ofPattern(pattern);
         LocalDateTime dateTime = LocalDateTime.parse(expiry, formatter);
-        return dateTime;
+        return dateTime.atZone(ZoneId.systemDefault()).toInstant();
       } catch (DateTimeParseException e) {
-        lastException = e;
+        parseException.addSuppressed(e);
       }
     }
-    throw lastException;
+
+    throw parseException;
   }
 
   private String getProcessStream(InputStream stream) throws IOException {
@@ -83,7 +105,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/main/java/com/databricks/sdk/core/utils/SystemClockSupplier.java

@@ -5,6 +5,6 @@
 public class SystemClockSupplier implements ClockSupplier {
   @Override
   public Clock getClock() {
-    return Clock.systemDefaultZone();
+    return Clock.systemUTC();
   }
 }

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

@@ -1,6 +1,7 @@
 package com.databricks.sdk.core;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.mockConstruction;
@@ -14,7 +15,9 @@
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.time.Duration;
+import java.time.Instant;
 import java.time.LocalDateTime;
+import java.time.ZoneId;
 import java.time.ZonedDateTime;
 import java.time.format.DateTimeFormatter;
 import java.time.format.DateTimeParseException;
@@ -27,7 +30,6 @@
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
@@ -157,30 +159,58 @@ public void testRefreshWithExpiry(
     }
   }
 
-  @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);
+  private static Stream<Arguments> expiryProvider() {
+    return Stream.of(
+        Arguments.of(
+            "2023-07-17T09:02:22.330612218Z",
+            Instant.parse("2023-07-17T09:02:22.330612218Z"),
+            "9-digit nanos"),
+        Arguments.of(
+            "2023-07-17T09:02:22.33061221Z",
+            Instant.parse("2023-07-17T09:02:22.330612210Z"),
+            "8-digit nanos"),
+        Arguments.of(
+            "2023-07-17T09:02:22.330612Z",
+            Instant.parse("2023-07-17T09:02:22.330612000Z"),
+            "6-digit nanos"),
+        Arguments.of(
+            "2023-07-17T10:02:22.330612218+01:00",
+            Instant.parse("2023-07-17T09:02:22.330612218Z"),
+            "+01:00 offset, 9-digit nanos"),
+        Arguments.of(
+            "2023-07-17T04:02:22.330612218-05:00",
+            Instant.parse("2023-07-17T09:02:22.330612218Z"),
+            "-05:00 offset, 9-digit nanos"),
+        Arguments.of(
+            "2023-07-17T10:02:22.330612+01:00",
+            Instant.parse("2023-07-17T09:02:22.330612000Z"),
+            "+01:00 offset, 6-digit nanos"),
+        Arguments.of("2023-07-17T09:02:22.33061221987Z", null, "Invalid: >9 nanos"),
+        Arguments.of("17-07-2023 09:02:22", null, "Invalid date format"),
+        Arguments.of(
+            "2023-07-17 09:02:22.330612218",
+            LocalDateTime.parse("2023-07-17T09:02:22.330612218")
+                .atZone(ZoneId.systemDefault())
+                .toInstant(),
+            "Space separator, 9-digit nanos"),
+        Arguments.of(
+            "2023-07-17 09:02:22.330612",
+            LocalDateTime.parse("2023-07-17T09:02:22.330612")
+                .atZone(ZoneId.systemDefault())
+                .toInstant(),
+            "Space separator, 6-digit nanos"),
+        Arguments.of(
+            "2023-07-17 09:02:22.33061221987", null, "Space separator, Invalid: >9 nanos"));
   }
 
-  @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);
-  }
-
-  @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);
-  }
-
-  @Test
-  public void testParseExpiryWithMoreThanNineNanoSecondDigits() {
-    try {
-      CliTokenSource.parseExpiry("2023-07-17T09:02:22.33061221987Z");
-    } catch (DateTimeParseException e) {
-      assert (e.getMessage().contains("could not be parsed"));
+  @ParameterizedTest(name = "{2}")
+  @MethodSource("expiryProvider")
+  public void testParseExpiry(String input, Instant expectedInstant, String description) {
+    if (expectedInstant == null) {
+      assertThrows(DateTimeParseException.class, () -> CliTokenSource.parseExpiry(input));
+    } else {
+      Instant parsedInstant = CliTokenSource.parseExpiry(input);
+      assertEquals(expectedInstant, parsedInstant);
     }
   }
 }

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");