diff --git a/app/data/__init__.py b/app/data/__init__.py index 60a75dac..dffc5ef9 100644 --- a/app/data/__init__.py +++ b/app/data/__init__.py @@ -3,19 +3,21 @@ from ..services.location.jhu import JhuLocationService from ..services.location.nyt import NYTLocationService -# Mapping of services to data-sources. -DATA_SOURCES = { - "jhu": JhuLocationService(), - "csbs": CSBSLocationService(), - "nyt": NYTLocationService(), -} +class Data: + def __init__(self, source) + # Mapping of services to data-sources. + DATA_SOURCES = { + "jhu": JhuLocationService(), + "csbs": CSBSLocationService(), + "nyt": NYTLocationService(), + } -def data_source(source): - """ - Retrieves the provided data-source service. + def data_source(self): + """ + Retrieves the provided data-source service. - :returns: The service. - :rtype: LocationService - """ - return DATA_SOURCES.get(source.lower()) + :returns: The service. + :rtype: LocationService + """ + return self.DATA_SOURCES.get(source.lower()) diff --git a/app/main.py b/app/main.py index b9aff949..ddcf6d3d 100644 --- a/app/main.py +++ b/app/main.py @@ -14,7 +14,7 @@ from sentry_sdk.integrations.asgi import SentryAsgiMiddleware from .config import get_settings -from .data import data_source +from .data import Data from .routers import V1, V2 from .utils.httputils import setup_client_session, teardown_client_session @@ -74,7 +74,7 @@ async def add_datasource(request: Request, call_next): Attach the data source to the request.state. """ # Retrieve the datas ource from query param. - source = data_source(request.query_params.get("source", default="jhu")) + source = Data().data_source(request.query_params.get("source", default="jhu")) # Abort with 404 if source cannot be found. if not source: