22from datetime import datetime
33from enum import Enum
44from pathlib import Path
5- from typing import Any , List , Optional , Set , Union
5+ from typing import Any , Callable , List , Optional , Set , Union
66
77from 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