acryl-datahub 0.15.0.1rc17__py3-none-any.whl → 0.15.0.2rc1__py3-none-any.whl

{acryl_datahub-0.15.0.1rc17.dist-info → acryl_datahub-0.15.0.2rc1.dist-info}/RECORD

@@ -1,4 +1,4 @@
-datahub/__init__.py,sha256=GvEPomdTJt9ZrIZUgZuwaVdBYiJA2qKcUyUKxDy3owo,577
+datahub/__init__.py,sha256=ubEB5EHYmuiGuDtVdhhKbNxGtfw-kzV0eBW81uVifQU,576
 datahub/__main__.py,sha256=pegIvQ9hzK7IhqVeUi1MeADSZ2QlP-D3K0OQdEg55RU,106
 datahub/entrypoints.py,sha256=3-qSfXAx3Z0FEkBV5tlO8fQr4xk4ySeDRMVTpS5Xd6A,7793
 datahub/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -59,7 +59,7 @@ datahub/api/graphql/base.py,sha256=9q637r6v-RGOd8Mk8HW2g0vt9zpqFexsQ5R6TPEHVbs,1
 datahub/api/graphql/operation.py,sha256=h7OXbVRrpJgoth1X4cgeIFhD5JY1MGKg2KjVlQK1gqE,5116
 datahub/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/cli/check_cli.py,sha256=9dXNyzZayHeoFjwFjLkMVyx6DiCZfeESyI-sYtGA6bE,12850
-datahub/cli/cli_utils.py,sha256=gFmcOGAT6IdrTwmpRFSwaqzGmoqS4dbWrxILB1uvlGk,13214
+datahub/cli/cli_utils.py,sha256=d_Q9vPZTPxO7XyyghD-i1Nkr4DX0M8cs2IWrMUQAu0c,13539
 datahub/cli/config_utils.py,sha256=yuXw7RzpRY5x_-MAoqWbv46qUkIeRNAJL4_OeJpYdBE,4879
 datahub/cli/delete_cli.py,sha256=VLeHi7MLFCtTk7MI4y8r_k_7aLcCUZIglU2MNLsXU6M,23051
 datahub/cli/docker_check.py,sha256=rED4wHXqxcQ_qNFyIgFEZ85BHT9ZTE5YC-oUKqbRqi0,9432
@@ -119,7 +119,7 @@ datahub/emitter/mcp.py,sha256=hAAYziDdkwjazQU0DtWMbQWY8wS09ACrKJbqxoWXdgc,9637
 datahub/emitter/mcp_builder.py,sha256=eOcuz41c4a3oTkNk39yYl9bTxpksxqATPHLcqyhPGT0,9856
 datahub/emitter/mcp_patch_builder.py,sha256=oonC8iGOvDzqj890CxOjWlBdDEF1RnwvbSZy1sivlTY,4572
 datahub/emitter/request_helper.py,sha256=33ORG3S3OVy97_jlWBRn7yUM5XCIkRN6WSdJvN7Ofcg,670
-datahub/emitter/rest_emitter.py,sha256=YpRQEyuDBq31Iw7bZtOe5arm4YviCQLpvaObzVwheBY,16759
+datahub/emitter/rest_emitter.py,sha256=O9IJ7r-AXL4Pi892pEFOygvUKTbD8V6ey8KObuqHqgk,17876
 datahub/emitter/serialization_helper.py,sha256=q12Avmf70Vy4ttQGMJoTKlE5EsybMKNg2w3MQeZiHvk,3652
 datahub/emitter/sql_parsing_builder.py,sha256=Cr5imZrm3dYDSCACt5MFscgHCtVbHTD6IjUmsvsKoEs,11991
 datahub/emitter/synchronized_file_emitter.py,sha256=s4ATuxalI4GDAkrZTaGSegxBdvvNPZ9jRSdtElU0kNs,1805
@@ -164,8 +164,8 @@ datahub/ingestion/glossary/classifier.py,sha256=zp8Fe3he80H5Zz1EwymKjThUPkTpw6Pg
 datahub/ingestion/glossary/classifier_registry.py,sha256=yFOYLQhDgCLqXYMG3L1BquXafeLcZDcmp8meyw6k9ts,307
 datahub/ingestion/glossary/datahub_classifier.py,sha256=8VhwuLDhyOqqOr0jqAPIgorb4eAOnvTr4m13Y2Wy1-E,7515
 datahub/ingestion/graph/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datahub/ingestion/graph/client.py,sha256=AYDFwP9a_M-fCZv-PcWMSr5tc53XWJl372SWKwdu37E,64651
-datahub/ingestion/graph/config.py,sha256=3b_Gxa5wcBnphP63bBiAFdWS7PJhUHRE1WZL_q4Cw8k,749
+datahub/ingestion/graph/client.py,sha256=R50K7NmE3TYgVXvdLnvLZn7N0fkiCXOK0MoJz9ueglA,64963
+datahub/ingestion/graph/config.py,sha256=_oha8Je7P80ZmrkZUAaRHyYbdMmTkMI5JkYjEP2Ri1Q,751
 datahub/ingestion/graph/connections.py,sha256=9462L0ZWGKURyypAln25eMPhK3pcufBar9tNDoqspXs,741
 datahub/ingestion/graph/filters.py,sha256=UeUZQHoimavIYx-jXLA0WGkOUe10TaO8uEZkfa-QgNE,6188
 datahub/ingestion/reporting/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -217,7 +217,7 @@ datahub/ingestion/source/abs/report.py,sha256=fzkTdTewYlWrTk4f2Cyl-e8RV4qw9wEVtm
 datahub/ingestion/source/abs/source.py,sha256=pzxW-R_cWGKPneEhX8JWdTZiX2k1kAZOPKgMxp9mAEI,24533
 datahub/ingestion/source/aws/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/aws/aws_common.py,sha256=DfdQgkJ_s2isFx8WvqKTlAcBk4KE8SgfpmA5BgC3fgY,17716
-datahub/ingestion/source/aws/glue.py,sha256=r7y1MPDK__BKX_mrJjVa_CEmSXM3Pa02gt19o0sSLE8,56815
+datahub/ingestion/source/aws/glue.py,sha256=lJW3QHHz1_SWqLEB-vUSTxSuL0EgUQ0ptdQns_NLNds,57343
 datahub/ingestion/source/aws/s3_boto_utils.py,sha256=Wyp9k9tapsCuw9dyH4FCXJr_wmeLaYFoCtKvrV6SEDk,3892
 datahub/ingestion/source/aws/s3_util.py,sha256=OFypcgmVC6jnZM90-gjcPpAMtTV1lbnreCaMhCzNlzs,2149
 datahub/ingestion/source/aws/sagemaker.py,sha256=Bl2tkBYnrindgx61VHYgNovUF_Kp_fXNcivQn28vC2w,5254
@@ -233,10 +233,10 @@ datahub/ingestion/source/azure/abs_folder_utils.py,sha256=7skXus-4fSIoKpqCeU-GG0
 datahub/ingestion/source/azure/abs_utils.py,sha256=KdAlCK-PMrn35kFHxz5vrsjajyx2PD5GRgoBKdoRvcg,2075
 datahub/ingestion/source/azure/azure_common.py,sha256=Zl0pPuE6L3QcM5B1P0LsPthZmD0h7fUUS0kg2okl6IY,4053
 datahub/ingestion/source/bigquery_v2/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datahub/ingestion/source/bigquery_v2/bigquery.py,sha256=YMsyj6s7fggzisWfDdbT4w1MKJ3eRdNERsCShnu0Zqo,13681
+datahub/ingestion/source/bigquery_v2/bigquery.py,sha256=c7g8sWuDIMhCSAX0D76P2arxZgTmzd-e0qlO7yt_zJY,13841
 datahub/ingestion/source/bigquery_v2/bigquery_audit.py,sha256=IlbHA8a-gNJvnubgBfxVHpUk8rFNIG80gk5HWXa2lyE,25108
 datahub/ingestion/source/bigquery_v2/bigquery_audit_log_api.py,sha256=LuGJ6LgPViLIfDQfylxlQ3CA7fZYM5MDt8M-7sfzm84,5096
-datahub/ingestion/source/bigquery_v2/bigquery_config.py,sha256=PqbYNqM4-KukCU1meuvsk0qbiWa7UFh5hqHrHsvOSWQ,25889
+datahub/ingestion/source/bigquery_v2/bigquery_config.py,sha256=sjCW997-Su14cVgWd1ZVx1E67yqfTIV5Wjp9Me0hfOw,26289
 datahub/ingestion/source/bigquery_v2/bigquery_data_reader.py,sha256=DeT3v_Z82__8En0FcZ0kavBAWQoRvSZ5Rppm9eeDAb8,2393
 datahub/ingestion/source/bigquery_v2/bigquery_helper.py,sha256=QER3gY8e_k1_eNVj7cBso7ZzrWl_vO5PYSa6CpvqNx8,1554
 datahub/ingestion/source/bigquery_v2/bigquery_platform_resource_helper.py,sha256=8nuQ8hMuJEswWDZtV2RjbK8RvDJUzT_S74dnyPpGFdQ,4857
@@ -305,7 +305,7 @@ datahub/ingestion/source/gc/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMp
 datahub/ingestion/source/gc/datahub_gc.py,sha256=W6uoeV7B4WIXdxT4tOEdDksdJm656WwwvkH79L7f_8Q,12969
 datahub/ingestion/source/gc/dataprocess_cleanup.py,sha256=86Tm3NNWMf0xM4TklNIEeNOjEingKpYy-XvCPeaAb4k,17125
 datahub/ingestion/source/gc/execution_request_cleanup.py,sha256=9jsyCIspWSSYSAVPHjKHr05885rXxM6FCH7KzTBceic,10139
-datahub/ingestion/source/gc/soft_deleted_entity_cleanup.py,sha256=2JpESfsqoJRdLskV3AHYU8nRj_NvNtIaLZ4_RRNIod4,11229
+datahub/ingestion/source/gc/soft_deleted_entity_cleanup.py,sha256=zRtgC_AcZui4qGf9jBASI3R-CrYZxNe3Pm-gNSLT3rw,11420
 datahub/ingestion/source/gcs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/gcs/gcs_source.py,sha256=iwvj4JwjyVWRP1Vq106sUtQhh0GuOYVSu9zCa1wCZN0,6189
 datahub/ingestion/source/gcs/gcs_utils.py,sha256=_78KM863XXgkVLmZLtYGF5PJNnZas1go-XRtOq-79lo,1047
@@ -429,7 +429,7 @@ datahub/ingestion/source/snowflake/constants.py,sha256=22n-0r04nuy-ImxWFFpmbrt_G
 datahub/ingestion/source/snowflake/oauth_config.py,sha256=ol9D3RmruGStJAeL8PYSQguSqcD2HfkjPkMF2AB_eZs,1277
 datahub/ingestion/source/snowflake/oauth_generator.py,sha256=fu2VnREGuJXeTqIV2jx4TwieVnznf83HQkrE0h2DGGM,3423
 datahub/ingestion/source/snowflake/snowflake_assertion.py,sha256=_l3k4aI9wvioE81xxdeizJn9nJCZ_nMIXgk9N6pEk5o,4803
-datahub/ingestion/source/snowflake/snowflake_config.py,sha256=jQGSa7ZQs3EsXB9ANShZ4xv9RqrhRfVHRSLeFiDwwxc,17974
+datahub/ingestion/source/snowflake/snowflake_config.py,sha256=UehWUvqTXRsWmE5bBS53IoLjUL06-wJq6K4O2MTT2R8,18374
 datahub/ingestion/source/snowflake/snowflake_connection.py,sha256=yzv-01FdmfDSCJY5rqKNNodXxzg3SS5DF7oA4WXArOA,17793
 datahub/ingestion/source/snowflake/snowflake_data_reader.py,sha256=ffR5E2uhD71FUMXd3XOg2rHwrp1rbbGEFTAbqKcmI2s,2195
 datahub/ingestion/source/snowflake/snowflake_lineage_v2.py,sha256=FBmiONx4EGHWV8RNJT6zHZyntKinPFFyd2oKbTUIbhE,21319
@@ -444,7 +444,7 @@ datahub/ingestion/source/snowflake/snowflake_summary.py,sha256=kTmuCtRnvHqM8WBYh
 datahub/ingestion/source/snowflake/snowflake_tag.py,sha256=fyfWmFVz2WZrpTJWNIe9m0WpDHgeFrGPf8diORJZUwo,6212
 datahub/ingestion/source/snowflake/snowflake_usage_v2.py,sha256=0rXgz8bvRiI9SYVMa0UGLeg_DcjqBy6kQsdq0Uq0HVk,24685
 datahub/ingestion/source/snowflake/snowflake_utils.py,sha256=MoI8-DR9tuMuHMBQcpDo4GFjvcoQZWLNkdFZsTkgK-M,12786
-datahub/ingestion/source/snowflake/snowflake_v2.py,sha256=hIWtzlxuSQ_3w48o4AF2l9CQOcWIe6AmD07I89sH2B0,31860
+datahub/ingestion/source/snowflake/snowflake_v2.py,sha256=0doZaPPMO64Qi9uN4w8ZYe3gKkkieGJKI5xntF7vS6w,32020
 datahub/ingestion/source/sql/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/sql/athena.py,sha256=G3cIY8H_76lIUAzQWW2kLnZOEsfbakmojxbiHb3dYZ8,24059
 datahub/ingestion/source/sql/clickhouse.py,sha256=jzvaXP5Wr0SMhj2rtuvVE821xnfpKiXhO3cm0xblgHs,27299
@@ -986,8 +986,8 @@ datahub_provider/operators/datahub_assertion_operator.py,sha256=uvTQ-jk2F0sbqqxp
 datahub_provider/operators/datahub_assertion_sensor.py,sha256=lCBj_3x1cf5GMNpHdfkpHuyHfVxsm6ff5x2Z5iizcAo,140
 datahub_provider/operators/datahub_operation_operator.py,sha256=aevDp2FzX7FxGlXrR0khoHNbxbhKR2qPEX5e8O2Jyzw,174
 datahub_provider/operators/datahub_operation_sensor.py,sha256=8fcdVBCEPgqy1etTXgLoiHoJrRt_nzFZQMdSzHqSG7M,168
-acryl_datahub-0.15.0.1rc17.dist-info/METADATA,sha256=dV3uL4e_h1lxLVGl_a9Cmnm5pBi8kgHTbR56ypsqjC8,173444
-acryl_datahub-0.15.0.1rc17.dist-info/WHEEL,sha256=A3WOREP4zgxI0fKrHUG8DC8013e3dK3n7a6HDbcEIwE,91
-acryl_datahub-0.15.0.1rc17.dist-info/entry_points.txt,sha256=xnPSPLK3bJGADxe4TDS4wL4u0FT_PGlahDa-ENYdYCQ,9512
-acryl_datahub-0.15.0.1rc17.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
-acryl_datahub-0.15.0.1rc17.dist-info/RECORD,,
+acryl_datahub-0.15.0.2rc1.dist-info/METADATA,sha256=cLoMIIavfob5z2bJbsj69AoG_J7tp489qvVvBf3W_Yo,173441
+acryl_datahub-0.15.0.2rc1.dist-info/WHEEL,sha256=A3WOREP4zgxI0fKrHUG8DC8013e3dK3n7a6HDbcEIwE,91
+acryl_datahub-0.15.0.2rc1.dist-info/entry_points.txt,sha256=xnPSPLK3bJGADxe4TDS4wL4u0FT_PGlahDa-ENYdYCQ,9512
+acryl_datahub-0.15.0.2rc1.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
+acryl_datahub-0.15.0.2rc1.dist-info/RECORD,,

datahub/__init__.py

@@ -3,7 +3,7 @@ import warnings
 
 # Published at https://pypi.org/project/acryl-datahub/.
 __package_name__ = "acryl-datahub"
-__version__ = "0.15.0.1rc17"
+__version__ = "0.15.0.2rc1"
 
 
 def is_dev_mode() -> bool:

datahub/cli/cli_utils.py

@@ -3,7 +3,7 @@ import logging
 import time
 import typing
 from datetime import datetime
-from typing import Any, Dict, List, Optional, Tuple, Type, Union
+from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar, Union
 
 import click
 import requests
@@ -33,6 +33,15 @@ def first_non_null(ls: List[Optional[str]]) -> Optional[str]:
     return next((el for el in ls if el is not None and el.strip() != ""), None)
 
 
+_T = TypeVar("_T")
+
+
+def get_or_else(value: Optional[_T], default: _T) -> _T:
+    # Normally we'd use `value or default`. However, that runs into issues
+    # when value is falsey but not None.
+    return value if value is not None else default
+
+
 def parse_run_restli_response(response: requests.Response) -> dict:
     response_json = response.json()
     if response.status_code != 200:
@@ -321,6 +330,8 @@ def get_frontend_session_login_as(
 def _ensure_valid_gms_url_acryl_cloud(url: str) -> str:
     if "acryl.io" not in url:
         return url
+    if url.endswith(":8080"):
+        url = url.replace(":8080", "")
     if url.startswith("http://"):
         url = url.replace("http://", "https://")
     if url.endswith("acryl.io"):

datahub/emitter/rest_emitter.py

@@ -1,9 +1,21 @@
+from __future__ import annotations
+
 import functools
 import json
 import logging
 import os
 from json.decoder import JSONDecodeError
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+)
 
 import requests
 from deprecated import deprecated
@@ -12,9 +24,13 @@ from requests.exceptions import HTTPError, RequestException
 
 from datahub import nice_version_name
 from datahub.cli import config_utils
-from datahub.cli.cli_utils import ensure_has_system_metadata, fixup_gms_url
+from datahub.cli.cli_utils import ensure_has_system_metadata, fixup_gms_url, get_or_else
 from datahub.cli.env_utils import get_boolean_env_variable
-from datahub.configuration.common import ConfigurationError, OperationalError
+from datahub.configuration.common import (
+    ConfigModel,
+    ConfigurationError,
+    OperationalError,
+)
 from datahub.emitter.generic_emitter import Emitter
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.request_helper import make_curl_command
@@ -31,10 +47,8 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
-_DEFAULT_CONNECT_TIMEOUT_SEC = 30  # 30 seconds should be plenty to connect
-_DEFAULT_READ_TIMEOUT_SEC = (
-    30  # Any ingest call taking longer than 30 seconds should be abandoned
-)
+_DEFAULT_TIMEOUT_SEC = 30  # 30 seconds should be plenty to connect
+_TIMEOUT_LOWER_BOUND_SEC = 1  # if below this, we log a warning
 _DEFAULT_RETRY_STATUS_CODES = [  # Additional status codes to retry on
     429,
     500,
@@ -63,15 +77,76 @@ BATCH_INGEST_MAX_PAYLOAD_LENGTH = int(
 )
 
 
+class RequestsSessionConfig(ConfigModel):
+    timeout: Union[float, Tuple[float, float], None] = _DEFAULT_TIMEOUT_SEC
+
+    retry_status_codes: List[int] = _DEFAULT_RETRY_STATUS_CODES
+    retry_methods: List[str] = _DEFAULT_RETRY_METHODS
+    retry_max_times: int = _DEFAULT_RETRY_MAX_TIMES
+
+    extra_headers: Dict[str, str] = {}
+
+    ca_certificate_path: Optional[str] = None
+    client_certificate_path: Optional[str] = None
+    disable_ssl_verification: bool = False
+
+    def build_session(self) -> requests.Session:
+        session = requests.Session()
+
+        if self.extra_headers:
+            session.headers.update(self.extra_headers)
+
+        if self.client_certificate_path:
+            session.cert = self.client_certificate_path
+
+        if self.ca_certificate_path:
+            session.verify = self.ca_certificate_path
+
+        if self.disable_ssl_verification:
+            session.verify = False
+
+        try:
+            # Set raise_on_status to False to propagate errors:
+            # https://stackoverflow.com/questions/70189330/determine-status-code-from-python-retry-exception
+            # Must call `raise_for_status` after making a request, which we do
+            retry_strategy = Retry(
+                total=self.retry_max_times,
+                status_forcelist=self.retry_status_codes,
+                backoff_factor=2,
+                allowed_methods=self.retry_methods,
+                raise_on_status=False,
+            )
+        except TypeError:
+            # Prior to urllib3 1.26, the Retry class used `method_whitelist` instead of `allowed_methods`.
+            retry_strategy = Retry(
+                total=self.retry_max_times,
+                status_forcelist=self.retry_status_codes,
+                backoff_factor=2,
+                method_whitelist=self.retry_methods,
+                raise_on_status=False,
+            )
+
+        adapter = HTTPAdapter(
+            pool_connections=100, pool_maxsize=100, max_retries=retry_strategy
+        )
+        session.mount("http://", adapter)
+        session.mount("https://", adapter)
+
+        if self.timeout is not None:
+            # Shim session.request to apply default timeout values.
+            # Via https://stackoverflow.com/a/59317604.
+            session.request = functools.partial(  # type: ignore
+                session.request,
+                timeout=self.timeout,
+            )
+
+        return session
+
+
 class DataHubRestEmitter(Closeable, Emitter):
     _gms_server: str
     _token: Optional[str]
     _session: requests.Session
-    _connect_timeout_sec: float = _DEFAULT_CONNECT_TIMEOUT_SEC
-    _read_timeout_sec: float = _DEFAULT_READ_TIMEOUT_SEC
-    _retry_status_codes: List[int] = _DEFAULT_RETRY_STATUS_CODES
-    _retry_methods: List[str] = _DEFAULT_RETRY_METHODS
-    _retry_max_times: int = _DEFAULT_RETRY_MAX_TIMES
 
     def __init__(
         self,
@@ -102,15 +177,13 @@ class DataHubRestEmitter(Closeable, Emitter):
 
         self._session = requests.Session()
 
-        self._session.headers.update(
-            {
-                "X-RestLi-Protocol-Version": "2.0.0",
-                "X-DataHub-Py-Cli-Version": nice_version_name(),
-                "Content-Type": "application/json",
-            }
-        )
+        headers = {
+            "X-RestLi-Protocol-Version": "2.0.0",
+            "X-DataHub-Py-Cli-Version": nice_version_name(),
+            "Content-Type": "application/json",
+        }
         if token:
-            self._session.headers.update({"Authorization": f"Bearer {token}"})
+            headers["Authorization"] = f"Bearer {token}"
         else:
             # HACK: When no token is provided but system auth env variables are set, we use them.
             # Ideally this should simply get passed in as config, instead of being sneakily injected
@@ -119,75 +192,43 @@ class DataHubRestEmitter(Closeable, Emitter):
             # rest emitter, and the rest sink uses the rest emitter under the hood.
             system_auth = config_utils.get_system_auth()
             if system_auth is not None:
-                self._session.headers.update({"Authorization": system_auth})
-
-        if extra_headers:
-            self._session.headers.update(extra_headers)
-
-        if client_certificate_path:
-            self._session.cert = client_certificate_path
-
-        if ca_certificate_path:
-            self._session.verify = ca_certificate_path
-
-        if disable_ssl_verification:
-            self._session.verify = False
-
-        self._connect_timeout_sec = (
-            connect_timeout_sec or timeout_sec or _DEFAULT_CONNECT_TIMEOUT_SEC
-        )
-        self._read_timeout_sec = (
-            read_timeout_sec or timeout_sec or _DEFAULT_READ_TIMEOUT_SEC
-        )
-
-        if self._connect_timeout_sec < 1 or self._read_timeout_sec < 1:
-            logger.warning(
-                f"Setting timeout values lower than 1 second is not recommended. Your configuration is connect_timeout:{self._connect_timeout_sec}s, read_timeout:{self._read_timeout_sec}s"
-            )
-
-        if retry_status_codes is not None:  # Only if missing. Empty list is allowed
-            self._retry_status_codes = retry_status_codes
-
-        if retry_methods is not None:
-            self._retry_methods = retry_methods
-
-        if retry_max_times:
-            self._retry_max_times = retry_max_times
+                headers["Authorization"] = system_auth
 
-        try:
-            # Set raise_on_status to False to propagate errors:
-            # https://stackoverflow.com/questions/70189330/determine-status-code-from-python-retry-exception
-            # Must call `raise_for_status` after making a request, which we do
-            retry_strategy = Retry(
-                total=self._retry_max_times,
-                status_forcelist=self._retry_status_codes,
-                backoff_factor=2,
-                allowed_methods=self._retry_methods,
-                raise_on_status=False,
-            )
-        except TypeError:
-            # Prior to urllib3 1.26, the Retry class used `method_whitelist` instead of `allowed_methods`.
-            retry_strategy = Retry(
-                total=self._retry_max_times,
-                status_forcelist=self._retry_status_codes,
-                backoff_factor=2,
-                method_whitelist=self._retry_methods,
-                raise_on_status=False,
+        timeout: float | tuple[float, float]
+        if connect_timeout_sec is not None or read_timeout_sec is not None:
+            timeout = (
+                connect_timeout_sec or timeout_sec or _DEFAULT_TIMEOUT_SEC,
+                read_timeout_sec or timeout_sec or _DEFAULT_TIMEOUT_SEC,
             )
+            if (
+                timeout[0] < _TIMEOUT_LOWER_BOUND_SEC
+                or timeout[1] < _TIMEOUT_LOWER_BOUND_SEC
+            ):
+                logger.warning(
+                    f"Setting timeout values lower than {_TIMEOUT_LOWER_BOUND_SEC} second is not recommended. Your configuration is (connect_timeout, read_timeout) = {timeout} seconds"
+                )
+        else:
+            timeout = get_or_else(timeout_sec, _DEFAULT_TIMEOUT_SEC)
+            if timeout < _TIMEOUT_LOWER_BOUND_SEC:
+                logger.warning(
+                    f"Setting timeout values lower than {_TIMEOUT_LOWER_BOUND_SEC} second is not recommended. Your configuration is timeout = {timeout} seconds"
+                )
 
-        adapter = HTTPAdapter(
-            pool_connections=100, pool_maxsize=100, max_retries=retry_strategy
-        )
-        self._session.mount("http://", adapter)
-        self._session.mount("https://", adapter)
-
-        # Shim session.request to apply default timeout values.
-        # Via https://stackoverflow.com/a/59317604.
-        self._session.request = functools.partial(  # type: ignore
-            self._session.request,
-            timeout=(self._connect_timeout_sec, self._read_timeout_sec),
+        self._session_config = RequestsSessionConfig(
+            timeout=timeout,
+            retry_status_codes=get_or_else(
+                retry_status_codes, _DEFAULT_RETRY_STATUS_CODES
+            ),
+            retry_methods=get_or_else(retry_methods, _DEFAULT_RETRY_METHODS),
+            retry_max_times=get_or_else(retry_max_times, _DEFAULT_RETRY_MAX_TIMES),
+            extra_headers={**headers, **(extra_headers or {})},
+            ca_certificate_path=ca_certificate_path,
+            client_certificate_path=client_certificate_path,
+            disable_ssl_verification=disable_ssl_verification,
         )
 
+        self._session = self._session_config.build_session()
+
     def test_connection(self) -> None:
         url = f"{self._gms_server}/config"
         response = self._session.get(url)

datahub/ingestion/graph/client.py

@@ -179,21 +179,24 @@ class DataHubGraph(DatahubRestEmitter):
 
     @classmethod
     def from_emitter(cls, emitter: DatahubRestEmitter) -> "DataHubGraph":
+        session_config = emitter._session_config
+        if isinstance(session_config.timeout, tuple):
+            # TODO: This is slightly lossy. Eventually, we want to modify the emitter
+            # to accept a tuple for timeout_sec, and then we'll be able to remove this.
+            timeout_sec: Optional[float] = session_config.timeout[0]
+        else:
+            timeout_sec = session_config.timeout
         return cls(
             DatahubClientConfig(
                 server=emitter._gms_server,
                 token=emitter._token,
-                timeout_sec=emitter._read_timeout_sec,
-                retry_status_codes=emitter._retry_status_codes,
-                retry_max_times=emitter._retry_max_times,
-                extra_headers=emitter._session.headers,
-                disable_ssl_verification=emitter._session.verify is False,
-                ca_certificate_path=(
-                    emitter._session.verify
-                    if isinstance(emitter._session.verify, str)
-                    else None
-                ),
-                client_certificate_path=emitter._session.cert,
+                timeout_sec=timeout_sec,
+                retry_status_codes=session_config.retry_status_codes,
+                retry_max_times=session_config.retry_max_times,
+                extra_headers=session_config.extra_headers,
+                disable_ssl_verification=session_config.disable_ssl_verification,
+                ca_certificate_path=session_config.ca_certificate_path,
+                client_certificate_path=session_config.client_certificate_path,
             )
         )
 

datahub/ingestion/graph/config.py

@@ -10,7 +10,7 @@ class DatahubClientConfig(ConfigModel):
     # by callers / the CLI, but the actual client should not have any magic.
     server: str
     token: Optional[str] = None
-    timeout_sec: Optional[int] = None
+    timeout_sec: Optional[float] = None
     retry_status_codes: Optional[List[int]] = None
     retry_max_times: Optional[int] = None
     extra_headers: Optional[Dict[str, str]] = None

datahub/ingestion/source/aws/glue.py

@@ -1054,49 +1054,66 @@ class GlueSource(StatefulIngestionSourceBase):
             yield from self.gen_database_containers(database)
 
         for table in tables:
-            database_name = table["DatabaseName"]
             table_name = table["Name"]
-            full_table_name = f"{database_name}.{table_name}"
-            self.report.report_table_scanned()
-            if not self.source_config.database_pattern.allowed(
-                database_name
-            ) or not self.source_config.table_pattern.allowed(full_table_name):
-                self.report.report_table_dropped(full_table_name)
-                continue
+            try:
+                yield from self._gen_table_wu(table=table)
+            except KeyError as e:
+                self.report.report_failure(
+                    message="Failed to extract workunit for table",
+                    context=f"Table: {table_name}",
+                    exc=e,
+                )
+        if self.extract_transforms:
+            yield from self._transform_extraction()
 
-            dataset_urn = make_dataset_urn_with_platform_instance(
-                platform=self.platform,
-                name=full_table_name,
-                env=self.env,
-                platform_instance=self.source_config.platform_instance,
-            )
+    def _gen_table_wu(self, table: Dict) -> Iterable[MetadataWorkUnit]:
+        database_name = table["DatabaseName"]
+        table_name = table["Name"]
+        full_table_name = f"{database_name}.{table_name}"
+        self.report.report_table_scanned()
+        if not self.source_config.database_pattern.allowed(
+            database_name
+        ) or not self.source_config.table_pattern.allowed(full_table_name):
+            self.report.report_table_dropped(full_table_name)
+            return
+
+        dataset_urn = make_dataset_urn_with_platform_instance(
+            platform=self.platform,
+            name=full_table_name,
+            env=self.env,
+            platform_instance=self.source_config.platform_instance,
+        )
 
-            mce = self._extract_record(dataset_urn, table, full_table_name)
-            yield MetadataWorkUnit(full_table_name, mce=mce)
+        mce = self._extract_record(dataset_urn, table, full_table_name)
+        yield MetadataWorkUnit(full_table_name, mce=mce)
 
-            # We also want to assign "table" subType to the dataset representing glue table - unfortunately it is not
-            # possible via Dataset snapshot embedded in a mce, so we have to generate a mcp.
-            yield MetadataChangeProposalWrapper(
-                entityUrn=dataset_urn,
-                aspect=SubTypes(typeNames=[DatasetSubTypes.TABLE]),
-            ).as_workunit()
+        # We also want to assign "table" subType to the dataset representing glue table - unfortunately it is not
+        # possible via Dataset snapshot embedded in a mce, so we have to generate a mcp.
+        yield MetadataChangeProposalWrapper(
+            entityUrn=dataset_urn,
+            aspect=SubTypes(typeNames=[DatasetSubTypes.TABLE]),
+        ).as_workunit()
 
-            yield from self._get_domain_wu(
-                dataset_name=full_table_name,
-                entity_urn=dataset_urn,
-            )
-            yield from self.add_table_to_database_container(
-                dataset_urn=dataset_urn, db_name=database_name
-            )
+        yield from self._get_domain_wu(
+            dataset_name=full_table_name,
+            entity_urn=dataset_urn,
+        )
+        yield from self.add_table_to_database_container(
+            dataset_urn=dataset_urn, db_name=database_name
+        )
 
-            wu = self.get_lineage_if_enabled(mce)
-            if wu:
-                yield wu
+        wu = self.get_lineage_if_enabled(mce)
+        if wu:
+            yield wu
 
+        try:
             yield from self.get_profile_if_enabled(mce, database_name, table_name)
-
-        if self.extract_transforms:
-            yield from self._transform_extraction()
+        except KeyError as e:
+            self.report.report_failure(
+                message="Failed to extract profile for table",
+                context=f"Table: {dataset_urn}",
+                exc=e,
+            )
 
     def _transform_extraction(self) -> Iterable[MetadataWorkUnit]:
         dags: Dict[str, Optional[Dict[str, Any]]] = {}

datahub/ingestion/source/bigquery_v2/bigquery.py

@@ -281,6 +281,8 @@ class BigqueryV2Source(StatefulIngestionSourceBase, TestableSource):
                     include_lineage=self.config.include_table_lineage,
                     include_usage_statistics=self.config.include_usage_statistics,
                     include_operations=self.config.usage.include_operational_stats,
+                    include_queries=self.config.include_queries,
+                    include_query_usage_statistics=self.config.include_query_usage_statistics,
                     top_n_queries=self.config.usage.top_n_queries,
                     region_qualifiers=self.config.region_qualifiers,
                 ),

datahub/ingestion/source/bigquery_v2/bigquery_config.py

@@ -447,6 +447,14 @@ class BigQueryV2Config(
         default=False,
         description="If enabled, uses the new queries extractor to extract queries from bigquery.",
     )
+    include_queries: bool = Field(
+        default=True,
+        description="If enabled, generate query entities associated with lineage edges. Only applicable if `use_queries_v2` is enabled.",
+    )
+    include_query_usage_statistics: bool = Field(
+        default=True,
+        description="If enabled, generate query popularity statistics. Only applicable if `use_queries_v2` is enabled.",
+    )
 
     @property
     def have_table_data_read_permission(self) -> bool:

datahub/ingestion/source/gc/soft_deleted_entity_cleanup.py

@@ -19,8 +19,8 @@ from datahub.utilities.urns._urn_base import Urn
 
 logger = logging.getLogger(__name__)
 
-QUERY_QUERY_ENTITY = """
-query listQueries($input: ScrollAcrossEntitiesInput!) {
+QUERY_ENTITIES = """
+query listEntities($input: ScrollAcrossEntitiesInput!) {
   scrollAcrossEntities(input: $input) {
     nextScrollId
     count
@@ -29,6 +29,9 @@ query listQueries($input: ScrollAcrossEntitiesInput!) {
         ... on QueryEntity {
           urn
         }
+        ... on DataProcessInstance {
+          urn
+        }
       }
     }
   }
@@ -225,16 +228,16 @@ class SoftDeletedEntitiesCleanup:
                     time.sleep(self.config.delay)
         return futures
 
-    def _get_soft_deleted_queries(self) -> Iterable[str]:
+    def _get_soft_deleted(self, graphql_query: str, entity_type: str) -> Iterable[str]:
         assert self.ctx.graph
         scroll_id: Optional[str] = None
         while True:
             try:
                 result = self.ctx.graph.execute_graphql(
-                    QUERY_QUERY_ENTITY,
+                    graphql_query,
                     {
                         "input": {
-                            "types": ["QUERY"],
+                            "types": [entity_type],
                             "query": "*",
                             "scrollId": scroll_id if scroll_id else None,
                             "count": self.config.batch_size,
@@ -254,7 +257,7 @@ class SoftDeletedEntitiesCleanup:
                 )
             except Exception as e:
                 self.report.failure(
-                    f"While trying to get queries with {scroll_id}", exc=e
+                    f"While trying to get {entity_type} with {scroll_id}", exc=e
                 )
                 break
             scroll_across_entities = result.get("scrollAcrossEntities")
@@ -275,7 +278,8 @@ class SoftDeletedEntitiesCleanup:
             status=RemovedStatusFilter.ONLY_SOFT_DELETED,
             batch_size=self.config.batch_size,
         )
-        yield from self._get_soft_deleted_queries()
+        yield from self._get_soft_deleted(QUERY_ENTITIES, "QUERY")
+        yield from self._get_soft_deleted(QUERY_ENTITIES, "DATA_PROCESS_INSTANCE")
 
     def _times_up(self) -> bool:
         if (