Skip to content

Commit bc863ae

Browse files
Adding support for serving endpoints (#990)
## Changes Assessment did not crawled permissions for serving endpoints, this PR aims to fix it ### Tests - [X] added integration tests
1 parent a628eeb commit bc863ae

File tree

4 files changed

+61
-0
lines changed

4 files changed

+61
-0
lines changed

src/databricks/labs/ucx/mixins/fixtures.py

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
from databricks.sdk.errors import NotFound, ResourceConflict
1919
from databricks.sdk.retries import retried
2020
from databricks.sdk.service import compute, iam, jobs, pipelines, sql, workspace
21+
from databricks.sdk.service._internal import Wait
2122
from databricks.sdk.service.catalog import (
2223
AzureServicePrincipal,
2324
CatalogInfo,
@@ -28,6 +29,12 @@
2829
TableInfo,
2930
TableType,
3031
)
32+
from databricks.sdk.service.serving import (
33+
EndpointCoreConfigInput,
34+
ServedModelInput,
35+
ServedModelInputWorkloadSize,
36+
ServingEndpointDetailed,
37+
)
3138
from databricks.sdk.service.sql import (
3239
CreateWarehouseRequestWarehouseType,
3340
GetResponse,
@@ -1095,3 +1102,27 @@ def remove(storage_credential: StorageCredentialInfo):
10951102
ws.storage_credentials.delete(storage_credential.name, force=True)
10961103

10971104
yield from factory("storage_credential_from_spn", create, remove)
1105+
1106+
1107+
@pytest.fixture
1108+
def make_serving_endpoint(ws, make_random, make_model):
1109+
def create() -> Wait[ServingEndpointDetailed]:
1110+
endpoint_name = make_random(4)
1111+
model = make_model()
1112+
endpoint = ws.serving_endpoints.create(
1113+
endpoint_name,
1114+
EndpointCoreConfigInput(
1115+
served_models=[
1116+
ServedModelInput(model.name, "1", ServedModelInputWorkloadSize.SMALL, scale_to_zero_enabled=True)
1117+
]
1118+
),
1119+
)
1120+
return endpoint
1121+
1122+
def remove(endpoint_name: str):
1123+
try:
1124+
ws.serving_endpoints.delete(endpoint_name)
1125+
except RuntimeError as e:
1126+
logger.info(f"Can't remove endpoint {e}")
1127+
1128+
yield from factory("Serving endpoint", create, remove)

src/databricks/labs/ucx/workspace_access/manager.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ def factory(
5151
generic.Listing(ws.warehouses.list, "id", "sql/warehouses"),
5252
generic.Listing(ws.jobs.list, "job_id", "jobs"),
5353
generic.Listing(ws.pipelines.list_pipelines, "pipeline_id", "pipelines"),
54+
generic.Listing(ws.serving_endpoints.list, "id", "serving-endpoints"),
5455
generic.Listing(generic.experiments_listing(ws), "experiment_id", "experiments"),
5556
generic.Listing(generic.models_listing(ws, num_threads), "id", "registered-models"),
5657
generic.Listing(generic.tokens_and_passwords, "object_id", "authorization"),

tests/integration/workspace_access/test_generic.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -411,3 +411,31 @@ def test_verify_permissions(ws, make_group, make_job, make_job_permissions):
411411
result = task()
412412

413413
assert result
414+
415+
416+
@retried(on=[NotFound], timeout=timedelta(minutes=3))
417+
def test_endpoints(
418+
ws, make_group, make_serving_endpoint, make_serving_endpoint_permissions
419+
): # pylint: disable=invalid-name
420+
group_a = make_group()
421+
group_b = make_group()
422+
endpoint = make_serving_endpoint()
423+
make_serving_endpoint_permissions(
424+
object_id=endpoint.response.id,
425+
permission_level=PermissionLevel.CAN_MANAGE,
426+
group_name=group_a.display_name,
427+
)
428+
429+
generic_permissions = GenericPermissionsSupport(ws, [Listing(ws.serving_endpoints.list, "id", "serving-endpoints")])
430+
before = generic_permissions.load_as_dict("serving-endpoints", endpoint.response.id)
431+
assert before[group_a.display_name] == PermissionLevel.CAN_MANAGE
432+
433+
apply_tasks(
434+
generic_permissions,
435+
[
436+
MigratedGroup.partial_info(group_a, group_b),
437+
],
438+
)
439+
440+
after = generic_permissions.load_as_dict("serving-endpoints", endpoint.response.id)
441+
assert after[group_b.display_name] == PermissionLevel.CAN_MANAGE

tests/unit/workspace_access/test_manager.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -214,6 +214,7 @@ def test_factory(mocker):
214214
"secrets",
215215
"entitlements",
216216
"roles",
217+
'serving-endpoints',
217218
"ANY FILE",
218219
"FUNCTION",
219220
"ANONYMOUS FUNCTION",

0 commit comments

Comments
 (0)