🎉 Initial commit. (#248)

amundsen-io · Jan 28, 2021 · 1e2e3c8 · 1e2e3c8
1 parent 7e63c1d
commit 1e2e3c8
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 3 deletions.
diff --git a/metadata_service/config.py b/metadata_service/config.py
@@ -17,6 +17,7 @@
 PROXY_ENCRYPTED = 'PROXY_ENCRYPTED'
 PROXY_VALIDATE_SSL = 'PROXY_VALIDATE_SSL'
 PROXY_CLIENT = 'PROXY_CLIENT'
+PROXY_CLIENT_KWARGS = 'PROXY_CLIENT_KWARGS'
 
 PROXY_CLIENTS = {
     'NEO4J': 'metadata_service.proxy.neo4j_proxy.Neo4jProxy',
@@ -84,6 +85,10 @@ class Config:
     # format to be used in tables
     WATERMARK_DATE_FORMATS = ['%Y%m%d']
 
+    # Custom kwargs that will be passed to proxy client. Can be used to fine-tune parameters like timeout
+    # or num of retries
+    PROXY_CLIENT_KWARGS: Dict = dict()
+
 
 # NB: If you're using the gremlin proxy, the appropriate GremlinConfig must be added to any other configs
 class LocalConfig(LocalGremlinConfig, Config):

diff --git a/metadata_service/proxy/__init__.py b/metadata_service/proxy/__init__.py
@@ -35,12 +35,15 @@ def get_proxy_client() -> BaseProxy:
             encrypted = current_app.config[config.PROXY_ENCRYPTED]
             validate_ssl = current_app.config[config.PROXY_VALIDATE_SSL]
 
+            client_kwargs = current_app.config[config.PROXY_CLIENT_KWARGS]
+
             client = import_string(current_app.config[config.PROXY_CLIENT])
             _proxy_client = client(host=host,
                                    port=port,
                                    user=user,
                                    password=password,
                                    encrypted=encrypted,
-                                   validate_ssl=validate_ssl)
+                                   validate_ssl=validate_ssl,
+                                   client_kwargs=client_kwargs)
 
     return _proxy_client
diff --git a/metadata_service/proxy/atlas_proxy.py b/metadata_service/proxy/atlas_proxy.py
@@ -70,7 +70,8 @@ def __init__(self, *,
                  user: str = 'admin',
                  password: str = '',
                  encrypted: bool = False,
-                 validate_ssl: bool = False) -> None:
+                 validate_ssl: bool = False,
+                 client_kwargs: dict = dict()) -> None:
         """
         Initiate the Apache Atlas client with the provided credentials
         """
@@ -80,7 +81,8 @@ def __init__(self, *,
                              username=user,
                              password=password,
                              protocol=protocol,
-                             validate_ssl=validate_ssl)
+                             validate_ssl=validate_ssl,
+                             **client_kwargs)
 
     def _get_ids_from_basic_search(self, *, params: Dict) -> List[str]:
         """