client.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License 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.

import threading
from typing import Optional, List, Type

from dubbo.bootstrap import Dubbo
from dubbo.classes import MethodDescriptor
from dubbo.configs import ReferenceConfig
from dubbo.constants import common_constants
from dubbo.extension import extensionLoader
from dubbo.protocol import Invoker, Protocol
from dubbo.proxy import RpcCallable, RpcCallableFactory
from dubbo.proxy.callables import DefaultRpcCallableFactory
from dubbo.registry.protocol import RegistryProtocol
from dubbo.types import (
    DeserializingFunction,
    RpcTypes,
    SerializingFunction,
)
from dubbo.url import URL
from dubbo.codec import DubboSerializationService

__all__ = ["Client"]


class Client:
    def __init__(self, reference: ReferenceConfig, dubbo: Optional[Dubbo] = None):
        self._initialized = False
        self._global_lock = threading.RLock()

        self._dubbo = dubbo or Dubbo()
        self._reference = reference

        self._url: Optional[URL] = None
        self._protocol: Optional[Protocol] = None
        self._invoker: Optional[Invoker] = None

        self._callable_factory: RpcCallableFactory = DefaultRpcCallableFactory()

        # initialize the invoker
        self._initialize()

    def _initialize(self):
        """
        Initialize the invoker.
        """
        with self._global_lock:
            if self._initialized:
                return

            # get the protocol
            protocol = extensionLoader.get_extension(Protocol, self._reference.protocol)()

            registry_config = self._dubbo.registry_config

            self._protocol = RegistryProtocol(registry_config, protocol) if registry_config else protocol

            # build url
            reference_url = self._reference.to_url()
            if registry_config:
                self._url = registry_config.to_url().copy()
                self._url.path = reference_url.path
                for k, v in reference_url.parameters.items():
                    self._url.parameters[k] = v
            else:
                self._url = reference_url

            # create invoker
            self._invoker = self._protocol.refer(self._url)

            self._initialized = True

    def _create_rpc_callable(
        self,
        rpc_type: str,
        method_name: str,
        params_types: List[Type],
        return_type: Type,
        codec: Optional[str] = None,
        request_serializer: Optional[SerializingFunction] = None,
        response_deserializer: Optional[DeserializingFunction] = None,
    ) -> RpcCallable:
        """
        Create RPC callable with the specified type.
        """
        print("2", params_types)
        # Determine serializers
        if request_serializer and response_deserializer:
            req_ser = request_serializer
            res_deser = response_deserializer
        else:
            req_ser, res_deser = DubboSerializationService.create_serialization_functions(
                codec or "json",
                parameter_types=params_types,
                return_type=return_type,
            )

        # Create MethodDescriptor
        descriptor = MethodDescriptor(
            method_name=method_name,
            arg_serialization=(req_ser, None),
            return_serialization=(None, res_deser),
            rpc_type=rpc_type,
        )

        return self._callable(descriptor)

    def unary(self, method_name: str, params_types: List[Type], return_type: Type, **kwargs) -> RpcCallable:
        return self._create_rpc_callable(
            rpc_type=RpcTypes.UNARY.value,
            method_name=method_name,
            params_types=params_types,
            return_type=return_type,
            **kwargs,
        )

    def client_stream(self, method_name: str, params_types: List[Type], return_type: Type, **kwargs) -> RpcCallable:
        return self._create_rpc_callable(
            rpc_type=RpcTypes.CLIENT_STREAM.value,
            method_name=method_name,
            params_types=params_types,
            return_type=return_type,
            **kwargs,
        )

    def server_stream(self, method_name: str, params_types: List[Type], return_type: Type, **kwargs) -> RpcCallable:
        return self._create_rpc_callable(
            rpc_type=RpcTypes.SERVER_STREAM.value,
            method_name=method_name,
            params_types=params_types,
            return_type=return_type,
            **kwargs,
        )

    def bi_stream(self, method_name: str, params_types: List[Type], return_type: Type, **kwargs) -> RpcCallable:
        return self._create_rpc_callable(
            rpc_type=RpcTypes.BI_STREAM.value,
            method_name=method_name,
            params_types=params_types,
            return_type=return_type,
            **kwargs,
        )

    def _callable(self, method_descriptor: MethodDescriptor) -> RpcCallable:
        """
        Generate a proxy for the given method.
        """
        url = self._invoker.get_url().copy()
        url.parameters[common_constants.METHOD_KEY] = method_descriptor.get_method_name()
        url.attributes[common_constants.METHOD_DESCRIPTOR_KEY] = method_descriptor
        return self._callable_factory.get_callable(self._invoker, url)