Support QLDB shell migration (#43)

byronlin13 · web-flow · commit 11c6b5d5afa2 · 2021-01-27T16:33:43.000-08:00
Add support for QLDB shell migration
diff --git a/.gitignore b/.gitignore
@@ -17,4 +17,5 @@ __pycache__/
 /doc/_apidoc/
 /docs/build/*
 /build
-/dist
+/dist
+/venv
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,6 +6,8 @@ Add support for obtaining basic server-side statistics on individual statement e
 * `get_consumed_ios` returns a dictionary containing the number of read IO requests for a statement execution.
 * `get_timing_information` returns a dictionary containing the server side processing time in milliseconds for a statement execution.
 * `get_consumed_ios` and `get_timing_information` methods in the `StreamCursor` class are stateful, meaning the statistics returned by them reflect the state at the time of method execution.
+* Add `transaction_id` property in `Executor` to provide the Transaction ID if needed.
+* The `Config` parameter of `QldbDriver` now appends the `user_agent_extra` value instead of overwriting it.
 
 ### Release 3.0.0 (August 20, 2020)
 This is a public and generally available(GA) release of the driver, and this version can be used in production applications.
diff --git a/README.md b/README.md
@@ -5,6 +5,7 @@ to write software that makes use of AmazonQLDB.
 
 [![Latest Version](https://img.shields.io/pypi/v/pyqldb.svg)](https://pypi.python.org/pypi/pyqldb)
 [![Documentation Status](https://readthedocs.org/projects/amazon-qldb-driver-python/badge/?version=latest)](https://amazon-qldb-driver-python.readthedocs.io/en/latest/?badge=latest)
+[![Build Status](https://github.com/awslabs/amazon-qldb-driver-python/workflows/Python%20application/badge.svg)](https://github.com/awslabs/amazon-qldb-driver-python/actions?query=workflow%3A%22Python+application%22)
 
 For our tutorial, see [Python and Amazon QLDB](https://docs.aws.amazon.com/qldb/latest/developerguide/getting-started.python.html).
 
diff --git a/pyqldb/driver/qldb_driver.py b/pyqldb/driver/qldb_driver.py
@@ -31,7 +31,6 @@
 SERVICE_DESCRIPTION = 'QLDB Driver for Python v{}'.format(__version__)
 SERVICE_NAME = 'qldb-session'
 SERVICE_RETRY = {'max_attempts': 0}
-DEFAULT_CONFIG = Config(user_agent_extra=SERVICE_DESCRIPTION, retries=SERVICE_RETRY)
 DEFAULT_RETRY_CONFIG = RetryConfig()
 
 
@@ -74,7 +73,7 @@ class QldbDriver:
     :param aws_session_token: See [1].
 
     :type config: :py:class:`botocore.config.Config`
-    :param config: See [2]. Note that parameter user_agent_extra will be overwritten.
+    :param config: See [2]. Note that parameter user_agent_extra will be appended and retries will be overwritten.
 
     :type boto3_session: :py:class:`boto3.session.Session`
     :param boto3_session: The boto3 session to create the client with (see [1]). The boto3 session is expected to be
@@ -116,9 +115,14 @@ def __init__(self, ledger_name, retry_config=None, read_ahead=0, executor=None,
             if not isinstance(config, Config):
                 raise TypeError('config must be of type botocore.config.Config. Found: {}'
                                 .format(type(config).__name__))
-            self._config = config.merge(DEFAULT_CONFIG)
+            self._config = config
+            self._config.retries = SERVICE_RETRY
+            if self._config.user_agent_extra:
+                self._config.user_agent_extra = ' '.join([SERVICE_DESCRIPTION, self._config.user_agent_extra])
+            else:
+                self._config.user_agent_extra = SERVICE_DESCRIPTION
         else:
-            self._config = DEFAULT_CONFIG
+            self._config = Config(user_agent_extra=SERVICE_DESCRIPTION, retries=SERVICE_RETRY)
 
         if retry_config is not None:
             if not isinstance(retry_config, RetryConfig):
diff --git a/pyqldb/execution/executor.py b/pyqldb/execution/executor.py
@@ -25,6 +25,13 @@ class Executor:
     def __init__(self, transaction):
         self._transaction = transaction
 
+    @property
+    def transaction_id(self):
+        """
+        The **read-only** ID of this transaction.
+        """
+        return self._transaction.transaction_id
+
     def abort(self):
         """
         Abort the transaction and roll back any changes.
diff --git a/tests/unit/test_executor.py b/tests/unit/test_executor.py
@@ -26,9 +26,11 @@
 class TestExecutor(TestCase):
 
     def test_Executor(self, mock_transaction):
+        mock_transaction.transaction_id = 'txnId'
         executor = Executor(mock_transaction)
 
         self.assertEqual(executor._transaction, mock_transaction)
+        self.assertEqual(executor.transaction_id, 'txnId')
 
     def test_abort(self, mock_transaction):
         executor = Executor(mock_transaction)
diff --git a/tests/unit/test_qldb_driver.py b/tests/unit/test_qldb_driver.py
@@ -16,7 +16,7 @@
 from botocore.config import Config
 from boto3.session import Session
 
-from pyqldb.driver.qldb_driver import QldbDriver
+from pyqldb.driver.qldb_driver import QldbDriver, SERVICE_DESCRIPTION
 from pyqldb.errors import DriverClosedError
 from pyqldb.session.qldb_session import QldbSession
 
@@ -27,7 +27,7 @@
 DEFAULT_BACKOFF_BASE = 10
 DEFAULT_TIMEOUT_SECONDS = 0.001
 EMPTY_STRING = ''
-MOCK_CONFIG = Config()
+MOCK_CONFIG = Config(user_agent_extra='user_agent')
 MOCK_LEDGER_NAME = 'QLDB'
 MOCK_MESSAGE = 'message'
 MOCK_BOTO3_SESSION = Session()
@@ -39,22 +39,18 @@ class TestQldbDriver(TestCase):
     @patch('pyqldb.driver.qldb_driver.AtomicInteger')
     @patch('pyqldb.driver.qldb_driver.BoundedSemaphore')
     @patch('pyqldb.driver.qldb_driver.client')
-    @patch('pyqldb.driver.qldb_driver.Config.merge')
-    def test_constructor_with_valid_config(self, mock_config_merge, mock_client, mock_bounded_semaphore,
+    def test_constructor_with_valid_config(self, mock_client, mock_bounded_semaphore,
                                            mock_atomic_integer, mock_queue):
         mock_queue.return_value = mock_queue
         mock_atomic_integer.return_value = mock_atomic_integer
         mock_bounded_semaphore.return_value = mock_bounded_semaphore
         mock_client.return_value = mock_client
-        mock_config_merge.return_value = mock_config_merge
-        mock_config_merge.max_pool_connections = DEFAULT_MAX_CONCURRENT_TRANSACTIONS
 
         qldb_driver = QldbDriver(MOCK_LEDGER_NAME, config=MOCK_CONFIG)
 
-        mock_config_merge.assert_called_once()
         mock_client.assert_called_once_with(DEFAULT_SESSION_NAME, aws_access_key_id=None,
                                             aws_secret_access_key=None, aws_session_token=None,
-                                            config=mock_config_merge, endpoint_url=None, region_name=None, verify=None)
+                                            config=MOCK_CONFIG, endpoint_url=None, region_name=None, verify=None)
         self.assertEqual(qldb_driver._ledger_name, MOCK_LEDGER_NAME)
         self.assertEqual(qldb_driver._retry_config.retry_limit, DEFAULT_RETRY_LIMIT)
         self.assertEqual(qldb_driver._retry_config.base, DEFAULT_BACKOFF_BASE)
@@ -77,26 +73,22 @@ def test_constructor_with_invalid_config(self, mock_client):
     @patch('pyqldb.driver.qldb_driver.AtomicInteger')
     @patch('pyqldb.driver.qldb_driver.BoundedSemaphore')
     @patch('pyqldb.driver.qldb_driver.client')
-    @patch('pyqldb.driver.qldb_driver.Config.merge')
-    def test_default_constructor_with_parameters(self, mock_config_merge, mock_client, mock_bounded_semaphore,
+    def test_default_constructor_with_parameters(self, mock_client, mock_bounded_semaphore,
                                                  mock_atomic_integer, mock_queue):
         mock_queue.return_value = mock_queue
         mock_atomic_integer.return_value = mock_atomic_integer
         mock_bounded_semaphore.return_value = mock_bounded_semaphore
         mock_client.return_value = mock_client
-        mock_config_merge.return_value = mock_config_merge
-        mock_config_merge.max_pool_connections = DEFAULT_MAX_CONCURRENT_TRANSACTIONS
 
         qldb_driver = QldbDriver(MOCK_LEDGER_NAME, region_name=EMPTY_STRING, verify=EMPTY_STRING,
                                  endpoint_url=EMPTY_STRING, aws_access_key_id=EMPTY_STRING,
                                  aws_secret_access_key=EMPTY_STRING, aws_session_token=EMPTY_STRING,
                                  config=MOCK_CONFIG)
 
-        mock_config_merge.assert_called_once()
         mock_client.assert_called_once_with(DEFAULT_SESSION_NAME, region_name=EMPTY_STRING, verify=EMPTY_STRING,
                                             endpoint_url=EMPTY_STRING, aws_access_key_id=EMPTY_STRING,
                                             aws_secret_access_key=EMPTY_STRING, aws_session_token=EMPTY_STRING,
-                                            config=mock_config_merge)
+                                            config=MOCK_CONFIG)
         self.assertEqual(qldb_driver._ledger_name, MOCK_LEDGER_NAME)
         self.assertEqual(qldb_driver._retry_config.retry_limit, DEFAULT_RETRY_LIMIT)
         self.assertEqual(qldb_driver._retry_config.base, DEFAULT_BACKOFF_BASE)
@@ -108,28 +100,22 @@ def test_default_constructor_with_parameters(self, mock_config_merge, mock_clien
         mock_atomic_integer.assert_called_once_with(DEFAULT_MAX_CONCURRENT_TRANSACTIONS)
         mock_queue.assert_called_once_with()
 
-    @patch('pyqldb.driver.qldb_driver.Config.merge')
-    def test_constructor_with_boto3_session(self, mock_config_merge):
+    def test_constructor_with_boto3_session(self):
         mock_session = Mock(spec=MOCK_BOTO3_SESSION)
-        mock_config_merge.return_value = mock_config_merge
-        mock_config_merge.max_pool_connections = DEFAULT_MAX_CONCURRENT_TRANSACTIONS
 
         qldb_driver = QldbDriver(MOCK_LEDGER_NAME, boto3_session=mock_session, config=MOCK_CONFIG)
-        mock_session.client.assert_called_once_with(DEFAULT_SESSION_NAME, config=mock_config_merge, endpoint_url=None,
+        mock_session.client.assert_called_once_with(DEFAULT_SESSION_NAME, config=MOCK_CONFIG, endpoint_url=None,
                                                     verify=None)
         self.assertEqual(qldb_driver._client, mock_session.client())
+        self.assertTrue(SERVICE_DESCRIPTION in qldb_driver._config.user_agent_extra)
 
     @patch('pyqldb.driver.qldb_driver.logger.warning')
-    @patch('pyqldb.driver.qldb_driver.Config.merge')
-    def test_constructor_with_boto3_session_and_parameters_that_may_overwrite(self, mock_config_merge,
-                                                                              mock_logger_warning):
+    def test_constructor_with_boto3_session_and_parameters_that_may_overwrite(self, mock_logger_warning):
         mock_session = Mock(spec=MOCK_BOTO3_SESSION)
-        mock_config_merge.return_value = mock_config_merge
-        mock_config_merge.max_pool_connections = DEFAULT_MAX_CONCURRENT_TRANSACTIONS
         region_name = 'region_name'
         qldb_driver = QldbDriver(MOCK_LEDGER_NAME, boto3_session=mock_session, config=MOCK_CONFIG,
                                  region_name=region_name)
-        mock_session.client.assert_called_once_with(DEFAULT_SESSION_NAME, config=mock_config_merge, endpoint_url=None,
+        mock_session.client.assert_called_once_with(DEFAULT_SESSION_NAME, config=MOCK_CONFIG, endpoint_url=None,
                                                     verify=None)
         self.assertEqual(qldb_driver._client, mock_session.client())
         mock_logger_warning.assert_called_once()
