chore: lint codebase

Author: Brooke-white

.pre-commit-config.yaml

@@ -19,4 +19,5 @@ repos:
     hooks:
       - id: mypy
         additional_dependencies: [types-setuptools, types-requests, types-python-dateutil]
+        args: [--ignore-missing-imports, --disable-error-code, "annotation-unchecked"]
         verbose: true

redshift_connector/__init__.py

@@ -56,13 +56,22 @@
 logging.getLogger(__name__).addHandler(logging.NullHandler())
 _logger: logging.Logger = logging.getLogger(__name__)
 
-IDC_PLUGINS_LIST = ("redshift_connector.plugin.BrowserIdcAuthPlugin", "BrowserIdcAuthPlugin",
-                    "redshift_connector.plugin.IdpTokenAuthPlugin", "IdpTokenAuthPlugin")
+IDC_PLUGINS_LIST = (
+    "redshift_connector.plugin.BrowserIdcAuthPlugin",
+    "BrowserIdcAuthPlugin",
+    "redshift_connector.plugin.IdpTokenAuthPlugin",
+    "IdpTokenAuthPlugin",
+)
 IDC_OR_NATIVE_IDP_PLUGINS_LIST = (
-    "redshift_connector.plugin.BrowserAzureOAuth2CredentialsProvider", "BrowserAzureOAuth2CredentialsProvider",
-    "redshift_connector.plugin.BasicJwtCredentialsProvider", "BasicJwtCredentialsProvider",
-    "redshift_connector.plugin.BrowserIdcAuthPlugin", "BrowserIdcAuthPlugin",
-    "redshift_connector.plugin.IdpTokenAuthPlugin", "IdpTokenAuthPlugin")
+    "redshift_connector.plugin.BrowserAzureOAuth2CredentialsProvider",
+    "BrowserAzureOAuth2CredentialsProvider",
+    "redshift_connector.plugin.BasicJwtCredentialsProvider",
+    "BasicJwtCredentialsProvider",
+    "redshift_connector.plugin.BrowserIdcAuthPlugin",
+    "BrowserIdcAuthPlugin",
+    "redshift_connector.plugin.IdpTokenAuthPlugin",
+    "IdpTokenAuthPlugin",
+)
 
 # Copyright (c) 2007-2009, Mathieu Fenniak
 # Copyright (c) The Contributors

redshift_connector/core.py

@@ -565,12 +565,13 @@ def get_calling_module() -> str:
                 init_params["idp_type"] = "AzureAD"
 
             if credentials_provider.split(".")[-1] in (
-                    "IdpTokenAuthPlugin",
-                    "BrowserIdcAuthPlugin",
+                "IdpTokenAuthPlugin",
+                "BrowserIdcAuthPlugin",
             ):
                 redshift_native_auth = True
-                self.set_idc_plugins_params(init_params, credentials_provider, identity_namespace, token_type,
-                                            idc_client_display_name)
+                self.set_idc_plugins_params(
+                    init_params, credentials_provider, identity_namespace, token_type, idc_client_display_name
+                )
 
             if redshift_native_auth and provider_name:
                 init_params["provider_name"] = provider_name
@@ -1808,8 +1809,10 @@ def execute(self: "Connection", cursor: Cursor, operation: str, vals) -> None:
             val: typing.Union[bytes, bytearray] = bytearray(statement_name_bin)
             typing.cast(bytearray, val).extend(statement.encode(_client_encoding) + NULL_BYTE)
             if len(params) > 32767:
-                raise DataError("Prepared statement exceeds bind parameter limit 32767. {} bind parameters were "
-                                "provided. Please retry with fewer bind parameters.".format(len(params)))
+                raise DataError(
+                    "Prepared statement exceeds bind parameter limit 32767. {} bind parameters were "
+                    "provided. Please retry with fewer bind parameters.".format(len(params))
+                )
             typing.cast(bytearray, val).extend(h_pack(len(params)))
             for oid, fc, send_func in params:  # type: ignore
                 # Parse message doesn't seem to handle the -1 type_oid for NULL
@@ -2551,12 +2554,15 @@ def tpc_recover(self: "Connection") -> typing.List[typing.Tuple[typing.Any, ...]
         finally:
             self.autocommit = previous_autocommit_mode
 
-    def set_idc_plugins_params(self: "Connection", init_params: typing.Dict[str, typing.Optional[typing.Union[str, bytes]]],
-                               credentials_provider: typing.Optional[str] = None,
-                               identity_namespace: typing.Optional[str] = None,
-                               token_type: typing.Optional[str] = None,
-                               idc_client_display_name: typing.Optional[str] = None) -> None:
-        plugin_name = credentials_provider.split(".")[-1]
+    def set_idc_plugins_params(
+        self: "Connection",
+        init_params: typing.Dict[str, typing.Optional[typing.Union[str, bytes]]],
+        credentials_provider: typing.Optional[str] = None,
+        identity_namespace: typing.Optional[str] = None,
+        token_type: typing.Optional[str] = None,
+        idc_client_display_name: typing.Optional[str] = None,
+    ) -> None:
+        plugin_name = typing.cast(str, credentials_provider).split(".")[-1]
         init_params["idp_type"] = "AwsIdc"
 
         if identity_namespace:
@@ -2569,5 +2575,3 @@ def set_idc_plugins_params(self: "Connection", init_params: typing.Dict[str, typ
 
         if idc_client_display_name:
             init_params["idc_client_display_name"] = idc_client_display_name
-
-

redshift_connector/plugin/__init__.py

@@ -1,6 +1,5 @@
 from .adfs_credentials_provider import AdfsCredentialsProvider
 from .azure_credentials_provider import AzureCredentialsProvider
-from .idp_token_auth_plugin import IdpTokenAuthPlugin
 from .browser_azure_credentials_provider import BrowserAzureCredentialsProvider
 from .browser_azure_oauth2_credentials_provider import (
     BrowserAzureOAuth2CredentialsProvider,
@@ -9,6 +8,7 @@
 from .browser_saml_credentials_provider import BrowserSamlCredentialsProvider
 from .common_credentials_provider import CommonCredentialsProvider
 from .idp_credentials_provider import IdpCredentialsProvider
+from .idp_token_auth_plugin import IdpTokenAuthPlugin
 from .jwt_credentials_provider import (
     BasicJwtCredentialsProvider,
     JwtCredentialsProvider,

redshift_connector/plugin/browser_idc_auth_plugin.py

@@ -1,12 +1,15 @@
-import boto3
 import logging
-import typing
 import time
+import typing
 import webbrowser
 
+import boto3
 from botocore.exceptions import ClientError
+
 from redshift_connector.error import InterfaceError
-from redshift_connector.plugin.common_credentials_provider import CommonCredentialsProvider
+from redshift_connector.plugin.common_credentials_provider import (
+    CommonCredentialsProvider,
+)
 from redshift_connector.redshift_property import RedshiftProperty
 
 logging.getLogger(__name__).addHandler(logging.NullHandler())
@@ -18,10 +21,10 @@ class BrowserIdcAuthPlugin(CommonCredentialsProvider):
     Class to get IdC Token using SSO OIDC APIs
     """
 
-    DEFAULT_IDC_CLIENT_DISPLAY_NAME = 'Amazon Redshift Python connector'
-    CLIENT_TYPE = 'public'
-    GRANT_TYPE = 'urn:ietf:params:oauth:grant-type:device_code'
-    IDC_SCOPE = 'redshift:connect'
+    DEFAULT_IDC_CLIENT_DISPLAY_NAME = "Amazon Redshift Python connector"
+    CLIENT_TYPE = "public"
+    GRANT_TYPE = "urn:ietf:params:oauth:grant-type:device_code"
+    IDC_SCOPE = "redshift:connect"
     DEFAULT_BROWSER_AUTH_VERIFY_TIMEOUT_IN_SEC = 120
     DEFAULT_CREATE_TOKEN_INTERVAL_IN_SEC = 1
 
@@ -32,11 +35,11 @@ def __init__(self: "BrowserIdcAuthPlugin") -> None:
         self.register_client_cache: typing.Dict[str, dict] = {}
         self.start_url: typing.Optional[str] = None
         self.idc_region: typing.Optional[str] = None
-        self.sso_oidc_client: "SSOOIDC.Client" = None
+        self.sso_oidc_client: boto3.client = None
 
     def add_parameter(
-            self: "BrowserIdcAuthPlugin",
-            info: RedshiftProperty,
+        self: "BrowserIdcAuthPlugin",
+        info: RedshiftProperty,
     ) -> None:
         """
         Adds parameters to the BrowserIdcAuthPlugin
@@ -65,11 +68,13 @@ def check_required_parameters(self: "BrowserIdcAuthPlugin") -> None:
         if not self.start_url:
             _logger.error("IdC authentication failed: start_url needs to be provided in connection params")
             raise InterfaceError(
-                "IdC authentication failed: The start URL must be included in the connection parameters.")
+                "IdC authentication failed: The start URL must be included in the connection parameters."
+            )
         if not self.idc_region:
             _logger.error("IdC authentication failed: idc_region needs to be provided in connection params")
             raise InterfaceError(
-                "IdC authentication failed: The IdC region must be included in the connection parameters.")
+                "IdC authentication failed: The IdC region must be included in the connection parameters."
+            )
 
     def get_cache_key(self: "BrowserIdcAuthPlugin") -> str:
         """
@@ -93,16 +98,21 @@ def get_idc_token(self: "BrowserIdcAuthPlugin") -> str:
         try:
             self.check_required_parameters()
 
-            self.sso_oidc_client = boto3.client('sso-oidc', region_name=self.idc_region)
+            self.sso_oidc_client = boto3.client("sso-oidc", region_name=self.idc_region)
             register_client_cache_key: str = f"{self.idc_client_display_name}:{self.idc_region}"
 
-            register_client_result: typing.Dict[str, typing.Any] = self.register_client(register_client_cache_key,
-                                                                                        self.idc_client_display_name,
-                                                                                        self.CLIENT_TYPE,
-                                                                                        self.IDC_SCOPE)
+            register_client_result: typing.Dict[str, typing.Any] = self.register_client(
+                register_client_cache_key,
+                self.idc_client_display_name,
+                typing.cast(str, self.CLIENT_TYPE),
+                self.IDC_SCOPE,
+            )
             start_device_auth_result: typing.Dict[str, typing.Any] = self.start_device_authorization(
-                register_client_result['clientId'], register_client_result['clientSecret'], self.start_url)
-            self.open_browser(start_device_auth_result['verificationUriComplete'])
+                register_client_result["clientId"],
+                register_client_result["clientSecret"],
+                typing.cast(str, self.start_url),
+            )
+            self.open_browser(start_device_auth_result["verificationUriComplete"])
 
             return self.poll_for_create_token(register_client_result, start_device_auth_result, self.GRANT_TYPE)
         except InterfaceError as e:
@@ -111,8 +121,9 @@ def get_idc_token(self: "BrowserIdcAuthPlugin") -> str:
             _logger.debug("An error occurred while trying to obtain an IdC token : {}".format(str(e)))
             raise InterfaceError("There was an error during authentication.")
 
-    def register_client(self: "BrowserIdcAuthPlugin", register_client_cache_key: str, client_name: str,
-                        client_type: str, scope: str) -> typing.Dict[str, typing.Any]:
+    def register_client(  # type: ignore
+        self: "BrowserIdcAuthPlugin", register_client_cache_key: str, client_name: str, client_type: str, scope: str
+    ) -> typing.Dict[str, typing.Any]:
         """
         Registers the client with IdC.
         :param register_client_cache_key: str
@@ -126,21 +137,25 @@ def register_client(self: "BrowserIdcAuthPlugin", register_client_cache_key: str
         :return: dict
             The register client result from IdC
         """
-        if register_client_cache_key in self.register_client_cache and \
-                self.register_client_cache[register_client_cache_key]['clientSecretExpiresAt'] > time.time():
+        if (
+            register_client_cache_key in self.register_client_cache
+            and self.register_client_cache[register_client_cache_key]["clientSecretExpiresAt"] > time.time()
+        ):
             _logger.debug("Valid registerClient result found from cache")
             return self.register_client_cache[register_client_cache_key]
 
         try:
             register_client_result: typing.Dict[str, typing.Any] = self.sso_oidc_client.register_client(
-                clientName=client_name, clientType=client_type, scopes=[scope])
+                clientName=client_name, clientType=client_type, scopes=[scope]
+            )
             self.register_client_cache[register_client_cache_key] = register_client_result
             return register_client_result
         except ClientError as e:
             self.handle_error(e, "registering client with IdC")
 
-    def start_device_authorization(self: "BrowserIdcAuthPlugin", client_id: str, client_secret: str,
-                                   start_url: str) -> typing.Dict[str, typing.Any]:
+    def start_device_authorization(  # type: ignore
+        self: "BrowserIdcAuthPlugin", client_id: str, client_secret: str, start_url: str
+    ) -> typing.Dict[str, typing.Any]:
         """
         Starts device authorization flow with IdC.
         :param client_id: str
@@ -154,9 +169,7 @@ def start_device_authorization(self: "BrowserIdcAuthPlugin", client_id: str, cli
         """
         try:
             response: typing.Dict[str, typing.Any] = self.sso_oidc_client.start_device_authorization(
-                clientId=client_id,
-                clientSecret=client_secret,
-                startUrl=start_url
+                clientId=client_id, clientSecret=client_secret, startUrl=start_url
             )
             return response
         except ClientError as e:
@@ -173,9 +186,12 @@ def open_browser(self: "BrowserIdcAuthPlugin", url: str) -> None:
         self.validate_url(url)
         webbrowser.open(url)
 
-    def poll_for_create_token(self: "BrowserIdcAuthPlugin",
-                              register_client_result: typing.Dict[str, typing.Any],
-                              start_device_auth_result: typing.Dict[str, typing.Any], grant_type: str) -> str:
+    def poll_for_create_token(
+        self: "BrowserIdcAuthPlugin",
+        register_client_result: typing.Dict[str, typing.Any],
+        start_device_auth_result: typing.Dict[str, typing.Any],
+        grant_type: str,
+    ) -> str:
         """
         Polls for IdC access token using SSO OIDC APIs.
         :param register_client_result: dict
@@ -191,22 +207,22 @@ def poll_for_create_token(self: "BrowserIdcAuthPlugin",
         polling_end_time: float = time.time() + self.idc_response_timeout
 
         polling_interval_in_sec: int = self.DEFAULT_CREATE_TOKEN_INTERVAL_IN_SEC
-        if start_device_auth_result['interval']:
-            polling_interval_in_sec = start_device_auth_result['interval']
+        if start_device_auth_result["interval"]:
+            polling_interval_in_sec = start_device_auth_result["interval"]
 
         while time.time() < polling_end_time:
             try:
                 response: typing.Dict[str, typing.Any] = self.sso_oidc_client.create_token(
-                    clientId=register_client_result['clientId'],
-                    clientSecret=register_client_result['clientSecret'],
+                    clientId=register_client_result["clientId"],
+                    clientSecret=register_client_result["clientSecret"],
                     grantType=grant_type,
-                    deviceCode=start_device_auth_result['deviceCode']
+                    deviceCode=start_device_auth_result["deviceCode"],
                 )
-                if not response['accessToken']:
+                if not response["accessToken"]:
                     raise InterfaceError("IdC authentication failed : The credential token couldn't be created.")
-                return response['accessToken']
+                return response["accessToken"]
             except ClientError as e:
-                if e.response['Error']['Code'] == 'AuthorizationPendingException':
+                if e.response["Error"]["Code"] == "AuthorizationPendingException":
                     _logger.debug("Browser authorization pending from user")
                     time.sleep(polling_interval_in_sec)
                 else:
@@ -225,11 +241,11 @@ def handle_error(self: "BrowserIdcAuthPlugin", e: ClientError, operation: str) -
         :raises InterfaceError: A client error to be returned to the user with appropriate error message
         """
         _logger.debug("Error response = {} ".format(e.response))
-        error_message = e.response['Error']['Message']
+        error_message = e.response["Error"]["Message"]
         if not error_message:
-            error_message = e.response['error_description'] if e.response[
-                'error_description'] else "Something unexpected happened"
-        error_code = e.response['Error']['Code']
-        _logger.debug(
-            "An error occurred while {}: ClientError = {} - {}".format(operation, error_code, error_message))
+            error_message = (
+                e.response["error_description"] if e.response["error_description"] else "Something unexpected happened"
+            )
+        error_code = e.response["Error"]["Code"]
+        _logger.debug("An error occurred while {}: ClientError = {} - {}".format(operation, error_code, error_message))
         raise InterfaceError("IdC authentication failed : An error occurred during the request.")

redshift_connector/plugin/common_credentials_provider.py

@@ -84,4 +84,4 @@ def set_group_federation(self: "CommonCredentialsProvider", group_federation: bo
         pass
 
     def get_sub_type(self: "CommonCredentialsProvider") -> int:
-        return IamHelper.IDC_PLUGIN
+        return IamHelper.IDC_PLUGIN

redshift_connector/plugin/idp_token_auth_plugin.py

@@ -2,7 +2,9 @@
 import typing
 
 from redshift_connector.error import InterfaceError
-from redshift_connector.plugin.common_credentials_provider import CommonCredentialsProvider
+from redshift_connector.plugin.common_credentials_provider import (
+    CommonCredentialsProvider,
+)
 from redshift_connector.redshift_property import RedshiftProperty
 
 logging.getLogger(__name__).addHandler(logging.NullHandler())
@@ -36,9 +38,10 @@ def check_required_parameters(self: "IdpTokenAuthPlugin") -> None:
         if not self.token_type:
             _logger.error("IdC authentication failed: token_type needs to be provided in connection params")
             raise InterfaceError(
-                "IdC authentication failed: The token type must be included in the connection parameters.")
+                "IdC authentication failed: The token type must be included in the connection parameters."
+            )
 
-    def get_cache_key(self: "IdpTokenAuthPlugin") -> str:
+    def get_cache_key(self: "IdpTokenAuthPlugin") -> str:  # type: ignore
         pass
 
     def get_auth_token(self: "IdpTokenAuthPlugin") -> str: