Fix Docker Import (#12455)

docker/validate_compose.py

@@ -1,10 +1,16 @@
+"""
+Helper functions used for ingestion of sample data into docker by calling airflow dags
+"""
+
+import time
+from pprint import pprint
 from typing import Tuple
 
-from pprint import pprint
 import requests
+from metadata.utils.logger import log_ansi_encoded_string
+from metadata.cli.docker import REQUESTS_TIMEOUT
 from requests.auth import HTTPBasicAuth
-import time
-from metadata.utils.ansi import print_ansi_encoded_string
+
 
 HEADER_JSON = {"Content-Type": "application/json"}
 BASIC_AUTH = HTTPBasicAuth("admin", "admin")
@@ -16,10 +22,10 @@ def get_last_run_info() -> Tuple[str, str]:
     """
     dag_runs = None
     while not dag_runs:
-        print_ansi_encoded_string(message="Waiting for DAG Run data...")
+        log_ansi_encoded_string(message="Waiting for DAG Run data...")
         time.sleep(5)
         runs = requests.get(
-            "http://localhost:8080/api/v1/dags/sample_data/dagRuns", auth=BASIC_AUTH
+            "http://localhost:8080/api/v1/dags/sample_data/dagRuns", auth=BASIC_AUTH, timeout=REQUESTS_TIMEOUT
         ).json()
         dag_runs = runs.get("dag_runs")
 
@@ -33,6 +39,7 @@ def print_last_run_logs() -> None:
     logs = requests.get(
         "http://localhost:8080/api/v1/openmetadata/last_dag_logs?dag_id=sample_data",
         auth=BASIC_AUTH,
+        timeout=REQUESTS_TIMEOUT
     ).text
     pprint(logs)
 
@@ -42,13 +49,11 @@ def main():
     state = None
     while state != "success":
 
-        print_ansi_encoded_string(
-            message="Waiting for sample data ingestion to be a success. We'll show some logs along the way."
+        log_ansi_encoded_string(
+            message="Waiting for sample data ingestion to be a success. We'll show some logs along the way.",
         )
-
         dag_run_id, state = get_last_run_info()
-        print_ansi_encoded_string(message=f"DAG run: [{dag_run_id}, {state}]")
-
+        log_ansi_encoded_string(message=f"DAG run: [{dag_run_id}, {state}]")
         print_last_run_logs()
         time.sleep(10)
 

ingestion/tests/unit/profiler/pandas/test_profiler.py

@@ -13,7 +13,6 @@
 Test Profiler behavior
 """
 import os
-from concurrent.futures import TimeoutError
 from datetime import datetime
 from unittest import TestCase, mock
 from uuid import uuid4