From b74e6d02c301fb437dedd0ddc55cf26d77a340cb Mon Sep 17 00:00:00 2001 From: Achal Shah Date: Fri, 5 Aug 2022 13:19:45 -0700 Subject: [PATCH 1/4] fix: Return an empty infra object from sql registry when it doesn't exist Signed-off-by: Achal Shah --- sdk/python/feast/infra/registry_stores/sql.py | 26 +++++++++++-------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/sdk/python/feast/infra/registry_stores/sql.py b/sdk/python/feast/infra/registry_stores/sql.py index 9c6b47a714..37eb70d561 100644 --- a/sdk/python/feast/infra/registry_stores/sql.py +++ b/sdk/python/feast/infra/registry_stores/sql.py @@ -560,16 +560,19 @@ def update_infra(self, infra: Infra, project: str, commit: bool = True): ) def get_infra(self, project: str, allow_cache: bool = False) -> Infra: - return self._get_object( - managed_infra, - "infra_obj", - project, - InfraProto, - Infra, - "infra_name", - "infra_proto", - None, - ) + try: + return self._get_object( + managed_infra, + "infra_obj", + project, + InfraProto, + Infra, + "infra_name", + "infra_proto", + None, + ) + except Exception: + return Infra.from_proto(InfraProto()) def apply_user_metadata( self, @@ -687,7 +690,8 @@ def _apply_object( ): self._maybe_init_project_metadata(project) - name = name or obj.name + name = name or obj.name if hasattr(obj, name) else None + assert name, f"name needs to be provided for {obj}" with self.engine.connect() as conn: update_datetime = datetime.utcnow() update_time = int(update_datetime.timestamp()) From 1d777497480de97505f84bf045f6a780a2b633dc Mon Sep 17 00:00:00 2001 From: Achal Shah Date: Fri, 5 Aug 2022 13:39:00 -0700 Subject: [PATCH 2/4] better Signed-off-by: Achal Shah --- sdk/python/feast/infra/registry_stores/sql.py | 30 ++++++++++--------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/sdk/python/feast/infra/registry_stores/sql.py b/sdk/python/feast/infra/registry_stores/sql.py index 37eb70d561..b0a502a8e3 100644 --- a/sdk/python/feast/infra/registry_stores/sql.py +++ b/sdk/python/feast/infra/registry_stores/sql.py @@ -560,19 +560,18 @@ def update_infra(self, infra: Infra, project: str, commit: bool = True): ) def get_infra(self, project: str, allow_cache: bool = False) -> Infra: - try: - return self._get_object( - managed_infra, - "infra_obj", - project, - InfraProto, - Infra, - "infra_name", - "infra_proto", - None, - ) - except Exception: - return Infra.from_proto(InfraProto()) + infra_object = self._get_object( + managed_infra, + "infra_obj", + project, + InfraProto, + Infra, + "infra_name", + "infra_proto", + None, + ) + infra_object = infra_object or InfraProto() + return Infra.from_proto(infra_object) def apply_user_metadata( self, @@ -786,7 +785,10 @@ def _get_object( if row: _proto = proto_class.FromString(row[proto_field_name]) return python_class.from_proto(_proto) - raise not_found_exception(name, project) + if not_found_exception: + raise not_found_exception(name, project) + else: + return None def _list_objects( self, table, project, proto_class, python_class, proto_field_name From 7e85fba8ce3aed9f8e420d4fb0568b948e6a35a6 Mon Sep 17 00:00:00 2001 From: Achal Shah Date: Fri, 5 Aug 2022 14:19:20 -0700 Subject: [PATCH 3/4] types Signed-off-by: Achal Shah --- sdk/python/feast/infra/registry_stores/sql.py | 37 +++++++++++++------ 1 file changed, 25 insertions(+), 12 deletions(-) diff --git a/sdk/python/feast/infra/registry_stores/sql.py b/sdk/python/feast/infra/registry_stores/sql.py index b0a502a8e3..af437b56d9 100644 --- a/sdk/python/feast/infra/registry_stores/sql.py +++ b/sdk/python/feast/infra/registry_stores/sql.py @@ -2,7 +2,7 @@ from datetime import datetime from enum import Enum from pathlib import Path -from typing import Any, List, Optional, Set, Union +from typing import Any, Callable, List, Optional, Set, Union from sqlalchemy import ( # type: ignore BigInteger, @@ -685,7 +685,13 @@ def commit(self): pass def _apply_object( - self, table, project: str, id_field_name, obj, proto_field_name, name=None + self, + table: Table, + project: str, + id_field_name, + obj, + proto_field_name, + name=None, ): self._maybe_init_project_metadata(project) @@ -752,7 +758,9 @@ def _maybe_init_project_metadata(self, project): conn.execute(insert_stmt) usage.set_current_project_uuid(new_project_uuid) - def _delete_object(self, table, name, project, id_field_name, not_found_exception): + def _delete_object( + self, table: Table, name: str, project: str, id_field_name: str, not_found_exception: Optional[Callable] + ): with self.engine.connect() as conn: stmt = delete(table).where( getattr(table.c, id_field_name) == name, table.c.project_id == project @@ -766,14 +774,14 @@ def _delete_object(self, table, name, project, id_field_name, not_found_exceptio def _get_object( self, - table, - name, - project, - proto_class, - python_class, - id_field_name, - proto_field_name, - not_found_exception, + table: Table, + name: str, + project: str, + proto_class: Any, + python_class: Any, + id_field_name: str, + proto_field_name: str, + not_found_exception: Optional[Callable], ): self._maybe_init_project_metadata(project) @@ -791,7 +799,12 @@ def _get_object( return None def _list_objects( - self, table, project, proto_class, python_class, proto_field_name + self, + table: Table, + project: str, + proto_class: Any, + python_class: Any, + proto_field_name: str, ): self._maybe_init_project_metadata(project) with self.engine.connect() as conn: From d9111ef8ae8baf59c071710628baf5dbc6dc284f Mon Sep 17 00:00:00 2001 From: Achal Shah Date: Fri, 5 Aug 2022 14:43:09 -0700 Subject: [PATCH 4/4] fix hasattr Signed-off-by: Achal Shah --- sdk/python/feast/infra/registry_stores/sql.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/sdk/python/feast/infra/registry_stores/sql.py b/sdk/python/feast/infra/registry_stores/sql.py index af437b56d9..e4f08f5413 100644 --- a/sdk/python/feast/infra/registry_stores/sql.py +++ b/sdk/python/feast/infra/registry_stores/sql.py @@ -695,7 +695,7 @@ def _apply_object( ): self._maybe_init_project_metadata(project) - name = name or obj.name if hasattr(obj, name) else None + name = name or obj.name if hasattr(obj, "name") else None assert name, f"name needs to be provided for {obj}" with self.engine.connect() as conn: update_datetime = datetime.utcnow() @@ -759,7 +759,12 @@ def _maybe_init_project_metadata(self, project): usage.set_current_project_uuid(new_project_uuid) def _delete_object( - self, table: Table, name: str, project: str, id_field_name: str, not_found_exception: Optional[Callable] + self, + table: Table, + name: str, + project: str, + id_field_name: str, + not_found_exception: Optional[Callable], ): with self.engine.connect() as conn: stmt = delete(table).where(