From e6ce4324c98d69d41a89db99bd467ce2fc39c86a Mon Sep 17 00:00:00 2001
From: Mayuri Nehate <33225191+mayurinehate@users.noreply.github.com>
Date: Mon, 18 Jul 2022 02:25:20 +0530
Subject: [PATCH] fix(ingest): looker - pass transport options to all api calls
 (#5417)

---
 .../src/datahub/ingestion/source/looker.py               | 9 ++++++++-
 .../src/datahub/ingestion/source/looker_common.py        | 6 +++++-
 .../tests/integration/looker/test_looker.py              | 6 +++++-
 3 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/metadata-ingestion/src/datahub/ingestion/source/looker.py b/metadata-ingestion/src/datahub/ingestion/source/looker.py
index 1712abaa54..f31b22aca1 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/looker.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/looker.py
@@ -1053,7 +1053,11 @@ class LookerDashboardSource(Source):
 
     def _populate_userwise_runs_counts(self, dashboard_usages):
         userwise_count_rows = LookerUtil.run_inline_query(
-            self.client, usage_queries["counts_per_day_per_user_per_dashboard"]
+            self.client,
+            usage_queries["counts_per_day_per_user_per_dashboard"],
+            transport_options=self.source_config.transport_options.get_transport_options()
+            if self.source_config.transport_options is not None
+            else None,
         )
 
         for row in userwise_count_rows:
@@ -1108,6 +1112,9 @@ class LookerDashboardSource(Source):
         count_rows = LookerUtil.run_inline_query(
             self.client,
             usage_queries["counts_per_day_per_dashboard"],
+            transport_options=self.source_config.transport_options.get_transport_options()
+            if self.source_config.transport_options is not None
+            else None,
         )
         for row in count_rows:
             dashboard_usages[
diff --git a/metadata-ingestion/src/datahub/ingestion/source/looker_common.py b/metadata-ingestion/src/datahub/ingestion/source/looker_common.py
index 2535ca2148..34b6424222 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/looker_common.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/looker_common.py
@@ -476,11 +476,14 @@ class LookerUtil:
         )
 
     @staticmethod
-    def run_inline_query(client: Looker31SDK, q: dict) -> List:
+    def run_inline_query(
+        client: Looker31SDK, q: dict, transport_options: Optional[TransportOptions]
+    ) -> List:
 
         response_sql = client.run_inline_query(
             result_format="sql",
             body=LookerUtil.create_query_request(q),
+            transport_options=transport_options,
         )
         logger.debug("=================Query=================")
         logger.debug(response_sql)
@@ -488,6 +491,7 @@ class LookerUtil:
         response_json = client.run_inline_query(
             result_format="json",
             body=LookerUtil.create_query_request(q),
+            transport_options=transport_options,
         )
 
         logger.debug("=================Response=================")
diff --git a/metadata-ingestion/tests/integration/looker/test_looker.py b/metadata-ingestion/tests/integration/looker/test_looker.py
index 598990c35d..53e2ec875f 100644
--- a/metadata-ingestion/tests/integration/looker/test_looker.py
+++ b/metadata-ingestion/tests/integration/looker/test_looker.py
@@ -1,9 +1,11 @@
 import json
 import time
 from datetime import datetime
+from typing import Optional
 from unittest import mock
 
 from freezegun import freeze_time
+from looker_sdk.rtl.transport import TransportOptions
 from looker_sdk.sdk.api31.models import (
     Dashboard,
     DashboardElement,
@@ -294,7 +296,9 @@ def setup_mock_user(mocked_client):
     mocked_client.user.return_value = User(id=1, email="test@looker.com")
 
 
-def side_effect_query_inline(result_format: str, body: WriteQuery) -> str:
+def side_effect_query_inline(
+    result_format: str, body: WriteQuery, transport_options: Optional[TransportOptions]
+) -> str:
     query_type = None
     if result_format == "sql":
         return ""  # Placeholder for sql text