diff --git a/databricks/sdk/core.py b/databricks/sdk/core.py index e3d4f8cc9..04a1b78ea 100644 --- a/databricks/sdk/core.py +++ b/databricks/sdk/core.py @@ -346,7 +346,8 @@ def __init__(self, cmd: List[str], token_type_field: str, access_token_field: st @staticmethod def _parse_expiry(expiry: str) -> datetime: - for fmt in ("%Y-%m-%d %H:%M:%S.%f", "%Y-%m-%d %H:%M:%S", "%Y-%m-%dT%H:%M:%S.%f%z"): + expiry = expiry.rstrip("Z").split(".")[0] + for fmt in ("%Y-%m-%d %H:%M:%S", "%Y-%m-%dT%H:%M:%S"): try: return datetime.strptime(expiry, fmt) except ValueError as e: diff --git a/tests/test_core.py b/tests/test_core.py index 2b865bbd3..d7e2c8f41 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -6,6 +6,7 @@ import random import string import typing +from datetime import datetime from http.server import BaseHTTPRequestHandler from typing import Iterator, List @@ -14,10 +15,10 @@ from databricks.sdk import WorkspaceClient from databricks.sdk.azure import ENVIRONMENTS, AzureEnvironment -from databricks.sdk.core import (ApiClient, Config, CredentialsProvider, - DatabricksCliTokenSource, DatabricksError, - HeaderFactory, StreamingResponse, - databricks_cli) +from databricks.sdk.core import (ApiClient, CliTokenSource, Config, + CredentialsProvider, DatabricksCliTokenSource, + DatabricksError, HeaderFactory, + StreamingResponse, databricks_cli) from databricks.sdk.service.catalog import PermissionsChange from databricks.sdk.service.iam import AccessControlRequest from databricks.sdk.version import __version__ @@ -52,6 +53,14 @@ def test_databricks_cli_token_source_not_installed(config, monkeypatch): DatabricksCliTokenSource(config) +@pytest.mark.parametrize("date_string,expected", + [("2023-12-01T15:19:48.007742617Z", datetime(2023, 12, 1, 15, 19, 48)), + ("2023-12-05T15:59:01.40081+11:00", datetime(2023, 12, 5, 15, 59, 1)), + ("2023-12-06 10:06:05", datetime(2023, 12, 6, 10, 6, 5))]) +def test_databricks_cli_token_parse_expiry(date_string, expected): + assert CliTokenSource._parse_expiry(date_string) == expected + + def write_small_dummy_executable(path: pathlib.Path): cli = path.joinpath('databricks') cli.write_text('#!/bin/sh\necho "hello world"\n')