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 all 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
18 changes: 13 additions & 5 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,10 +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 +155,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_pandas().values.tolist()
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
20 changes: 19 additions & 1 deletion 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 @@ -67,6 +67,24 @@ def connect(
runtime = runtime or DEFAULT_RUNTIME
region = region or DEFAULT_REGION

if ws_url:
logging.info(
"Using existing %s/%s runtime in %s from %s ...",
runtime.name,
runtime.value,
region.value,
host,
)
session_uri = ws_url
return connect_direct(
uri=http_to_ws(session_uri),
headers=headers,
read_timeout=read_timeout,
results_format=results_format,
data_compression=data_compression,
geometry_representation=geometry_representation,
)

logging.info(
"Requesting %s/%s runtime in %s from %s ...",
runtime.name,
Expand Down