-
Notifications
You must be signed in to change notification settings - Fork 96
Added export CLI functionality for assessment results #2553
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 35 commits
4f0c6dc
d288e0a
5aee6b6
afc714c
ad01abb
f26c522
522d43e
0c7083d
66dbab1
32adee0
77a88cd
9db2d70
33f1f11
898eeef
2c8e92b
d6243ad
99bf9a1
ee446f1
e57cbfe
1240c63
5cdeb3c
c63904e
8b71454
a8a325e
6ebc90c
f47f125
cd67ee4
3e8d342
92fd1c0
e5e2e10
a71b086
21785e4
077a03f
14fa408
8d607c8
20c19f2
cf6fa3c
6a33aca
7b798aa
66c281f
a00ecf8
92937ce
b4d58d5
4599dfa
dce6f86
d61a5ce
8c7dd67
d94d581
007891d
598c596
82146db
f4b7170
90a8189
5eb417c
9069f3d
44b4ac4
bcfd6e1
f90d2c1
6e3ddf9
26799fc
b2e3c85
9fb3f49
3d567b4
ae5832d
204ab4a
0a6e72f
93d496e
ee20112
2f62a0f
8fee9d8
21da7cc
2a09a8f
d1dd0c5
7cba9b0
04a4956
15c5536
2611b11
555e83a
d744465
04918be
0a03def
be955ea
788c273
1d12391
d3a0e39
4d2e9a6
b163801
188bc24
cb89541
571fc8b
77786a2
84830fd
2975950
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change | ||||
---|---|---|---|---|---|---|
|
@@ -10,8 +10,10 @@ so that you'll be able to [scope the migration](docs/assessment.md) and execute | |||||
The [README notebook](#readme-notebook), which can be found in the installation folder contains further instructions and explanations of the different ucx workflows & dashboards. | ||||||
Once the migration is scoped, you can start with the [table migration process](#Table-Migration). | ||||||
|
||||||
|
||||||
More workflows, like notebook code migration are coming in future releases. | ||||||
|
||||||
|
||||||
UCX also provides a number of command line utilities accessible via `databricks labs ucx`. | ||||||
|
||||||
For questions, troubleshooting or bug fixes, please see our [troubleshooting guide](docs/troubleshooting.md) or submit [an issue](https://github.com/databrickslabs/ucx/issues). | ||||||
|
@@ -88,6 +90,7 @@ See [contributing instructions](CONTRIBUTING.md) to help improve this project. | |||||
* [`open-remote-config` command](#open-remote-config-command) | ||||||
* [`installations` command](#installations-command) | ||||||
* [`report-account-compatibility` command](#report-account-compatibility-command) | ||||||
* [`export-assessment` command](#export-assessment-command) | ||||||
* [Metastore related commands](#metastore-related-commands) | ||||||
* [`show-all-metastores` command](#show-all-metastores-command) | ||||||
* [`assign-metastore` command](#assign-metastore-command) | ||||||
|
@@ -1109,6 +1112,29 @@ databricks labs ucx report-account-compatibility --profile labs-azure-account | |||||
|
||||||
[[back to top](#databricks-labs-ucx)] | ||||||
|
||||||
## `export-assessment` command | ||||||
|
||||||
```commandline | ||||||
databricks labs ucx export-assessment | ||||||
``` | ||||||
The `export-assessment` command exports UCX results to a specified location and prompts for details on the path and report type. If not specified, it defaults to exporting the main results to the current directory as `export_to_zipped_csv.zip`. | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I would prefer that the "specified location" becomes the output zip, instead of the hardcoded
Suggested change
|
||||||
|
||||||
### Prompts | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. why do we need this heading in ToC? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. will remove. |
||||||
|
||||||
- **Choose a path to save the UCX Assessment results:** | ||||||
- **Description:** Specify the path where the results should be saved. If not provided, results will be saved in the current directory. | ||||||
|
||||||
- **Choose which assessment results to export:** | ||||||
- **Description:** Select the type of results to export. Options include: | ||||||
- `azure` | ||||||
- `estimates` | ||||||
- `interactive` | ||||||
- `main` | ||||||
- **Default:** `main` | ||||||
|
||||||
[[back to top](#databricks-labs-ucx)] | ||||||
|
||||||
|
||||||
# Metastore related commands | ||||||
|
||||||
These commands are used to assign a Unity Catalog metastore to a workspace. The metastore assignment is a pre-requisite | ||||||
|
Original file line number | Diff line number | Diff line change | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
@@ -0,0 +1,51 @@ | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
import logging | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
from pathlib import Path | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
from databricks.labs.blueprint.tui import Prompts | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
from databricks.labs.ucx.config import WorkspaceConfig | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
from databricks.labs.lsql.backends import SqlBackend | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
from databricks.labs.lsql.dashboards import DashboardMetadata | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
logger = logging.getLogger(__name__) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
class AssessmentExporter: | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
# File and Path Constants | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
_EXPORT_FILE_NAME = "ucx_assessment_results.zip" | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
def __init__(self, sql_backend: SqlBackend, config: WorkspaceConfig): | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
self._sql_backend = sql_backend | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
self._config = config | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
def export_results(self, prompts: Prompts): | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
"""Main method to export results to CSV files inside a ZIP archive.""" | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
project_root = Path(__file__).resolve().parents[3] | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
queries_path_root = project_root / f"labs/ucx/queries/assessment" | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
valid_queries = {subdir.name for subdir in queries_path_root.iterdir() if subdir.is_dir()} | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. this makes no sense. dashboard is a folder, not a query. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Could you please clarify if you want me to rename the variable name from valid_queries to something more descriptive, like valid_directories? |
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
export_path = Path( | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
prompts.question( | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
"Choose a path to save the UCX Assessment results", | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
default=Path.cwd().as_posix(), | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
validate=lambda p_: Path(p_).exists(), | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
query_choice = prompts.question( | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. there's There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The code defaults to running queries available in Main. But we also wanted to extend the capability to run other queries in folders 'estimates', 'interactive', 'azure', 'main' under labs/ucx/queries/assessment There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. @JCZuurmond suggested adding this capability. PFA for previous comments suggesting this need. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. do we really need this? can we simplify? |
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
"Choose which assessment results to export", | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
default="main", | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
validate=lambda q: q in valid_queries, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
queries_path = queries_path_root / query_choice | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
display_name = "UCX Assessment Results" | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
dashboard_metadata = DashboardMetadata(display_name) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
assessment_results = ( | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
DashboardMetadata.from_path(queries_path) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
.replace_database(catalog="hive_metastore", database=self._config.inventory_database) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
logger.info("Exporting assessment results....") | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
results_path = assessment_results.export_to_zipped_csv(self._sql_backend, export_path) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
please simplify this entire method. |
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
logger.info(f"Results exported to {results_path}") |
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -19,6 +19,8 @@ | |
from databricks.labs.ucx.install import AccountInstaller | ||
from databricks.labs.ucx.source_code.linters.files import LocalCodeLinter | ||
|
||
from databricks.labs.ucx.assessment.export import AssessmentExporter | ||
|
||
ucx = App(__file__) | ||
logger = get_logger(__file__) | ||
|
||
|
@@ -566,7 +568,6 @@ def join_collection(a: AccountClient, workspace_ids: str): | |
w_ids = [int(_.strip()) for _ in workspace_ids.split(",") if _] | ||
account_installer.join_collection(w_ids) | ||
|
||
|
||
@ucx.command | ||
def upload( | ||
file: Path | str, | ||
|
@@ -637,5 +638,13 @@ def lint_local_code( | |
linter.lint(prompts, None if path is None else Path(path)) | ||
|
||
|
||
@ucx.command | ||
def export_assessment(w: WorkspaceClient, prompts: Prompts): | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Looks like a duplicate of the above command? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Removed now. |
||
"""Export the UCX assessment queries to a zip file.""" | ||
ctx: WorkspaceContext = WorkspaceContext(w) | ||
exporter : AssessmentExporter = ctx.assessment_exporter | ||
exporter.export_results(prompts) | ||
rportilla-databricks marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
||
|
||
if __name__ == "__main__": | ||
ucx() |
Original file line number | Diff line number | Diff line change | ||||||
---|---|---|---|---|---|---|---|---|
|
@@ -112,6 +112,111 @@ | |||||||
f'--parent_run_id=' + dbutils.widgets.get('parent_run_id')) | ||||||||
""" | ||||||||
|
||||||||
EXPORT_TO_EXCEL_NOTEBOOK = """ | ||||||||
nfx marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||||
# Databricks notebook source | ||||||||
# MAGIC %md | ||||||||
# MAGIC ##### Exporter of UCX assessment results | ||||||||
# MAGIC ##### Instructions: | ||||||||
# MAGIC 1. Execute using an all-purpose cluster with Databricks Runtime 14 or higher. | ||||||||
# MAGIC 1. Hit **Run all** button and wait for completion. | ||||||||
# MAGIC 1. Go to the bottom of the notebook and click the Download UCX Results button. | ||||||||
# MAGIC | ||||||||
# MAGIC ##### Important: | ||||||||
# MAGIC Please note that this is only meant to serve as example code. | ||||||||
# MAGIC | ||||||||
# MAGIC Example code developed by **Databricks Shared Technical Services team**. | ||||||||
|
||||||||
# COMMAND ---------- | ||||||||
|
||||||||
# DBTITLE 1,Installing Packages | ||||||||
# MAGIC %pip install {remote_wheel} -qqq | ||||||||
# MAGIC %pip install xlsxwriter -qqq | ||||||||
# MAGIC dbutils.library.restartPython() | ||||||||
|
||||||||
# COMMAND ---------- | ||||||||
|
||||||||
# DBTITLE 1,Libraries Import and Setting UCX | ||||||||
import os | ||||||||
nfx marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||||
import shutil | ||||||||
import logging, threading | ||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
|
||||||||
from functools import partial | ||||||||
from threading import Lock | ||||||||
|
||||||||
# third party Libraries imports | ||||||||
import pandas as pd | ||||||||
import xlsxwriter | ||||||||
|
||||||||
# Databricks imports | ||||||||
from databricks.sdk.config import with_user_agent_extra | ||||||||
from databricks.labs.blueprint.logger import install_logger | ||||||||
from databricks.labs.blueprint.parallel import Threads | ||||||||
from databricks.labs.lsql.dashboards import QueryTile | ||||||||
from databricks.labs.ucx.contexts.workflow_task import RuntimeContext | ||||||||
|
||||||||
# ctx | ||||||||
install_logger() | ||||||||
with_user_agent_extra("cmd", "export-assessment") | ||||||||
named_parameters = dict(config="/Workspace{config_file}") | ||||||||
ctx = RuntimeContext(named_parameters) | ||||||||
lock = Lock() | ||||||||
|
||||||||
# COMMAND ---------- | ||||||||
|
||||||||
# DBTITLE 1,Assessment Export | ||||||||
|
||||||||
# File and Path Constants | ||||||||
FILE_NAME = "assessment_results.xlsx" | ||||||||
TMP_PATH = f'/tmp' + ctx.installation.install_folder() + '/excel-export' | ||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. use |
||||||||
DOWNLOAD_PATH = "/dbfs/FileStore/excel-export" | ||||||||
|
||||||||
def _cleanup() -> None: | ||||||||
\"\"\"Move the temporary results file to the download path and clean up the temp directory.\"\"\" | ||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. use There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We had doc strings here and the guidance from JCZuurmond was to use double quotes (which we need to escape inside the contents). Getting conflicting feedback here. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. use |
||||||||
shutil.move( | ||||||||
os.path.join(TMP_PATH, FILE_NAME), | ||||||||
os.path.join(DOWNLOAD_PATH, FILE_NAME), | ||||||||
) | ||||||||
shutil.rmtree(TMP_PATH) | ||||||||
|
||||||||
def _prepare_directories() -> None: | ||||||||
'''Ensure that the necessary directories exist.''' | ||||||||
os.makedirs(TMP_PATH, exist_ok=True) | ||||||||
os.makedirs(DOWNLOAD_PATH, exist_ok=True) | ||||||||
|
||||||||
def _to_excel(tile: QueryTile, writer: ...) -> None: | ||||||||
'''Execute a SQL query and write the result to an Excel sheet.''' | ||||||||
sdf = spark.sql(tile.content) | ||||||||
df = sdf.toPandas() | ||||||||
with lock: | ||||||||
df.to_excel(writer, sheet_name=tile.metadata.title, index=False) | ||||||||
|
||||||||
def _render_export() -> None: | ||||||||
'''Render an HTML link for downloading the results.''' | ||||||||
html_content = f''' | ||||||||
<style>@font-face{{font-family:'DM Sans';src:url(https://cdn.bfldr.com/9AYANS2F/at/p9qfs3vgsvnp5c7txz583vgs/dm-sans-regular.ttf?auto=webp&format=ttf) format('truetype');font-weight:400;font-style:normal}}body{{font-family:'DM Sans',Arial,sans-serif}}.export-container{{text-align:center;margin-top:20px}}.export-container h2{{color:#1B3139;font-size:24px;margin-bottom:20px}}.export-container a{{display:inline-block;padding:12px 25px;background-color:#1B3139;color:#fff;text-decoration:none;border-radius:4px;font-size:18px;font-weight:500;transition:background-color 0.3s ease,transform 0.3s ease}}.export-container a:hover{{background-color:#FF3621;transform:translateY(-2px)}}</style><div class="export-container"><h2>Export Results</h2><a href='{workspace_host}files/excel-export/assessment_results.xlsx?o={workspace_id}' target='_blank' download>Download UCX Results </a></div> | ||||||||
''' | ||||||||
displayHTML(html_content) | ||||||||
|
||||||||
def export_results() -> None: | ||||||||
'''Main method to export results to an Excel file.''' | ||||||||
_prepare_directories() | ||||||||
try: | ||||||||
target = TMP_PATH + '/assessment_results.xlsx' | ||||||||
with pd.ExcelWriter(target, engine="xlsxwriter") as writer: | ||||||||
tasks = [] | ||||||||
for query_tile in ctx.assessment_exporter.queries(): | ||||||||
tasks.append(partial(_to_excel, query_tile, writer)) | ||||||||
Threads.strict("exporting", tasks) | ||||||||
_cleanup() | ||||||||
_render_export() | ||||||||
except Exception as e: | ||||||||
print(f"Error exporting results ", e) | ||||||||
|
||||||||
# COMMAND ---------- | ||||||||
|
||||||||
# DBTITLE 1,Data Export | ||||||||
export_results() | ||||||||
""" | ||||||||
|
||||||||
|
||||||||
class DeployedWorkflows: | ||||||||
def __init__(self, ws: WorkspaceClient, install_state: InstallState): | ||||||||
|
@@ -792,6 +897,15 @@ def _create_debug(self, remote_wheels: list[str]): | |||||||
).encode("utf8") | ||||||||
self._installation.upload('DEBUG.py', content) | ||||||||
|
||||||||
def _create_export(self, remote_wheels: list[str]): | ||||||||
content = EXPORT_TO_EXCEL_NOTEBOOK.format( | ||||||||
remote_wheel=remote_wheels, | ||||||||
config_file=self._config_file, | ||||||||
workspace_host=self._ws.config.host, | ||||||||
workspace_id=self._ws.get_workspace_id(), | ||||||||
).encode("utf8") | ||||||||
self._installation.upload('EXPORT_TO_EXCEL_NOTEBOOK.py', content) | ||||||||
rportilla-databricks marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||||
|
||||||||
|
||||||||
class MaxedStreamHandler(logging.StreamHandler): | ||||||||
|
||||||||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,39 @@ | ||
from databricks.labs.ucx.config import WorkspaceConfig | ||
from databricks.labs.ucx.assessment.export import AssessmentExport | ||
from databricks.labs.lsql.backends import MockBackend | ||
from databricks.labs.blueprint.tui import MockPrompts | ||
from databricks.labs.lsql.core import Row | ||
|
||
|
||
def test_export(tmp_path): | ||
"""Test the export_results method of the AssessmentExport class.""" | ||
query = { | ||
"SELECT\n one\nFROM ucx.external_locations": [ | ||
Row(location="s3://bucket1/folder1", table_count=1), | ||
Row(location="abfss://container1@storage1.dfs.core.windows.net/folder1", table_count=1), | ||
Row(location="gcp://folder1", table_count=2), | ||
] | ||
} | ||
|
||
# Setup workspace configuration | ||
config = WorkspaceConfig(inventory_database="ucx") | ||
|
||
# Prepare temporary paths and files | ||
export_path = tmp_path / "export" | ||
export_path.mkdir(parents=True, exist_ok=True) | ||
|
||
# Mock backend and prompts | ||
mock_backend = MockBackend(rows=query) | ||
mock_prompts = MockPrompts( | ||
{ | ||
"Choose a path to save the UCX Assessment results": export_path.as_posix(), | ||
"Choose which assessment results to export": "main", | ||
} | ||
) | ||
|
||
# Execute export process | ||
export = AssessmentExporter(mock_backend, config) | ||
exported = export.export_results(mock_prompts) | ||
|
||
# Assertions | ||
assert exported == export_path / "export_to_zipped_csv.zip" |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Remove redundant newlines