1313class Check :
1414 """Represents a check in the "checks" DB table"""
1515
16+ @staticmethod
17+ def _convert_to_dict_if_needed (result : Record | None , as_dict : bool ) -> Record | dict | None :
18+ if as_dict and result :
19+ return dict (result )
20+ return result
21+
22+ @staticmethod
23+ def _convert_list_to_dict_if_needed (
24+ results : list [Record ], as_dict : bool
25+ ) -> list [Record ] | list [dict ]:
26+ if as_dict :
27+ return [dict (result ) for result in results ]
28+ return results
29+
1630 @classmethod
17- async def get_by_id (cls , check_id : int , with_deleted : bool = False ) -> Record | None :
31+ async def get_by_id (
32+ cls , check_id : int , with_deleted : bool = False , as_dict : bool = False
33+ ) -> Record | dict | None :
1834 pool = await context .pool ()
1935 async with pool .acquire () as connection :
2036 q = """
@@ -24,12 +40,13 @@ async def get_by_id(cls, check_id: int, with_deleted: bool = False) -> Record |
2440 """
2541 if not with_deleted :
2642 q += " AND catalog.deleted = FALSE"
27- return await connection .fetchrow (q , check_id )
43+ result = await connection .fetchrow (q , check_id )
44+ return cls ._convert_to_dict_if_needed (result , as_dict )
2845
2946 @classmethod
3047 async def get_by_resource_id (
31- cls , resource_id : str , with_deleted : bool = False
32- ) -> Record | None :
48+ cls , resource_id : str , with_deleted : bool = False , as_dict : bool = False
49+ ) -> Record | dict | None :
3350 pool = await context .pool ()
3451 async with pool .acquire () as connection :
3552 q = """
@@ -39,23 +56,25 @@ async def get_by_resource_id(
3956 """
4057 if not with_deleted :
4158 q += " AND catalog.deleted = FALSE"
42- return await connection .fetchrow (q , resource_id )
59+ result = await connection .fetchrow (q , resource_id )
60+ return cls ._convert_to_dict_if_needed (result , as_dict )
4361
4462 @classmethod
45- async def get_by_url (cls , url : str ) -> list [Record ]:
63+ async def get_by_url (cls , url : str , as_dict : bool = False ) -> list [Record ] | list [ dict ]:
4664 pool = await context .pool ()
4765 async with pool .acquire () as connection :
4866 q = """
4967 SELECT * FROM checks
5068 WHERE url = $1
5169 ORDER BY created_at DESC
5270 """
53- return await connection .fetch (q , url )
71+ results = await connection .fetch (q , url )
72+ return cls ._convert_list_to_dict_if_needed (results , as_dict )
5473
5574 @classmethod
5675 async def get_latest (
57- cls , url : str | None = None , resource_id : str | None = None
58- ) -> Record | None :
76+ cls , url : str | None = None , resource_id : str | None = None , as_dict : bool = False
77+ ) -> Record | dict | None :
5978 column : str = "url" if url else "resource_id"
6079 pool = await context .pool ()
6180 async with pool .acquire () as connection :
@@ -66,10 +85,13 @@ async def get_latest(
6685 WHERE catalog.{ column }
6786 AND checks.id = catalog.last_check
6887 """
69- return await connection .fetchrow (q , url or resource_id )
88+ result = await connection .fetchrow (q , url or resource_id )
89+ return cls ._convert_to_dict_if_needed (result , as_dict )
7090
7191 @classmethod
72- async def get_all (cls , url : str | None = None , resource_id : str | None = None ) -> list [Record ]:
92+ async def get_all (
93+ cls , url : str | None = None , resource_id : str | None = None , as_dict : bool = False
94+ ) -> list [Record ] | list [dict ]:
7395 column : str = "url" if url else "resource_id"
7496 pool = await context .pool ()
7597 async with pool .acquire () as connection :
@@ -81,12 +103,13 @@ async def get_all(cls, url: str | None = None, resource_id: str | None = None) -
81103 AND catalog.{ column } { column }
82104 ORDER BY created_at DESC
83105 """
84- return await connection .fetch (q , url or resource_id )
106+ results = await connection .fetch (q , url or resource_id )
107+ return cls ._convert_list_to_dict_if_needed (results , as_dict )
85108
86109 @classmethod
87110 async def get_group_by_for_date (
88- cls , column : str , date : date , page_size : int = 20
89- ) -> list [Record ]:
111+ cls , column : str , date : date , page_size : int = 20 , as_dict : bool = False
112+ ) -> list [Record ] | list [ dict ] :
90113 pool = await context .pool ()
91114 async with pool .acquire () as connection :
92115 q = f"""
@@ -97,22 +120,20 @@ async def get_group_by_for_date(
97120 ORDER BY count desc
98121 LIMIT $2
99122 """
100- return await connection .fetch (q , date , page_size )
123+ results = await connection .fetch (q , date , page_size )
124+ return cls ._convert_list_to_dict_if_needed (results , as_dict )
101125
102126 @classmethod
103- async def insert (cls , data : dict , returning : str = "id" ) -> dict :
127+ async def insert (cls , data : dict , returning : str = "id" , as_dict : bool = True ) -> Record | dict :
104128 """
105129 Insert a new check in DB, associate it with the resource and return the check dict, optionally associated with the resource dataset_id.
106130 This uses the info from the last check of the same resource.
107-
108- Note: Returns dict instead of Record because this method performs additional operations beyond simple insertion (joins with catalog table, adds dataset_id).
109131 """
110132 json_data = convert_dict_values_to_json (data )
111133 q1 : str = compute_insert_query (table_name = "checks" , data = json_data , returning = returning )
112134 pool = await context .pool ()
113135 async with pool .acquire () as connection :
114136 last_check : Record = await connection .fetchrow (q1 , * json_data .values ())
115- last_check_dict = dict (last_check )
116137 q2 = (
117138 """UPDATE catalog SET last_check = $1 WHERE resource_id = $2 RETURNING dataset_id"""
118139 )
@@ -121,17 +142,17 @@ async def insert(cls, data: dict, returning: str = "id") -> dict:
121142 )
122143 # Add the dataset_id arg to the check response, if we can, and if it's asked
123144 if returning in ["*" , "dataset_id" ] and updated_resource :
145+ last_check_dict = dict (last_check )
124146 last_check_dict ["dataset_id" ] = updated_resource ["dataset_id" ]
125- return last_check_dict
147+ return last_check_dict if as_dict else last_check
148+ return dict (last_check ) if as_dict else last_check
126149
127150 @classmethod
128- async def update (cls , check_id : int , data : dict , return_as_dict : bool = False ) -> Record | dict :
129- check : Record = await update_table_record (
151+ async def update (cls , check_id : int , data : dict , as_dict : bool = False ) -> Record | dict | None :
152+ check : Record | None = await update_table_record (
130153 table_name = "checks" , record_id = check_id , data = data
131154 )
132- if return_as_dict :
133- return dict (check )
134- return check
155+ return cls ._convert_to_dict_if_needed (check , as_dict )
135156
136157 @classmethod
137158 async def delete (cls , check_id : int ) -> None :
0 commit comments