-
Notifications
You must be signed in to change notification settings - Fork 208
Expand file tree
/
Copy pathclient.py
More file actions
166 lines (129 loc) · 4.58 KB
/
Copy pathclient.py
File metadata and controls
166 lines (129 loc) · 4.58 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
"""RPC client using SDK Core. (unstable)
Nothing in this module should be considered stable. The API may change.
"""
from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import timedelta
from typing import TypeVar
import google.protobuf.message
import temporalio.bridge.runtime
import temporalio.bridge.temporal_sdk_bridge
from temporalio.bridge.temporal_sdk_bridge import (
RPCError, # type:ignore[reportUnusedImport]
)
@dataclass
class ClientTlsConfig:
"""Python representation of the Rust struct for configuring TLS."""
server_root_ca_cert: bytes | None
domain: str | None
client_cert: bytes | None
client_private_key: bytes | None
@dataclass
class ClientRetryConfig:
"""Python representation of the Rust struct for configuring retry."""
initial_interval_millis: int
randomization_factor: float
multiplier: float
max_interval_millis: int
max_elapsed_time_millis: int | None
max_retries: int
@dataclass
class ClientKeepAliveConfig:
"""Python representation of the Rust struct for configuring keep alive."""
interval_millis: int
timeout_millis: int
@dataclass
class ClientHttpConnectProxyConfig:
"""Python representation of the Rust struct for configuring HTTP proxy."""
target_host: str
basic_auth: tuple[str, str] | None
@dataclass
class ClientDnsLoadBalancingConfig:
"""Python representation of the Rust struct for configuring DNS load
balancing.
"""
resolution_interval_millis: int
@dataclass
class ClientConfig:
"""Python representation of the Rust struct for configuring the client."""
target_url: str
metadata: Mapping[str, str | bytes]
api_key: str | None
identity: str
tls_config: ClientTlsConfig | None
retry_config: ClientRetryConfig | None
keep_alive_config: ClientKeepAliveConfig | None
client_name: str
client_version: str
http_connect_proxy_config: ClientHttpConnectProxyConfig | None
dns_load_balancing_config: ClientDnsLoadBalancingConfig | None
grpc_compression: str
@dataclass
class RpcCall:
"""Python representation of the Rust struct for an RPC call."""
rpc: str
req: bytes
retry: bool
metadata: Mapping[str, str | bytes]
timeout_millis: int | None
ProtoMessage = TypeVar("ProtoMessage", bound=google.protobuf.message.Message)
class Client:
"""RPC client using SDK Core."""
@staticmethod
async def connect(
runtime: temporalio.bridge.runtime.Runtime, config: ClientConfig
) -> Client:
"""Establish connection with server."""
return Client(
runtime,
await temporalio.bridge.temporal_sdk_bridge.connect_client(
runtime._ref, config
),
)
def __init__(
self,
runtime: temporalio.bridge.runtime.Runtime,
ref: temporalio.bridge.temporal_sdk_bridge.ClientRef,
):
"""Initialize client with underlying SDK Core reference."""
self._runtime = runtime
self._ref = ref
def update_metadata(self, metadata: Mapping[str, str | bytes]) -> None:
"""Update underlying metadata on Core client."""
self._ref.update_metadata(metadata)
def update_api_key(self, api_key: str | None) -> None:
"""Update underlying API key on Core client."""
self._ref.update_api_key(api_key)
async def call(
self,
*,
service: str,
rpc: str,
req: google.protobuf.message.Message,
resp_type: type[ProtoMessage],
retry: bool,
metadata: Mapping[str, str | bytes],
timeout: timedelta | None,
) -> ProtoMessage:
"""Make RPC call using SDK Core."""
# Prepare call
timeout_millis = round(timeout.total_seconds() * 1000) if timeout else None
call = RpcCall(rpc, req.SerializeToString(), retry, metadata, timeout_millis)
# Do call (this throws an RPCError on failure)
if service == "workflow":
resp_fut = self._ref.call_workflow_service(call)
elif service == "operator":
resp_fut = self._ref.call_operator_service(call)
elif service == "cloud":
resp_fut = self._ref.call_cloud_service(call)
elif service == "test":
resp_fut = self._ref.call_test_service(call)
elif service == "health":
resp_fut = self._ref.call_health_service(call)
else:
raise ValueError(f"Unrecognized service {service}")
# Convert response
resp = resp_type()
resp.ParseFromString(await resp_fut)
return resp