Skip to content

Commit

Permalink
Handle nested query parameters in ApiClient.do() (#249)
Browse files Browse the repository at this point in the history
## Changes
The Query History list API filter_by query parameter is modeled by a
dictionary, rather than a primitive type, defining the allowed filters
including `query_start_time_range`, `statuses`, `user_ids`, and
`warehouse_ids`. To be compatible with gRPC transcoding, query
parameters modeled by message types as opposed to primitives need to be
separated into one query parameter per nested field, where the key is
the path to that query parameter. For example:

```
{'filter_by': {'user_ids': [123, 456]}}
```
becomes
```
filter_by.user_ids=123&filter_by.user_ids=456
```
For this to be compatible with the requests library we use today, we
need to convert the first dictionary to
```
{'filter_by.user_ids': [123, 456]}
```

This resolves one of the problems from #99. The issue with the conflict
between filter_by and next_page_token will be resolved by the backend
service.

## Tests
Added an integration test that covers query history listing with a
filter_by parameter.

- [x] `make test` run locally
- [x] `make fmt` applied
- [x] relevant integration tests applied
  • Loading branch information
mgyucht authored Jul 25, 2023
1 parent aa9d048 commit 2e07e0a
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 2 deletions.
24 changes: 22 additions & 2 deletions databricks/sdk/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
import urllib.parse
from datetime import datetime
from json import JSONDecodeError
from typing import Callable, Dict, Iterable, List, Optional, Union
from typing import Any, Callable, Dict, Iterable, List, Optional, Union

import requests
import requests.auth
Expand Down Expand Up @@ -911,7 +911,27 @@ def _fix_query_string(query: Optional[dict] = None) -> Optional[dict]:
# See: https://github.com/databricks/databricks-sdk-py/issues/142
if query is None:
return None
return {k: v if type(v) != bool else ('true' if v else 'false') for k, v in query.items()}
with_fixed_bools = {k: v if type(v) != bool else ('true' if v else 'false') for k, v in query.items()}

# Query parameters may be nested, e.g.
# {'filter_by': {'user_ids': [123, 456]}}
# The HTTP-compatible representation of this is
# filter_by.user_ids=123&filter_by.user_ids=456
# To achieve this, we convert the above dictionary to
# {'filter_by.user_ids': [123, 456]}
# See the following for more information:
# https://cloud.google.com/endpoints/docs/grpc-service-config/reference/rpc/google.api#google.api.HttpRule
def flatten_dict(d: Dict[str, Any]) -> Dict[str, Any]:
for k1, v1 in d.items():
if isinstance(v1, dict):
v1 = dict(flatten_dict(v1))
for k2, v2 in v1.items():
yield f"{k1}.{k2}", v2
else:
yield k1, v1

flattened = dict(flatten_dict(with_fixed_bools))
return flattened

def do(self,
method: str,
Expand Down
15 changes: 15 additions & 0 deletions tests/integration/test_sql.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
from datetime import datetime

from databricks.sdk.service.sql import QueryFilter, TimeRange


def test_query_history_list_with_filter(w):

def date_to_ms(date):
return int(datetime.strptime(date, '%Y-%m-%d').timestamp() * 1000)

filter = QueryFilter(query_start_time_range=TimeRange(start_time_ms=date_to_ms('2023-01-01'),
end_time_ms=date_to_ms('2023-01-02')))
queries = w.query_history.list(filter_by=filter)
for q in queries:
print(q)

0 comments on commit 2e07e0a

Please sign in to comment.