From 5c4740ac53e984a671a22ec77dd5fbd680f56b75 Mon Sep 17 00:00:00 2001
From: Manthan Gupta <manthangupta109@gmail.com>
Date: Fri, 27 Sep 2024 15:24:30 +0530
Subject: [PATCH] Fix: create api key when phi ws create is invoked

---
 phi/constants.py          |  1 +
 phi/workspace/operator.py | 21 +++++++++++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/phi/constants.py b/phi/constants.py
index c20f76db34..d6be4e0193 100644
--- a/phi/constants.py
+++ b/phi/constants.py
@@ -2,6 +2,7 @@
 PHI_RUNTIME_ENV_VAR: str = "PHI_RUNTIME"
 PHI_API_KEY_ENV_VAR: str = "PHI_API_KEY"
 PHI_WS_KEY_ENV_VAR: str = "PHI_WS_KEY"
+PHI_API_ENDPOINT_ENV_VAR: str = "PHI_API_ENDPOINT"
 
 SCRIPTS_DIR_ENV_VAR: str = "PHI_SCRIPTS_DIR"
 STORAGE_DIR_ENV_VAR: str = "PHI_STORAGE_DIR"
diff --git a/phi/workspace/operator.py b/phi/workspace/operator.py
index 7bf63f3ca9..43043f00be 100644
--- a/phi/workspace/operator.py
+++ b/phi/workspace/operator.py
@@ -1,6 +1,8 @@
+from os import getenv
 from pathlib import Path
 from typing import Optional, Dict, List
 
+import requests
 
 from phi.api.workspace import log_workspace_event
 from phi.api.schemas.workspace import (
@@ -17,6 +19,8 @@
     print_subheading,
     log_config_not_available_msg,
 )
+from phi.cli.credentials import read_auth_token
+from phi.constants import PHI_API_ENDPOINT_ENV_VAR
 from phi.infra.type import InfraType
 from phi.infra.resources import InfraResources
 from phi.workspace.config import WorkspaceConfig
@@ -344,6 +348,23 @@ def setup_workspace(ws_root_path: Path, team: Optional[str] = None) -> bool:
             else:
                 logger.debug("Failed to sync workspace with api. Please setup again")
 
+        ######################################################
+        # 2.3 Create API Key for User Workspace
+        ######################################################
+        if ws_schema:
+            url = getenv(PHI_API_ENDPOINT_ENV_VAR, "") + "v1/user-keys/refresh"
+            headers = {
+                "X-PHIDATA-AUTH-TOKEN": read_auth_token(),
+            }
+            params = {"workspace_url": ws_schema.ws_name}
+            data = {"user": {"id_user": phi_config.user.id_user, "email": phi_config.user.email}}
+            response = requests.post(url, params=params, json=data, headers=headers)
+            logger.info(response.json())
+            if response.status_code == 200:
+                print_info(f"New API Key created for workspace: {ws_schema.ws_name}")
+            else:
+                logger.error(f"Failed to create API Key for workspace: {ws_schema.ws_name}")
+
     if ws_config is not None:
         # logger.debug("Workspace Config: {}".format(ws_config.model_dump_json(indent=2)))
         print_subheading("Setup complete! Next steps:")