Skip to content

Commit e2a919b

Browse files
JCZuurmondmicheledaddetta-databricksgueniai
authored
Implement disposition field in SQL backend (#3477)
### Description This feature aims to the chance of running SQL statement for assessment results export in case of large workspaces with a large amount of findings. Introducing the `query_statement_disposition` value during the UCX installation can allow users to select the proper disposition method to use when running large queries. This parameter is added in the `config.yml` file and used for the SqlBackend definition. Using this approach, large queries will not fail. ### Changes Added SQL query statement disposition choice during the UCX installation Added `query_statement_disposition` value in the config file ### Linked issues Resolves #3447 From other PR #3455 ### Functionality - [X] modified existing commands: `databricks labs install ucx`, `databricks labs ucx export-assessment` ### Tests - [X] manually tested --------- Co-authored-by: micheledaddetta-databricks <michele.daddetta@databricks.com> Co-authored-by: Guenia Izquierdo Delgado <guenia.izquierdo@databricks.com>
1 parent 3e96d4f commit e2a919b

File tree

5 files changed

+19
-7
lines changed

5 files changed

+19
-7
lines changed

src/databricks/labs/ucx/config.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from dataclasses import dataclass
22

33
from databricks.sdk.core import Config
4+
from databricks.sdk.service.sql import Disposition
45

56
__all__ = ["WorkspaceConfig"]
67

@@ -92,6 +93,9 @@ class WorkspaceConfig: # pylint: disable=too-many-instance-attributes
9293
# Skip TACL migration during table migration
9394
skip_tacl_migration: bool = False
9495

96+
# Select SQL query statement disposition
97+
query_statement_disposition: Disposition = Disposition.INLINE
98+
9599
def replace_inventory_variable(self, text: str) -> str:
96100
return text.replace("$inventory", f"hive_metastore.{self.inventory_database}")
97101

src/databricks/labs/ucx/contexts/workspace_cli.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,9 @@ def workspace_client(self) -> WorkspaceClient:
4343

4444
@cached_property
4545
def sql_backend(self) -> SqlBackend:
46-
return StatementExecutionBackend(self.workspace_client, self.config.warehouse_id)
46+
return StatementExecutionBackend(
47+
self.workspace_client, self.config.warehouse_id, disposition=self.config.query_statement_disposition
48+
)
4749

4850
@cached_property
4951
def cluster_access(self) -> ClusterAccess:

src/databricks/labs/ucx/install.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -257,6 +257,7 @@ def _prompt_for_new_installation(self) -> WorkspaceConfig:
257257
recon_tolerance_percent = int(
258258
self.prompts.question("Reconciliation threshold, in percentage", default="5", valid_number=True)
259259
)
260+
260261
return WorkspaceConfig(
261262
inventory_database=inventory_database,
262263
ucx_catalog=ucx_catalog,

tests/unit/install/test_install.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -377,6 +377,7 @@ def test_configure_sets_expected_workspace_configuration_values(
377377
"warehouse_id": "abc",
378378
"workspace_start_path": "/",
379379
"num_days_submit_runs_history": 30,
380+
'query_statement_disposition': 'INLINE',
380381
"recon_tolerance_percent": 5,
381382
"managed_table_external_storage": "CLONE",
382383
}
@@ -412,6 +413,7 @@ def test_configure_with_default_owner_group(
412413
"min_workers": 1,
413414
"max_workers": 10,
414415
"policy_id": "foo",
416+
"query_statement_disposition": "INLINE",
415417
"renamed_group_prefix": "db-temp-",
416418
"warehouse_id": "abc",
417419
"workspace_start_path": "/",
@@ -506,6 +508,7 @@ def test_create_cluster_policy(ws, mock_installation) -> None:
506508
'renamed_group_prefix': 'db-temp-',
507509
'warehouse_id': 'abc',
508510
'workspace_start_path': '/',
511+
'query_statement_disposition': 'INLINE',
509512
'recon_tolerance_percent': 5,
510513
'managed_table_external_storage': "CLONE",
511514
},
@@ -1751,6 +1754,7 @@ def test_save_config_ext_hms(ws, mock_installation) -> None:
17511754
'warehouse_id': 'abc',
17521755
'workspace_start_path': '/',
17531756
'num_days_submit_runs_history': 30,
1757+
'query_statement_disposition': 'INLINE',
17541758
'recon_tolerance_percent': 5,
17551759
'managed_table_external_storage': "CLONE",
17561760
},

tests/unit/test_cli.py

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import pytest
99
import yaml
1010
from databricks.labs.blueprint.tui import MockPrompts
11+
from databricks.sdk.service.sql import Disposition
1112

1213
from databricks.labs.ucx.aws.credentials import IamRoleCreation
1314
from databricks.sdk import AccountClient, WorkspaceClient
@@ -207,7 +208,7 @@ def test_skip_with_schema(ws) -> None:
207208
byte_limit=None,
208209
catalog=None,
209210
schema=None,
210-
disposition=None,
211+
disposition=Disposition.INLINE,
211212
format=sql.Format.JSON_ARRAY,
212213
wait_timeout=None,
213214
)
@@ -222,7 +223,7 @@ def test_skip_with_table(ws) -> None:
222223
byte_limit=None,
223224
catalog=None,
224225
schema=None,
225-
disposition=None,
226+
disposition=Disposition.INLINE,
226227
format=sql.Format.JSON_ARRAY,
227228
wait_timeout=None,
228229
)
@@ -237,7 +238,7 @@ def test_skip_with_view(ws) -> None:
237238
byte_limit=None,
238239
catalog=None,
239240
schema=None,
240-
disposition=None,
241+
disposition=Disposition.INLINE,
241242
format=sql.Format.JSON_ARRAY,
242243
wait_timeout=None,
243244
)
@@ -264,7 +265,7 @@ def test_unskip_with_schema(ws) -> None:
264265
byte_limit=None,
265266
catalog=None,
266267
schema=None,
267-
disposition=None,
268+
disposition=Disposition.INLINE,
268269
format=sql.Format.JSON_ARRAY,
269270
wait_timeout=None,
270271
)
@@ -279,7 +280,7 @@ def test_unskip_with_table(ws) -> None:
279280
byte_limit=None,
280281
catalog=None,
281282
schema=None,
282-
disposition=None,
283+
disposition=Disposition.INLINE,
283284
format=sql.Format.JSON_ARRAY,
284285
wait_timeout=None,
285286
)
@@ -294,7 +295,7 @@ def test_unskip_with_view(ws) -> None:
294295
byte_limit=None,
295296
catalog=None,
296297
schema=None,
297-
disposition=None,
298+
disposition=Disposition.INLINE,
298299
format=sql.Format.JSON_ARRAY,
299300
wait_timeout=None,
300301
)

0 commit comments

Comments
 (0)