Coverage for chalicelib/core/db_request_handler.py: 0%

1import logging 

2 

3from chalicelib.utils import pg_client 

4from chalicelib.utils.log import sanitize 

5 

6logger = logging.getLogger(__name__) 

7 

8 

9class DatabaseRequestHandler: 

10 def __init__(self, table_name): 

11 self.table_name = table_name 

12 self.constraints = [] 

13 self.params = {} 

14 self.order_clause = "" 

15 self.sort_clause = "" 

16 self.select_columns = [] 

17 self.sub_queries = [] 

18 self.joins = [] 

19 self.group_by_clause = "" 

20 self.client = pg_client 

21 self.logger = logging.getLogger(__name__) 

22 self.pagination = {} 

23 

24 def add_constraint(self, constraint, param=None): 

25 self.constraints.append(constraint) 

26 if param: 

27 self.params.update(param) 

28 

29 def add_subquery(self, subquery, alias, param=None): 

30 self.sub_queries.append((subquery, alias)) 

31 if param: 

32 self.params.update(param) 

33 

34 def add_join(self, join_clause): 

35 self.joins.append(join_clause) 

36 

37 def add_param(self, key, value): 

38 self.params[key] = value 

39 

40 def set_order_by(self, order_by): 

41 self.order_clause = order_by 

42 

43 def set_sort_by(self, sort_by): 

44 self.sort_clause = sort_by 

45 

46 def set_select_columns(self, columns): 

47 self.select_columns = columns 

48 

49 def set_group_by(self, group_by_clause): 

50 self.group_by_clause = group_by_clause 

51 

52 def set_pagination(self, page, page_size): 

53 """ 

54 Set pagination parameters for the query. 

55 :param page: The page number (1-indexed) 

56 :param page_size: Number of items per page 

57 """ 

58 self.pagination = { 

59 'offset': (page - 1) * page_size, 

60 'limit': page_size 

61 } 

62 

63 def build_query(self, action="select", additional_clauses=None, data=None): 

64 

65 if action == "select": 

66 query = f"SELECT {', '.join(self.select_columns)} FROM {self.table_name}" 

67 elif action == "insert": 

68 columns = ', '.join(data.keys()) 

69 placeholders = ', '.join(f'%({k})s' for k in data.keys()) 

70 query = f"INSERT INTO {self.table_name} ({columns}) VALUES ({placeholders})" 

71 elif action == "update": 

72 set_clause = ', '.join(f"{k} = %({k})s" for k in data.keys()) 

73 query = f"UPDATE {self.table_name} SET {set_clause}" 

74 elif action == "delete": 

75 query = f"DELETE FROM {self.table_name}" 

76 

77 for join in self.joins: 

78 query += f" {join}" 

79 for subquery, alias in self.sub_queries: 

80 query += f", ({subquery}) AS {alias}" 

81 if self.constraints: 

82 query += " WHERE " + " AND ".join(self.constraints) 

83 if action == "select": 

84 if self.group_by_clause: 

85 query += " GROUP BY " + self.group_by_clause 

86 if self.sort_clause: 

87 query += " ORDER BY " + self.sort_clause 

88 if self.order_clause: 

89 query += " " + self.order_clause 

90 if hasattr(self, 'pagination') and self.pagination: 

91 query += " LIMIT %(limit)s OFFSET %(offset)s" 

92 self.params.update(self.pagination) 

93 

94 if additional_clauses: 

95 query += " " + additional_clauses 

96 

97 logger.debug(f"Query: {query}") 

98 return query 

99 

100 def execute_query(self, query, data=None): 

101 try: 

102 with self.client.PostgresClient() as cur: 

103 mogrified_query = cur.mogrify(query, {**data, **self.params} if data else self.params) 

104 cur.execute(mogrified_query) 

105 return cur.fetchall() if cur.description else None 

106 except Exception as e: 

107 self.logger.error(f"Database operation failed: {sanitize(str(e))}") 

108 raise 

109 

110 def fetchall(self): 

111 query = self.build_query() 

112 return self.execute_query(query) 

113 

114 def fetchone(self): 

115 query = self.build_query() 

116 result = self.execute_query(query) 

117 return result[0] if result else None 

118 

119 def insert(self, data): 

120 query = self.build_query(action="insert", data=data) 

121 query += " RETURNING *;" 

122 

123 result = self.execute_query(query, data) 

124 return result[0] if result else None 

125 

126 def update(self, data): 

127 query = self.build_query(action="update", data=data) 

128 query += " RETURNING *;" 

129 

130 result = self.execute_query(query, data) 

131 return result[0] if result else None 

132 

133 def delete(self): 

134 query = self.build_query(action="delete") 

135 return self.execute_query(query) 

136 

137 def batch_insert(self, items): 

138 if not items: 

139 return None 

140 

141 columns = ', '.join(items[0].keys()) 

142 

143 # Building a values string with unique parameter names for each item 

144 all_values_query = ', '.join( 

145 '(' + ', '.join([f"%({key}_{i})s" for key in item]) + ')' 

146 for i, item in enumerate(items) 

147 ) 

148 

149 query = f"INSERT INTO {self.table_name} ({columns}) VALUES {all_values_query} RETURNING *;" 

150 

151 try: 

152 with self.client.PostgresClient() as cur: 

153 # Flatten items into a single dictionary with unique keys 

154 combined_params = {f"{k}_{i}": v for i, item in enumerate(items) for k, v in item.items()} 

155 mogrified_query = cur.mogrify(query, combined_params) 

156 cur.execute(mogrified_query) 

157 return cur.fetchall() 

158 except Exception as e: 

159 self.logger.error(f"Database batch insert operation failed: {sanitize(str(e))}") 

160 raise 

161 

162 def raw_query(self, query, params=None): 

163 try: 

164 with self.client.PostgresClient() as cur: 

165 mogrified_query = cur.mogrify(query, params) 

166 cur.execute(mogrified_query) 

167 return cur.fetchall() if cur.description else None 

168 except Exception as e: 

169 self.logger.error(f"Database operation failed: {sanitize(str(e))}") 

170 raise 

171 

172 def batch_update(self, items): 

173 if not items: 

174 return None 

175 

176 id_column = list(items[0])[0] 

177 

178 # Building the set clause for the update statement 

179 update_columns = list(items[0].keys()) 

180 update_columns.remove(id_column) 

181 set_clause = ', '.join([f"{col} = v.{col}" for col in update_columns]) 

182 

183 # Building the values part for the 'VALUES' section 

184 values_rows = [] 

185 for item in items: 

186 values = ', '.join([f"%({key})s" for key in item.keys()]) 

187 values_rows.append(f"({values})") 

188 values_query = ', '.join(values_rows) 

189 

190 # Constructing the full update query 

191 query = f""" 

192 UPDATE {self.table_name} AS t  

193 SET {set_clause}  

194 FROM (VALUES {values_query}) AS v ({', '.join(items[0].keys())})  

195 WHERE t.{id_column} = v.{id_column}; 

196 """ 

197 

198 try: 

199 with self.client.PostgresClient() as cur: 

200 # Flatten items into a single dictionary for mogrify 

201 combined_params = {k: v for item in items for k, v in item.items()} 

202 mogrified_query = cur.mogrify(query, combined_params) 

203 cur.execute(mogrified_query) 

204 except Exception as e: 

205 self.logger.error(f"Database batch update operation failed: {sanitize(str(e))}") 

206 raise