formatting (black)

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

src/databricks/sql/__init__.py

@@ -38,8 +38,7 @@ def filter(self, record):
                 )
         else:
             record.args = tuple(
-                (self.redact(arg) if isinstance(arg, str) else arg)
-                for arg in record.args
+                (self.redact(arg) if isinstance(arg, str) else arg) for arg in record.args
             )
 
         return True

src/databricks/sql/auth/authenticators.py

@@ -68,9 +68,7 @@ def __init__(
         try:
             idp_endpoint = get_oauth_endpoints(hostname, auth_type == "azure-oauth")
             if not idp_endpoint:
-                raise NotImplementedError(
-                    f"OAuth is not supported for host ${hostname}"
-                )
+                raise NotImplementedError(f"OAuth is not supported for host ${hostname}")
 
             # Convert to the corresponding scopes in the corresponding IdP
             cloud_scopes = idp_endpoint.get_scopes_mapping(scopes)
@@ -179,9 +177,7 @@ class AzureServicePrincipalCredentialProvider(CredentialsProvider):
     AZURE_MANAGED_RESOURCE = "https://management.core.windows.net/"
 
     DATABRICKS_AZURE_SP_TOKEN_HEADER = "X-Databricks-Azure-SP-Management-Token"
-    DATABRICKS_AZURE_WORKSPACE_RESOURCE_ID_HEADER = (
-        "X-Databricks-Azure-Workspace-Resource-Id"
-    )
+    DATABRICKS_AZURE_WORKSPACE_RESOURCE_ID_HEADER = "X-Databricks-Azure-Workspace-Resource-Id"
 
     def __init__(
         self,
@@ -195,9 +191,7 @@ def __init__(
         self.azure_client_id = azure_client_id
         self.azure_client_secret = azure_client_secret
         self.azure_workspace_resource_id = azure_workspace_resource_id
-        self.azure_tenant_id = azure_tenant_id or get_azure_tenant_id_from_host(
-            hostname
-        )
+        self.azure_tenant_id = azure_tenant_id or get_azure_tenant_id_from_host(hostname)
 
     def auth_type(self) -> str:
         return AuthType.AZURE_SP_M2M.value
@@ -211,9 +205,7 @@ def get_token_source(self, resource: str) -> RefreshableTokenSource:
         )
 
     def __call__(self, *args, **kwargs) -> HeaderFactory:
-        inner = self.get_token_source(
-            resource=get_effective_azure_login_app_id(self.hostname)
-        )
+        inner = self.get_token_source(resource=get_effective_azure_login_app_id(self.hostname))
         cloud = self.get_token_source(resource=self.AZURE_MANAGED_RESOURCE)
 
         def header_factory() -> Dict[str, str]:

src/databricks/sql/auth/endpoint.py

@@ -58,9 +58,7 @@ def infer_cloud_from_host(hostname: str) -> Optional[CloudType]:
 
 def is_supported_databricks_oauth_host(hostname: str) -> bool:
     host = hostname.lower().replace("https://", "").split("/")[0]
-    domains = (
-        DATABRICKS_AWS_DOMAINS + DATABRICKS_GCP_DOMAINS + DATABRICKS_OAUTH_AZURE_DOMAINS
-    )
+    domains = DATABRICKS_AWS_DOMAINS + DATABRICKS_GCP_DOMAINS + DATABRICKS_OAUTH_AZURE_DOMAINS
     return any(e for e in domains if host.endswith(e))
 
 
@@ -106,7 +104,9 @@ def get_authorization_url(self, hostname: str):
         return f"{get_databricks_oidc_url(hostname)}/oauth2/v2.0/authorize"
 
     def get_openid_config_url(self, hostname: str):
-        return "https://login.microsoftonline.com/organizations/v2.0/.well-known/openid-configuration"
+        return (
+            "https://login.microsoftonline.com/organizations/v2.0/.well-known/openid-configuration"
+        )
 
 
 class InHouseOAuthEndpointCollection(OAuthEndpointCollection):
@@ -123,9 +123,7 @@ def get_openid_config_url(self, hostname: str):
         return f"{idp_url}/.well-known/oauth-authorization-server"
 
 
-def get_oauth_endpoints(
-    hostname: str, use_azure_auth: bool
-) -> Optional[OAuthEndpointCollection]:
+def get_oauth_endpoints(hostname: str, use_azure_auth: bool) -> Optional[OAuthEndpointCollection]:
     cloud = infer_cloud_from_host(hostname)
 
     if cloud in [CloudType.AWS, CloudType.GCP]:

src/databricks/sql/auth/oauth.py

@@ -41,9 +41,7 @@ def __init__(self, access_token: str, token_type: str, refresh_token: str):
 
     def is_expired(self) -> bool:
         try:
-            decoded_token = jwt.decode(
-                self.access_token, options={"verify_signature": False}
-            )
+            decoded_token = jwt.decode(self.access_token, options={"verify_signature": False})
             exp_time = decoded_token.get("exp")
             current_time = time.time()
             buffer_time = 30  # 30 seconds buffer
@@ -134,9 +132,7 @@ def __fetch_well_known_config(self, hostname: str):
     def __get_challenge():
         verifier_string = OAuthManager.__token_urlsafe(32)
         digest = hashlib.sha256(verifier_string.encode("UTF-8")).digest()
-        challenge_string = (
-            base64.urlsafe_b64encode(digest).decode("UTF-8").replace("=", "")
-        )
+        challenge_string = base64.urlsafe_b64encode(digest).decode("UTF-8").replace("=", "")
         return verifier_string, challenge_string
 
     def __get_authorization_code(self, client, auth_url, scope, state, challenge):
@@ -158,9 +154,7 @@ def __get_authorization_code(self, client, auth_url, scope, state, challenge):
                     logger.info(f"Opening {auth_req_uri}")
 
                     webbrowser.open_new(auth_req_uri)
-                    logger.info(
-                        f"Listening for OAuth authorization callback at {redirect_url}"
-                    )
+                    logger.info(f"Listening for OAuth authorization callback at {redirect_url}")
                     httpd.handle_request()
                 self.redirect_port = port
                 break
@@ -182,9 +176,7 @@ def __get_authorization_code(self, client, auth_url, scope, state, challenge):
             raise RuntimeError(msg)
         # This is a kludge because the parsing library expects https callbacks
         # We should probably set it up using https
-        full_redirect_url = (
-            f"https://localhost:{self.redirect_port}/{handler.request_path}"
-        )
+        full_redirect_url = f"https://localhost:{self.redirect_port}/{handler.request_path}"
         try:
             authorization_code_response = client.parse_request_uri_response(
                 full_redirect_url, state=state
@@ -197,9 +189,7 @@ def __get_authorization_code(self, client, auth_url, scope, state, challenge):
     def __send_auth_code_token_request(
         self, client, token_request_url, redirect_url, code, verifier
     ):
-        token_request_body = client.prepare_request_body(
-            code=code, redirect_uri=redirect_url
-        )
+        token_request_body = client.prepare_request_body(code=code, redirect_uri=redirect_url)
         data = f"{token_request_body}&code_verifier={verifier}"
         return self.__send_token_request(token_request_url, data)
 
@@ -227,15 +217,11 @@ def __send_refresh_token_request(self, hostname, refresh_token):
     def __get_tokens_from_response(oauth_response):
         access_token = oauth_response["access_token"]
         refresh_token = (
-            oauth_response["refresh_token"]
-            if "refresh_token" in oauth_response
-            else None
+            oauth_response["refresh_token"] if "refresh_token" in oauth_response else None
         )
         return access_token, refresh_token
 
-    def check_and_refresh_access_token(
-        self, hostname: str, access_token: str, refresh_token: str
-    ):
+    def check_and_refresh_access_token(self, hostname: str, access_token: str, refresh_token: str):
         now = datetime.now(tz=timezone.utc)
         # If we can't decode an expiration time, this will be expired by default.
         expiration_time = now
@@ -246,9 +232,7 @@ def check_and_refresh_access_token(
             # an unnecessary signature verification.
             access_token_payload = access_token.split(".")[1]
             # add padding
-            access_token_payload = access_token_payload + "=" * (
-                -len(access_token_payload) % 4
-            )
+            access_token_payload = access_token_payload + "=" * (-len(access_token_payload) % 4)
             decoded = json.loads(base64.standard_b64decode(access_token_payload))
             expiration_time = datetime.fromtimestamp(decoded["exp"], tz=timezone.utc)
         except Exception as e:
@@ -265,13 +249,9 @@ def check_and_refresh_access_token(
             raise RuntimeError(msg)
 
         # Try to refresh using the refresh token
-        logger.debug(
-            f"Attempting to refresh OAuth access token that expired on {expiration_time}"
-        )
+        logger.debug(f"Attempting to refresh OAuth access token that expired on {expiration_time}")
         oauth_response = self.__send_refresh_token_request(hostname, refresh_token)
-        fresh_access_token, fresh_refresh_token = self.__get_tokens_from_response(
-            oauth_response
-        )
+        fresh_access_token, fresh_refresh_token = self.__get_tokens_from_response(oauth_response)
         return fresh_access_token, fresh_refresh_token, True
 
     def get_tokens(self, hostname: str, scope=None):
@@ -285,9 +265,7 @@ def get_tokens(self, hostname: str, scope=None):
         client = oauthlib.oauth2.WebApplicationClient(self.client_id)
 
         try:
-            auth_response = self.__get_authorization_code(
-                client, auth_url, scope, state, challenge
-            )
+            auth_response = self.__get_authorization_code(client, auth_url, scope, state, challenge)
         except OAuth2Error as e:
             msg = f"OAuth Authorization Error: {e.description}"
             logger.error(msg)
@@ -359,6 +337,4 @@ def refresh(self) -> Token:
                     oauth_response.refresh_token,
                 )
             else:
-                raise Exception(
-                    f"Failed to get token: {response.status_code} {response.text}"
-                )
+                raise Exception(f"Failed to get token: {response.status_code} {response.text}")

src/databricks/sql/auth/retry.py

@@ -167,9 +167,7 @@ def __private_init__(
         new_object.command_type = command_type
         return new_object
 
-    def new(
-        self, **urllib3_incremented_counters: typing.Any
-    ) -> "DatabricksRetryPolicy":
+    def new(self, **urllib3_incremented_counters: typing.Any) -> "DatabricksRetryPolicy":
         """This method is responsible for passing the entire Retry state to its next iteration.
 
         urllib3 calls Retry.new() between successive requests as part of its `.increment()` method
@@ -435,9 +433,7 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
             "Failed requests are retried by default per configured DatabricksRetryPolicy",
         )
 
-    def is_retry(
-        self, method: str, status_code: int, has_retry_after: bool = False
-    ) -> bool:
+    def is_retry(self, method: str, status_code: int, has_retry_after: bool = False) -> bool:
         """
         Called by urllib3 when determining whether or not to retry
 

src/databricks/sql/auth/thrift_http_client.py

@@ -199,9 +199,7 @@ def flush(self):
         self.headers = self.__resp.headers
 
         logger.info(
-            "HTTP Response with status code {}, message: {}".format(
-                self.code, self.message
-            )
+            "HTTP Response with status code {}, message: {}".format(self.code, self.message)
         )
 
     @staticmethod