Skip to content

Commit cc0df2e

Browse files
authored
Fix oss client create function (unitycatalog#1012)
**PR Checklist** - [ ] A description of the changes is added to the description of this PR. - [ ] If there is a related issue, make sure it is linked to this PR. - [ ] If you've fixed a bug or added code that should be tested, add tests! - [ ] If you've added or modified a feature, documentation in `docs` is updated **Description of changes** get_function_async returns None if no function existing, fix the corresponding error in create_function_async <!-- Please state what you've changed and how it might affect the users. --> --------- Signed-off-by: serena-ruan <serena.rxy@gmail.com>
1 parent ad0d465 commit cc0df2e

2 files changed

Lines changed: 11 additions & 12 deletions

File tree

ai/core/src/unitycatalog/ai/core/client.py

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -464,14 +464,14 @@ async def create_function_async(
464464
)
465465

466466
try:
467-
await self.get_function_async(str(function_name), timeout=timeout)
468-
if replace:
469-
_logger.info(f"Function {function_name} already exists, replacing it.")
470-
await self.delete_function_async(str(function_name), timeout=timeout)
471-
else:
472-
raise ValueError(
473-
f"Function {function_name} already exists. Set replace=True to overwrite it."
474-
)
467+
if await self.get_function_async(str(function_name), timeout=timeout):
468+
if replace:
469+
_logger.info(f"Function {function_name} already exists, replacing it.")
470+
await self.delete_function_async(str(function_name), timeout=timeout)
471+
else:
472+
raise ValueError(
473+
f"Function {function_name} already exists. Set replace=True to overwrite it."
474+
)
475475
except ServiceException:
476476
pass
477477

@@ -739,7 +739,7 @@ def wrapper(x: int, y:int) -> int:
739739

740740
async def get_function_async(
741741
self, function_name: str, timeout: Optional[float] = None, **kwargs: Any
742-
) -> FunctionInfo:
742+
) -> Optional[FunctionInfo]:
743743
"""
744744
Retrieve a function by its full name asynchronously.
745745

ai/core/tests/core/oss/test_oss_client.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@
4040
FunctionsApi,
4141
SchemasApi,
4242
)
43-
from unitycatalog.client.exceptions import NotFoundException, ServiceException
43+
from unitycatalog.client.exceptions import NotFoundException
4444

4545
SCHEMA = "uc_test"
4646
_logger = logging.getLogger(__name__)
@@ -763,8 +763,7 @@ async def test_create_function_invalid_data_type(uc_client):
763763
async def test_get_nonexistent_function(uc_client):
764764
function_name = f"{CATALOG}.{SCHEMA}.nonexistent_function"
765765

766-
with pytest.raises(ServiceException, match="(500)"):
767-
uc_client.get_function(function_name=function_name)
766+
assert uc_client.get_function(function_name=function_name) is None
768767

769768

770769
@pytest.mark.asyncio

0 commit comments

Comments
 (0)