Skip to content

Commit 1809956

Browse files
house constants in enums for readability and immutability
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent 94879c0 commit 1809956

File tree

2 files changed

+45
-4
lines changed

2 files changed

+45
-4
lines changed

src/databricks/sql/backend/sea/backend.py

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,10 @@
55

66
from databricks.sql.backend.sea.utils.constants import (
77
ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP,
8+
ResultFormat,
9+
ResultDisposition,
10+
ResultCompression,
11+
WaitTimeout,
812
)
913

1014
if TYPE_CHECKING:
@@ -405,17 +409,25 @@ def execute_command(
405409
)
406410
)
407411

408-
format = "ARROW_STREAM" if use_cloud_fetch else "JSON_ARRAY"
409-
disposition = "EXTERNAL_LINKS" if use_cloud_fetch else "INLINE"
410-
result_compression = "LZ4_FRAME" if lz4_compression else None
412+
format = (
413+
ResultFormat.ARROW_STREAM if use_cloud_fetch else ResultFormat.JSON_ARRAY
414+
).value
415+
disposition = (
416+
ResultDisposition.EXTERNAL_LINKS
417+
if use_cloud_fetch
418+
else ResultDisposition.INLINE
419+
).value
420+
result_compression = (
421+
ResultCompression.LZ4_FRAME if lz4_compression else ResultCompression.NONE
422+
).value
411423

412424
request = ExecuteStatementRequest(
413425
warehouse_id=self.warehouse_id,
414426
session_id=sea_session_id,
415427
statement=operation,
416428
disposition=disposition,
417429
format=format,
418-
wait_timeout="0s" if async_op else "10s",
430+
wait_timeout=(WaitTimeout.ASYNC if async_op else WaitTimeout.SYNC).value,
419431
on_wait_timeout="CONTINUE",
420432
row_limit=max_rows,
421433
parameters=sea_parameters if sea_parameters else None,

src/databricks/sql/backend/sea/utils/constants.py

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
"""
44

55
from typing import Dict
6+
from enum import Enum
67

78
# from https://docs.databricks.com/aws/en/sql/language-manual/sql-ref-parameters
89
ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP: Dict[str, str] = {
@@ -15,3 +16,31 @@
1516
"TIMEZONE": "UTC",
1617
"USE_CACHED_RESULT": "true",
1718
}
19+
20+
21+
class ResultFormat(Enum):
22+
"""Enum for result format values."""
23+
24+
ARROW_STREAM = "ARROW_STREAM"
25+
JSON_ARRAY = "JSON_ARRAY"
26+
27+
28+
class ResultDisposition(Enum):
29+
"""Enum for result disposition values."""
30+
31+
EXTERNAL_LINKS = "EXTERNAL_LINKS"
32+
INLINE = "INLINE"
33+
34+
35+
class ResultCompression(Enum):
36+
"""Enum for result compression values."""
37+
38+
LZ4_FRAME = "LZ4_FRAME"
39+
NONE = None
40+
41+
42+
class WaitTimeout(Enum):
43+
"""Enum for wait timeout values."""
44+
45+
ASYNC = "0s"
46+
SYNC = "10s"

0 commit comments

Comments
 (0)