Merge pull request #4 from MeltanoLabs/refactor_to_boto_authenticator

refactor: refactor to boto authenticator class

tap_dynamodb/aws_auth.py → tap_dynamodb/aws_authenticators.py

@@ -4,8 +4,13 @@
 import boto3
 
 
-class AWSAuth:
-    def __init__(self, config):
+class AWSBotoAuthenticator:
+    def __init__(self, config, service_name):
+        self._service_name = service_name
+        self._config = config
+        self._client = None
+        self._resource = None
+        self.logger = logging.getLogger(__name__)
         # config for use environment variables
         if config.get("use_aws_env_vars"):
             self.aws_access_key_id = os.environ.get("AWS_ACCESS_KEY_ID")
@@ -23,7 +28,23 @@ def __init__(self, config):
         self.aws_endpoint_url = config.get("aws_endpoint_url")
         self.aws_assume_role_arn = config.get("aws_assume_role_arn")
 
-        self.logger = logging.getLogger(__name__)
+    @property
+    def client(self):
+        if self._client:
+            return self._client
+        else:
+            session = self.get_session()
+            self._client = self.get_client(session, self._service_name)
+            return self._client
+
+    @property
+    def resource(self):
+        if self._resource:
+            return self._resource
+        else:
+            session = self.get_session()
+            self._resource = self.get_resource(session, self._service_name)
+            return self._resource
 
     def get_session(self):
         session = None

tap_dynamodb/dynamo.py

@@ -2,11 +2,11 @@
 import orjson
 from botocore.exceptions import ClientError
 
-from tap_dynamodb.aws_base import AWSBase
+from tap_dynamodb.aws_authenticators import AWSBotoAuthenticator
 from tap_dynamodb.exception import EmptyTableException
 
 
-class DynamoDB(AWSBase):
+class DynamoDB(AWSBotoAuthenticator):
     def __init__(self, config):
         super().__init__(config, "dynamodb")
 

tests/test_aws_auth.py → tests/test_aws_authenticators.py

@@ -1,21 +1,22 @@
 import pytest
-from contextlib import nullcontext as does_not_raise
-from unittest import mock
 from moto import mock_dynamodb, mock_sts
-from tap_dynamodb.aws_auth import AWSAuth
-import os
+from tap_dynamodb.aws_authenticators import AWSBotoAuthenticator
 from unittest.mock import patch
-import botocore
 
-@patch("tap_dynamodb.aws_auth.boto3.Session", return_value="mock_session")
+
+@patch(
+    "tap_dynamodb.aws_authenticators.boto3.Session",
+    return_value="mock_session",
+)
 @mock_dynamodb
 def test_get_session_base(patch):
-    auth = AWSAuth(
+    auth = AWSBotoAuthenticator(
         {
             "aws_access_key_id": "foo",
             "aws_secret_access_key": "bar",
             "aws_default_region": "baz",
-        }
+        },
+        "dynamodb",
     )
     session = auth.get_session()
     patch.assert_called_with(
@@ -25,16 +26,18 @@ def test_get_session_base(patch):
     )
     assert session == "mock_session"
 
-@patch("tap_dynamodb.aws_auth.boto3.Session", return_value="mock_session")
+
+@patch("tap_dynamodb.aws_authenticators.boto3.Session", return_value="mock_session")
 @mock_dynamodb
 def test_get_session_w_token(patch):
-    auth = AWSAuth(
+    auth = AWSBotoAuthenticator(
         {
             "aws_access_key_id": "foo",
             "aws_secret_access_key": "bar",
             "aws_session_token": "abc",
             "aws_default_region": "baz",
-        }
+        },
+        "dynamodb",
     )
     session = auth.get_session()
     patch.assert_called_with(
@@ -46,13 +49,14 @@ def test_get_session_w_token(patch):
     assert session == "mock_session"
 
 
-@patch("tap_dynamodb.aws_auth.boto3.Session", return_value="mock_session")
+@patch("tap_dynamodb.aws_authenticators.boto3.Session", return_value="mock_session")
 @mock_dynamodb
 def test_get_session_w_profile(patch):
-    auth = AWSAuth(
+    auth = AWSBotoAuthenticator(
         {
             "aws_profile": "foo",
-        }
+        },
+        "dynamodb",
     )
     session = auth.get_session()
     patch.assert_called_with(
@@ -64,44 +68,48 @@ def test_get_session_w_profile(patch):
 @mock_dynamodb
 def test_get_session_empty_fail():
     with pytest.raises(Exception):
-        auth = AWSAuth({})
+        auth = AWSBotoAuthenticator({})
         auth.get_session()
 
 
 @mock_dynamodb
 @mock_sts
 def test_get_session_assume_role():
-    auth = AWSAuth(
+    auth = AWSBotoAuthenticator(
         {
             "aws_access_key_id": "foo",
             "aws_secret_access_key": "bar",
             "aws_default_region": "baz",
-            "aws_assume_role_arn": "arn:aws:iam::123456778910:role/my-role-name"
-        }
+            "aws_assume_role_arn": "arn:aws:iam::123456778910:role/my-role-name",
+        },
+        "dynamodb",
     )
     session = auth.get_session()
 
 
 @mock_dynamodb
 def test_get_client():
-    auth = AWSAuth(
+    auth = AWSBotoAuthenticator(
         {
             "aws_access_key_id": "foo",
             "aws_secret_access_key": "bar",
             "aws_default_region": "baz",
-        }
+        },
+        "dynamodb",
     )
     session = auth.get_session()
     client = auth.get_client(session, "dynamodb")
 
+
 @mock_dynamodb
 def test_get_resource():
-    auth = AWSAuth(
+    auth = AWSBotoAuthenticator(
         {
             "aws_access_key_id": "foo",
             "aws_secret_access_key": "bar",
             "aws_default_region": "baz",
-        }
+        },
+        "dynamodb",
     )
     session = auth.get_session()
     resource = auth.get_resource(session, "dynamodb")