qwak-core 0.4.246__py3-none-any.whl → 0.4.248__py3-none-any.whl

qwak/inner/model_loggers_utils.py

@@ -4,7 +4,6 @@ from typing import Optional
 
 import requests
 
-from _qwak_proto.qwak.builds.builds_orchestrator_service_pb2 import AuthenticationDetail
 from qwak.clients.model_management.client import ModelsManagementClient
 from qwak.exceptions import QwakException
 
@@ -69,33 +68,22 @@ def fetch_build_id() -> Optional[str]:
 def upload_data(
     upload_url: str,
     data: bytes,
-    authentication_details: Optional[AuthenticationDetail],
+    headers: dict,
     content_type: str = "text/plain",
-) -> None:
+):
     """
     Upload data
-    Args:
-        upload_url: the url to upload to.
-        data: the data to upload
-        authentication_details: authentication details for upload data
-        content_type: Uploaded content-type
+    :param upload_url: the url to upload to.
+    :param data: the data to upload
+    :param headers: authentication details for upload data
+    :param content_type: Uploaded content-type
     """
     try:
-        auth = None
-        if (
-            authentication_details.WhichOneof("integration_type")
-            == "jfrog_authentication_detail"
-        ):
-            auth = (
-                authentication_details.jfrog_authentication_detail.username,
-                authentication_details.jfrog_authentication_detail.token,
-            )
-
+        headers["Content-Type"] = content_type
         http_response = requests.put(  # nosec B113
             upload_url,
             data=data,
-            headers={"content-type": content_type},
-            auth=auth,
+            headers=headers,
         )
 
         if http_response.status_code not in [200, 201]:

qwak/model_loggers/artifact_logger.py

@@ -2,6 +2,9 @@ import urllib.request
 from typing import Optional
 
 from _qwak_proto.qwak.builds.build_url_pb2 import BuildVersioningTagsType
+from _qwak_proto.qwak.builds.builds_orchestrator_service_pb2 import (
+    GetBuildVersioningUploadURLResponse,
+)
 from qwak.clients.build_orchestrator.client import BuildOrchestratorClient
 from qwak.clients.file_versioning.client import FileVersioningManagementClient
 from qwak.exceptions import QwakException
@@ -38,7 +41,9 @@ def log_file(
         # Checking if called inside a model - then build id saved as environment variable or stays
         build_id = fetch_build_id()
 
-    upload_url_response = BuildOrchestratorClient().get_build_versioning_upload_url(
+    upload_url_response: (
+        GetBuildVersioningUploadURLResponse
+    ) = BuildOrchestratorClient().get_build_versioning_upload_url(
         build_id=build_id,
         model_id=model_id,
         tag=tag,
@@ -53,7 +58,7 @@ def log_file(
         upload_data(
             upload_url_response.upload_url,
             f.read(),
-            upload_url_response.authentication_detail,
+            upload_url_response.headers,
         )
 
 

qwak/model_loggers/data_logger.py

@@ -12,6 +12,9 @@ from _qwak_proto.qwak.builds.build_pb2 import (
     DataTableDefinition,
 )
 from _qwak_proto.qwak.builds.build_url_pb2 import BuildVersioningTagsType
+from _qwak_proto.qwak.builds.builds_orchestrator_service_pb2 import (
+    GetBuildVersioningUploadURLResponse,
+)
 from qwak.clients.build_orchestrator.client import BuildOrchestratorClient
 from qwak.clients.data_versioning.client import DataVersioningManagementClient
 from qwak.exceptions import QwakException
@@ -63,11 +66,13 @@ def log_data(
         build_id = fetch_build_id()
 
     client = BuildOrchestratorClient()
-    upload_url_response = client.get_build_versioning_upload_url(
-        build_id=build_id,
-        model_id=model_id,
-        tag=str(data_tag),
-        tag_type=BuildVersioningTagsType.DATA_TAG_TYPE,
+    upload_url_response: GetBuildVersioningUploadURLResponse = (
+        client.get_build_versioning_upload_url(
+            build_id=build_id,
+            model_id=model_id,
+            tag=str(data_tag),
+            tag_type=BuildVersioningTagsType.DATA_TAG_TYPE,
+        )
     )
 
     string_buffer = StringIO()
@@ -76,7 +81,7 @@ def log_data(
         upload_url=upload_url_response.upload_url,
         data=gzip.compress(bytes(string_buffer.getvalue(), "utf-8")),
         content_type="text/plain",
-        authentication_details=upload_url_response.authentication_detail,
+        headers=upload_url_response.headers,
     )
 
     dataframe_definition = DataTableDefinition(

{qwak_core-0.4.246.dist-info → qwak_core-0.4.248.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: qwak-core
-Version: 0.4.246
+Version: 0.4.248
 Summary: Qwak Core contains the necessary objects and communication tools for using the Qwak Platform
 License: Apache-2.0
 Keywords: mlops,ml,deployment,serving,model

{qwak_core-0.4.246.dist-info → qwak_core-0.4.248.dist-info}/RECORD

@@ -540,6 +540,12 @@ _qwak_proto/qwak/self_service/user/v1/user_pb2_grpc.py,sha256=1oboBPFxaTEXt9Aw7E
 _qwak_proto/qwak/self_service/user/v1/user_service_pb2.py,sha256=31GMuClpxHWHnFPXKxKlxLby-NofL9yHj4GULLJQ8ng,10389
 _qwak_proto/qwak/self_service/user/v1/user_service_pb2.pyi,sha256=Y4BlejbLGPraMkujz9Du_uC4WwJdCUgWSPv1IEFTYKk,8148
 _qwak_proto/qwak/self_service/user/v1/user_service_pb2_grpc.py,sha256=HPNgP5ejFRwdgF_eVoDqu2-wpG1odK_p5nyZ9iXTIL8,10512
+_qwak_proto/qwak/service_discovery/service_discovery_location_pb2.py,sha256=KnfvXbH36EvzBtG0aWgYKy761JEsRkvgCdMhG5-Zq9Q,3329
+_qwak_proto/qwak/service_discovery/service_discovery_location_pb2.pyi,sha256=33gciRiYabLn_gmN5Ksw5reequ4GkfUbGxFW4gionf4,2406
+_qwak_proto/qwak/service_discovery/service_discovery_location_pb2_grpc.py,sha256=1oboBPFxaTEXt9Aw7EAj8gXHDCNMhZD2VXqocC9l_gk,159
+_qwak_proto/qwak/service_discovery/service_discovery_location_service_pb2.py,sha256=NkkenQmPrFSm9T9bji2vJtD5difVTdFE8pN5-2Z5gmM,3813
+_qwak_proto/qwak/service_discovery/service_discovery_location_service_pb2.pyi,sha256=7lfas7n0SXrqcrjCjivcqPVEV8SOQhW8vEXbhvdVFeg,1357
+_qwak_proto/qwak/service_discovery/service_discovery_location_service_pb2_grpc.py,sha256=LfpjLTdJ8lTJjErhDHgoM9ew-7c_FfXcnSqFEL2plVU,13684
 _qwak_proto/qwak/traffic/v1/traffic_api_pb2.py,sha256=4KmYUxK7Gku8IddNOc7VCOjS89H5qtAv4eOMqAJ4v2o,6909
 _qwak_proto/qwak/traffic/v1/traffic_api_pb2.pyi,sha256=tsh59fnAmLLsKzc5-dj9Jx8MDXJWukoNzJtvGjGQjfI,5186
 _qwak_proto/qwak/traffic/v1/traffic_api_pb2_grpc.py,sha256=46oskdNbo8xUzu8QGNtqtvUL93zYuOsSHCVFY1e8TLM,8228
@@ -576,7 +582,7 @@ _qwak_proto/qwak/workspace/workspace_pb2_grpc.py,sha256=1oboBPFxaTEXt9Aw7EAj8gXH
 _qwak_proto/qwak/workspace/workspace_service_pb2.py,sha256=AB3C9S_AbOD7Nx1Ni4j1rW6PNtYTV1zjiqFQk-goQ74,21429
 _qwak_proto/qwak/workspace/workspace_service_pb2.pyi,sha256=nKKCHwnovZhsy8TSVmdz-Vtl0nviOOoX56HD-41Xo08,13726
 _qwak_proto/qwak/workspace/workspace_service_pb2_grpc.py,sha256=yKGuexxTBza99Ihe0DSTniV2ZSd_AG47inHenqfi890,27193
-qwak/__init__.py,sha256=NI8Z0miVc9RaYzA7QH7AlrhoEjQOmUcaFlF7EcOmxO0,587
+qwak/__init__.py,sha256=5N6mBzO1DFCpFv6gezURZqEuGMsOSAcu57c6E9MIo7g,587
 qwak/automations/__init__.py,sha256=qFZRvCxUUn8gcxkJR0v19ulHW2oJ0x6-Rif7HiheDP4,1522
 qwak/automations/automation_executions.py,sha256=5MeH_epYYWb8NKXgAozwT_jPyyUDednBHG7izloi7RY,3228
 qwak/automations/automations.py,sha256=3yx8e2v0uSKDnXbqyknasyEoQ5vxGni6K40Hbi1_zkk,12599
@@ -630,7 +636,7 @@ qwak/clients/feature_store/__init__.py,sha256=mMCPBHDga6Y7dtJfNoHvfOvCyjNUHrVDX5
 qwak/clients/feature_store/execution_management_client.py,sha256=tKLIML8wsvgsl1hBfx74izFL0rHYdNM69sEstomDfYM,4051
 qwak/clients/feature_store/job_registry_client.py,sha256=7VMtEj7aofKcABrYpldgdxyv8Vkdq_mocjns1O0uCqA,2635
 qwak/clients/feature_store/management_client.py,sha256=UqFx_wVrDx0N7ArMmaLiBqAzpekqZgkmQlC9UHkze_M,20498
-qwak/clients/feature_store/offline_serving_client.py,sha256=0E9LzBcz2NOcxzAhtUZDnoX3LR5maXQBJ8d5vH8rpaY,8473
+qwak/clients/feature_store/offline_serving_client.py,sha256=gz8hqaboPA1Und8leOf1O0dXa9xorHDTU3b7-Ne9YSE,9344
 qwak/clients/feature_store/operator_client.py,sha256=qChgvX8m7A_tffMIaRx22pLSt5YxAbWeHPO-67LrRKE,5811
 qwak/clients/file_versioning/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 qwak/clients/file_versioning/client.py,sha256=ywOy9olB4mekjQQrKUufAbaC-3sWzQnqd3ChcwrAJzc,2505
@@ -644,6 +650,8 @@ qwak/clients/integration_management/openai/__init__.py,sha256=47DEQpj8HBSa-_TImW
 qwak/clients/integration_management/openai/openai_system_secret.py,sha256=y7NT-F-S50ol7CUcnrSK522xZr7k0lD7bYo9VlyZ1j8,2119
 qwak/clients/kube_deployment_captain/__init__.py,sha256=rJUEEy3zNH0aTFyuO_UBexzaUKdjvwU9P2vV1MDj684,41
 qwak/clients/kube_deployment_captain/client.py,sha256=g5nLf4l7A4SFpdfF8Q5KwEincSFiWTHKKEFKztE5QVw,9308
+qwak/clients/location_discovery/__init__.py,sha256=sqGQ75YHFE6nvOcir38fykUUmAa6cFEIze8PJYgYWRc,44
+qwak/clients/location_discovery/client.py,sha256=O-wjMGiSTJ5qGjD0VhofE6Mv7NqvmwCpRnJgzyxKg6I,2867
 qwak/clients/logging_client/__init__.py,sha256=1OCHnigQBYThBwGbxCreYA0BgP0HcuLFzNEWd3Yxh-c,34
 qwak/clients/logging_client/client.py,sha256=IdPa93JjlVm0B_Gdc4lH9LYHoVn5DjT_gIPaZ-W3Hdw,4906
 qwak/clients/model_management/__init__.py,sha256=vjWVP8MjmK4_A70WOgJqa6x24AeLK-ABjGJtogGzw9w,43
@@ -685,8 +693,7 @@ qwak/feature_store/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuF
 qwak/feature_store/_common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 qwak/feature_store/_common/artifact_utils.py,sha256=iIdz6EfFuE_yub6g5KXlOOMt7TA7pEuuLwjlzVSRpO8,1977
 qwak/feature_store/_common/feature_set_utils.py,sha256=dhtPWnwv1nB3h3EA8FAYklx1qGkO_Dj1jecaBGNNpew,8634
-qwak/feature_store/_common/featureset_asterisk_handler.py,sha256=82X-h26hGDMM4dvgcXyNRsvoaGW8Ef6iJTovBLHDE3M,4707
-qwak/feature_store/_common/functions.py,sha256=z033VNAqGZak2oB7pUSS9iJrh9aMRblVGP3iLWCChko,1298
+qwak/feature_store/_common/functions.py,sha256=kSNYJ7dy48NN09HG9asm4ibQh0JaCGcZYsRDjRWlUHE,659
 qwak/feature_store/_common/packaging.py,sha256=eyTyjIB_C6wGHcdU7xR5u7-N2Ld98s4LPU_iZtcbRO0,8122
 qwak/feature_store/_common/source_code_spec.py,sha256=OqlrGVtLvZNtJjycEyC5v5VCIfD12JkYcM2JbnR-x_k,1804
 qwak/feature_store/_common/source_code_spec_factory.py,sha256=zDif0nlWs43T87-k2VCUYyQyR6O3vbCHYCOtVfLUI8g,1760
@@ -748,12 +755,8 @@ qwak/feature_store/feature_sets/transformations/functions/qwak_pandas.py,sha256=
 qwak/feature_store/feature_sets/transformations/functions/schema.py,sha256=kuu8MZ3d2Y9DkmgPZCRgDb0ecsc8isFHQG1lFLPBR3Y,1156
 qwak/feature_store/feature_sets/transformations/transformations.py,sha256=k10Tg64XOfvpL9pQPMEfkuXtMT5EYxZ9-xkTOoD1fpM,15328
 qwak/feature_store/feature_sets/transformations/validations/validations_util.py,sha256=Nr7MyWTLxe9J5VwQYCPRpaQdsv5VmWsWLwPC3C-cwOk,3179
-qwak/feature_store/offline/__init__.py,sha256=on-3Twm3-THTeYXEuVHMZVd2sgVGVChtzvqzSJyO_4I,173
+qwak/feature_store/offline/__init__.py,sha256=v6cf_Ne02UxZpsxPCWna0cZBLyF2czN_iL4eO3XIPOI,96
 qwak/feature_store/offline/_offline_serving_validations.py,sha256=a3Ga9-IPTS0TPf_ghjuQ9qumBCMKFaQBkWxdoawUr-s,1216
-qwak/feature_store/offline/_query_engine.py,sha256=LBVejISK9daYYWUYclXuMifezfsuzF4JgBhp2YlHZf0,738
-qwak/feature_store/offline/athena/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-qwak/feature_store/offline/athena/athena_query_engine.py,sha256=Z8Y46nf9-_TI7AXpS5aO57UTDQFnFvARislMQcwQb8A,5182
-qwak/feature_store/offline/client.py,sha256=dGtCnqlhuL4y5T7AtbTj1Pzu78-8JR_O1wnrnKmozHE,28651
 qwak/feature_store/offline/client_v2.py,sha256=MKIlVGohW7aMmCpBUKE9wP9X6gtU_pLVfOgnBn2Bp98,14393
 qwak/feature_store/offline/feature_set_features.py,sha256=Eh_rPz_m12cAQvrhIF-xtA-ZGVKy1HlgKJdjNvsmPoo,739
 qwak/feature_store/online/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -828,7 +831,7 @@ qwak/inner/di_configuration/containers.py,sha256=1Vvs_p5laNcduDWH47dZpJJa2xu3v5P
 qwak/inner/di_configuration/session.py,sha256=Pwqapcu-0dsbb5dDn7ZGewVk4bYRN9sWk6ity72IDqY,452
 qwak/inner/instance_template/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 qwak/inner/instance_template/verify_template_id.py,sha256=-c0Fh89XAfphKCpkurYaD-YyX3M3NkkpSLH4QsNPyko,1916
-qwak/inner/model_loggers_utils.py,sha256=egsXQMfBTqsHZxNJq8cTGrFCS_Gxobni6xEfNTeupBE,2933
+qwak/inner/model_loggers_utils.py,sha256=irVEuNYlMefINR8N5UiD9IcO8fzdcuc8jBbBzNzKC7M,2441
 qwak/inner/provider.py,sha256=3evQnyp0v0enpvGGDyaZziusO4BGi-U9j1sno8DAHo4,70
 qwak/inner/runtime_di/__init__.py,sha256=ylaHwo2Hlq3pmeziA0Ip87DMnkcuBffKKi0tTH7gYVI,281
 qwak/inner/runtime_di/containers.py,sha256=-AREM-IUAd9JhrhQablqixxpj4L9Tk1iQkK1fA5YbrA,545
@@ -958,8 +961,8 @@ qwak/model/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,
 qwak/model/utils/extract_wrapped_function.py,sha256=uIle1zL8vbmeS3PGAuaNFLIUQAsvpuzk3LlH-Teba94,320
 qwak/model/utils/feature_utils.py,sha256=jB4pfmhqN2_B4THR7YADlCWbymU3cm_8u2RzTegrv_w,2389
 qwak/model_loggers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-qwak/model_loggers/artifact_logger.py,sha256=CZHTe-WO6Ov4WRYR6yVw5BuNSI0NFi-TUS117c4JQ1Y,2938
-qwak/model_loggers/data_logger.py,sha256=4eSfcsfnG_EvIRRX6Zj1tI-SfSOYUWCnR9naW7AzVT4,5472
+qwak/model_loggers/artifact_logger.py,sha256=E5A9MaGD4bDAWlj7cuc5f9Pzet9_p5jgkGuoZiVn1B0,3090
+qwak/model_loggers/data_logger.py,sha256=aTRnRwKmKUFZq9SRg4sY-b1yGDY2XB-pExlQ0WlcH0k,5629
 qwak/model_loggers/model_logger.py,sha256=WxKLEpfzDaRn3kanrI0kj4zqqKQrwT2sin1IS1WNEA8,852
 qwak/qwak_client/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 qwak/qwak_client/batch_jobs/execution.py,sha256=L0CjS6BuEsPGJbQ4_yV75FsIx-jvdUg8DirYdH6ddFQ,1757
@@ -1035,11 +1038,12 @@ qwak_services_mock/mocks/integration_management_service.py,sha256=j5ew7epeG0-n-q
 qwak_services_mock/mocks/internal_build_orchestrator_service.py,sha256=deupu0eGAepqwh8bbebCY3i1BavrXv_mTpRdktaENPk,1046
 qwak_services_mock/mocks/job_registry_service_api.py,sha256=BEU8eH9yE0_Jrc792rdp_M2KtrXO97uB3hklmQZsUtI,2696
 qwak_services_mock/mocks/kube_captain_service_api.py,sha256=raZEWAP3aOk-DeM9qVHMGl97AK-8_smiCYCo0Euj0zg,1583
+qwak_services_mock/mocks/location_discovery_service_api.py,sha256=sEy-FEZg1rQnGqmNiewH4azkxbj-riaoiQ4h8XjeA2o,3679
 qwak_services_mock/mocks/logging_service.py,sha256=rkOud7CFAvb0xfP6WXf3kP541eXdYJK5Lv_Q_2I1Fek,7381
 qwak_services_mock/mocks/model_management_service.py,sha256=H3bMTr7VLrGjZjIrjgrjBk3kNDy1UjQXOeUyL5eAL_o,4153
 qwak_services_mock/mocks/project_manager_service.py,sha256=WzcYzN3O04kSpvuTwX0QIF0HCh_ggsV9-gMel_vQVuc,3090
 qwak_services_mock/mocks/prompt_manager_service.py,sha256=tEhVUVwkvRy68r7MNX9V3V8ZBZ1ynqzr6v1zREMm0_s,9895
-qwak_services_mock/mocks/qwak_mocks.py,sha256=vfn1f608tvTcYjo8hioiXR8_Zxg7paCkPN7abukOJdw,5969
+qwak_services_mock/mocks/qwak_mocks.py,sha256=abayoevvX2gzUftDef_SYlL2t44P0r_7IaVh4GM0Hx0,6142
 qwak_services_mock/mocks/secret_service.py,sha256=BwNyhNSTmxHDACzrr7NHjZ0YchX3wHWuJkl1ioi0Aig,1406
 qwak_services_mock/mocks/self_service_user_service.py,sha256=UK22V-wAc0zHScXCE0QKDO0TdbrHf_uk7tEEIM47MpI,1205
 qwak_services_mock/mocks/system_secret_service.py,sha256=S8MfGGEex7bs0AvnAf8SbPOzE-YDfCvePT1TxXWzHRQ,1896
@@ -1049,9 +1053,9 @@ qwak_services_mock/mocks/utils/exception_handlers.py,sha256=Fhi9cx_qNytDnl4BtZmE
 qwak_services_mock/mocks/vector_serving_api.py,sha256=7ZSmDJVdw7_ne7BxLLrmy8eZkzVCeZ508gfL51wpKeY,5722
 qwak_services_mock/mocks/vectors_management_api.py,sha256=ePh8l3MlDGzs7E4NuPZzF5xK-bC9AdMgDaNZrq3FkHc,3594
 qwak_services_mock/mocks/workspace_manager_service_mock.py,sha256=O9ZSwln4T4kHVkR_usXnDQtarTeNfffSMON0P6wbT4g,7740
-qwak_services_mock/services_mock.py,sha256=CABZ2teXoSJ-X7CD3iRo5ht-w3jhW-Cvpb47APUHXEA,19089
+qwak_services_mock/services_mock.py,sha256=zXtHcX8a_acz7ynxuCBxxVpHpde7aAGjIn6Uw52LY1s,19593
 qwak_services_mock/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 qwak_services_mock/utils/service_utils.py,sha256=ZlB0CnB1J6oBn6_m7fQO2U8tKoboHdUa6ljjkRMYNXU,265
-qwak_core-0.4.246.dist-info/METADATA,sha256=HOEV8pRohcWvk7huZnH3R9NpzwIT5pLppjfhfr2bc2I,2150
-qwak_core-0.4.246.dist-info/WHEEL,sha256=IYZQI976HJqqOpQU6PHkJ8fb3tMNBFjg-Cn-pwAbaFM,88
-qwak_core-0.4.246.dist-info/RECORD,,
+qwak_core-0.4.248.dist-info/METADATA,sha256=V-53JLVMprm3wnTnnPBpn2DWt5N4PpJ33JR0muYz4Sc,2150
+qwak_core-0.4.248.dist-info/WHEEL,sha256=IYZQI976HJqqOpQU6PHkJ8fb3tMNBFjg-Cn-pwAbaFM,88
+qwak_core-0.4.248.dist-info/RECORD,,

qwak_services_mock/mocks/location_discovery_service_api.py

@@ -0,0 +1,104 @@
+from typing import Dict, Optional
+
+import grpc
+from _qwak_proto.qwak.service_discovery.service_discovery_location_pb2 import (
+    ServiceLocationDescriptor,
+)
+from _qwak_proto.qwak.service_discovery.service_discovery_location_service_pb2 import (
+    GetServingUrlRequestResponse,
+)
+from _qwak_proto.qwak.service_discovery.service_discovery_location_service_pb2_grpc import (
+    LocationDiscoveryServiceServicer,
+)
+
+
+class LocationDiscoveryServiceApiMock(LocationDiscoveryServiceServicer):
+    """
+    Mock implementation of the LocationDiscoveryService for testing SDK behavior.
+    Allows setting mock responses and optional error codes for each endpoint.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self._responses: Dict[str, Optional[ServiceLocationDescriptor]] = {}
+        self._error_codes: Dict[str, grpc.StatusCode] = {}
+
+    def _set_mock(
+        self,
+        key: str,
+        response: Optional[ServiceLocationDescriptor],
+        error_code: grpc.StatusCode = grpc.StatusCode.NOT_FOUND,
+    ):
+        self._responses[key] = response
+        self._error_codes[key] = error_code
+
+    def _handle(
+        self, key: str, context: grpc.ServicerContext
+    ) -> GetServingUrlRequestResponse:
+        response = self._responses.get(key)
+        if response:
+            return GetServingUrlRequestResponse(location=response)
+        context.set_code(self._error_codes.get(key, grpc.StatusCode.NOT_FOUND))
+        context.set_details(f"No mock response set for {key}")
+        return GetServingUrlRequestResponse()
+
+    # Setters
+    def set_get_offline_serving_url_response(
+        self,
+        response: Optional[ServiceLocationDescriptor],
+        error_code: grpc.StatusCode = grpc.StatusCode.NOT_FOUND,
+    ):
+        self._set_mock("offline", response, error_code)
+
+    def set_get_distribution_manager_url_response(
+        self,
+        response: Optional[ServiceLocationDescriptor],
+        error_code: grpc.StatusCode = grpc.StatusCode.NOT_FOUND,
+    ):
+        self._set_mock("distribution", response, error_code)
+
+    def set_get_analytics_engine_url_response(
+        self,
+        response: Optional[ServiceLocationDescriptor],
+        error_code: grpc.StatusCode = grpc.StatusCode.NOT_FOUND,
+    ):
+        self._set_mock("analytics", response, error_code)
+
+    def set_get_metrics_gateway_url_response(
+        self,
+        response: Optional[ServiceLocationDescriptor],
+        error_code: grpc.StatusCode = grpc.StatusCode.NOT_FOUND,
+    ):
+        self._set_mock("metrics", response, error_code)
+
+    def set_get_features_operator_url_response(
+        self,
+        response: Optional[ServiceLocationDescriptor],
+        error_code: grpc.StatusCode = grpc.StatusCode.NOT_FOUND,
+    ):
+        self._set_mock("features", response, error_code)
+
+    def set_get_hosting_gateway_url_response(
+        self,
+        response: Optional[ServiceLocationDescriptor],
+        error_code: grpc.StatusCode = grpc.StatusCode.NOT_FOUND,
+    ):
+        self._set_mock("hosting", response, error_code)
+
+    def GetOfflineServingUrl(self, request, context):
+        return self._handle("offline", context)
+
+    def GetDistributionManagerUrl(self, request, context):
+        return self._handle("distribution", context)
+
+    def GetAnalyticsEngineUrl(self, request, context):
+        return self._handle("analytics", context)
+
+    def GetMetricsGatewayUrl(self, request, context):
+        return self._handle("metrics", context)
+
+    def GetFeaturesOperatorUrl(self, request, context):
+        return self._handle("features", context)
+
+    def GetHostingGatewayUrl(self, request, context):
+        return self._handle("hosting", context)

qwak_services_mock/mocks/qwak_mocks.py

@@ -60,6 +60,9 @@ from qwak_services_mock.mocks.internal_build_orchestrator_service import (
 )
 from qwak_services_mock.mocks.job_registry_service_api import JobRegistryServiceApiMock
 from qwak_services_mock.mocks.kube_captain_service_api import KubeCaptainServiceApiMock
+from qwak_services_mock.mocks.location_discovery_service_api import (
+    LocationDiscoveryServiceApiMock,
+)
 from qwak_services_mock.mocks.logging_service import LoggingServiceApiMock
 from qwak_services_mock.mocks.model_management_service import (
     ModelsManagementServiceMock,
@@ -121,3 +124,4 @@ class QwakMocks:
     system_secret_service: SystemSecretServiceMock
     integration_management_service: IntegrationManagementServiceMock
     prompt_manager_service: PromptManagerServiceMock
+    location_discovery_service: LocationDiscoveryServiceApiMock

qwak_services_mock/services_mock.py

@@ -106,6 +106,9 @@ from _qwak_proto.qwak.secret_service.secret_service_pb2_grpc import (
 from _qwak_proto.qwak.self_service.user.v1.user_service_pb2_grpc import (
     add_UserServiceServicer_to_server,
 )
+from _qwak_proto.qwak.service_discovery.service_discovery_location_service_pb2_grpc import (
+    add_LocationDiscoveryServiceServicer_to_server,
+)
 from _qwak_proto.qwak.vectors.v1.collection.collection_service_pb2_grpc import (
     add_VectorCollectionServiceServicer_to_server,
 )
@@ -175,6 +178,9 @@ from qwak_services_mock.mocks.internal_build_orchestrator_service import (
 )
 from qwak_services_mock.mocks.job_registry_service_api import JobRegistryServiceApiMock
 from qwak_services_mock.mocks.kube_captain_service_api import KubeCaptainServiceApiMock
+from qwak_services_mock.mocks.location_discovery_service_api import (
+    LocationDiscoveryServiceApiMock,
+)
 from qwak_services_mock.mocks.logging_service import LoggingServiceApiMock
 from qwak_services_mock.mocks.model_management_service import (
     ModelsManagementServiceMock,
@@ -235,6 +241,7 @@ def qwak_container():
         file_versioning,
         instance_template,
         kube_deployment_captain,
+        location_discovery,
         logging_client,
         model_management,
         project,
@@ -280,6 +287,7 @@ def qwak_container():
             integration_manager_client,
             system_secret_client,
             prompt_manager_client,
+            location_discovery,
         ]
     )
 
@@ -492,6 +500,11 @@ def attach_servicers(free_port, server):
                 PromptManagerServiceMock,
                 add_PromptManagerServiceServicer_to_server,
             ),
+            (
+                "location_discovery_service",
+                LocationDiscoveryServiceApiMock,
+                add_LocationDiscoveryServiceServicer_to_server,
+            ),
             ("port", free_port, None),
         ],
     )

qwak/feature_store/_common/featureset_asterisk_handler.py

@@ -1,115 +0,0 @@
-import dataclasses
-import itertools
-import re
-from typing import Callable, Dict, List
-
-from qwak.feature_store.offline import OfflineClient
-from qwak.model.schema_entities import BaseFeature, Entity, FeatureStoreInput
-
-
-def is_asterisk(feature: BaseFeature) -> bool:
-    """
-    Checks whether the given Feature is a FeatureStoreInput with an '*'
-    Args:
-        feature: input feature
-    Return
-    """
-    p = "^[^.]+.\\*$"
-    # if this is not a FeatureStoreInput, no need to check
-    if isinstance(feature, FeatureStoreInput):
-        return bool(re.match(p, feature.name))
-    return False
-
-
-def inflate_feature(feature: BaseFeature, offline_fs: OfflineClient) -> List:
-    """
-    Inflates a Feature into a list. if the feature contains an '*', a list
-        a feature for each feature in the associate FeatureSet is returned. else -
-        a list containing the original feature is returned.
-    Args:
-        feature: input feature
-        offline_fs: OfflineFeatureStore to query
-    Return:
-        if an '*' exists, returns a List of FeatureStoreInuput for each feature in the feature set
-        else returns a list containing the feature.
-
-    """
-    if is_asterisk(feature):
-        featureset_name = feature.name.lower().split(".")[0]
-
-        p = f"^{featureset_name}+.*$"
-        feature_names = [
-            feature.lower()
-            for feature in offline_fs.get_columns_by_feature_set(featureset_name)
-            if re.match(p, feature.lower())
-        ]
-        return [dataclasses.replace(feature, name=name) for name in feature_names]
-
-    return [feature]
-
-
-def unpack_asterisk_features(
-    features: List[BaseFeature],
-    feature_store_generator: Callable[[], OfflineClient] = lambda: OfflineClient(),
-) -> List[BaseFeature]:
-    """
-    Handles features with an '*'.
-    If a feature of type FeatureStoreInput that matches '<featureset_name>.*' exists,
-    transforms this feature into the list of features actually present in the FeatureSet.
-    Other features remain unchanged.
-    the original order is maintained, and the internal order (inside a FeatureStoreInput
-     that has an '*') is inferred from the order of columns returned from
-     OfflineFeatureStore.get_columns_by_feature_set(...)
-     Args:
-         features: List of features
-         feature_store_generator: A function that generates an instance of OfflineFeatureStore
-                                  If no feature with an '*' is found, the function is not called.
-
-    Returns:
-        the unpacked features
-    """
-    asterisk_found = any([is_asterisk(f) for f in features])
-    if asterisk_found:
-        # at least 1 feature contains an asterisk, map each feature to
-        # a list of features (features with no '*' are mapped to a list of size 1), then flatten it.
-        offline_feature_store = feature_store_generator()
-        inflated_features = [
-            inflate_feature(feature, offline_feature_store) for feature in features
-        ]
-        return list(itertools.chain.from_iterable(inflated_features))
-
-    return features
-
-
-def unpack_asterisk_features_from_key_mapping(
-    key_to_features: Dict[str, List[str]],
-    feature_store_generator: Callable[[], OfflineClient] = lambda: OfflineClient(),
-) -> dict:
-    """
-    Handles features with an '*'.
-    If a feature name that matches '<featureset_name>.*' exists in one of the lists,
-    transforms this feature into the list of feature names actually present in the FeatureSet.
-    Other features remain unchanged.
-    the original order is maintained, and the internal order (inside a feature
-     that has an '*') is inferred from the order of columns returned from
-     OfflineFeatureStore.get_columns_by_feature_set(...)
-     Args:
-         key_to_features: dictionary of entity keys to requested features (same as in the offline store)
-             >>> key_to_features = {'uuid': ['user_purchases.*',
-             >>>                             'user_purchases.avg_purchase_amount']}
-         feature_store_generator: A function that generates an instance of OfflineFeatureStore
-                                  If no feature with an '*' is found, the function is not called.
-    Returns:
-        the unpacked features
-    """
-    inflated_key_to_features = {}
-
-    for entity, features in key_to_features.items():
-        feature_inputs = [
-            FeatureStoreInput(name=feature, entity=Entity(name=entity, type=str))
-            for feature in features
-        ]
-        unpacked = unpack_asterisk_features(feature_inputs, feature_store_generator)
-        feature_list = [feature.name for feature in unpacked]
-        inflated_key_to_features[entity] = feature_list
-    return inflated_key_to_features

qwak/feature_store/offline/_query_engine.py

@@ -1,32 +0,0 @@
-import uuid
-from abc import ABC, abstractmethod
-
-import pandas as pd
-
-
-class BaseQueryEngine(ABC):
-    @abstractmethod
-    def upload_table(self, df: pd.DataFrame):
-        pass
-
-    @abstractmethod
-    def cleanup(self):
-        pass
-
-    @abstractmethod
-    def run_query(self, query: str):
-        pass
-
-    @abstractmethod
-    def read_pandas_from_query(self, query: str, parse_dates=None):
-        pass
-
-    @staticmethod
-    @abstractmethod
-    def get_quotes():
-        pass
-
-    class JoinTableSpec:
-        def __init__(self, join_tables_db_name: str, quotes: str):
-            self.table_name = str(uuid.uuid4())
-            self.join_table_full_path = f"{quotes}{join_tables_db_name}{quotes}.{quotes}{self.table_name}{quotes}"