RDBC-700 Subscription negotiation

Author: poissoncorp

ravendb/documents/commands/subscriptions.py

@@ -63,12 +63,14 @@ def __init__(
         certificate: Optional[str] = None,
         urls: Optional[List[str]] = None,
         node_tag: Optional[str] = None,
+        server_id: Optional[str] = None,
     ):
         self.port = port
         self.url = url
         self.certificate = certificate
         self.urls = urls
         self.node_tag = node_tag
+        self.server_id = server_id
 
     @classmethod
     def from_json(cls, json_dict: Dict) -> TcpConnectionInfo:
@@ -78,6 +80,7 @@ def from_json(cls, json_dict: Dict) -> TcpConnectionInfo:
             json_dict.get("Certificate", None),
             json_dict.get("Urls", None),
             json_dict.get("NodeTag", None),
+            json_dict.get("ServerId", None),
         )
 
 

ravendb/documents/subscriptions/worker.py

@@ -14,7 +14,7 @@
 
 from ravendb import constants
 from ravendb.documents.session.entity_to_json import EntityToJson
-from ravendb.documents.commands.subscriptions import GetTcpInfoForRemoteTaskCommand
+from ravendb.documents.commands.subscriptions import GetTcpInfoForRemoteTaskCommand, TcpConnectionInfo
 from ravendb.documents.session.document_session_operations.in_memory_document_session_operations import (
     InMemoryDocumentSessionOperations,
 )
@@ -319,21 +319,16 @@ def _connect_to_server(self) -> socket:
             except ClientVersionMismatchException:
                 tcp_info = self._legacy_try_get_tcp_info(request_executor)
 
-        self._tcp_client, chosen_url = TcpUtils.connect_with_priority(
-            tcp_info, command.result.certificate, self._store.certificate_pem_path
+        result = TcpUtils.connect_secured_tcp_socket(
+            tcp_info,
+            command.result.certificate,
+            self._store.certificate_pem_path,
+            None,
+            TcpConnectionHeaderMessage.OperationTypes.SUBSCRIPTION,
+            self.__negotiate_protocol_version_for_subscription,
         )
-
-        database_name = self._store.get_effective_database(self._db_name)
-
-        parameters = TcpNegotiateParameters()
-        parameters.database = database_name
-        parameters.operation = TcpConnectionHeaderMessage.OperationTypes.SUBSCRIPTION
-        parameters.version = TcpConnectionHeaderMessage.SUBSCRIPTION_TCP_VERSION
-        parameters.read_response_and_get_version_callback = self._read_server_response_and_get_version
-        parameters.destination_node_tag = self.current_node_tag
-        parameters.destination_url = chosen_url
-
-        self._supported_features = TcpNegotiation.negotiate_protocol_version(self._tcp_client, parameters)
+        self._tcp_client = result.socket
+        self._supported_features = result.supported_features
 
         if self._supported_features.protocol_version <= 0:
             raise RuntimeError(
@@ -363,6 +358,22 @@ def _connect_to_server(self) -> socket:
 
         return self._tcp_client
 
+    def __negotiate_protocol_version_for_subscription(
+        self, chosen_url: str, tcp_info: TcpConnectionInfo, s: socket
+    ) -> TcpConnectionHeaderMessage.SupportedFeatures:
+        database_name = self._store.get_effective_database(self._db_name)
+
+        parameters = TcpNegotiateParameters()
+        parameters.database = database_name
+        parameters.operation = TcpConnectionHeaderMessage.OperationTypes.SUBSCRIPTION
+        parameters.version = TcpConnectionHeaderMessage.SUBSCRIPTION_TCP_VERSION
+        parameters.read_response_and_get_version_callback = self._read_server_response_and_get_version
+        parameters.destination_node_tag = self.current_node_tag
+        parameters.destination_url = chosen_url
+        parameters.destination_server_id = tcp_info.server_id
+
+        return TcpNegotiation.negotiate_protocol_version(s, parameters)
+
     def _legacy_try_get_tcp_info(self, request_executor: RequestExecutor, node: Optional[ServerNode] = None):
         tcp_command = GetTcpInfoCommand(f"Subscription/{self._db_name}", self._db_name)
 
@@ -378,10 +389,10 @@ def _ensure_parser(self) -> None:
         # python doesn't use parsers
         pass
 
-    def _read_server_response_and_get_version(self, url: str) -> int:
+    def _read_server_response_and_get_version(self, url: str, sock: socket) -> int:
         # reading reply from server
         self._ensure_parser()
-        response = self._tcp_client.recv(self._options.receive_buffer_size)
+        response = sock.recv(self._options.receive_buffer_size)
         reply = TcpConnectionHeaderResponse.from_json(json.loads(response.decode("utf-8")))
 
         if reply.status == TcpConnectionStatus.OK:
@@ -897,6 +908,10 @@ def items(self) -> List[Item[_T]]:
     def number_of_items_in_batch(self) -> int:
         return 0 if self.items is None else len(self.items)
 
+    @property
+    def number_of_includes(self) -> int:
+        return len(self._includes) if self._includes is not None else 0
+
     def open_session(self, options: Optional[SessionOptions] = None) -> DocumentSession:
         if not options:
             options = SessionOptions()

ravendb/serverwide/tcp.py

@@ -253,7 +253,8 @@ def __init__(
         source_node_tag: Optional[str] = None,
         destination_node_tag: Optional[str] = None,
         destination_url: Optional[str] = None,
-        read_response_and_get_version_callback: Optional[Callable[[str], int]] = None,
+        read_response_and_get_version_callback: Optional[Callable[[str, socket.socket], int]] = None,
+        destination_server_id: Optional[str] = None,
     ):
         self.operation = operation
         self.authorize_info = authorize_info
@@ -263,6 +264,7 @@ def __init__(
         self.destination_node_tag = destination_node_tag
         self.destination_url = destination_url
         self.read_response_and_get_version_callback = read_response_and_get_version_callback
+        self.destination_server_id = destination_server_id
 
 
 class TcpNegotiation:
@@ -280,7 +282,7 @@ def negotiate_protocol_version(
         current = parameters.version
         while True:
             cls._send_tcp_version_info(sock, parameters, current)
-            version = parameters.read_response_and_get_version_callback(parameters.destination_url)
+            version = parameters.read_response_and_get_version_callback(parameters.destination_url, sock)
 
             cls.logger.info(
                 f"Read response from {parameters.source_node_tag or parameters.destination_url} "

ravendb/util/tcp_utils.py

@@ -1,9 +1,10 @@
 import base64
 import socket
 import ssl
-from typing import Tuple, Optional
+from typing import Tuple, Optional, Callable
 
 from ravendb.documents.commands.subscriptions import TcpConnectionInfo
+from ravendb.serverwide.tcp import TcpConnectionHeaderMessage
 
 
 class TcpUtils:
@@ -55,3 +56,55 @@ def connect_with_priority(
         )
 
         return s, info.url
+
+    @staticmethod
+    def invoke_negotiation(
+        info: TcpConnectionInfo,
+        operation_type: TcpConnectionHeaderMessage.OperationTypes,
+        negotiation_callback,
+        url: str,
+        socket: socket.socket,
+    ) -> TcpConnectionHeaderMessage.SupportedFeatures:
+        if operation_type == TcpConnectionHeaderMessage.OperationTypes.SUBSCRIPTION:
+            return negotiation_callback(url, info, socket)
+        else:
+            raise NotImplementedError(f"Operation type '{operation_type}' not supported")
+
+    class ConnectSecuredTcpSocketResult:
+        def __init__(
+            self,
+            url: str = None,
+            socket: socket.socket = None,
+            supported_features: TcpConnectionHeaderMessage.SupportedFeatures = None,
+        ):
+            self.url = url
+            self.socket = socket
+            self.supported_features = supported_features
+
+    @staticmethod
+    def connect_secured_tcp_socket(
+        info: TcpConnectionInfo,
+        server_certificate: str,
+        client_certificate_pem_path: str,
+        certificate_private_key_password: Optional[str],
+        operation_type: TcpConnectionHeaderMessage.OperationTypes,
+        negotiation_callback: Callable,
+    ) -> ConnectSecuredTcpSocketResult:
+        if info.urls:
+            for url in info.urls:
+                try:
+                    s = TcpUtils.connect(
+                        url, server_certificate, client_certificate_pem_path, certificate_private_key_password
+                    )
+                    supported_features = TcpUtils.invoke_negotiation(info, operation_type, negotiation_callback, url, s)
+
+                    return TcpUtils.ConnectSecuredTcpSocketResult(url, s, supported_features)
+                except Exception as e:
+                    pass
+                    # ignored
+        s = TcpUtils.connect(
+            info.url, server_certificate, client_certificate_pem_path, certificate_private_key_password
+        )
+
+        supported_features = TcpUtils.invoke_negotiation(info, operation_type, negotiation_callback, info.url, s)
+        return TcpUtils.ConnectSecuredTcpSocketResult(info.url, s, supported_features)