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 <[email protected]>
Co-authored-by: Guenia Izquierdo Delgado <[email protected]>

src/databricks/labs/ucx/config.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 
 from databricks.sdk.core import Config
+from databricks.sdk.service.sql import Disposition
 
 __all__ = ["WorkspaceConfig"]
 
@@ -92,6 +93,9 @@ class WorkspaceConfig:  # pylint: disable=too-many-instance-attributes
     # Skip TACL migration during table migration
     skip_tacl_migration: bool = False
 
+    # Select SQL query statement disposition
+    query_statement_disposition: Disposition = Disposition.INLINE
+
     def replace_inventory_variable(self, text: str) -> str:
         return text.replace("$inventory", f"hive_metastore.{self.inventory_database}")
 

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

@@ -43,7 +43,9 @@ def workspace_client(self) -> WorkspaceClient:
 
     @cached_property
     def sql_backend(self) -> SqlBackend:
-        return StatementExecutionBackend(self.workspace_client, self.config.warehouse_id)
+        return StatementExecutionBackend(
+            self.workspace_client, self.config.warehouse_id, disposition=self.config.query_statement_disposition
+        )
 
     @cached_property
     def cluster_access(self) -> ClusterAccess:

src/databricks/labs/ucx/install.py

@@ -257,6 +257,7 @@ def _prompt_for_new_installation(self) -> WorkspaceConfig:
         recon_tolerance_percent = int(
             self.prompts.question("Reconciliation threshold, in percentage", default="5", valid_number=True)
         )
+
         return WorkspaceConfig(
             inventory_database=inventory_database,
             ucx_catalog=ucx_catalog,

tests/unit/install/test_install.py

@@ -377,6 +377,7 @@ def test_configure_sets_expected_workspace_configuration_values(
         "warehouse_id": "abc",
         "workspace_start_path": "/",
         "num_days_submit_runs_history": 30,
+        'query_statement_disposition': 'INLINE',
         "recon_tolerance_percent": 5,
         "managed_table_external_storage": "CLONE",
     }
@@ -412,6 +413,7 @@ def test_configure_with_default_owner_group(
         "min_workers": 1,
         "max_workers": 10,
         "policy_id": "foo",
+        "query_statement_disposition": "INLINE",
         "renamed_group_prefix": "db-temp-",
         "warehouse_id": "abc",
         "workspace_start_path": "/",
@@ -506,6 +508,7 @@ def test_create_cluster_policy(ws, mock_installation) -> None:
             'renamed_group_prefix': 'db-temp-',
             'warehouse_id': 'abc',
             'workspace_start_path': '/',
+            'query_statement_disposition': 'INLINE',
             'recon_tolerance_percent': 5,
             'managed_table_external_storage': "CLONE",
         },
@@ -1751,6 +1754,7 @@ def test_save_config_ext_hms(ws, mock_installation) -> None:
             'warehouse_id': 'abc',
             'workspace_start_path': '/',
             'num_days_submit_runs_history': 30,
+            'query_statement_disposition': 'INLINE',
             'recon_tolerance_percent': 5,
             'managed_table_external_storage': "CLONE",
         },

tests/unit/test_cli.py

@@ -8,6 +8,7 @@
 import pytest
 import yaml
 from databricks.labs.blueprint.tui import MockPrompts
+from databricks.sdk.service.sql import Disposition
 
 from databricks.labs.ucx.aws.credentials import IamRoleCreation
 from databricks.sdk import AccountClient, WorkspaceClient
@@ -207,7 +208,7 @@ def test_skip_with_schema(ws) -> None:
         byte_limit=None,
         catalog=None,
         schema=None,
-        disposition=None,
+        disposition=Disposition.INLINE,
         format=sql.Format.JSON_ARRAY,
         wait_timeout=None,
     )
@@ -222,7 +223,7 @@ def test_skip_with_table(ws) -> None:
         byte_limit=None,
         catalog=None,
         schema=None,
-        disposition=None,
+        disposition=Disposition.INLINE,
         format=sql.Format.JSON_ARRAY,
         wait_timeout=None,
     )
@@ -237,7 +238,7 @@ def test_skip_with_view(ws) -> None:
         byte_limit=None,
         catalog=None,
         schema=None,
-        disposition=None,
+        disposition=Disposition.INLINE,
         format=sql.Format.JSON_ARRAY,
         wait_timeout=None,
     )
@@ -264,7 +265,7 @@ def test_unskip_with_schema(ws) -> None:
         byte_limit=None,
         catalog=None,
         schema=None,
-        disposition=None,
+        disposition=Disposition.INLINE,
         format=sql.Format.JSON_ARRAY,
         wait_timeout=None,
     )
@@ -279,7 +280,7 @@ def test_unskip_with_table(ws) -> None:
         byte_limit=None,
         catalog=None,
         schema=None,
-        disposition=None,
+        disposition=Disposition.INLINE,
         format=sql.Format.JSON_ARRAY,
         wait_timeout=None,
     )
@@ -294,7 +295,7 @@ def test_unskip_with_view(ws) -> None:
         byte_limit=None,
         catalog=None,
         schema=None,
-        disposition=None,
+        disposition=Disposition.INLINE,
         format=sql.Format.JSON_ARRAY,
         wait_timeout=None,
     )