Enhance CI (add black as pre-commit hook and add pylint checks)

.github/workflows/code-quality-checks.yml

@@ -156,6 +156,11 @@ jobs:
       #----------------------------------------------
       - name: Black
         run: poetry run black --check src
+      #----------------------------------------------
+      #              pylint the code
+      #----------------------------------------------
+      - name: Pylint
+        run: poetry run pylint --rcfile=pylintrc src
 
   check-types:
     runs-on: ubuntu-latest

.pre-commit-config.yaml

@@ -0,0 +1,13 @@
+repos:
+-   repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.5.0
+    hooks:
+    -   id: trailing-whitespace
+    -   id: end-of-file-fixer
+    -   id: check-added-large-files
+
+-   repo: https://github.com/psf/black
+    rev: 22.3.0
+    hooks:
+    -   id: black
+        exclude: '/(\.eggs|\.git|\.hg|\.mypy_cache|\.nox|\.tox|\.venv|\.svn|_build|buck-out|build|dist|thrift_api)/'

CONTRIBUTING.md

@@ -145,9 +145,11 @@ The `PySQLStagingIngestionTestSuite` namespace requires a cluster running DBR ve
 The suites marked `[not documented]` require additional configuration which will be documented at a later time.
 
 
-### Code formatting
+### Code formatting and linting
 
-This project uses [Black](https://pypi.org/project/black/).
+This project uses [Black](https://pypi.org/project/black/) for code formatting and [Pylint](https://pylint.org/) for linting.
+
+#### Black
 
 ```
 poetry run python3 -m black src --check
@@ -157,6 +159,25 @@ Remove the `--check` flag to write reformatted files to disk.
 
 To simplify reviews you can format your changes in a separate commit.
 
+#### Pylint
+
+```
+poetry run pylint --rcfile=pylintrc src
+```
+
+#### Pre-commit hooks
+
+We use [pre-commit](https://pre-commit.com/) to automatically run Black and other checks before each commit.
+
+To set up pre-commit hooks:
+
+```bash
+# Set up the git hooks
+poetry run pre-commit install
+```
+
+This will set up the hooks defined in `.pre-commit-config.yaml` to run automatically on each commit.
+
 ### Change a pinned dependency version
 
 Modify the dependency specification (syntax can be found [here](https://python-poetry.org/docs/dependency-specification/)) in `pyproject.toml` and run one of the following in your terminal:

pylintrc

@@ -0,0 +1,22 @@
+[MASTER]
+ignore=thrift_api
+
+[MESSAGES CONTROL]
+disable=too-many-arguments,
+        too-many-locals,
+        too-many-public-methods,
+        too-many-branches,
+        too-many-statements,
+        fixme,
+        missing-docstring,
+        line-too-long,
+        too-few-public-methods,
+        too-many-instance-attributes,
+        too-many-lines
+
+[FORMAT]
+max-line-length=100
+
+[REPORTS]
+output-format=text
+reports=yes

pyproject.toml

@@ -37,6 +37,7 @@ mypy = "^1.10.1"
 pylint = ">=2.12.0"
 black = "^22.3.0"
 pytest-dotenv = "^0.5.2"
+pre-commit = {version = "^3.5.0"}
 numpy = [
     { version = ">=1.16.6", python = ">=3.8,<3.11" },
     { version = ">=1.23.4", python = ">=3.11" },
@@ -55,6 +56,7 @@ ignore_missing_imports = "true"
 exclude = ['ttypes\.py$', 'TCLIService\.py$']
 
 [tool.black]
+line-length = 100
 exclude = '/(\.eggs|\.git|\.hg|\.mypy_cache|\.nox|\.tox|\.venv|\.svn|_build|buck-out|build|dist|thrift_api)/'
 
 [tool.pytest.ini_options]

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