Skip to content

Commit 8e3d78f

Browse files
authored
Add x-restate-server to the response headers (#20)
This commit adds to the HTTP response headers (only seen by restate-server) the x-restate-server header, that contains the sdk name and version.
1 parent 77115b5 commit 8e3d78f

File tree

4 files changed

+23
-12
lines changed

4 files changed

+23
-12
lines changed

Cargo.lock

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

python/restate/server.py

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,14 @@
1919
from restate.server_types import Receive, Scope, Send, binary_to_header, header_to_binary
2020
from restate.vm import VMWrapper
2121
from restate._internal import PyIdentityVerifier, IdentityVerificationException # pylint: disable=import-error,no-name-in-module
22+
from restate._internal import SDK_VERSION # pylint: disable=import-error,no-name-in-module
2223
from restate.aws_lambda import is_running_on_lambda, wrap_asgi_as_lambda_handler
2324

25+
X_RESTATE_SERVER = header_to_binary([("x-restate-server", f"restate-sdk-python/{SDK_VERSION}")])
2426

2527
async def send_status(send, receive, status_code: int):
2628
"""respond with a status code"""
27-
await send({'type': 'http.response.start', 'status': status_code})
29+
await send({'type': 'http.response.start', 'status': status_code, "headers": X_RESTATE_SERVER})
2830
# For more info on why this loop, see ServerInvocationContext.leave()
2931
# pylint: disable=R0801
3032
while True:
@@ -49,10 +51,12 @@ async def send_discovery(scope: Scope, send: Send, endpoint: Endpoint):
4951
else:
5052
discovered_as = "bidi"
5153
headers, js = compute_discovery_json(endpoint, 1, discovered_as)
54+
bin_headers = header_to_binary(headers.items())
55+
bin_headers.extend(X_RESTATE_SERVER)
5256
await send({
5357
'type': 'http.response.start',
5458
'status': 200,
55-
'headers': header_to_binary(headers.items()),
59+
'headers': bin_headers,
5660
'trailers': False
5761
})
5862
await send({
@@ -68,10 +72,12 @@ async def process_invocation_to_completion(vm: VMWrapper,
6872
send: Send):
6973
"""Invoke the user code."""
7074
status, res_headers = vm.get_response_head()
75+
res_bin_headers = header_to_binary(res_headers)
76+
res_bin_headers.extend(X_RESTATE_SERVER)
7177
await send({
7278
'type': 'http.response.start',
7379
'status': status,
74-
'headers': header_to_binary(res_headers),
80+
'headers': res_bin_headers,
7581
'trailers': False
7682
})
7783
assert status == 200

python/restate/vm.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515

1616
from dataclasses import dataclass
1717
import typing
18-
from restate._internal import PyVM, PyFailure, PySuspended, PyVoid, PyStateKeys # pylint: disable=import-error,no-name-in-module
18+
from restate._internal import PyVM, PyFailure, PySuspended, PyVoid, PyStateKeys # pylint: disable=import-error,no-name-in-module
1919

2020
@dataclass
2121
class Invocation:
@@ -59,7 +59,7 @@ def __init__(self, *args: object) -> None:
5959
class VMWrapper:
6060
"""
6161
A wrapper class for the restate_sdk._internal.PyVM class.
62-
It provides a type-friendly interface to our shared vm.
62+
It provides a type-friendly interface to our shared vm.
6363
"""
6464

6565
def __init__(self, headers: typing.List[typing.Tuple[str, str]]):

src/lib.rs

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,16 @@
11
use pyo3::create_exception;
22
use pyo3::prelude::*;
33
use pyo3::types::{PyBytes, PyNone};
4-
use restate_sdk_shared_core::{AsyncResultHandle, CoreVM, Failure, Header, IdentityVerifier, Input, NonEmptyValue, ResponseHead, RunEnterResult, SuspendedOrVMError, TakeOutputResult, Target, VMError, Value, VM};
4+
use restate_sdk_shared_core::{
5+
AsyncResultHandle, CoreVM, Failure, Header, IdentityVerifier, Input, NonEmptyValue,
6+
ResponseHead, RunEnterResult, SuspendedOrVMError, TakeOutputResult, Target, VMError, Value, VM,
7+
};
58
use std::borrow::Cow;
69
use std::time::Duration;
710

11+
// Current crate version
12+
const CURRENT_VERSION: &str = env!("CARGO_PKG_VERSION");
13+
814
// Data model
915

1016
#[pyclass]
@@ -103,7 +109,7 @@ impl From<PyFailure> for Failure {
103109
#[derive(Clone)]
104110
struct PyStateKeys {
105111
#[pyo3(get, set)]
106-
keys: Vec<String>
112+
keys: Vec<String>,
107113
}
108114

109115
#[pyclass]
@@ -237,7 +243,7 @@ impl PyVM {
237243
Ok(PyFailure::from(f).into_py(py).into_bound(py).into_any())
238244
}
239245
Ok(Some(Value::StateKeys(keys))) => {
240-
Ok(PyStateKeys {keys}.into_py(py).into_bound(py).into_any())
246+
Ok(PyStateKeys { keys }.into_py(py).into_bound(py).into_any())
241247
}
242248
}
243249
}
@@ -259,9 +265,7 @@ impl PyVM {
259265
.map_err(Into::into)
260266
}
261267

262-
fn sys_get_state_keys(
263-
mut self_: PyRefMut<'_, Self>,
264-
) -> Result<PyAsyncResultHandle, PyVMError> {
268+
fn sys_get_state_keys(mut self_: PyRefMut<'_, Self>) -> Result<PyAsyncResultHandle, PyVMError> {
265269
self_
266270
.vm
267271
.sys_state_get_keys()
@@ -563,5 +567,6 @@ fn _internal(m: &Bound<'_, PyModule>) -> PyResult<()> {
563567
"IdentityVerificationException",
564568
m.py().get_type_bound::<IdentityVerificationException>(),
565569
)?;
570+
m.add("SDK_VERSION", CURRENT_VERSION)?;
566571
Ok(())
567572
}

0 commit comments

Comments
 (0)