refactor: use importlib instead of pkg_resources

Author: Brooke-white

redshift_connector/iam_helper.py

@@ -3,7 +3,6 @@
 import logging
 import typing
 
-import pkg_resources
 from dateutil.tz import tzutc
 from packaging.version import Version
 
@@ -58,7 +57,7 @@ def can_support_v2(provider_type: "IamHelper.IAMAuthenticationType") -> bool:
                     IamHelper.IAMAuthenticationType.IAM_KEYS,
                     IamHelper.IAMAuthenticationType.IAM_KEYS_WITH_SESSION,
                 )
-            ) and Version(pkg_resources.get_distribution("boto3").version) >= Version("1.24.5")
+            ) and IdpAuthHelper.get_pkg_version("boto3") >= Version("1.24.5")
 
     credentials_cache: typing.Dict[str, dict] = {}
 
@@ -94,8 +93,8 @@ def set_iam_properties(info: RedshiftProperty) -> RedshiftProperty:
         IamHelper.set_auth_properties(info)
 
         if info._is_serverless and info.iam:
-            if Version(pkg_resources.get_distribution("boto3").version) < Version("1.24.11"):
-                raise pkg_resources.VersionConflict(
+            if IdpAuthHelper.get_pkg_version("boto3") < Version("1.24.11"):
+                raise ModuleNotFoundError(
                     "boto3 >= 1.24.11 required for authentication with Amazon Redshift serverless. "
                     "Please upgrade the installed version of boto3 to use this functionality."
                 )

redshift_connector/idp_auth_helper.py

@@ -2,6 +2,8 @@
 import typing
 from enum import Enum
 
+from packaging.version import Version
+
 from redshift_connector.error import InterfaceError, ProgrammingError
 from redshift_connector.plugin.i_plugin import IPlugin
 from redshift_connector.redshift_property import RedshiftProperty
@@ -32,14 +34,26 @@ class IdpAuthHelper:
     SAML_PLUGIN: int = 1
     JWT_PLUGIN: int = 2
 
+    @staticmethod
+    def get_pkg_version(module_name: str) -> Version:
+        """
+        Returns a Version object pertaining to the module name provided.
+        """
+        try:
+            from importlib.metadata import version as version
+        except ModuleNotFoundError:  # if importlib is not present, fallback to pkg_resources
+            import pkg_resources
+
+            return Version(pkg_resources.get_distribution(module_name).version)
+
+        return Version(version(module_name))
+
     @staticmethod
     def set_auth_properties(info: RedshiftProperty):
         """
         Helper function to handle IAM and Native Auth connection properties and ensure required parameters are specified.
         Parameters
         """
-        import pkg_resources
-        from packaging.version import Version
 
         if info is None:
             raise InterfaceError("Invalid connection property setting. info must be specified")
@@ -65,13 +79,13 @@ def set_auth_properties(info: RedshiftProperty):
         #         "AWS credentials, Amazon Redshift authentication profile, or AWS profile"
         #     )
         if info.iam is True:
-            _logger.debug("boto3 version: {}".format(Version(pkg_resources.get_distribution("boto3").version)))
-            _logger.debug("botocore version: {}".format(Version(pkg_resources.get_distribution("botocore").version)))
+            _logger.debug("boto3 version: {}".format(IdpAuthHelper.get_pkg_version("boto3")))
+            _logger.debug("botocore version: {}".format(IdpAuthHelper.get_pkg_version("botocore")))
 
             # Check for IAM keys and AuthProfile first
             if info.auth_profile is not None:
-                if Version(pkg_resources.get_distribution("boto3").version) < Version("1.17.111"):
-                    raise pkg_resources.VersionConflict(
+                if IdpAuthHelper.get_pkg_version("boto3") < Version("1.17.111"):
+                    raise ModuleNotFoundError(
                         "boto3 >= 1.17.111 required for authentication via Amazon Redshift authentication profile. "
                         "Please upgrade the installed version of boto3 to use this functionality."
                     )

test/manual/auth/test_aws_credentials.py

@@ -38,7 +38,7 @@ def test_use_aws_credentials_default_profile():
 0) Generate credentials using instructions:  https://docs.aws.amazon.com/sdk-for-javascript/v2/developer-guide/getting-your-credentials.html
 1) In the connect method below,  specify the connection parameters
 3) Specify the AWS IAM credentials in the variables above
-4) Update iam_helper.py to include correct min version. line `Version(pkg_resources.get_distribution("boto3").version) > Version("9.99.9999"):`
+4) Update iam_helper.py to include correct min version.
 5) Manually execute this test
 """
 

test/unit/test_iam_helper.py

@@ -781,13 +781,13 @@ def test_get_cluster_credentials_api_type_will_use_correct_api(conn_params, prov
     ),
 )
 def test_set_iam_properties_raises_exception_when_insufficient_boto3_version(mocker, boto3_version):
-    mock_boto3_dist_obj = MagicMock()
-    mock_boto3_dist_obj.version = boto3_version
+    from packaging.version import Version
 
-    mocker.patch("pkg_resources.get_distribution", return_value=mock_boto3_dist_obj)
-    import pkg_resources
+    mock_boto3_dist_obj = Version(boto3_version)
 
-    with pytest.raises(pkg_resources.VersionConflict) as excinfo:
+    mocker.patch("redshift_connector.idp_auth_helper.IdpAuthHelper.get_pkg_version", return_value=mock_boto3_dist_obj)
+
+    with pytest.raises(ModuleNotFoundError) as excinfo:
         IamHelper.set_iam_properties(
             make_basic_redshift_property(
                 **{"iam": True, "ssl": True, "auth_profile": "SomeTestProfile", "cluster_identifier": "my_cluster"}