Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commitd9bcdbe

Browse files
remove irrelevant changes
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent8985c62 commitd9bcdbe

File tree

8 files changed

+107
-774
lines changed

8 files changed

+107
-774
lines changed

‎examples/experimental/sea_connector_test.py‎

Lines changed: 14 additions & 110 deletions
Original file line numberDiff line numberDiff line change
@@ -6,122 +6,34 @@
66
logging.basicConfig(level=logging.DEBUG)
77
logger=logging.getLogger(__name__)
88

9-
10-
deftest_sea_query_exec():
11-
"""
12-
Test executing a query using the SEA backend with result compression.
13-
14-
This function connects to a Databricks SQL endpoint using the SEA backend,
15-
executes a simple query with result compression enabled and disabled,
16-
and verifies that execution completes successfully.
17-
"""
18-
server_hostname=os.environ.get("DATABRICKS_SERVER_HOSTNAME")
19-
http_path=os.environ.get("DATABRICKS_HTTP_PATH")
20-
access_token=os.environ.get("DATABRICKS_TOKEN")
21-
catalog=os.environ.get("DATABRICKS_CATALOG")
22-
23-
ifnotall([server_hostname,http_path,access_token]):
24-
logger.error("Missing required environment variables.")
25-
logger.error(
26-
"Please set DATABRICKS_SERVER_HOSTNAME, DATABRICKS_HTTP_PATH, and DATABRICKS_TOKEN."
27-
)
28-
sys.exit(1)
29-
30-
try:
31-
# Test with compression enabled
32-
logger.info("Creating connection with LZ4 compression enabled")
33-
connection=Connection(
34-
server_hostname=server_hostname,
35-
http_path=http_path,
36-
access_token=access_token,
37-
catalog=catalog,
38-
schema="default",
39-
use_sea=True,
40-
user_agent_entry="SEA-Test-Client",
41-
use_cloud_fetch=True,# Enable cloud fetch to use compression
42-
enable_query_result_lz4_compression=True,# Enable LZ4 compression
43-
)
44-
45-
logger.info(
46-
f"Successfully opened SEA session with ID:{connection.get_session_id_hex()}"
47-
)
48-
logger.info(f"backend type:{type(connection.session.backend)}")
49-
50-
# Execute a simple query with compression enabled
51-
cursor=connection.cursor(arraysize=0,buffer_size_bytes=0)
52-
logger.info("Executing query with LZ4 compression: SELECT 1 as test_value")
53-
cursor.execute("SELECT 1 as test_value")
54-
logger.info("Query with compression executed successfully")
55-
cursor.close()
56-
connection.close()
57-
logger.info("Successfully closed SEA session with compression enabled")
58-
59-
# Test with compression disabled
60-
logger.info("Creating connection with LZ4 compression disabled")
61-
connection=Connection(
62-
server_hostname=server_hostname,
63-
http_path=http_path,
64-
access_token=access_token,
65-
catalog=catalog,
66-
schema="default",
67-
use_sea=True,
68-
user_agent_entry="SEA-Test-Client",
69-
use_cloud_fetch=False,# Enable cloud fetch
70-
enable_query_result_lz4_compression=False,# Disable LZ4 compression
71-
)
72-
73-
logger.info(
74-
f"Successfully opened SEA session with ID:{connection.get_session_id_hex()}"
75-
)
76-
77-
# Execute a simple query with compression disabled
78-
cursor=connection.cursor(arraysize=0,buffer_size_bytes=0)
79-
logger.info("Executing query without compression: SELECT 1 as test_value")
80-
cursor.execute("SELECT 1 as test_value")
81-
logger.info("Query without compression executed successfully")
82-
cursor.close()
83-
connection.close()
84-
logger.info("Successfully closed SEA session with compression disabled")
85-
86-
exceptExceptionase:
87-
logger.error(f"Error during SEA query execution test:{str(e)}")
88-
importtraceback
89-
90-
logger.error(traceback.format_exc())
91-
sys.exit(1)
92-
93-
logger.info("SEA query execution test with compression completed successfully")
94-
95-
969
deftest_sea_session():
9710
"""
9811
Test opening and closing a SEA session using the connector.
99-
12+
10013
This function connects to a Databricks SQL endpoint using the SEA backend,
10114
opens a session, and then closes it.
102-
15+
10316
Required environment variables:
10417
- DATABRICKS_SERVER_HOSTNAME: Databricks server hostname
10518
- DATABRICKS_HTTP_PATH: HTTP path for the SQL endpoint
10619
- DATABRICKS_TOKEN: Personal access token for authentication
10720
"""
21+
10822
server_hostname=os.environ.get("DATABRICKS_SERVER_HOSTNAME")
10923
http_path=os.environ.get("DATABRICKS_HTTP_PATH")
11024
access_token=os.environ.get("DATABRICKS_TOKEN")
11125
catalog=os.environ.get("DATABRICKS_CATALOG")
112-
26+
11327
ifnotall([server_hostname,http_path,access_token]):
11428
logger.error("Missing required environment variables.")
115-
logger.error(
116-
"Please set DATABRICKS_SERVER_HOSTNAME, DATABRICKS_HTTP_PATH, and DATABRICKS_TOKEN."
117-
)
29+
logger.error("Please set DATABRICKS_SERVER_HOSTNAME, DATABRICKS_HTTP_PATH, and DATABRICKS_TOKEN.")
11830
sys.exit(1)
119-
31+
12032
logger.info(f"Connecting to{server_hostname}")
12133
logger.info(f"HTTP Path:{http_path}")
12234
ifcatalog:
12335
logger.info(f"Using catalog:{catalog}")
124-
36+
12537
try:
12638
logger.info("Creating connection with SEA backend...")
12739
connection=Connection(
@@ -130,33 +42,25 @@ def test_sea_session():
13042
access_token=access_token,
13143
catalog=catalog,
13244
schema="default",
133-
use_sea=True,
134-
user_agent_entry="SEA-Test-Client",# add custom user agent
135-
)
136-
137-
logger.info(
138-
f"Successfully opened SEA session with ID:{connection.get_session_id_hex()}"
45+
use_sea=True,
46+
user_agent_entry="SEA-Test-Client"# add custom user agent
13947
)
48+
49+
logger.info(f"Successfully opened SEA session with ID:{connection.get_session_id_hex()}")
14050
logger.info(f"backend type:{type(connection.session.backend)}")
141-
51+
14252
# Close the connection
14353
logger.info("Closing the SEA session...")
14454
connection.close()
14555
logger.info("Successfully closed SEA session")
146-
56+
14757
exceptExceptionase:
14858
logger.error(f"Error testing SEA session:{str(e)}")
14959
importtraceback
150-
15160
logger.error(traceback.format_exc())
15261
sys.exit(1)
153-
62+
15463
logger.info("SEA session test completed successfully")
15564

156-
15765
if__name__=="__main__":
158-
# Test session management
15966
test_sea_session()
160-
161-
# Test query execution with compression
162-
test_sea_query_exec()

‎src/databricks/sql/backend/databricks_client.py‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616

1717
fromdatabricks.sql.thrift_api.TCLIServiceimportttypes
1818
fromdatabricks.sql.backend.typesimportSessionId,CommandId,CommandState
19+
fromdatabricks.sql.utilsimportExecuteResponse
20+
fromdatabricks.sql.typesimportSSLOptions
1921

2022
# Forward reference for type hints
2123
fromtypingimportTYPE_CHECKING
@@ -86,6 +88,34 @@ def execute_command(
8688
async_op:bool,
8789
enforce_embedded_schema_correctness:bool,
8890
)->Union["ResultSet",None]:
91+
"""
92+
Executes a SQL command or query within the specified session.
93+
94+
This method sends a SQL command to the server for execution and handles
95+
the response. It can operate in both synchronous and asynchronous modes.
96+
97+
Args:
98+
operation: The SQL command or query to execute
99+
session_id: The session identifier in which to execute the command
100+
max_rows: Maximum number of rows to fetch in a single fetch batch
101+
max_bytes: Maximum number of bytes to fetch in a single fetch batch
102+
lz4_compression: Whether to use LZ4 compression for result data
103+
cursor: The cursor object that will handle the results
104+
use_cloud_fetch: Whether to use cloud fetch for retrieving large result sets
105+
parameters: List of parameters to bind to the query
106+
async_op: Whether to execute the command asynchronously
107+
enforce_embedded_schema_correctness: Whether to enforce schema correctness
108+
109+
Returns:
110+
If async_op is False, returns a ResultSet object containing the
111+
query results and metadata. If async_op is True, returns None and the
112+
results must be fetched later using get_execution_result().
113+
114+
Raises:
115+
ValueError: If the session ID is invalid
116+
OperationalError: If there's an error executing the command
117+
ServerOperationError: If the server encounters an error during execution
118+
"""
89119
pass
90120

91121
@abstractmethod

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp