Skip to content

Commit 329bc47

Browse files
achalsadchia
authored andcommitted
fix: Return an empty infra object from sql registry when it doesn't exist (#3022)
* fix: Return an empty infra object from sql registry when it doesn't exist Signed-off-by: Achal Shah <[email protected]> * better Signed-off-by: Achal Shah <[email protected]> * types Signed-off-by: Achal Shah <[email protected]> * fix hasattr Signed-off-by: Achal Shah <[email protected]>
1 parent 1db9e3c commit 329bc47

1 file changed

Lines changed: 39 additions & 15 deletions

File tree

  • sdk/python/feast/infra/registry_stores

‎sdk/python/feast/infra/registry_stores/sql.py‎

Lines changed: 39 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
from datetime import datetime
33
from enum import Enum
44
from pathlib import Path
5-
from typing import Any, List, Optional, Set, Union
5+
from typing import Any, Callable, List, Optional, Set, Union
66

77
from sqlalchemy import ( # type: ignore
88
BigInteger,
@@ -555,7 +555,7 @@ def update_infra(self, infra: Infra, project: str, commit: bool = True):
555555
)
556556

557557
def get_infra(self, project: str, allow_cache: bool = False) -> Infra:
558-
return self._get_object(
558+
infra_object = self._get_object(
559559
managed_infra,
560560
"infra_obj",
561561
project,
@@ -565,6 +565,8 @@ def get_infra(self, project: str, allow_cache: bool = False) -> Infra:
565565
"infra_proto",
566566
None,
567567
)
568+
infra_object = infra_object or InfraProto()
569+
return Infra.from_proto(infra_object)
568570

569571
def apply_user_metadata(
570572
self,
@@ -676,11 +678,18 @@ def commit(self):
676678
pass
677679

678680
def _apply_object(
679-
self, table, project: str, id_field_name, obj, proto_field_name, name=None
681+
self,
682+
table: Table,
683+
project: str,
684+
id_field_name,
685+
obj,
686+
proto_field_name,
687+
name=None,
680688
):
681689
self._maybe_init_project_metadata(project)
682690

683-
name = name or obj.name
691+
name = name or obj.name if hasattr(obj, "name") else None
692+
assert name, f"name needs to be provided for {obj}"
684693
with self.engine.connect() as conn:
685694
update_datetime = datetime.utcnow()
686695
update_time = int(update_datetime.timestamp())
@@ -738,7 +747,14 @@ def _maybe_init_project_metadata(self, project):
738747
conn.execute(insert_stmt)
739748
usage.set_current_project_uuid(new_project_uuid)
740749

741-
def _delete_object(self, table, name, project, id_field_name, not_found_exception):
750+
def _delete_object(
751+
self,
752+
table: Table,
753+
name: str,
754+
project: str,
755+
id_field_name: str,
756+
not_found_exception: Optional[Callable],
757+
):
742758
with self.engine.connect() as conn:
743759
stmt = delete(table).where(
744760
getattr(table.c, id_field_name) == name, table.c.project_id == project
@@ -752,14 +768,14 @@ def _delete_object(self, table, name, project, id_field_name, not_found_exceptio
752768

753769
def _get_object(
754770
self,
755-
table,
756-
name,
757-
project,
758-
proto_class,
759-
python_class,
760-
id_field_name,
761-
proto_field_name,
762-
not_found_exception,
771+
table: Table,
772+
name: str,
773+
project: str,
774+
proto_class: Any,
775+
python_class: Any,
776+
id_field_name: str,
777+
proto_field_name: str,
778+
not_found_exception: Optional[Callable],
763779
):
764780
self._maybe_init_project_metadata(project)
765781

@@ -771,10 +787,18 @@ def _get_object(
771787
if row:
772788
_proto = proto_class.FromString(row[proto_field_name])
773789
return python_class.from_proto(_proto)
774-
raise not_found_exception(name, project)
790+
if not_found_exception:
791+
raise not_found_exception(name, project)
792+
else:
793+
return None
775794

776795
def _list_objects(
777-
self, table, project, proto_class, python_class, proto_field_name
796+
self,
797+
table: Table,
798+
project: str,
799+
proto_class: Any,
800+
python_class: Any,
801+
proto_field_name: str,
778802
):
779803
self._maybe_init_project_metadata(project)
780804
with self.engine.connect() as conn:

0 commit comments

Comments
 (0)