From 841d83b91e29f39fe3a2375b4a7f32b0a8964b7c Mon Sep 17 00:00:00 2001
From: Mike Presman <mikepresman@gmail.com>
Date: Thu, 15 Jul 2021 12:05:02 -0400
Subject: [PATCH] Aggregate Pattern for DataSource functionallity implemented.
 Allows for a deeper level of abstraction while encapsulating data sources in
 such a way that the data sources and the boundary can only be accessed
 through the root

---
 app/data/__init__.py | 31 ++++++++++++++++++-------------
 app/main.py          |  6 ++++--
 app/routers/v2.py    | 11 ++++++-----
 3 files changed, 28 insertions(+), 20 deletions(-)

diff --git a/app/data/__init__.py b/app/data/__init__.py
index 60a75dac..6a66b1fb 100644
--- a/app/data/__init__.py
+++ b/app/data/__init__.py
@@ -2,20 +2,25 @@
 from ..services.location.csbs import CSBSLocationService
 from ..services.location.jhu import JhuLocationService
 from ..services.location.nyt import NYTLocationService
+from ..services.location import LocationService
 
-# Mapping of services to data-sources.
-DATA_SOURCES = {
-    "jhu": JhuLocationService(),
-    "csbs": CSBSLocationService(),
-    "nyt": NYTLocationService(),
-}
+class DataSource:
+    __DATA_SOURCES = {}
 
+    def __init__(self):
+        self.__DATA_SOURCES['jhu'] = JhuLocationService()
+        self.__DATA_SOURCES['csbs'] = CSBSLocationService()
+        self.__DATA_SOURCES['nyt'] = NYTLocationService()
 
-def data_source(source):
-    """
-    Retrieves the provided data-source service.
+    # Mapping of services to data-sources.
+    @classmethod
+    def get_data_source(self, source: str) -> LocationService:
+        return self.__DATA_SOURCES.get(source.lower())
 
-    :returns: The service.
-    :rtype: LocationService
-    """
-    return DATA_SOURCES.get(source.lower())
+    @classmethod
+    def add_data_source(self, source: str, reference_to_source: LocationService) -> None:
+        self.__DATA_SOURCES[source] = reference_to_source
+    
+    @classmethod
+    def get_all_sources(self) -> dict:
+        return self.__DATA_SOURCES
\ No newline at end of file
diff --git a/app/main.py b/app/main.py
index b9aff949..15660f56 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 DataSource
 from .routers import V1, V2
 from .utils.httputils import setup_client_session, teardown_client_session
 
@@ -74,7 +74,9 @@ 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 = DataSource()
+    source.get_data_source(request.query_params.get("source", default = "jhu"))
+    
 
     # Abort with 404 if source cannot be found.
     if not source:
diff --git a/app/routers/v2.py b/app/routers/v2.py
index 31eb408c..d9042623 100644
--- a/app/routers/v2.py
+++ b/app/routers/v2.py
@@ -3,7 +3,7 @@
 
 from fastapi import APIRouter, HTTPException, Request
 
-from ..data import DATA_SOURCES
+from ..data import DataSource
 from ..models import LatestResponse, LocationResponse, LocationsResponse
 
 V2 = APIRouter()
@@ -26,7 +26,7 @@ async def get_latest(
     """
     Getting latest amount of total confirmed cases, deaths, and recoveries.
     """
-    locations = await request.state.source.get_all()
+    locations = await request.state.source.get_all_sources()[source].get_all()
     return {
         "latest": {
             "confirmed": sum(map(lambda location: location.confirmed, locations)),
@@ -57,7 +57,7 @@ async def get_locations(
     params.pop("timelines", None)
 
     # Retrieve all the locations.
-    locations = await request.state.source.get_all()
+    locations = await request.state.source.get_all_sources()[source].get_all()
 
     # Attempt to filter out locations with properties matching the provided query params.
     for key, value in params.items():
@@ -98,7 +98,7 @@ async def get_location_by_id(
     """
     Getting specific location by id.
     """
-    location = await request.state.source.get(id)
+    location = await request.state.source.get_all_sources()[source].get(id)
     return {"location": location.serialize(timelines)}
 
 
@@ -107,4 +107,5 @@ async def sources():
     """
     Retrieves a list of data-sources that are availble to use.
     """
-    return {"sources": list(DATA_SOURCES.keys())}
+    data_source = DataSource()
+    return {"sources": list(data_source.get_all_sources().keys())}
\ No newline at end of file