Coverage for src/app/repositories/conversation_postgres_repository.py: 100%

1import psycopg2 

2from entities.conversation_entity import ConversationEntity 

3 

4class ConversationPostgresRepository: 

5 def __init__(self, db_config: dict): 

6 ''' 

7 Initializes the PostgresRepository with the given database configuration. 

8 Args: 

9 db_config (dict): The configuration dictionary for the PostgreSQL database. 

10 ''' 

11 self.__db_config = db_config 

12 

13 def __connect(self): 

14 ''' 

15 Establishes a new connection to the PostgreSQL database. 

16 Returns: 

17 psycopg2.extensions.connection: The connection object to the PostgreSQL database. 

18 ''' 

19 return psycopg2.connect(**self.__db_config) 

20 

21 def get_conversation(self, conversation: ConversationEntity) -> ConversationEntity: 

22 ''' 

23 Retrieves a conversation from the PostgreSQL database by its ID. 

24 Args: 

25 conversation (ConversationEntity): The conversation entity containing the ID to retrieve. 

26 Returns: 

27 ConversationEntity: The retrieved conversation, or None if not found. 

28 Raises: 

29 psycopg2.Error: If an error occurs while retrieving the conversation from the PostgreSQL database. 

30 ''' 

31 id = conversation.get_id() 

32 

33 query = "SELECT id, title, user_id FROM Conversations WHERE id = %s;" 

34 with self.__connect() as conn: 

35 with conn.cursor() as cursor: 

36 cursor.execute(query, (id,)) 

37 result = cursor.fetchone() 

38 if result: 

39 return ConversationEntity(id=result[0], title=result[1], user_id=result[2]) 

40 else: 

41 return None 

42 

43 

44 def get_conversations(self, conversation: ConversationEntity) -> list[ConversationEntity]: 

45 ''' 

46 Retrieves all conversations associated with a specific user from the PostgreSQL database. 

47 Args: 

48 user_id (int): The ID of the user whose conversations are to be retrieved. 

49 Returns: 

50 list[ConversationEntity]: A list of all retrieved conversations. 

51 Raises: 

52 psycopg2.Error: If an error occurs while retrieving the conversations from the PostgreSQL database. 

53 ''' 

54 

55 query = """ 

56 SELECT id, title, user_id 

57 FROM Conversations 

58 WHERE user_id = %s; 

59 """ 

60 with self.__connect() as conn: 

61 with conn.cursor() as cursor: 

62 cursor.execute(query, (conversation.get_user_id(),)) 

63 results = cursor.fetchall() 

64 return [ConversationEntity(id=row[0], title=row[1], user_id=row[2]) for row in results] 

65 

66 def save_conversation_title(self, conversation: ConversationEntity) -> int: 

67 ''' 

68 Saves the title of a conversation in the PostgreSQL database. 

69 If the conversation does not exist, it creates a new one. 

70 Args: 

71 conversation (ConversationEntity): The conversation entity containing the user ID and title. 

72 Returns: 

73 int: The ID of the saved conversation. 

74 Raises: 

75 psycopg2.Error: If an error occurs while saving the conversation title in the PostgreSQL database. 

76 ''' 

77 insert_query = """ 

78 INSERT INTO Conversations (title, user_id)  

79 VALUES (%s, %s)  

80 RETURNING id; 

81 """ 

82 with self.__connect() as conn: 

83 with conn.cursor() as cursor: 

84 cursor.execute(insert_query, (conversation.get_title(), conversation.get_user_id())) 

85 saved_id = cursor.fetchone()[0] 

86 conn.commit() 

87 return saved_id 

88 

89 def delete_conversation(self, conversation: ConversationEntity)-> bool: 

90 ''' 

91 Deletes a conversation from the PostgreSQL database. 

92 Args: 

93 conversation (ConversationEntity): The conversation entity to delete. 

94 Returns: 

95 bool: True if the conversation was successfully deleted, False otherwise. 

96 Raises: 

97 psycopg2.Error: If an error occurs while deleting the conversation from the PostgreSQL database. 

98 ''' 

99 id = conversation.get_id() 

100 delete_query = "DELETE FROM Conversations WHERE id = %s;" 

101 with self.__connect() as conn: 

102 with conn.cursor() as cursor: 

103 cursor.execute(delete_query, (id,)) 

104 conn.commit() 

105 return cursor.rowcount > 0