From 0af3d0c3026205198565a34eaeba61305bbb22ed Mon Sep 17 00:00:00 2001
From: Thiago Castro Ferreira <thiago@aixplain.com>
Date: Thu, 6 Jun 2024 17:54:18 -0300
Subject: [PATCH] Create and get Pipelines with api key as input parameter

---
 aixplain/factories/pipeline_factory.py | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/aixplain/factories/pipeline_factory.py b/aixplain/factories/pipeline_factory.py
index 404a5556..4ebdc439 100644
--- a/aixplain/factories/pipeline_factory.py
+++ b/aixplain/factories/pipeline_factory.py
@@ -73,7 +73,9 @@ def get(cls, pipeline_id: Text, api_key: Optional[Text] = None) -> Pipeline:
         resp = None
         try:
             url = urljoin(cls.backend_url, f"sdk/pipelines/{pipeline_id}")
-            if cls.aixplain_key != "":
+            if api_key is not None:
+                headers = {"Authorization": f"Token {api_key}", "Content-Type": "application/json"}
+            elif cls.aixplain_key != "":
                 headers = {"x-aixplain-key": f"{cls.aixplain_key}", "Content-Type": "application/json"}
             else:
                 headers = {"Authorization": f"Token {config.TEAM_API_KEY}", "Content-Type": "application/json"}
@@ -86,7 +88,7 @@ def get(cls, pipeline_id: Text, api_key: Optional[Text] = None) -> Pipeline:
                 resp["api_key"] = api_key
             pipeline = cls.__from_response(resp)
             return pipeline
-        except Exception as e:
+        except Exception:
             status_code = 400
             if resp is not None and "statusCode" in resp:
                 status_code = resp["statusCode"]
@@ -172,7 +174,7 @@ def list(
         else:
             headers = {"Authorization": f"Token {config.TEAM_API_KEY}", "Content-Type": "application/json"}
 
-        assert 0 < page_size <= 100, f"Pipeline List Error: Page size must be greater than 0 and not exceed 100."
+        assert 0 < page_size <= 100, "Pipeline List Error: Page size must be greater than 0 and not exceed 100."
         payload = {
             "pageSize": page_size,
             "pageNumber": page_number,
@@ -223,13 +225,16 @@ def list(
         return {"results": pipelines, "page_total": page_total, "page_number": page_number, "total": total}
 
     @classmethod
-    def create(cls, name: Text, pipeline: Union[Text, Dict], status: Text = "draft") -> Pipeline:
+    def create(
+        cls, name: Text, pipeline: Union[Text, Dict], status: Text = "draft", api_key: Optional[Text] = None
+    ) -> Pipeline:
         """Pipeline Creation
 
         Args:
             name (Text): Pipeline Name
             pipeline (Union[Text, Dict]): Pipeline as a Python dictionary or in a JSON file
             status (Text, optional): Status of the pipeline. Currently only draft pipelines can be saved. Defaults to "draft".
+            api_key (Optional[Text], optional): _description_. Defaults to None.
 
         Raises:
             Exception: Currently just the creation of draft pipelines are supported
@@ -250,11 +255,12 @@ def create(cls, name: Text, pipeline: Union[Text, Dict], status: Text = "draft")
             # prepare payload
             payload = {"name": name, "status": "draft", "architecture": pipeline}
             url = urljoin(cls.backend_url, "sdk/pipelines")
-            headers = {"Authorization": f"Token {config.TEAM_API_KEY}", "Content-Type": "application/json"}
+            api_key = api_key if api_key is not None else config.TEAM_API_KEY
+            headers = {"Authorization": f"Token {api_key}", "Content-Type": "application/json"}
             logging.info(f"Start service for POST Create Pipeline - {url} - {headers} - {json.dumps(payload)}")
             r = _request_with_retry("post", url, headers=headers, json=payload)
             response = r.json()
 
-            return Pipeline(response["id"], name, config.TEAM_API_KEY)
+            return Pipeline(response["id"], name, api_key)
         except Exception as e:
             raise Exception(e)