Coverage for pepdbagent/pepdbagent.py: 94%

1from pepdbagent.const import POSTGRES_DIALECT 

2from pepdbagent.db_utils import BaseEngine 

3from pepdbagent.modules.annotation import PEPDatabaseAnnotation 

4from pepdbagent.modules.namespace import PEPDatabaseNamespace 

5from pepdbagent.modules.project import PEPDatabaseProject 

6from pepdbagent.modules.sample import PEPDatabaseSample 

7from pepdbagent.modules.schema import PEPDatabaseSchema 

8from pepdbagent.modules.user import PEPDatabaseUser 

9from pepdbagent.modules.view import PEPDatabaseView 

10 

11 

12class PEPDatabaseAgent(object): 

13 def __init__( 

14 self, 

15 host: str = "localhost", 

16 port: int = 5432, 

17 database: str = "pep-db", 

18 user: str | None = None, 

19 password: str | None = None, 

20 drivername: str = POSTGRES_DIALECT, 

21 dsn: str | None = None, 

22 echo: bool = False, 

23 run_migrations: bool = False, 

24 ): 

25 """Initialize connection to the pep_db database. 

26 

27 Args: 

28 host: Database server address, e.g., localhost or an IP address. 

29 port: Port number (default: 5432). 

30 database: Name of the database to connect to. 

31 user: Username for authentication. 

32 password: Password for authentication. 

33 drivername: Database driver (default: postgresql). 

34 dsn: libpq connection string, e.g., "localhost://username:password@pdp_db:5432". 

35 echo: Log all SQL statements if True. 

36 run_migrations: Run migrations on the database if True. 

37 """ 

38 

39 pep_db_engine = BaseEngine( 

40 host=host, 

41 port=port, 

42 database=database, 

43 user=user, 

44 password=password, 

45 drivername=drivername, 

46 dsn=dsn, 

47 echo=echo, 

48 run_migrations=run_migrations, 

49 ) 

50 sa_engine = pep_db_engine.engine 

51 

52 self.pep_db_engine = pep_db_engine 

53 self._sa_engine = sa_engine 

54 

55 self._project = PEPDatabaseProject(pep_db_engine) 

56 self._annotation = PEPDatabaseAnnotation(pep_db_engine) 

57 self._namespace = PEPDatabaseNamespace(pep_db_engine) 

58 self._sample = PEPDatabaseSample(pep_db_engine) 

59 self._user = PEPDatabaseUser(pep_db_engine) 

60 self._view = PEPDatabaseView(pep_db_engine) 

61 self._schema = PEPDatabaseSchema(pep_db_engine) 

62 

63 self._db_name = database 

64 

65 @property 

66 def project(self) -> PEPDatabaseProject: 

67 return self._project 

68 

69 @property 

70 def annotation(self) -> PEPDatabaseAnnotation: 

71 return self._annotation 

72 

73 @property 

74 def namespace(self) -> PEPDatabaseNamespace: 

75 return self._namespace 

76 

77 @property 

78 def user(self) -> PEPDatabaseUser: 

79 return self._user 

80 

81 @property 

82 def sample(self) -> PEPDatabaseSample: 

83 return self._sample 

84 

85 @property 

86 def view(self) -> PEPDatabaseView: 

87 return self._view 

88 

89 @property 

90 def schema(self) -> PEPDatabaseSchema: 

91 return self._schema 

92 

93 def __str__(self): 

94 return f"Connection to the database: '{self.__db_name}' is set!" 

95 

96 def __exit__(self): 

97 self._sa_engine.__exit__() 

98 

99 @property 

100 def connection(self) -> BaseEngine: 

101 return self._sa_engine