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

1import psycopg2 

2from entities.template_entity import TemplateEntity 

3 

4class TemplatePostgresRepository: 

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_template(self, template: TemplateEntity) -> TemplateEntity: 

22 ''' 

23 Retrieves a template from the PostgreSQL database using the provided template entity. 

24 Args: 

25 template (TemplateEntity): The template entity containing the ID of the template to retrieve. 

26 Returns: 

27 TemplateEntity: The retrieved template, or None if no template is found. 

28 Raises: 

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

30 ''' 

31 

32 query = "SELECT id, question, answer, author, last_modified FROM Templates WHERE id = %s;" 

33 with self.__connect() as conn: 

34 with conn.cursor() as cursor: 

35 cursor.execute(query, (template.get_id(),)) 

36 result = cursor.fetchone() 

37 if result: 

38 return TemplateEntity(id=result[0], question=result[1], answer=result[2], author_id=result[3], last_modified=result[4]) 

39 else: 

40 raise ValueError(f"No template found for template ID {template.get_id()}.") 

41 

42 

43 def get_template_list(self) -> list[TemplateEntity]: 

44 ''' 

45 Retrieves all templates from the Templates table. 

46 Returns: 

47 list[TemplateEntity]: A list of TemplateEntity objects. 

48 ''' 

49 

50 query = "SELECT id, question, answer, author, last_modified FROM Templates;" 

51 with self.__connect() as conn: 

52 with conn.cursor() as cursor: 

53 cursor.execute(query) 

54 rows = cursor.fetchall() 

55 return [TemplateEntity(id=row[0], question=row[1], answer=row[2], author_id=row[3], last_modified=row[4]) for row in rows] 

56 

57 def save_template(self, template: TemplateEntity) -> int: 

58 ''' 

59 Saves a new template to the PostgreSQL database. 

60 Args: 

61 template (TemplateEntity): The template entity to save. 

62 Returns: 

63 int: The ID of the newly created template. 

64 Raises: 

65 psycopg2.Error: If an error occurs while saving the template in the PostgreSQL database. 

66 ''' 

67 query = """ 

68 INSERT INTO Templates (question, answer, author, last_modified) 

69 VALUES (%s, %s, %s, %s) 

70 RETURNING id; 

71 """ 

72 with self.__connect() as conn: 

73 with conn.cursor() as cursor: 

74 cursor.execute(query, ( 

75 template.get_question(), 

76 template.get_answer(), 

77 template.get_author_id(), 

78 template.get_last_modified() 

79 )) 

80 conn.commit() 

81 return cursor.fetchone()[0] 

82 

83 def delete_template(self, template: TemplateEntity) -> bool: 

84 ''' 

85 Deletes a template from the PostgreSQL database based on its ID. 

86 Args: 

87 template (TemplateEntity): The template entity containing the ID of the template to delete. 

88 Returns: 

89 bool: True if the template was deleted, False otherwise. 

90 Raises: 

91 psycopg2.Error: If an error occurs while deleting the template from the PostgreSQL database. 

92 ''' 

93 

94 delete_query = "DELETE FROM Templates WHERE id = %s;" 

95 with self.__connect() as conn: 

96 with conn.cursor() as cursor: 

97 cursor.execute(delete_query, (template.get_id(),)) 

98 if cursor.rowcount > 0: 

99 conn.commit() 

100 return True 

101 else: 

102 return False