Skip to content

Commit

Permalink
Properly parse ISO date strings (#473)
Browse files Browse the repository at this point in the history
Pre Python 3.11 the `datetime.fromisoformat` isn't a generic date time
parser and can only be used to parse ISO dates generated by the python
`datatime` module. This PR normalized the date string the same way the
google oauth2 library does by discarding the timezone and millisecond
part of the date string.

## Changes
Make date parsing for tokens from the CLI more robust

## Tests
* Added unit test
* Manually tested with CLI auth provider
  • Loading branch information
fjakobs authored Dec 7, 2023
1 parent 5955333 commit 741238b
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 5 deletions.
3 changes: 2 additions & 1 deletion databricks/sdk/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
17 changes: 13 additions & 4 deletions tests/test_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import random
import string
import typing
from datetime import datetime
from http.server import BaseHTTPRequestHandler
from typing import Iterator, List

Expand All @@ -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__
Expand Down Expand Up @@ -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')
Expand Down

0 comments on commit 741238b

Please sign in to comment.