Merge pull request #884 from DalgoT4D/877-github-pat-not-saved-to-pre…

…fect-block

update the prefect secret block when the github PAT is updated

ddpui/api/dbt_api.py

@@ -15,8 +15,8 @@
 from ddpui.models.tasks import TaskProgressHashPrefix
 from ddpui.utils.taskprogress import TaskProgress
 from ddpui.ddpdbt import dbt_service
-from ddpui.ddpprefect import DBTCLIPROFILE, prefect_service
-from ddpui.ddpprefect.schema import OrgDbtGitHub, OrgDbtSchema, OrgDbtTarget
+from ddpui.ddpprefect import DBTCLIPROFILE, SECRET, prefect_service
+from ddpui.ddpprefect.schema import OrgDbtGitHub, OrgDbtSchema, OrgDbtTarget, PrefectSecretBlockEdit
 from ddpui.models.org import OrgPrefectBlockv1, Org
 from ddpui.models.org_user import OrgUser, OrgUserResponse
 from ddpui.core.orgdbt_manager import DbtProjectManager
@@ -70,6 +70,43 @@ def put_dbt_github(request, payload: OrgDbtGitHub):
     org.dbt.gitrepo_access_token_secret = payload.gitrepoAccessToken
     org.dbt.save()
 
+    # ignore if token is *******
+    if set(payload.gitrepoAccessToken.strip()) == set("*"):
+        pass
+
+    # if token is empty, delete the secret block
+    elif payload.gitrepoAccessToken in [None, ""]:
+        block_name = f"{orguser.org.slug}-git-pull-url"
+        secret_block = OrgPrefectBlockv1.objects.filter(
+            org=orguser.org, block_type=SECRET, block_name=block_name
+        ).first()
+        if secret_block:
+            prefect_service.delete_secret_block(secret_block.block_id)
+            secret_block.delete()
+
+    # else create / update the secret block
+    else:
+        gitrepo_url = payload.gitrepoUrl.replace(
+            "github.com", "oauth2:" + payload.gitrepoAccessToken + "@github.com"
+        )
+
+        # update the git oauth endpoint with token in the prefect secret block
+        secret_block_edit_params = PrefectSecretBlockEdit(
+            block_name=f"{orguser.org.slug}-git-pull-url",
+            secret=gitrepo_url,
+        )
+
+        response = prefect_service.upsert_secret_block(secret_block_edit_params)
+        if not OrgPrefectBlockv1.objects.filter(
+            org=orguser.org, block_type=SECRET, block_name=secret_block_edit_params.block_name
+        ).exists():
+            OrgPrefectBlockv1.objects.create(
+                org=orguser.org,
+                block_type=SECRET,
+                block_name=secret_block_edit_params.block_name,
+                block_id=response["block_id"],
+            )
+
     org_dir = DbtProjectManager.get_org_dir(org)
 
     task = clone_github_repo.delay(
@@ -104,8 +141,13 @@ def get_dbt_workspace(request):
     if orguser.org.dbt is None:
         return {"error": "no dbt workspace has been configured"}
 
+    secret_block_exists = OrgPrefectBlockv1.objects.filter(
+        org=orguser.org, block_type=SECRET, block_name=f"{orguser.org.slug}-git-pull-url"
+    ).exists()
+
     return {
         "gitrepo_url": orguser.org.dbt.gitrepo_url,
+        "gitrepo_access_token": "*********" if secret_block_exists else None,
         "target_type": orguser.org.dbt.target_type,
         "default_schema": orguser.org.dbt.default_schema,
     }

ddpui/ddpprefect/prefect_service.py

@@ -11,6 +11,7 @@
 from ddpui.ddpprefect.schema import (
     PrefectDataFlowCreateSchema3,
     PrefectSecretBlockCreate,
+    PrefectSecretBlockEdit,
     PrefectShellTaskSetup,
     PrefectDbtTaskSetup,
     PrefectDataFlowUpdateSchema3,
@@ -306,6 +307,15 @@ def create_secret_block(secret_block: PrefectSecretBlockCreate):
     return response
 
 
+def upsert_secret_block(secret_block: PrefectSecretBlockEdit):
+    """This will create a secret block in the prefect to store any password like string"""
+    response = prefect_put(
+        "blocks/secret/",
+        {"blockName": secret_block.block_name, "secret": secret_block.secret},
+    )
+    return response
+
+
 def delete_secret_block(block_id) -> None:
     """Delete secret block in prefect"""
     prefect_delete_a_block(block_id)

ddpui/ddpprefect/schema.py

@@ -212,6 +212,13 @@ class PrefectSecretBlockCreate(Schema):
     block_name: str
 
 
+class PrefectSecretBlockEdit(Schema):
+    """Docstring"""
+
+    secret: str
+    block_name: str
+
+
 class OrgDbtSchema(Schema):
     """Docstring"""
 

ddpui/tests/api_tests/test_dbt_api.py

@@ -20,8 +20,9 @@
     put_dbt_github,
 )
 from ddpui.auth import ACCOUNT_MANAGER_ROLE
+from ddpui.ddpprefect import SECRET
 from ddpui.ddpprefect.schema import DbtProfile, OrgDbtGitHub, OrgDbtSchema
-from ddpui.models.org import Org, OrgDbt
+from ddpui.models.org import Org, OrgDbt, OrgPrefectBlockv1
 from ddpui.models.org_user import OrgUser, OrgUserRole
 from ddpui.models.role_based_access import Permission, Role, RolePermission
 from ddpui.tests.api_tests.test_user_org_api import mock_request, seed_db
@@ -115,22 +116,29 @@ def test_put_dbt_github(orguser):
 
     payload = OrgDbtGitHub(gitrepoUrl="new-url", gitrepoAccessToken="new-access-token")
 
+    OrgPrefectBlockv1.objects.create(
+        org=request.orguser.org,
+        block_type=SECRET,
+        block_name=f"{request.orguser.org.slug}-git-pull-url",
+    )
+
     mocked_task = Mock()
     mocked_task.id = "task-id"
     with patch(
         "ddpui.celeryworkers.tasks.clone_github_repo.delay", return_value=mocked_task
     ) as delay:
         with patch("ddpui.api.dbt_api.dbt_service.check_repo_exists", return_value=True):
-            put_dbt_github(request, payload)
-            delay.assert_called_once_with(
-                "org-slug",
-                "new-url",
-                "new-access-token",
-                os.getenv("CLIENTDBT_ROOT") + "/org-slug",
-                None,
-            )
-            assert request.orguser.org.dbt.gitrepo_url == "new-url"
-            assert request.orguser.org.dbt.gitrepo_access_token_secret == "new-access-token"
+            with patch("ddpui.ddpprefect.prefect_service.upsert_secret_block"):
+                put_dbt_github(request, payload)
+                delay.assert_called_once_with(
+                    "org-slug",
+                    "new-url",
+                    "new-access-token",
+                    os.getenv("CLIENTDBT_ROOT") + "/org-slug",
+                    None,
+                )
+                assert request.orguser.org.dbt.gitrepo_url == "new-url"
+                assert request.orguser.org.dbt.gitrepo_access_token_secret == "new-access-token"
 
 
 def test_dbt_delete_no_org(orguser):

ddpui/tests/services/test_prefect_service.py

@@ -24,7 +24,9 @@
     delete_airbyte_connection_block,
     delete_dbt_core_block,
     PrefectSecretBlockCreate,
+    PrefectSecretBlockEdit,
     create_secret_block,
+    upsert_secret_block,
     delete_secret_block,
     update_dbt_core_block_schema,
     get_flow_runs_by_deployment_id,
@@ -339,6 +341,18 @@ def test_create_secret_block(mock_post: Mock):
     )
 
 
+@patch("ddpui.ddpprefect.prefect_service.prefect_put")
+def test_upsert_secret_block(mock_put: Mock):
+    mock_put.return_value = {"block_id": "block-id"}
+    secret_block = PrefectSecretBlockEdit(block_name="bname", secret="secret")
+    response = upsert_secret_block(secret_block)
+    assert response == {"block_id": "block-id"}
+    mock_put.assert_called_once_with(
+        "blocks/secret/",
+        {"blockName": "bname", "secret": "secret"},
+    )
+
+
 @patch("ddpui.ddpprefect.prefect_service.prefect_delete_a_block")
 def test_delete_secret_block(mock_delete: Mock):
     delete_secret_block("blockid")