[Story]: Oauth - Retreiving oauth keys from db if exists

valmi-io · Feb 7, 2024 · 2338314 · 2338314
1 parent caa45ae
commit 2338314
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 21 deletions.
diff --git a/core/api.py b/core/api.py
@@ -40,6 +40,8 @@
 
 from .models import Account, Connector, Credential, Destination, Source, Sync, User, Workspace, OAuthApiKeys
 
+from valmi_app_backend.utils import replace_values_in_json
+
 
 router = Router()
 
@@ -73,7 +75,7 @@ def list_spaces(request):
 
 @router.get("/workspaces/{workspace_id}/connectors/{connector_type}/spec", response=Json)
 def connector_spec(request, workspace_id, connector_type):
-    workspace = Workspace.objects.get(id=workspace_id)
+    # workspace = Workspace.objects.get(id=workspace_id)
     connector = Connector.objects.get(type=connector_type)
 
     return requests.post(
@@ -88,13 +90,20 @@ def connector_check(request, workspace_id, connector_type, payload: ConnectorCon
     workspace = Workspace.objects.get(id=workspace_id)
     connector = Connector.objects.get(type=connector_type)
 
-    # Replacing Oauth keys
-    oauth_proxy_keys = config("OAUTH_SECRETS", default="", cast=Csv(str))
-    if len(oauth_proxy_keys) > 0:
-        config_str = json.dumps(payload.config)
-        for key in oauth_proxy_keys:
-            config_str = config_str.replace(key, config(key))
-        payload.config = json.loads(config_str)
+    queryset = OAuthApiKeys.objects.filter(workspace=workspace, type=connector_type)
+
+    if queryset.exists():
+        keys = queryset.first()
+        # Replacing oauth keys with db values
+        replace_values_in_json(payload.config, keys.oauth_config)
+    else:
+        # Replacing oauth keys with .env values
+        oauth_proxy_keys = config("OAUTH_SECRETS", default="", cast=Csv(str))
+        if len(oauth_proxy_keys) > 0:
+            config_str = json.dumps(payload.config)
+            for key in oauth_proxy_keys:
+                config_str = config_str.replace(key, config(key))
+            payload.config = json.loads(config_str)
 
     return requests.post(
         f"{CONNECTOR_PREFIX_URL}/{connector.type}/check",
@@ -108,13 +117,22 @@ def connector_discover(request, workspace_id, connector_type, payload: Connector
     workspace = Workspace.objects.get(id=workspace_id)
     connector = Connector.objects.get(type=connector_type)
 
-    # Replacing Oauth keys
-    oauth_proxy_keys = config("OAUTH_SECRETS", default="", cast=Csv(str))
-    if len(oauth_proxy_keys) > 0:
-        config_str = json.dumps(payload.config)
-        for key in oauth_proxy_keys:
-            config_str = config_str.replace(key, config(key))
-        payload.config = json.loads(config_str)
+    queryset = OAuthApiKeys.objects.filter(workspace=workspace, type=connector_type)
+
+    if queryset.exists():
+        keys = queryset.first()
+
+        # Replacing oauth keys with db values
+        replace_values_in_json(payload.config, keys.oauth_config)
+
+    else:
+        # Replacing oauth keys with .env values
+        oauth_proxy_keys = config("OAUTH_SECRETS", default="", cast=Csv(str))
+        if len(oauth_proxy_keys) > 0:
+            config_str = json.dumps(payload.config)
+            for key in oauth_proxy_keys:
+                config_str = config_str.replace(key, config(key))
+            payload.config = json.loads(config_str)
 
     return requests.post(
         f"{CONNECTOR_PREFIX_URL}/{connector.type}/discover",

diff --git a/core/engine_api.py b/core/engine_api.py
@@ -14,11 +14,17 @@
 
 from core.schemas import ConnectorSchema, DetailSchema, SyncSchema
 
-from .models import Connector, Sync
+from .models import (
+    Connector,
+    Sync,
+    OAuthApiKeys
+)
 import json
 
 from opentelemetry.metrics import get_meter_provider
-from opentelemetry import trace
+# from opentelemetry import trace
+
+from valmi_app_backend.utils import replace_values_in_json
 
 router = Router()
 
@@ -47,10 +53,24 @@ def get_all_syncs(request):
             for sync in syncs:
                 src_dst = [sync.source, sync.destination]
                 for obj in src_dst:
-                    config_str = json.dumps(obj.credential.connector_config)
-                    for key in oauth_proxy_keys:
-                        config_str = config_str.replace(key, config(key))
-                    obj.credential.connector_config = json.loads(config_str)
+                    workspace = obj.credential.workspace
+                    connector_type = obj.credential.connector.type
+
+                    queryset = OAuthApiKeys.objects.filter(workspace=workspace, type=connector_type)
+
+                    if queryset.exists():
+
+                        keys = queryset.first()
+
+                        # Replacing oauth keys with db values
+                        replace_values_in_json(obj.credential.connector_config, keys.oauth_config)
+                    else:
+
+                        # Replacing oauth keys with .env values
+                        config_str = json.dumps(obj.credential.connector_config)
+                        for key in oauth_proxy_keys:
+                            config_str = config_str.replace(key, config(key))
+                        obj.credential.connector_config = json.loads(config_str)
 
         return syncs
     except Exception:

diff --git a/valmi_app_backend/utils.py b/valmi_app_backend/utils.py
@@ -12,3 +12,13 @@ class BearerAuthentication(authentication.TokenAuthentication):
     """
 
     keyword = "Bearer"
+
+
+def replace_values_in_json(json_obj, replacements):
+    for key, value in json_obj.items():
+        if isinstance(value, dict):
+            replace_values_in_json(value, replacements)
+        elif key in replacements:
+            replacement_value = replacements.get(key)
+            if replacement_value is not None:
+                json_obj[key] = replacement_value