Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[EWT-1250] Sqlalchemy/Superset expectations from python-DBAPI #17

Draft
wants to merge 8 commits into
base: main
Choose a base branch
from
Draft

[EWT-1250] Sqlalchemy/Superset expectations from python-DBAPI #17

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
c43c8a4
feat: add ws_url parameter to connect
prantogg Aug 1, 2024
c21261a
feat: change return type to List of Tuples
prantogg Aug 1, 2024
fa5d1b1
feat: allow passing commit() and rollback()
prantogg Aug 1, 2024
df3fff5
add temporary debug logging
prantogg Aug 1, 2024
b802b80
remove temporary debug logging
prantogg Aug 1, 2024
65120fb
add types to type map
prantogg Aug 1, 2024
2c9b5f8
address comments
prantogg Aug 1, 2024
96e44dc
fix typo
prantogg Aug 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 13 additions & 6 deletions wherobots/db/connection.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
GeometryRepresentation,
)
from wherobots.db.cursor import Cursor
from wherobots.db.errors import NotSupportedError, OperationalError
from wherobots.db.errors import OperationalError


@dataclass
Expand Down Expand Up @@ -78,11 +78,10 @@ def close(self):
self.__ws.close()

def commit(self):
raise NotSupportedError
pass

def rollback(self):
raise NotSupportedError

pass
peterfoldes marked this conversation as resolved.
Show resolved Hide resolved
def cursor(self) -> Cursor:
return Cursor(self.__execute_sql, self.__cancel_query)

Expand Down Expand Up @@ -155,12 +154,20 @@ def __listen(self):

query.state = ExecutionState.COMPLETED
if result_format == ResultsFormat.JSON:
query.handler(json.loads(result_bytes.decode("utf-8")))
data = json.loads(result_bytes.decode("utf-8"))
columns = data["columns"]
column_types = data.get("column_types")
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is column_types optional? If so, then it's good that you're using data.get() here, but then in Cursor.__get_results you expect column_types to be non-None. You either need to ensure column_types is always provided, or change __get_results to be more defensive.

rows = data["rows"]
query.handler((columns, column_types, rows))
elif result_format == ResultsFormat.ARROW:
buffer = pyarrow.py_buffer(result_bytes)
stream = pyarrow.input_stream(buffer, result_compression)
with pyarrow.ipc.open_stream(stream) as reader:
query.handler(reader.read_pandas())
schema = reader.schema
columns = schema.names
column_types = [field.type for field in schema]
rows = reader.read_all().to_pandas().values.tolist()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you need both read_all().to_pandas() or can you read_pandas() directly? (to be fair according to the docs it does the same thing)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch! Using read_pandas() now for better readability.

query.handler((columns, column_types, rows))
else:
query.handler(
OperationalError(f"Unsupported results format {result_format}")
Expand Down
20 changes: 14 additions & 6 deletions wherobots/db/cursor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,19 @@

_TYPE_MAP = {
"object": "STRING",
"string": "STRING",
"int32": "NUMBER",
"int64": "NUMBER",
"float32": "NUMBER",
"float64": "NUMBER",
"datetime64[ns]": "DATETIME",
"timedelta[ns]": "DATETIME",
"double": "NUMBER",
"bool": "NUMBER", # Assuming boolean is stored as number
"bytes": "BINARY",
"struct": "STRUCT",
"list": "LIST",
"geometry": "GEOMETRY"
}


Expand Down Expand Up @@ -54,20 +61,21 @@ def __get_results(self) -> Optional[List[Tuple[Any, ...]]]:
if isinstance(result, DatabaseError):
raise result

self.__rowcount = len(result)
self.__results = result
if not result.empty:
columns, column_types, rows = result
self.__rowcount = len(rows)
self.__results = rows
if rows:
self.__description = [
(
col_name, # name
_TYPE_MAP.get(str(result[col_name].dtype), "STRING"), # type_code
_TYPE_MAP.get(str(column_types[i]), "STRING"), # type_code
None, # display_size
result[col_name].memory_usage(), # internal_size
None, # internal_size
peterfoldes marked this conversation as resolved.
Show resolved Hide resolved
None, # precision
None, # scale
True, # null_ok; Assuming all columns can accept NULL values
)
for col_name in result.columns
for i, col_name in enumerate(columns)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Use https://docs.python.org/3/library/functions.html#zip to avoid jumping hoops with an index (it's much nicer to read, and also more efficient):

self.__description = [
  (
    col_name,
    _TYPE_MAP.get(col_type, 'STRING'),
    ...
  )
  for (col_name, col_type) in zip(columns, column_types)
]

]

return self.__results
Expand Down
99 changes: 51 additions & 48 deletions wherobots/db/driver.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@
"""

import logging
import os
import urllib.parse
import queue
import requests
Expand Down Expand Up @@ -51,6 +50,7 @@ def connect(
results_format: Union[ResultsFormat, None] = None,
data_compression: Union[DataCompression, None] = None,
geometry_representation: Union[GeometryRepresentation, None] = None,
ws_url: str = None,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this required, instead of calling connect_direct() directly?

The reason this parameter wasn't included in connect() is that it creates ambiguity between the rest of the parameters (like runtime/region) and the runtime you'd actually connect to when providing a ws_url, which may not match those choices.

) -> Connection:
if not token and not api_key:
raise ValueError("At least one of `token` or `api_key` is required")
Expand All @@ -76,54 +76,57 @@ def connect(
)

# Default to HTTPS if the hostname doesn't explicitly specify a scheme.
if not host.startswith("http:"):
host = f"https://{host}"

try:
resp = requests.post(
url=f"{host}/sql/session",
params={"region": region.value},
json={
"runtimeId": runtime.value,
"shutdownAfterInactiveSeconds": shutdown_after_inactive_seconds,
},
headers=headers,
if ws_url:
session_uri = ws_url
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Couple of things here:

  1. This can be done earlier, so if we feel confident in using the ws_url we can do that right after checking the token/api-keys.
  2. This needs a little bit different logging, since we're not requesting a new runtime (see line 70).
  3. This should early return so that we don't have to indent everything afterwards. Makes it more readable.

else:
if not host.startswith("http:"):
host = f"https://{host}"

try:
resp = requests.post(
url=f"{host}/sql/session",
params={"region": region.value},
json={
"runtimeId": runtime.value,
"shutdownAfterInactiveSeconds": shutdown_after_inactive_seconds,
},
headers=headers,
)
resp.raise_for_status()
except requests.HTTPError as e:
raise InterfaceError("Failed to create SQL session!", e)

# At this point we've been redirected to /sql/session/{session_id}, which we'll need to keep polling until the
# session is in READY state.
session_id_url = resp.url

@tenacity.retry(
stop=tenacity.stop_after_delay(wait_timeout),
wait=tenacity.wait_exponential(multiplier=1, min=1, max=5),
retry=tenacity.retry_if_not_exception_type(
(requests.HTTPError, OperationalError)
),
)
resp.raise_for_status()
except requests.HTTPError as e:
raise InterfaceError("Failed to create SQL session!", e)

# At this point we've been redirected to /sql/session/{session_id}, which we'll need to keep polling until the
# session is in READY state.
session_id_url = resp.url

@tenacity.retry(
stop=tenacity.stop_after_delay(wait_timeout),
wait=tenacity.wait_exponential(multiplier=1, min=1, max=5),
retry=tenacity.retry_if_not_exception_type(
(requests.HTTPError, OperationalError)
),
)
def get_session_uri() -> str:
r = requests.get(session_id_url, headers=headers)
r.raise_for_status()
payload = r.json()
status = AppStatus(payload.get("status"))
logging.info(" ... %s", status)
if status.is_starting():
raise tenacity.TryAgain("SQL Session is not ready yet")
elif status == AppStatus.READY:
return payload["appMeta"]["url"]
else:
logging.error("SQL session creation failed: %s; should not retry.", status)
raise OperationalError(f"Failed to create SQL session: {status}")

try:
logging.info("Getting SQL session status from %s ...", session_id_url)
session_uri = get_session_uri()
logging.debug("SQL session URI from app status: %s", session_uri)
except Exception as e:
raise InterfaceError("Could not acquire SQL session!", e)
def get_session_uri() -> str:
r = requests.get(session_id_url, headers=headers)
r.raise_for_status()
payload = r.json()
status = AppStatus(payload.get("status"))
logging.info(" ... %s", status)
if status.is_starting():
raise tenacity.TryAgain("SQL Session is not ready yet")
elif status == AppStatus.READY:
return payload["appMeta"]["url"]
else:
logging.error("SQL session creation failed: %s; should not retry.", status)
raise OperationalError(f"Failed to create SQL session: {status}")

try:
logging.info("Getting SQL session status from %s ...", session_id_url)
session_uri = get_session_uri()
logging.debug("SQL session URI from app status: %s", session_uri)
except Exception as e:
raise InterfaceError("Could not acquire SQL session!", e)

return connect_direct(
uri=http_to_ws(session_uri),
Expand Down