From cd6b28df1a65abff7332ca72792a4ef29dfdf116 Mon Sep 17 00:00:00 2001
From: Mayur Singal <39544459+ulixius9@users.noreply.github.com>
Date: Sat, 5 Mar 2022 01:32:51 +0530
Subject: [PATCH] Fix #3066: Superset Auth Fixed (#3156)

---
 .../metadata/ingestion/ometa/superset_rest.py | 17 ++++++++++++---
 .../src/metadata/ingestion/source/superset.py | 21 ++++++++++---------
 2 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/ingestion/src/metadata/ingestion/ometa/superset_rest.py b/ingestion/src/metadata/ingestion/ometa/superset_rest.py
index 204670a19dc..0aaf161fc16 100644
--- a/ingestion/src/metadata/ingestion/ometa/superset_rest.py
+++ b/ingestion/src/metadata/ingestion/ometa/superset_rest.py
@@ -58,7 +58,13 @@ class SupersetAuthenticationProvider(AuthenticationProvider):
 
     def __init__(self, config: SupersetConfig):
         self.config = config
-        client_config = ClientConfig(base_url=config.url, api_version="api/v1")
+        client_config = ClientConfig(
+            base_url=config.url,
+            api_version="api/v1",
+            auth_token=lambda: ("no_token", 0),
+            auth_header="Authorization",
+            allow_redirects=True,
+        )
         self.client = REST(client_config)
         super().__init__()
 
@@ -69,7 +75,8 @@ class SupersetAuthenticationProvider(AuthenticationProvider):
     def auth_token(self) -> str:
         login_request = self._login_request()
         login_response = self.client.post("/security/login", login_request)
-        return login_response["access_token"]
+        self.generated_auth_token = login_response["access_token"]
+        self.expiry = 0
 
     def _login_request(self) -> str:
         auth_request = {
@@ -80,6 +87,10 @@ class SupersetAuthenticationProvider(AuthenticationProvider):
         }
         return json.dumps(auth_request)
 
+    def get_access_token(self):
+        self.auth_token()
+        return (self.generated_auth_token, self.expiry)
+
 
 class SupersetAPIClient:
     """
@@ -95,7 +106,7 @@ class SupersetAPIClient:
         client_config = ClientConfig(
             base_url=config.url,
             api_version="api/v1",
-            auth_token=f"{self._auth_provider.auth_token()}",
+            auth_token=lambda: self._auth_provider.get_access_token(),
             auth_header="Authorization",
             allow_redirects=True,
         )
diff --git a/ingestion/src/metadata/ingestion/source/superset.py b/ingestion/src/metadata/ingestion/source/superset.py
index 3f7ae60ef87..f4795ac5d12 100644
--- a/ingestion/src/metadata/ingestion/source/superset.py
+++ b/ingestion/src/metadata/ingestion/source/superset.py
@@ -278,17 +278,18 @@ class SupersetSource(Source[Entity]):
                         entity=Lineage_Dashboard,
                         fqdn=f"{self.config.service_name}.{dashboard['id']}",
                     )
-                    lineage = AddLineageRequest(
-                        edge=EntitiesEdge(
-                            fromEntity=EntityReference(
-                                id=from_entity.id.__root__, type="table"
-                            ),
-                            toEntity=EntityReference(
-                                id=to_entity.id.__root__, type="dashboard"
-                            ),
+                    if from_entity and to_entity:
+                        lineage = AddLineageRequest(
+                            edge=EntitiesEdge(
+                                fromEntity=EntityReference(
+                                    id=from_entity.id.__root__, type="table"
+                                ),
+                                toEntity=EntityReference(
+                                    id=to_entity.id.__root__, type="dashboard"
+                                ),
+                            )
                         )
-                    )
-                    yield lineage
+                        yield lineage
 
                 except Exception as err:
                     logger.debug(traceback.print_exc())