Add query support for elasticity changes in the system

Author: gmfeinberg

CHANGELOG.md

@@ -19,7 +19,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/).
     23.3 and higher and will not be immediately available in the cloud service
 -   Cloud only:  added support for Cloud feature Global Active Tables. This
     includes new requests and classes:
-    
+
     - AddReplicaRequest
     - DropReplicaRequest
     - ReplicaStatsRequest/Result
@@ -29,6 +29,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/).
     as well as additional replica-related information
     and interfaces in TableResult
 
+## Changed
+
+- Modified internal query processing to better support elasticity operations
 
 ## Fixed
 

src/borneo/client.py

@@ -22,7 +22,8 @@
                         OperationNotSupportedException, RequestSizeLimitException)
 from .http import RateLimiterMap, RequestUtils
 from .kv import StoreAccessTokenProvider
-from .operations import GetTableRequest, QueryResult, TableRequest, WriteRequest
+from .operations import (
+    GetTableRequest, QueryRequest, QueryResult, TableRequest, WriteRequest)
 from .query import QueryDriver
 from .serdeutil import SerdeUtil
 from .stats import StatsControl
@@ -94,6 +95,8 @@ def __init__(self, config, logger):
         self._sess.mount(self._url.scheme + '://', adapter)
         if self._proxy_host is not None:
             self._check_and_set_proxy(self._sess)
+        self.query_version = QueryDriver.QUERY_VERSION
+        self._topology_info = None
         self.serial_version = config.get_serial_version()
         # StoreAccessTokenProvider means onprem
         self._is_cloud = not isinstance(self._auth_provider, StoreAccessTokenProvider)
@@ -202,7 +205,6 @@ def execute(self, request):
                     'QueryRequest has no QueryDriver, but is prepared', 2)
                 driver = QueryDriver(request)
                 driver.set_client(self)
-                driver.set_topology_info(request.topology_info())
                 return QueryResult(request, False)
             """
             If we are here, then this is either (a) a simple query or (b) an
@@ -216,7 +218,11 @@ def execute(self, request):
             """
             self._trace(
                 'QueryRequest has no QueryDriver and is not prepared', 2)
+            request.incr_batch_counter()
+
         timeout_ms = request.get_timeout()
+        if request is QueryRequest or request.is_query_request():
+            request.set_topo_seq_num(self.get_topo_seq_num())
         headers = {'Host': self._url.hostname,
                    'Content-Type': 'application/octet-stream',
                    'Connection': 'keep-alive',
@@ -256,9 +262,27 @@ def execute(self, request):
         request_utils = RequestUtils(
             self._sess, self._logutils, request, self._retry_handler, self,
             self._rate_limiter_map)
-        return request_utils.do_post_request(self._request_uri, headers,
-                                             content, timeout_ms,
-                                             self._stats_control)
+        res = request_utils.do_post_request(self._request_uri, headers,
+                                            content, timeout_ms,
+                                            self._stats_control)
+        self._set_topology_info(res.get_topology_info())
+        if res is QueryResult and request.is_query_request():
+            request.set_query_traces(res.get_query_traces())
+        return res
+
+    @synchronized
+    def _set_topology_info(self, topo):
+        if topo is None:
+            return
+        if self.get_topo_seq_num() < topo.get_seq_num():
+            self._topology_info = topo
+
+    def get_topo_seq_num(self):
+        return -1 if self._topology_info is None else \
+          self._topology_info.get_seq_num()
+
+    def get_topology(self):
+        return self._topology_info
 
     # set the session cookie if in return headers (see RequestUtils in http.py)
     @synchronized
@@ -484,6 +508,24 @@ def _update_table_limiters(self, table_name):
                 'Background thread added limiters for table "' + table_name +
                 '"')
 
+    def decrement_query_version(self):
+        """
+        Decrements the query version, if it is greater than the minimum.
+        For internal use only.
+
+        The current minimum value is V3.
+        :returns: true if the version was decremented, false otherwise.
+        :rtype: bool
+        """
+        if self.query_version > QueryDriver.QUERY_V3:
+            self.query_version -= 1
+            msg = ('Unsupported query version error, decrementing ' +
+                   'version to ' + str(self.query_version) +
+                   ' and retrying')
+            self._logutils.log_info(msg)
+            return True
+        return False
+
     def decrement_serial_version(self):
         """
         Decrements the serial version, if it is greater than the minimum.
@@ -534,6 +576,8 @@ def serialize_request(self, request, headers):
         :returns: the bytearray that contains the content.
         :rtype: bytearray
         """
+        # in case it's a query or prepare
+        request.set_query_version(self.query_version)
         content = self._write_content(request)
         headers.update({'Content-Length': str(len(content))})
         return content

src/borneo/common.py

@@ -1075,7 +1075,7 @@ class PreparedStatement(object):
     """
     OPCODE_SELECT = 5
 
-    def __init__(self, sql_text, query_plan, query_schema, topology_info,
+    def __init__(self, sql_text, query_plan, query_schema,
                  proxy_statement, driver_plan, num_iterators, num_registers,
                  external_vars, namespace, table_name, operation):
         """
@@ -1091,8 +1091,6 @@ def __init__(self, sql_text, query_plan, query_schema, topology_info,
         self._sql_text = sql_text
         self._query_plan = query_plan
         self._query_schema = query_schema
-        # Applicable to advanced queries only.
-        self._topology_info = topology_info
         # The serialized PreparedStatement created at the backend store. It is
         # opaque for the driver. It is received from the proxy and sent back to
         # the proxy every time a new batch of results is needed.
@@ -1140,7 +1138,7 @@ def copy_statement(self):
         :rtype: PreparedStatement
         """
         return PreparedStatement(
-            self._sql_text, self._query_plan, self._query_schema, self._topology_info,
+            self._sql_text, self._query_plan, self._query_schema,
             self._proxy_statement, self._driver_query_plan, self._num_iterators,
             self._num_registers, self._variables, self._namespace,
             self._table_name, self._operation)
@@ -1227,16 +1225,6 @@ def print_driver_plan(self):
             return self._driver_query_plan.display()
         return None
 
-    @synchronized
-    def set_topology_info(self, topology_info):
-        if topology_info is None:
-            return
-        if self._topology_info is None:
-            self._topology_info = topology_info
-            return
-        if self._topology_info.get_seq_num() < topology_info.get_seq_num():
-            self._topology_info = topology_info
-
     def set_variable(self, variable, value):
         """
         Binds an external variable to a given value. The variable is identified
@@ -1277,14 +1265,6 @@ def set_variable(self, variable, value):
             raise IllegalArgumentException(
                 'There is no external variable at position ' + str(variable))
 
-    def topology_info(self):
-        return self._topology_info
-
-    @synchronized
-    def topology_seq_num(self):
-        return (-1 if self._topology_info is None else
-                self._topology_info.get_seq_num())
-
 
 class PutOption(object):
     """

src/borneo/exception.py

@@ -301,7 +301,7 @@ class UnsupportedQueryVersionException(NoSQLException):
     """
 
     def __init__(self, message):
-        super(UnauthorizedException, self).__init__(message)
+        super(UnsupportedQueryVersionException, self).__init__(message)
 
 
 class IndexExistsException(ResourceExistsException):

src/borneo/http.py

@@ -17,8 +17,8 @@
 from .exception import (
     IllegalStateException, NoSQLException,
     ReadThrottlingException, RequestTimeoutException, RetryableException,
-    SecurityInfoNotReadyException, UnsupportedProtocolException,
-    WriteThrottlingException)
+    SecurityInfoNotReadyException, UnsupportedQueryVersionException,
+    UnsupportedProtocolException, WriteThrottlingException)
 from .serdeutil import SerdeUtil
 
 try:
@@ -323,7 +323,6 @@ def _do_request(self, method, uri, headers, payload, timeout_ms,
                         stats_config.observe(self._request, req_size,
                                              len(response.content),
                                              network_time)
-
                     # check for a Set-Cookie header
                     cookie = response.headers.get('Set-Cookie', None)
                     if cookie is not None and cookie.startswith('session='):
@@ -405,6 +404,19 @@ def _do_request(self, method, uri, headers, payload, timeout_ms,
                 self._request.increment_retries()
                 exception = re
                 continue
+            except UnsupportedQueryVersionException as uqve:
+                if self._client.decrement_query_version():
+                    if self._request is not None:
+                        payload = self._client.serialize_request(self._request,
+                                                                 headers)
+                    self._request.increment_retries()
+                    # don't set exception for this case -- it is misleading
+                    # exception = uqve
+                    continue
+                self._logutils.log_error(
+                    'Client execution UnsupportedQueryVersionException: ' +
+                    str(uqve))
+                raise uqve
             except UnsupportedProtocolException as upe:
                 if self._client.decrement_serial_version():
                     if self._request is not None: