Coverage for src/app/repositories/message_postgres_repository.py: 100%
60 statements
« prev ^ index » next coverage.py v7.7.0, created at 2025-04-03 00:42 +0200
1import psycopg2
2from entities.message_entity import MessageEntity
3from entities.conversation_entity import ConversationEntity
5class MessagePostgresRepository:
6 def __init__(self, db_config: dict):
7 '''
8 Initializes the repository with the given database configuration.
9 Args:
10 db_config (dict): A dictionary containing the PostgreSQL database configuration parameters.
11 '''
12 self.__db_config = db_config
14 def __connect(self):
15 '''
16 Creates and returns a new connection to the PostgreSQL database.
17 Returns:
18 psycopg2.extensions.connection: A connection object to interact with the PostgreSQL database.
19 '''
20 return psycopg2.connect(**self.__db_config)
22 def get_message(self, message: MessageEntity) -> MessageEntity:
23 '''
24 Retrieves a message from the PostgreSQL database by its ID.
25 Args:
26 message (MessageEntity): The message entity containing the ID of the message to retrieve.
27 Returns:
28 MessageEntity: The retrieved message data.
29 Raises:
30 psycopg2.Error: If an error occurs while retrieving the message from the PostgreSQL database.
31 '''
33 select_message_query = """
34 SELECT id, text, created_at, is_bot, conversation_id, rating
35 FROM Messages
36 WHERE id = %s;
37 """
38 with self.__connect() as connection: # Call the method to get the connection object
39 with connection.cursor() as cursor:
40 cursor.execute(select_message_query, (message.get_id(),))
41 result = cursor.fetchone()
42 if result:
43 message = MessageEntity(
44 id=result[0],
45 text=result[1],
46 created_at=result[2],
47 is_bot=result[3],
48 conversation_id=result[4],
49 rating=result[5]
50 )
51 return message
52 else:
53 raise ValueError(f"Message with ID {message.get_id()} not found.")
55 def get_messages_by_conversation(self, conversation: MessageEntity) -> list[MessageEntity]:
56 '''
57 Retrieves all messages associated with a specific conversation from the PostgreSQL database.
58 Args:
59 conversation (ConversationEntity): The conversation entity containing the ID of the conversation.
60 Returns:
61 list[MessageEntity]: A list of MessageEntity objects.
62 Raises:
63 psycopg2.Error: If an error occurs while retrieving the messages from the PostgreSQL database.
64 '''
66 select_messages_query = """
67 SELECT id, text, created_at, is_bot, conversation_id, rating
68 FROM Messages
69 WHERE conversation_id = %s;
70 """
71 with self.__connect() as connection: # Call the method to get the connection object
72 with connection.cursor() as cursor:
73 cursor.execute(select_messages_query, (conversation.get_conversation_id(),))
74 rows = cursor.fetchall()
75 if rows:
76 return [
77 MessageEntity(
78 id=row[0],
79 text=row[1],
80 created_at=row[2],
81 is_bot=row[3],
82 conversation_id=row[4],
83 rating=row[5]
84 ) for row in rows
85 ]
86 else:
87 return []
89 def save_message(self, message: MessageEntity) -> int:
90 '''
91 Saves a message into the PostgreSQL database and returns the ID of the created message.
92 Args:
93 message (MessageEntity): The message data to be saved.
94 Returns:
95 int: The ID of the created message.
96 Raises:
97 psycopg2.Error: If an error occurs while saving the message in the PostgreSQL database.
98 '''
100 insert_message_query = """
101 INSERT INTO Messages (text, created_at, is_bot, conversation_id, rating)
102 VALUES (%s, %s, %s, %s, %s)
103 RETURNING id;
104 """
105 params = (message.get_text(), message.get_created_at(), message.get_is_bot(), message.get_conversation_id(), message.get_rating())
106 with self.__connect() as connection:
107 with connection.cursor() as cursor:
108 cursor.execute(insert_message_query, params)
109 created_id = cursor.fetchone()[0]
110 connection.commit()
111 return created_id
113 def delete_message(self, message: MessageEntity) -> bool:
114 '''
115 Deletes a message from the PostgreSQL database.
116 Args:
117 message (MessageEntity): The message to be deleted.
118 Returns:
119 bool: True if the message was deleted successfully, False otherwise.
120 Raises:
121 psycopg2.Error: If an error occurs while deleting the message from the PostgreSQL database.
122 '''
124 delete_message_query = """
125 DELETE FROM Messages
126 WHERE id = %s;
127 """
128 with self.__connect() as conn:
129 with conn.cursor() as cursor:
130 cursor.execute(delete_message_query, (message.get_id(),))
131 conn.commit()
132 return cursor.rowcount > 0
135 def update_message_rating(self, message: MessageEntity) -> bool:
136 '''
137 Updates the rating of a message in the PostgreSQL database by its ID.
138 Args:
139 message (MessageEntity): The message entity containing the ID and the new rating value.
140 Returns:
141 bool: True if the rating was successfully updated, False otherwise.
142 Raises:
143 psycopg2.Error: If an error occurs while updating the rating in the PostgreSQL database.
144 '''
146 update_rating_query = """
147 UPDATE Messages
148 SET rating = %s
149 WHERE id = %s;
150 """
151 params = (message.get_rating(), message.get_id())
152 with self.__connect() as connection:
153 with connection.cursor() as cursor:
154 cursor.execute(update_rating_query, params)
155 connection.commit()
156 return cursor.rowcount > 0
158 def fetch_messages(self) -> list[MessageEntity]:
159 """
160 Fetch the dashboard metrics data.
161 Returns:
162 list[MessageEntity]: A list of MessageEntity objects containing the messages data.
163 """
165 select_messages_query = """
166 SELECT id, text, created_at, is_bot, conversation_id, rating
167 FROM Messages;
168 """
169 with self.__connect() as connection:
170 with connection.cursor() as cursor:
171 cursor.execute(select_messages_query)
172 rows = cursor.fetchall()
173 if rows:
174 return [
175 MessageEntity(
176 id=row[0],
177 text=row[1],
178 created_at=row[2],
179 is_bot=row[3],
180 conversation_id=row[4],
181 rating=row[5]
182 ) for row in rows
183 ]
184 else:
185 return []