Add execution stats

Dominitio · alpian · commit 06c8146be081 · 2021-01-26T12:22:40.000-08:00
Add get_consumed_ios and get_timing_information methods in BufferedCursor and StreamCursor class.
diff --git a/.github/workflows/pythonapp.yml b/.github/workflows/pythonapp.yml
@@ -23,6 +23,7 @@ jobs:
       with:
         aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
         aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+        aws-session-token: ${{ secrets.AWS_SESSION_TOKEN }}
         aws-region: us-east-1
     - uses: actions/checkout@v2
     - name: Set up Python ${{ matrix.python-version }}
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,12 @@
+### Release 3.1.0
+Add support for obtaining basic server-side statistics on individual statement executions.
+
+#### :tada: Enhancements
+* Added `get_consumed_ios` and `get_timing_information` methods in `BufferedCursor` and `StreamCursor` classes to provide server-side execution statistics.
+* `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.
+
 ### 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/pyqldb/__init__.py b/pyqldb/__init__.py
@@ -9,4 +9,4 @@
 # CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
 # and limitations under the License.
 
-__version__ = '3.0.0'
+__version__ = '3.1.0'
diff --git a/pyqldb/cursor/buffered_cursor.py b/pyqldb/cursor/buffered_cursor.py
@@ -23,6 +23,8 @@ def __init__(self, cursor):
             self._buffered_values.append(item)
 
         self._buffered_values_iterator = iter(self._buffered_values)
+        self._consumed_ios = cursor.get_consumed_ios()
+        self._timing_information = cursor.get_timing_information()
 
     def __iter__(self):
         """
@@ -35,3 +37,21 @@ def __next__(self):
         Iterator function to implement the iterator protocol. Get next value in _buffered_values_iterator.
         """
         return next(self._buffered_values_iterator)
+
+    def get_consumed_ios(self):
+        """
+        Return a dictionary containing the total amount of IO requests for a statement's execution.
+
+        :rtype: dict
+        :return: The amount of read IO requests for a statement's execution.
+        """
+        return self._consumed_ios
+
+    def get_timing_information(self):
+        """
+        Return a dictionary containing the total amount of processing time for a statement's execution.
+
+        :rtype: dict
+        :return: The amount of processing time in milliseconds for a statement's execution.
+        """
+        return self._timing_information
diff --git a/pyqldb/cursor/read_ahead_cursor.py b/pyqldb/cursor/read_ahead_cursor.py
@@ -25,8 +25,8 @@ class ReadAheadCursor(StreamCursor):
     An iterable class representing a read ahead cursor on a statement's result set. This class will create a queue of
     size `read_ahead` and fetch results asynchronously to fill the queue.
 
-    :type page: dict
-    :param page: The page containing the initial result set data dictionary of the statement's execution.
+    :type statement_result: dict
+    :param statement_result: The initial result set data dictionary of the statement's execution.
 
     :type session: :py:class:`pyqldb.communication.session_client.SessionClient`
     :param session: The parent session that represents the communication channel to QLDB.
@@ -40,8 +40,8 @@ class ReadAheadCursor(StreamCursor):
     :type executor: :py:class:`concurrent.futures.thread.ThreadPoolExecutor`
     :param executor: The optional executor for asynchronous retrieval. If none specified, a new thread is created.
     """
-    def __init__(self, page, session, transaction_id, read_ahead, executor):
-        super().__init__(page, session, transaction_id)
+    def __init__(self, statement_result, session, transaction_id, read_ahead, executor):
+        super().__init__(statement_result, session, transaction_id)
         self._queue = Queue(read_ahead - 1)
         if executor is None:
             thread = Thread(target=self._populate_queue)
@@ -63,23 +63,24 @@ def _next_page(self):
         queue_result = self._queue.get()
         if isinstance(queue_result, Exception):
             raise queue_result
-        self._page = queue_result
+        super()._accumulate_query_stats(queue_result)
+        self._page = queue_result.get('Page')
         self._index = 0
 
     def _populate_queue(self):
         """
-        Fill the buffer queue with pages. If ClientError is received, it is put in the queue and execution stops.
-        If the parent transaction is closed, stop fetching results.
+        Fill the buffer queue with the statement_result fetched. If ClientError is received, it is put in the queue and
+        execution stops. If the parent transaction is closed, stop fetching results.
         """
         try:
             next_page_token = self._page.get('NextPageToken')
             while next_page_token is not None:
                 statement_result = self._session._fetch_page(self._transaction_id, next_page_token)
-                page = statement_result.get('Page')
                 while True:
                     try:
                         # Timeout of 50ms.
-                        self._queue.put(page, timeout=0.05)
+                        self._queue.put(statement_result, timeout=0.05)
+                        page = statement_result.get('Page')
                         next_page_token = page.get('NextPageToken')
                         break
                     except Full:
diff --git a/pyqldb/cursor/stream_cursor.py b/pyqldb/cursor/stream_cursor.py
@@ -17,21 +17,26 @@ class StreamCursor:
     """
     An iterable class representing a stream cursor on a statement's result set.
 
-    :type page: dict
-    :param page: The page containing the initial result set data dictionary of the statement's execution.
+    :type statement_result: dict
+    :param statement_result: The initial result set data dictionary of the statement's execution.
 
     :type session: :py:class:`pyqldb.communication.session_client.SessionClient`
     :param session: The parent session that represents the communication channel to QLDB.
 
     :type transaction_id: str
     :param transaction_id: The ID of this cursor's parent transaction, required to fetch pages.
     """
-    def __init__(self, page, session, transaction_id):
-        self._page = page
+    def __init__(self, statement_result, session, transaction_id):
+        self._page = statement_result.get('FirstPage')
         self._session = session
         self._index = 0
         self._is_open = True
         self._transaction_id = transaction_id
+        self._read_ios = None
+        self._write_ios = None
+        self._processing_time_milliseconds = None
+
+        self._accumulate_query_stats(statement_result)
 
     def __iter__(self):
         """
@@ -71,6 +76,24 @@ def close(self):
         """
         self._is_open = False
 
+    def get_consumed_ios(self):
+        """
+        Return a dictionary containing the accumulated amount of IO requests for a statement's execution.
+
+        :rtype: dict
+        :return: The amount of read IO requests for a statement's execution.
+        """
+        return {'ReadIOs': self._read_ios}
+
+    def get_timing_information(self):
+        """
+        Return a dictionary containing the accumulated amount of processing time for a statement's execution.
+
+        :rtype: dict
+        :return: The amount of processing time in milliseconds for a statement's execution.
+        """
+        return {'ProcessingTimeMilliseconds': self._processing_time_milliseconds}
+
     def _are_there_more_results(self):
         """
         Check if there are more results.
@@ -82,10 +105,33 @@ def _next_page(self):
         Get the next page using this cursor's session.
         """
         statement_result = self._session._fetch_page(self._transaction_id, self._page.get('NextPageToken'))
+        self._accumulate_query_stats(statement_result)
+
         page = statement_result.get('Page')
         self._page = page
         self._index = 0
 
+    def _accumulate_query_stats(self, statement_result):
+        """
+        From the statement_result, get the query stats and accumulate them.
+        """
+        self._processing_time_milliseconds = self._accumulate(statement_result, 'TimingInformation',
+                                                              'ProcessingTimeMilliseconds',
+                                                              self._processing_time_milliseconds)
+        self._read_ios = self._accumulate(statement_result, 'ConsumedIOs', 'ReadIOs', self._read_ios)
+        self._write_ios = self._accumulate(statement_result, 'ConsumedIOs', 'WriteIOs', self._write_ios)
+
+    @staticmethod
+    def _accumulate(statement_result, query_statistics_key, metric_key, metric_to_accumulate):
+        query_statistics = statement_result.get(query_statistics_key)
+        if query_statistics:
+            metric = query_statistics.get(metric_key)
+            if metric:
+                if metric_to_accumulate is None:
+                    metric_to_accumulate = 0
+                metric_to_accumulate += metric
+        return metric_to_accumulate
+
     @staticmethod
     def _value_holder_to_ion_value(value):
         """
diff --git a/pyqldb/driver/qldb_driver.py b/pyqldb/driver/qldb_driver.py
@@ -186,7 +186,7 @@ def list_tables(self):
         Get the list of table names in the ledger.
 
         :rtype: :py:class:`pyqldb.cursor.buffered_cursor.BufferedCursor`
-        :return: Iterable of table names in :py:class:`amazon.ion.simple_types.IonPyText`.
+        :return: Iterable of table names in :py:class:`amazon.ion.simple_types.IonPyText` format found in the ledger.
 
         :raises DriverClosedError: When this driver is closed.
         """
@@ -202,9 +202,12 @@ def execute_lambda(self, query_lambda, retry_config=None):
         This is the primary method to execute a transaction against Amazon QLDB ledger.
 
         :type query_lambda: function
-        :param query_lambda: The lambda function to execute. A lambda function cannot have any side effects as
-                             it may be invoked multiple times, and the result cannot be trusted until the transaction is
-                             committed.
+        :param query_lambda: The lambda function to execute. The function receives an instance of
+                             :py:class:`pyqldb.execution.executor.Executor` which can be used to execute statements.
+                             The instance of :py:class:`pyqldb.execution.executor.Executor` wraps an implicitly created
+                             transaction. The transaction will be implicitly committed when the passed function returns.
+                             The lambda function cannot have any side effects as it may be invoked multiple
+                             times, and the result cannot be trusted until the transaction is committed.
 
         :type retry_config: :py:class:`pyqldb.config.retry_config.RetryConfig`
         :param retry_config: Config to specify max number of retries, base and custom backoff strategy for retries.
diff --git a/pyqldb/transaction/transaction.py b/pyqldb/transaction/transaction.py
@@ -164,11 +164,11 @@ def _execute_statement(self, statement, *parameters):
         parameters = tuple(map(self._to_ion, parameters))
         self._update_hash(statement, parameters)
         statement_result = self._session._execute_statement(self._id, statement, parameters)
-        first_page = statement_result.get('FirstPage')
+
         if self._read_ahead > 0:
-            cursor = ReadAheadCursor(first_page, self._session, self._id, self._read_ahead, self._executor)
+            cursor = ReadAheadCursor(statement_result, self._session, self._id, self._read_ahead, self._executor)
         else:
-            cursor = StreamCursor(first_page, self._session, self._id)
+            cursor = StreamCursor(statement_result, self._session, self._id)
 
         self._cursors.append(cursor)
         return cursor
diff --git a/requirements.txt b/requirements.txt
@@ -1,6 +1,6 @@
-amazon.ion~=0.5.0
-boto3~=1.9.237
-botocore~=1.12.237
+amazon.ion~=0.7.0
+boto3~=1.16.56
+botocore~=1.19.56
 ionhash~=1.1.0
 pytest~=4.6.3
 pytest-cov~=2.7.1
diff --git a/setup.py b/setup.py
@@ -14,9 +14,9 @@
 
 ROOT = os.path.join(os.path.dirname(__file__), 'pyqldb')
 VERSION_RE = re.compile(r'''__version__ = ['"]([0-9.a-z\-]+)['"]''')
-requires = ['amazon.ion>=0.5.0,<0.6',
-            'boto3>=1.9.237,<2',
-            'botocore>=1.12.237,<2',
+requires = ['amazon.ion>=0.7.0,<1',
+            'boto3>=1.16.56,<2',
+            'botocore>=1.19.56,<2',
             'ionhash>=1.1.0,<2']
 
 
diff --git a/tests/integration/constants.py b/tests/integration/constants.py
@@ -19,5 +19,6 @@
 SINGLE_DOCUMENT_VALUE = "SingleDocumentValue"
 MULTIPLE_DOCUMENT_VALUE_1 = "MultipleDocumentValue1"
 MULTIPLE_DOCUMENT_VALUE_2 = "MultipleDocumentValue2"
+MULTIPLE_DOCUMENT_VALUE_3 = "MultipleDocumentValue3"
 
 LEDGER_NAME = "Pytest"
diff --git a/tests/integration/test_statement_execution.py b/tests/integration/test_statement_execution.py
@@ -9,7 +9,6 @@
 # CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
 # and limitations under the License.
 from time import sleep
-
 import pytest
 from unittest import TestCase
 
@@ -19,7 +18,7 @@
 from pyqldb.errors import is_occ_conflict_exception
 
 from .constants import CREATE_TABLE_NAME, COLUMN_NAME, MULTIPLE_DOCUMENT_VALUE_1, MULTIPLE_DOCUMENT_VALUE_2, \
-    INDEX_ATTRIBUTE, LEDGER_NAME, SINGLE_DOCUMENT_VALUE, TABLE_NAME
+    MULTIPLE_DOCUMENT_VALUE_3, INDEX_ATTRIBUTE, LEDGER_NAME, SINGLE_DOCUMENT_VALUE, TABLE_NAME
 from .integration_test_base import IntegrationTestBase
 
 
@@ -266,6 +265,60 @@ def execute_statement_and_return_value(txn, query, *parameters):
         # Then.
         self.assertEqual(SINGLE_DOCUMENT_VALUE, value)
 
+    def test_query_stats(self):
+        def fill_table():
+            # Create Ion structs to insert to populate table
+            parameter_1 = loads(dumps({COLUMN_NAME: MULTIPLE_DOCUMENT_VALUE_1}))
+            parameter_2 = loads(dumps({COLUMN_NAME: MULTIPLE_DOCUMENT_VALUE_2}))
+            parameter_3 = loads(dumps({COLUMN_NAME: MULTIPLE_DOCUMENT_VALUE_3}))
+            query = "INSERT INTO {} <<?, ?, ?>>".format(TABLE_NAME)
+
+            self.qldb_driver.execute_lambda(
+                lambda txn: txn.execute_statement(query, parameter_1, parameter_2, parameter_3))
+
+        def assert_io_usage_and_timing_information(cursor):
+            # Consumed IO test
+            consumed_ios = cursor.get_consumed_ios()
+            self.assertIsNotNone(consumed_ios)
+            read_ios = consumed_ios.get('ReadIOs')
+            self.assertTrue(read_ios >= 0)
+
+            # Timing information test
+            timing_information = cursor.get_timing_information()
+            self.assertIsNotNone(timing_information)
+            processing_time_milliseconds = timing_information.get('ProcessingTimeMilliseconds')
+            self.assertTrue(processing_time_milliseconds >= 0)
+
+        def stream_cursor_query_stats(query):
+            def execute_statement_and_test_stream_cursor(txn, query):
+                cursor = txn.execute_statement(query)
+                for row in cursor:
+                    assert_io_usage_and_timing_information(cursor)
+
+            self.qldb_driver.execute_lambda(
+                lambda txn: execute_statement_and_test_stream_cursor(txn, query))
+
+        def buffered_cursor_query_stats(query):
+            # Test buffered cursor's accumulated values of execution stats
+            def execute_statement_and_test_buffered_cursor(txn, query):
+                cursor = txn.execute_statement(query)
+                return cursor
+
+            buffered_cursor = self.qldb_driver.execute_lambda(
+                lambda txn: execute_statement_and_test_buffered_cursor(txn, query))
+            assert_io_usage_and_timing_information(buffered_cursor)
+
+        search_query = "SELECT * FROM {} as n1, {} as n2, {} as n3, {} as n4, {} as n5, {} as n6".format(TABLE_NAME,
+                                                                                                         TABLE_NAME,
+                                                                                                         TABLE_NAME,
+                                                                                                         TABLE_NAME,
+                                                                                                         TABLE_NAME,
+                                                                                                         TABLE_NAME)
+        # Call the routines
+        fill_table()
+        stream_cursor_query_stats(search_query)
+        buffered_cursor_query_stats(search_query)
+
     def test_insert_multiple_documents(self):
         # Given.
         # Create Ion structs to insert.
diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py
@@ -0,0 +1,10 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance with
+# the License. A copy of the License is located at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+# and limitations under the License.
diff --git a/tests/unit/helper_functions.py b/tests/unit/helper_functions.py
diff --git a/tests/unit/test_buffered_cursor.py b/tests/unit/test_buffered_cursor.py
diff --git a/tests/unit/test_read_ahead_cursor.py b/tests/unit/test_read_ahead_cursor.py
diff --git a/tests/unit/test_stream_cursor.py b/tests/unit/test_stream_cursor.py
diff --git a/tests/unit/test_transaction.py b/tests/unit/test_transaction.py
