Skip to content

Commit fab0e87

Browse files
sundarshankar89sundarshankar89
authored
Merge branch 'feature/switch-installer-integration' into feature/llm-transpile
2 parents c54e68f + c49c5b3 commit fab0e87

File tree

2 files changed

+154
-16
lines changed

2 files changed

+154
-16
lines changed

‎tests/unit/deployment/test_configurator.py‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,12 @@
99
SchemaInfo,
1010
VolumeInfo,
1111
)
12+
from databricks.sdk.service.serving import (
13+
EndpointCoreConfigSummary,
14+
FoundationModel,
15+
ServedEntitySpec,
16+
ServingEndpoint,
17+
)
1218
from databricks.sdk.service.sql import EndpointInfo, EndpointInfoWarehouseType, GetWarehouseResponse, State
1319

1420
from databricks.labs.lakebridge.deployment.configurator import ResourceConfigurator
@@ -314,3 +320,63 @@ def test_prompt_for_warehouse_setup_new(ws):
314320
catalog_operations = create_autospec(CatalogOperations)
315321
configurator = ResourceConfigurator(ws, prompts, catalog_operations)
316322
assert configurator.prompt_for_warehouse_setup("Test") == "new_w_id"
323+
324+
325+
def test_prompt_for_foundation_model_default_choice(ws):
326+
ws.serving_endpoints.list.return_value = [
327+
ServingEndpoint(
328+
name="databricks-claude-sonnet-4-5",
329+
config=EndpointCoreConfigSummary(
330+
served_entities=[
331+
ServedEntitySpec(
332+
foundation_model=FoundationModel(name="claude-sonnet-4.5"),
333+
)
334+
]
335+
),
336+
),
337+
ServingEndpoint(
338+
name="databricks-gpt-4",
339+
config=EndpointCoreConfigSummary(
340+
served_entities=[
341+
ServedEntitySpec(
342+
foundation_model=FoundationModel(name="gpt-4"),
343+
)
344+
]
345+
),
346+
),
347+
]
348+
prompts = MockPrompts({r"Select a Foundation Model serving endpoint:": "0"})
349+
catalog_operations = create_autospec(CatalogOperations)
350+
configurator = ResourceConfigurator(ws, prompts, catalog_operations)
351+
result = configurator.prompt_for_foundation_model_choice()
352+
assert result == "databricks-claude-sonnet-4-5"
353+
354+
355+
def test_prompt_for_foundation_model_non_default_choice(ws):
356+
ws.serving_endpoints.list.return_value = [
357+
ServingEndpoint(
358+
name="databricks-claude-sonnet-4-5",
359+
config=EndpointCoreConfigSummary(
360+
served_entities=[
361+
ServedEntitySpec(
362+
foundation_model=FoundationModel(name="claude-sonnet-4.5"),
363+
)
364+
]
365+
),
366+
),
367+
ServingEndpoint(
368+
name="databricks-gpt-4",
369+
config=EndpointCoreConfigSummary(
370+
served_entities=[
371+
ServedEntitySpec(
372+
foundation_model=FoundationModel(name="gpt-4"),
373+
)
374+
]
375+
),
376+
),
377+
]
378+
prompts = MockPrompts({r"Select a Foundation Model serving endpoint:": "1"})
379+
catalog_operations = create_autospec(CatalogOperations)
380+
configurator = ResourceConfigurator(ws, prompts, catalog_operations)
381+
result = configurator.prompt_for_foundation_model_choice()
382+
assert result == "databricks-gpt-4"

‎tests/unit/test_install.py‎

Lines changed: 88 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from databricks.sdk.service import iam
1010
from databricks.labs.blueprint.tui import MockPrompts
1111
from databricks.labs.blueprint.wheels import ProductInfo, WheelsV2
12+
1213
from databricks.labs.lakebridge.config import (
1314
DatabaseConfig,
1415
LSPConfigOptionV1,
@@ -44,10 +45,10 @@ def ws() -> WorkspaceClient:
4445

4546

4647
SET_IT_LATER = ["Set it later"]
47-
ALL_INSTALLED_DIALECTS_NO_LATER = sorted(["tsql", "snowflake"])
48+
ALL_INSTALLED_DIALECTS_NO_LATER = sorted(["tsql", "snowflake", "mssql"])
4849
ALL_INSTALLED_DIALECTS = SET_IT_LATER + ALL_INSTALLED_DIALECTS_NO_LATER
49-
TRANSPILERS_FOR_SNOWFLAKE_NO_LATER = sorted(["Remorph Community Transpiler", "Morpheus"])
50-
TRANSPILERS_FOR_SNOWFLAKE = SET_IT_LATER + TRANSPILERS_FOR_SNOWFLAKE_NO_LATER
50+
TRANSPILERS_LIST_NO_LATER = sorted(["Bladebridge", "Morpheus", "Remorph Community Transpiler", "Switch"])
51+
TRANSPILERS_LIST = SET_IT_LATER + TRANSPILERS_LIST_NO_LATER
5152
PATH_TO_TRANSPILER_CONFIG = "/some/path/to/config.yml"
5253

5354

@@ -69,7 +70,7 @@ def _all_installed_dialects(self):
6970
return ALL_INSTALLED_DIALECTS_NO_LATER
7071

7172
def _transpilers_with_dialect(self, dialect):
72-
return TRANSPILERS_FOR_SNOWFLAKE_NO_LATER
73+
return TRANSPILERS_LIST_NO_LATER
7374

7475
def _transpiler_config_path(self, transpiler):
7576
return PATH_TO_TRANSPILER_CONFIG
@@ -180,7 +181,7 @@ def test_workspace_installer_run_install_called_with_generated_config(
180181
{
181182
r"Do you want to override the existing installation?": "no",
182183
r"Select the source dialect": str(ALL_INSTALLED_DIALECTS.index("snowflake")),
183-
r"Select the transpiler": str(TRANSPILERS_FOR_SNOWFLAKE.index("Morpheus")),
184+
r"Select the transpiler": str(TRANSPILERS_LIST.index("Morpheus")),
184185
r"Enter input SQL path.*": "/tmp/queries/snow",
185186
r"Enter output directory.*": "/tmp/queries/databricks",
186187
r"Enter error file path.*": "/tmp/queries/errors.log",
@@ -231,7 +232,7 @@ def test_configure_transpile_no_existing_installation(
231232
{
232233
r"Do you want to override the existing installation?": "no",
233234
r"Select the source dialect": str(ALL_INSTALLED_DIALECTS.index("snowflake")),
234-
r"Select the transpiler": str(TRANSPILERS_FOR_SNOWFLAKE.index("Morpheus")),
235+
r"Select the transpiler": str(TRANSPILERS_LIST.index("Morpheus")),
235236
r"Enter input SQL path.*": "/tmp/queries/snow",
236237
r"Enter output directory.*": "/tmp/queries/databricks",
237238
r"Enter error file path.*": "/tmp/queries/errors.log",
@@ -345,7 +346,7 @@ def test_configure_transpile_installation_config_error_continue_install(
345346
{
346347
r"Do you want to override the existing installation?": "yes",
347348
r"Select the source dialect": str(ALL_INSTALLED_DIALECTS.index("snowflake")),
348-
r"Select the transpiler": str(TRANSPILERS_FOR_SNOWFLAKE.index("Morpheus")),
349+
r"Select the transpiler": str(TRANSPILERS_LIST.index("Morpheus")),
349350
r"Enter input SQL path.*": "/tmp/queries/snow",
350351
r"Enter output directory.*": "/tmp/queries/databricks",
351352
r"Enter error file path.*": "/tmp/queries/errors.log",
@@ -423,7 +424,7 @@ def test_configure_transpile_installation_with_no_validation(ws, ws_installer):
423424
prompts = MockPrompts(
424425
{
425426
r"Select the source dialect": ALL_INSTALLED_DIALECTS.index("snowflake"),
426-
r"Select the transpiler": TRANSPILERS_FOR_SNOWFLAKE.index("Morpheus"),
427+
r"Select the transpiler": TRANSPILERS_LIST.index("Morpheus"),
427428
r"Enter input SQL path.*": "/tmp/queries/snow",
428429
r"Enter output directory.*": "/tmp/queries/databricks",
429430
r"Enter error file path.*": "/tmp/queries/errors.log",
@@ -488,7 +489,7 @@ def test_configure_transpile_installation_with_validation_and_warehouse_id_from_
488489
prompts = MockPrompts(
489490
{
490491
r"Select the source dialect": str(ALL_INSTALLED_DIALECTS.index("snowflake")),
491-
r"Select the transpiler": str(TRANSPILERS_FOR_SNOWFLAKE.index("Morpheus")),
492+
r"Select the transpiler": str(TRANSPILERS_LIST.index("Morpheus")),
492493
r"Enter input SQL path.*": "/tmp/queries/snow",
493494
r"Enter output directory.*": "/tmp/queries/databricks",
494495
r"Enter error file path.*": "/tmp/queries/errors.log",
@@ -786,7 +787,7 @@ def test_configure_all_override_installation(
786787
{
787788
r"Do you want to override the existing installation?": "yes",
788789
r"Select the source dialect": str(ALL_INSTALLED_DIALECTS.index("snowflake")),
789-
r"Select the transpiler": str(TRANSPILERS_FOR_SNOWFLAKE.index("Morpheus")),
790+
r"Select the transpiler": str(TRANSPILERS_LIST.index("Morpheus")),
790791
r"Enter input SQL path.*": "/tmp/queries/snow",
791792
r"Enter output directory.*": "/tmp/queries/databricks",
792793
r"Enter error file path.*": "/tmp/queries/errors.log",
@@ -962,7 +963,7 @@ def test_runs_upgrades_on_more_recent_version(
962963
{
963964
r"Do you want to override the existing installation?": "yes",
964965
r"Select the source dialect": str(ALL_INSTALLED_DIALECTS.index("snowflake")),
965-
r"Select the transpiler": str(TRANSPILERS_FOR_SNOWFLAKE.index("Morpheus")),
966+
r"Select the transpiler": str(TRANSPILERS_LIST.index("Morpheus")),
966967
r"Enter input SQL path.*": "/tmp/queries/snow",
967968
r"Enter output directory.*": "/tmp/queries/databricks",
968969
r"Enter error file.*": "/tmp/queries/errors.log",
@@ -1020,7 +1021,7 @@ def test_runs_and_stores_confirm_config_option(
10201021
prompts = MockPrompts(
10211022
{
10221023
r"Select the source dialect": str(ALL_INSTALLED_DIALECTS.index("snowflake")),
1023-
r"Select the transpiler": str(TRANSPILERS_FOR_SNOWFLAKE.index("Remorph Community Transpiler")),
1024+
r"Select the transpiler": str(TRANSPILERS_LIST.index("Remorph Community Transpiler")),
10241025
r"Do you want to use the experimental Databricks generator ?": "yes",
10251026
r"Enter input SQL path.*": "/tmp/queries/snow",
10261027
r"Enter output directory.*": "/tmp/queries/databricks",
@@ -1114,7 +1115,7 @@ def test_runs_and_stores_force_config_option(
11141115
prompts = MockPrompts(
11151116
{
11161117
r"Select the source dialect": str(ALL_INSTALLED_DIALECTS.index("snowflake")),
1117-
r"Select the transpiler": str(TRANSPILERS_FOR_SNOWFLAKE.index("Remorph Community Transpiler")),
1118+
r"Select the transpiler": str(TRANSPILERS_LIST.index("Remorph Community Transpiler")),
11181119
r"Enter input SQL path.*": "/tmp/queries/snow",
11191120
r"Enter output directory.*": "/tmp/queries/databricks",
11201121
r"Enter error file path.*": "/tmp/queries/errors.log",
@@ -1194,7 +1195,7 @@ def test_runs_and_stores_question_config_option(
11941195
prompts = MockPrompts(
11951196
{
11961197
r"Select the source dialect": str(ALL_INSTALLED_DIALECTS.index("snowflake")),
1197-
r"Select the transpiler": str(TRANSPILERS_FOR_SNOWFLAKE.index("Remorph Community Transpiler")),
1198+
r"Select the transpiler": str(TRANSPILERS_LIST.index("Remorph Community Transpiler")),
11981199
r"Max number of heaps:": "1254",
11991200
r"Enter input SQL path.*": "/tmp/queries/snow",
12001201
r"Enter output directory.*": "/tmp/queries/databricks",
@@ -1276,7 +1277,7 @@ def test_runs_and_stores_choice_config_option(
12761277
prompts = MockPrompts(
12771278
{
12781279
r"Select the source dialect": str(ALL_INSTALLED_DIALECTS.index("snowflake")),
1279-
r"Select the transpiler": str(TRANSPILERS_FOR_SNOWFLAKE.index("Remorph Community Transpiler")),
1280+
r"Select the transpiler": str(TRANSPILERS_LIST.index("Remorph Community Transpiler")),
12801281
r"Select currency:": "2",
12811282
r"Enter input SQL path.*": "/tmp/queries/snow",
12821283
r"Enter output directory.*": "/tmp/queries/databricks",
@@ -1477,7 +1478,7 @@ def test_installer_upgrade_configure_if_changed(
14771478
prompts=MockPrompts(
14781479
{
14791480
r"Do you want to override the existing installation?": "yes",
1480-
r"Select the source dialect": "2",
1481+
r"Select the source dialect": str(ALL_INSTALLED_DIALECTS.index("tsql")),
14811482
r"Select the transpiler": "1",
14821483
r"Enter .*": "/tmp/updated",
14831484
r"Would you like to validate.*": "no",
@@ -1647,3 +1648,74 @@ def test_transpiler_installers_llm_flag(
16471648
)
16481649
assert installer.configure("transpile").include_switch == should_include_switch
16491650
assert installer.configure("all").include_switch == should_include_switch
1651+
1652+
1653+
def test_workspace_installer_run_install_called_with_generated_config_switch(
1654+
ws_installer: Callable[..., WorkspaceInstaller],
1655+
ws: WorkspaceClient,
1656+
) -> None:
1657+
prompts = MockPrompts(
1658+
{
1659+
r"Do you want to override the existing installation?": "no",
1660+
r"Select the source dialect": str(ALL_INSTALLED_DIALECTS.index("mssql")),
1661+
r"Select the transpiler": str(TRANSPILERS_LIST.index("Switch")),
1662+
r"Enter input SQL path.*": "/tmp/queries/mssql",
1663+
r"Enter output directory.*": "/tmp/queries/databricks",
1664+
r"Enter error file path.*": "/tmp/queries/errors.log",
1665+
r"Would you like to validate.*": "no",
1666+
r"Open .* in the browser?": "no",
1667+
r"Enter catalog name.*": "lakebridge",
1668+
r"Enter schema name.*": "switch",
1669+
r"Enter volume name.*": "switch_volume",
1670+
r"Select a Foundation Model serving endpoint.*": "0",
1671+
}
1672+
)
1673+
1674+
installation = MockInstallation()
1675+
resource_configurator = create_autospec(ResourceConfigurator)
1676+
resource_configurator.prompt_for_catalog_setup.return_value = "lakebridge"
1677+
resource_configurator.prompt_for_schema_setup.return_value = "switch"
1678+
resource_configurator.prompt_for_volume_setup.return_value = "switch_volume"
1679+
resource_configurator.prompt_for_foundation_model_choice.return_value = "databricks-claude-sonnet-4-5"
1680+
ctx = ApplicationContext(ws)
1681+
1682+
ctx.replace(
1683+
prompts=prompts,
1684+
installation=installation,
1685+
resource_configurator=resource_configurator,
1686+
workspace_installation=create_autospec(WorkspaceInstallation),
1687+
)
1688+
1689+
workspace_installer = ws_installer(
1690+
ctx.workspace_client,
1691+
ctx.prompts,
1692+
ctx.installation,
1693+
ctx.install_state,
1694+
ctx.product_info,
1695+
ctx.resource_configurator,
1696+
ctx.workspace_installation,
1697+
is_interactive=True,
1698+
include_llm=True,
1699+
)
1700+
workspace_installer.run("transpile")
1701+
installation.assert_file_written(
1702+
"config.yml",
1703+
{
1704+
"catalog_name": "remorph",
1705+
"transpiler_config_path": PATH_TO_TRANSPILER_CONFIG,
1706+
"source_dialect": "mssql",
1707+
"input_source": "/tmp/queries/mssql",
1708+
"output_folder": "/tmp/queries/databricks",
1709+
"error_file_path": "/tmp/queries/errors.log",
1710+
"schema_name": "transpiler",
1711+
"skip_validation": True,
1712+
"transpiler_options": {
1713+
"catalog": "lakebridge",
1714+
"foundation_model": "databricks-claude-sonnet-4-5",
1715+
"schema": "switch",
1716+
"transpiler_name": "Switch",
1717+
"volume": "switch_volume",
1718+
},
1719+
"version": 3,
1720+
},
1721+
)

0 commit comments

Comments
 (0)