From e6704037bc8bbc0641b9da7b358fff364986cf42 Mon Sep 17 00:00:00 2001 From: Abram Date: Thu, 7 Sep 2023 10:17:08 +0100 Subject: [PATCH 1/5] Feat - implemented database engine --- .../agenta_backend/models/db_engine.py | 44 +++++++++++++++++++ 1 file changed, 44 insertions(+) create mode 100644 agenta-backend/agenta_backend/models/db_engine.py diff --git a/agenta-backend/agenta_backend/models/db_engine.py b/agenta-backend/agenta_backend/models/db_engine.py new file mode 100644 index 0000000000..14f53ae905 --- /dev/null +++ b/agenta-backend/agenta_backend/models/db_engine.py @@ -0,0 +1,44 @@ +import os +import logging + +from odmantic import AIOEngine +from pymongo import MongoClient +from motor.motor_asyncio import AsyncIOMotorClient + + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) + + +class DBEngine(object): + """Database engine to initialize client and return engine based on mode""" + + def __init__(self, mode: str) -> None: + self.mode = mode + self.db_url = os.environ["MONGODB_URI"] + + @property + def initialize_client(self) -> AsyncIOMotorClient: + """Returns an instance of `AsyncIOMotorClient` initialized \ + with the provided `db_url`. + """ + + client = AsyncIOMotorClient(self.db_url) + return client + + def engine(self) -> AIOEngine: + """Returns an AIOEngine object with a specified database name based on the mode. + """ + + aio_engine = AIOEngine(client=self.initialize_client) + if self.mode == "test": + aio_engine.database = "agenta_test" + logger.info("Using test database...") + elif self.mode == "default": + aio_engine.database = "agenta" + logger.info("Using default database") + return aio_engine + + def remove_db(self) -> None: + client = MongoClient(self.db_url) + client.drop_database("agenta_test") From a50fe933c095507149c1d98ebd58b88001869698 Mon Sep 17 00:00:00 2001 From: Abram Date: Thu, 7 Sep 2023 10:44:06 +0100 Subject: [PATCH 2/5] Update - added method to remove database based on mode --- .../agenta_backend/models/db_engine.py | 23 ++++++++++++------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/agenta-backend/agenta_backend/models/db_engine.py b/agenta-backend/agenta_backend/models/db_engine.py index 14f53ae905..a084d3deda 100644 --- a/agenta-backend/agenta_backend/models/db_engine.py +++ b/agenta-backend/agenta_backend/models/db_engine.py @@ -11,7 +11,9 @@ class DBEngine(object): - """Database engine to initialize client and return engine based on mode""" + """ + Database engine to initialize client and return engine based on mode + """ def __init__(self, mode: str) -> None: self.mode = mode @@ -19,7 +21,8 @@ def __init__(self, mode: str) -> None: @property def initialize_client(self) -> AsyncIOMotorClient: - """Returns an instance of `AsyncIOMotorClient` initialized \ + """ + Returns an instance of `AsyncIOMotorClient` initialized \ with the provided `db_url`. """ @@ -27,18 +30,22 @@ def initialize_client(self) -> AsyncIOMotorClient: return client def engine(self) -> AIOEngine: - """Returns an AIOEngine object with a specified database name based on the mode. + """ + Returns an AIOEngine object with a specified database name based on the mode. """ - aio_engine = AIOEngine(client=self.initialize_client) if self.mode == "test": - aio_engine.database = "agenta_test" + aio_engine = AIOEngine(client=self.initialize_client, database="agenta_test") logger.info("Using test database...") + return aio_engine elif self.mode == "default": - aio_engine.database = "agenta" + aio_engine = AIOEngine(client=self.initialize_client, database="agenta") logger.info("Using default database") - return aio_engine + return aio_engine def remove_db(self) -> None: client = MongoClient(self.db_url) - client.drop_database("agenta_test") + if self.mode == "default": + client.drop_database("agenta") + elif self.mode == "test": + client.drop_database("agenta_test") From 3dfc5ca0dac6b460d78194988af41a0b5f37dd1d Mon Sep 17 00:00:00 2001 From: Abram Date: Thu, 7 Sep 2023 10:45:03 +0100 Subject: [PATCH 3/5] Refactor - use custom database enigine --- .../agenta_backend/services/db_manager.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/agenta-backend/agenta_backend/services/db_manager.py b/agenta-backend/agenta_backend/services/db_manager.py index 7c9a52e78c..93e28bac3b 100644 --- a/agenta-backend/agenta_backend/services/db_manager.py +++ b/agenta-backend/agenta_backend/services/db_manager.py @@ -1,4 +1,5 @@ import os +import logging from bson import ObjectId from typing import Dict, List, Any @@ -23,17 +24,13 @@ OrganizationDB, ) from agenta_backend.services import helpers +from agenta_backend.models.db_engine import DBEngine -from odmantic import AIOEngine, query -from motor.motor_asyncio import AsyncIOMotorClient - -import logging +from odmantic import query -# SQLite database connection -DATABASE_URL = os.environ["MONGODB_URI"] -client = AsyncIOMotorClient(DATABASE_URL) -engine = AIOEngine(client=client, database="agenta") +# Initialize database engine +engine = DBEngine(mode="default").engine() logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) From 613037639b63c8a72dddfc835949423cb8970fe5 Mon Sep 17 00:00:00 2001 From: Abram Date: Thu, 7 Sep 2023 11:04:30 +0100 Subject: [PATCH 4/5] Feat - created test db engine --- agenta-backend/tests/conftest.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/agenta-backend/tests/conftest.py b/agenta-backend/tests/conftest.py index c66872f036..e2fa4833af 100644 --- a/agenta-backend/tests/conftest.py +++ b/agenta-backend/tests/conftest.py @@ -1,5 +1,7 @@ import pytest from fastapi.testclient import TestClient + +from agenta_backend.models.db_engine import DBEngine from agenta_backend.main import app @@ -8,3 +10,11 @@ def test_app(): client = TestClient(app) yield client # provide the test client to the tests # teardown code goes here + +@pytest.fixture(scope="function") +def test_db_engine(): + # Initialize the DBEngine in 'test' mode + db_engine = DBEngine(mode="test") + test_engine = db_engine.engine() + yield test_engine + db_engine.remove_db() \ No newline at end of file From 7a5a1eba78238b10c843082a594695239485c96a Mon Sep 17 00:00:00 2001 From: Abram Date: Thu, 7 Sep 2023 11:05:24 +0100 Subject: [PATCH 5/5] :art: Format - ran black and foramt-fix --- agenta-backend/agenta_backend/models/db_engine.py | 10 ++++++---- agenta-backend/tests/conftest.py | 3 ++- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/agenta-backend/agenta_backend/models/db_engine.py b/agenta-backend/agenta_backend/models/db_engine.py index a084d3deda..26a309d980 100644 --- a/agenta-backend/agenta_backend/models/db_engine.py +++ b/agenta-backend/agenta_backend/models/db_engine.py @@ -25,7 +25,7 @@ def initialize_client(self) -> AsyncIOMotorClient: Returns an instance of `AsyncIOMotorClient` initialized \ with the provided `db_url`. """ - + client = AsyncIOMotorClient(self.db_url) return client @@ -33,16 +33,18 @@ def engine(self) -> AIOEngine: """ Returns an AIOEngine object with a specified database name based on the mode. """ - + if self.mode == "test": - aio_engine = AIOEngine(client=self.initialize_client, database="agenta_test") + aio_engine = AIOEngine( + client=self.initialize_client, database="agenta_test" + ) logger.info("Using test database...") return aio_engine elif self.mode == "default": aio_engine = AIOEngine(client=self.initialize_client, database="agenta") logger.info("Using default database") return aio_engine - + def remove_db(self) -> None: client = MongoClient(self.db_url) if self.mode == "default": diff --git a/agenta-backend/tests/conftest.py b/agenta-backend/tests/conftest.py index e2fa4833af..e7c5de5b25 100644 --- a/agenta-backend/tests/conftest.py +++ b/agenta-backend/tests/conftest.py @@ -11,10 +11,11 @@ def test_app(): yield client # provide the test client to the tests # teardown code goes here + @pytest.fixture(scope="function") def test_db_engine(): # Initialize the DBEngine in 'test' mode db_engine = DBEngine(mode="test") test_engine = db_engine.engine() yield test_engine - db_engine.remove_db() \ No newline at end of file + db_engine.remove_db()