Don't let .netrc take precedence in ApiClient (#526)

Author: kartikgupta-db

databricks_cli/sdk/api_client.py

@@ -39,6 +39,8 @@
 from . import version
 
 from requests.adapters import HTTPAdapter
+from requests.utils import get_netrc_auth
+from requests.auth import HTTPBasicAuth
 from six.moves.urllib.parse import urlparse
 
 try:
@@ -63,6 +65,30 @@ class TlsV1HttpAdapter(HTTPAdapter):
     def init_poolmanager(self, connections, maxsize, block=False):
         self.poolmanager = PoolManager(num_pools=connections, maxsize=maxsize, block=block, ssl_version=ssl.PROTOCOL_TLSv1_2)
 
+# https://github.com/psf/requests/issues/2773#issuecomment-174312831
+class FallbackNetrcAuth(requests.auth.AuthBase):
+    '''Force requests to ignore the ``.netrc`` if other authentication 
+    methods have been setup. Fallback to ``.netrc`` if not. 
+
+    Use with::
+
+        requests.get(url, auth=FallbackNetrcAuth())
+        
+        s = requests.Session()
+        s.auth = FallbackNetrcAuth()
+    '''
+
+    def __call__(self, r):
+        if "Authorization" in r.headers:
+            return r
+        
+        netrc_tuple = get_netrc_auth(r.url)
+
+        if netrc_tuple is None or not any(netrc_tuple):
+            return r
+        
+        return HTTPBasicAuth(*netrc_tuple)(r)
+
 class ApiClient(object):
     """
     A partial Python implementation of dbc rest api
@@ -82,6 +108,7 @@ def __init__(self, user=None, password=None, host=None, token=None,
             raise_on_status=False # return original response when retries have been exhausted
         )
         self.session = requests.Session()
+        self.session.auth = FallbackNetrcAuth()
         self.session.mount('https://', TlsV1HttpAdapter(max_retries=retries))
 
         parsed_url = urlparse(host)

tests/sdk/test_api_client.py

@@ -21,6 +21,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import json
+import os
 
 import pytest
 import requests
@@ -44,6 +45,12 @@ def m():
     with requests_mock.Mocker() as m:
         yield m
 
+@pytest.fixture(autouse=False)
+def netrc_file(tmp_path, monkeypatch):
+    netrc_file_path = os.path.join(str(tmp_path), ".netrc")
+    monkeypatch.setenv("NETRC", netrc_file_path)
+    return netrc_file_path
+
 def test_simple_request(m):
     data = {'cucumber': 'dade'}
     m.get('https://databricks.com/api/2.0/endpoint', text=json.dumps(data))
@@ -130,3 +137,38 @@ def test_api_client_url_parsing():
     # databricks_cli.configure.cli
     client = ApiClient(host='http://databricks.com')
     assert client.get_url('') == 'http://databricks.com/api/2.0'
+
+def test_api_client_auth_netrc_and_user_password(m, netrc_file):
+    with open(netrc_file, "w+") as netrc:
+        #generates header Authorization: 'Basic bmV0cmM6cGFzc3dvcmQ='
+        netrc.write("machine databricks.com login netrc password password")
+
+    m.get('https://databricks.com/api/2.0/endpoint', text=json.dumps({}))
+    client = ApiClient(user="apple", password="banana", host="https://databricks.com")
+    client.perform_query("GET", "/endpoint")
+    assert m.request_history[0].headers['Authorization'] == "Basic YXBwbGU6YmFuYW5h"
+
+def test_api_client_auth_only_valid_netrc(m, netrc_file):
+    with open(netrc_file, "w+") as netrc:
+        #generates header Authorization: 'Basic bmV0cmM6cGFzc3dvcmQ='
+        netrc.write("machine databricks.com login netrc password password")
+
+    m.get('https://databricks.com/api/2.0/endpoint', text=json.dumps({}))
+    client = ApiClient(host="https://databricks.com")
+    client.perform_query("GET", "/endpoint")
+    assert m.request_history[0].headers['Authorization'] == "Basic bmV0cmM6cGFzc3dvcmQ="
+
+def test_api_client_auth_invalid_netrc(m, netrc_file):
+    with open(netrc_file, "w+") as netrc:
+        netrc.write("garbage")
+
+    m.get('https://databricks.com/api/2.0/endpoint', text=json.dumps({}))
+    client = ApiClient(host="https://databricks.com")
+    client.perform_query("GET", "/endpoint")
+    assert "Authorization" not in m.request_history[0].headers
+
+def test_api_client_auth_no_netrc(m):
+    m.get('https://databricks.com/api/2.0/endpoint', text=json.dumps({}))
+    client = ApiClient(host="https://databricks.com")
+    client.perform_query("GET", "/endpoint")
+    assert "Authorization" not in m.request_history[0].headers