Skip to content

Commit

Permalink
split management plane and data plane operations.
Browse files Browse the repository at this point in the history
  • Loading branch information
marrobi committed Dec 21, 2023
1 parent 5e7a417 commit 97018a2
Show file tree
Hide file tree
Showing 8 changed files with 73 additions and 20 deletions.
2 changes: 1 addition & 1 deletion api_app/_version.py
Original file line number Diff line number Diff line change
@@ -1 +1 @@
__version__ = "0.17.8"
__version__ = "0.17.12"
6 changes: 3 additions & 3 deletions api_app/api/routes/health.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import asyncio
from fastapi import APIRouter
from fastapi import APIRouter, Request
from core import credentials
from models.schemas.status import HealthCheck, ServiceStatus, StatusEnum
from resources import strings
Expand All @@ -10,13 +10,13 @@


@router.get("/health", name=strings.API_GET_HEALTH_STATUS)
async def health_check() -> HealthCheck:
async def health_check(request: Request) -> HealthCheck:
# The health endpoint checks the status of key components of the system.
# Note that Resource Processor checks incur Azure management calls, so
# calling this endpoint frequently may result in API throttling.
async with credentials.get_credential_async() as credential:
cosmos, sb, rp = await asyncio.gather(
create_state_store_status(credential),
create_state_store_status(request),
create_service_bus_status(credential),
create_resource_processor_status(credential)
)
Expand Down
46 changes: 37 additions & 9 deletions api_app/db/events.py
Original file line number Diff line number Diff line change
@@ -1,17 +1,45 @@
from azure.cosmos.aio import CosmosClient
from azure.mgmt.cosmosdb import CosmosDBManagementClient

from api.dependencies.database import get_db_client
from db.repositories.resources import ResourceRepository
from core.config import SUBSCRIPTION_ID, RESOURCE_GROUP_NAME, RESOURCE_LOCATION, COSMOSDB_ACCOUNT_NAME, STATE_STORE_DATABASE, STATE_STORE_RESOURCES_CONTAINER, STATE_STORE_RESOURCE_TEMPLATES_CONTAINER, STATE_STORE_RESOURCES_HISTORY_CONTAINER, STATE_STORE_OPERATIONS_CONTAINER, STATE_STORE_AIRLOCK_REQUESTS_CONTAINER
from core.credentials import get_credential
from services.logging import logger


async def bootstrap_database(app) -> bool:
async def bootstrap_database() -> bool:
try:
client: CosmosClient = await get_db_client(app)
if client:
# Test access to database
await ResourceRepository.create(client)
return True
credential = get_credential()
db_mgmt_client = CosmosDBManagementClient(credential=credential, subscription_id=SUBSCRIPTION_ID)

repository_containers = [
STATE_STORE_RESOURCES_CONTAINER,
STATE_STORE_RESOURCE_TEMPLATES_CONTAINER,
STATE_STORE_RESOURCES_HISTORY_CONTAINER,
STATE_STORE_OPERATIONS_CONTAINER,
STATE_STORE_AIRLOCK_REQUESTS_CONTAINER
]

for container in repository_containers:
# create container if it doesn't exist
db_mgmt_client.sql_resources.begin_create_update_sql_container(
resource_group_name=RESOURCE_GROUP_NAME,
account_name=COSMOSDB_ACCOUNT_NAME,
database_name=STATE_STORE_DATABASE,
container_name=container,
create_update_sql_container_parameters={
"location": RESOURCE_LOCATION,
"resource": {
"id": container,
"partition_key": {
"paths": [
"/id"
],
"kind": "Hash"
}
}
}
)
return True

except Exception as e:
logger.exception("Could not bootstrap database")
logger.debug(e)
Expand Down
2 changes: 1 addition & 1 deletion api_app/db/repositories/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ def container(self) -> ContainerProxy:
async def _get_container(cls, container_name, partition_key_obj) -> ContainerProxy:
try:
database = cls._client.get_database_client(config.STATE_STORE_DATABASE)
container = await database.create_container_if_not_exists(id=container_name, partition_key=partition_key_obj)
container = database.get_container_client(container=container_name)
return container
except Exception:
raise UnableToAccessDatabase
Expand Down
2 changes: 1 addition & 1 deletion api_app/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@
async def lifespan(app: FastAPI):
app.state.cosmos_client = None

while not await bootstrap_database(app):
while not await bootstrap_database():
await asyncio.sleep(5)
logger.warning("Database connection could not be established")

Expand Down
6 changes: 3 additions & 3 deletions api_app/services/health_checker.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,19 +4,19 @@
from azure.mgmt.compute.aio import ComputeManagementClient
from azure.cosmos.exceptions import CosmosHttpResponseError
from azure.servicebus.exceptions import ServiceBusConnectionError, ServiceBusAuthenticationError
from api.dependencies.database import connect_to_db
from api.dependencies.database import get_db_client_from_request

from core import config
from models.schemas.status import StatusEnum
from resources import strings
from services.logging import logger


async def create_state_store_status(credential) -> Tuple[StatusEnum, str]:
async def create_state_store_status(request) -> Tuple[StatusEnum, str]:
status = StatusEnum.ok
message = ""
try:
cosmos_client = await connect_to_db()
cosmos_client = await get_db_client_from_request(request)
async with cosmos_client:
list_databases_response = cosmos_client.list_databases()
[database async for database in list_databases_response]
Expand Down
25 changes: 25 additions & 0 deletions api_app/tests_ma/test_db/test_events.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
from unittest.mock import AsyncMock, MagicMock, patch
from azure.core.exceptions import AzureError
from api_app.db import events


@patch("api_app.db.events.get_credential_async")
@patch("api_app.db.events.CosmosDBManagementClient")
async def test_bootstrap_database_success(cosmos_db_mgmt_client_mock, get_credential_async_mock):
get_credential_async_mock.return_value = AsyncMock()
cosmos_db_mgmt_client_mock.return_value = MagicMock()

result = await events.bootstrap_database()

assert result is True


@patch("api_app.db.events.get_credential_async")
@patch("api_app.db.events.CosmosDBManagementClient")
async def test_bootstrap_database_failure(cosmos_db_mgmt_client_mock, get_credential_async_mock):
get_credential_async_mock.return_value = AsyncMock()
cosmos_db_mgmt_client_mock.side_effect = AzureError()

result = await events.bootstrap_database()

assert result is False
4 changes: 2 additions & 2 deletions api_app/tests_ma/test_services/test_health_checker.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ async def test_get_state_store_status_responding(_, get_store_key_mock, get_cred

@patch("core.credentials.get_credential_async")
@patch("api.dependencies.database.get_store_key")
@patch("api.dependencies.database.CosmosClient")
@patch("api.dependencies.database.get_db_client")
async def test_get_state_store_status_not_responding(cosmos_client_mock, get_store_key_mock, get_credential_async) -> None:
get_credential_async.return_value = AsyncMock()
get_store_key_mock.return_value = None
Expand All @@ -38,7 +38,7 @@ async def test_get_state_store_status_not_responding(cosmos_client_mock, get_sto

@patch("core.credentials.get_credential_async")
@patch("api.dependencies.database.get_store_key")
@patch("api.dependencies.database.CosmosClient")
@patch("api.dependencies.database.get_db_client")
async def test_get_state_store_status_other_exception(cosmos_client_mock, get_store_key_mock, get_credential_async) -> None:
get_credential_async.return_value = AsyncMock()
get_store_key_mock.return_value = None
Expand Down

0 comments on commit 97018a2

Please sign in to comment.