File tree Expand file tree Collapse file tree 2 files changed +45
-4
lines changed
src/databricks/sql/backend/sea Expand file tree Collapse file tree 2 files changed +45
-4
lines changed Original file line number Diff line number Diff line change 5
5
6
6
from databricks .sql .backend .sea .utils .constants import (
7
7
ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP ,
8
+ ResultFormat ,
9
+ ResultDisposition ,
10
+ ResultCompression ,
11
+ WaitTimeout ,
8
12
)
9
13
10
14
if TYPE_CHECKING :
@@ -405,17 +409,25 @@ def execute_command(
405
409
)
406
410
)
407
411
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
411
423
412
424
request = ExecuteStatementRequest (
413
425
warehouse_id = self .warehouse_id ,
414
426
session_id = sea_session_id ,
415
427
statement = operation ,
416
428
disposition = disposition ,
417
429
format = format ,
418
- wait_timeout = "0s" if async_op else "10s" ,
430
+ wait_timeout = ( WaitTimeout . ASYNC if async_op else WaitTimeout . SYNC ). value ,
419
431
on_wait_timeout = "CONTINUE" ,
420
432
row_limit = max_rows ,
421
433
parameters = sea_parameters if sea_parameters else None ,
Original file line number Diff line number Diff line change 3
3
"""
4
4
5
5
from typing import Dict
6
+ from enum import Enum
6
7
7
8
# from https://docs.databricks.com/aws/en/sql/language-manual/sql-ref-parameters
8
9
ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP : Dict [str , str ] = {
15
16
"TIMEZONE" : "UTC" ,
16
17
"USE_CACHED_RESULT" : "true" ,
17
18
}
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"
You can’t perform that action at this time.
0 commit comments