autopush_common/db/bigtable/bigtable_client/
metadata.rs

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
/// gRPC metadata Resource prefix header
///
/// Generic across Google APIs. This "improves routing by the backend" as
/// described by other clients
const PREFIX_KEY: &str = "google-cloud-resource-prefix";

/// gRPC metadata Client information header
///
/// A `User-Agent` like header, likely its main use is for GCP's metrics
const METRICS_KEY: &str = "x-goog-api-client";

/// gRPC metadata Dynamic Routing header:
/// https://google.aip.dev/client-libraries/4222
///
/// See the googleapis protobuf for which routing header params are used for
/// each Spanner operation (under the `google.api.http` option).
///
/// https://github.com/googleapis/googleapis/blob/master/google/spanner/v1/spanner.proto
const ROUTING_KEY: &str = "x-goog-request-params";

/// gRPC metadata Leader Aware Routing header
///
/// Not well documented. Added to clients in early 2023 defaulting to disabled.
/// Clients have began defaulting it to enabled in late 2023.
///
/// "Enabling leader aware routing would route all requests in RW/PDML
/// transactions to the leader region." as described by other Spanner clients
const LEADER_AWARE_KEY: &str = "x-goog-spanner-route-to-leader";

/// The USER_AGENT string is a static value specified by Google.
/// Its meaning is not to be known to the uninitiated.
const USER_AGENT: &str = "gl-external/1.0 gccl/1.0";

/// Builds the [grpcio::Metadata] for all db operations
#[derive(Default)]
pub struct MetadataBuilder<'a> {
    prefix: &'a str,
    routing_params: Vec<(&'a str, &'a str)>,
    route_to_leader: bool,
}

impl<'a> MetadataBuilder<'a> {
    /// Initialize a new builder with a [PREFIX_KEY] header for the given
    /// resource
    pub fn with_prefix(prefix: &'a str) -> Self {
        Self {
            prefix,
            ..Default::default()
        }
    }

    /// Add a [ROUTING_KEY] header
    /// This normally specifies the session name, but unlike spanner, bigtable does not appear to have one of those?
    pub fn routing_param(mut self, key: &'a str, value: &'a str) -> Self {
        self.routing_params.push((key, value));
        self
    }

    /// Toggle the [LEADER_AWARE_KEY] header
    pub fn route_to_leader(mut self, route_to_leader: bool) -> Self {
        self.route_to_leader = route_to_leader;
        self
    }

    /// Build the [grpcio::Metadata]
    pub fn build(self) -> Result<grpcio::Metadata, grpcio::Error> {
        let mut meta = grpcio::MetadataBuilder::new();

        meta.add_str(PREFIX_KEY, self.prefix)?;
        meta.add_str(METRICS_KEY, USER_AGENT)?;
        if self.route_to_leader {
            meta.add_str(LEADER_AWARE_KEY, "true")?;
        }
        if !self.routing_params.is_empty() {
            meta.add_str(ROUTING_KEY, &self.routing_header())?;
        }
        Ok(meta.build())
    }

    fn routing_header(self) -> String {
        let mut ser = form_urlencoded::Serializer::new(String::new());
        for (key, val) in self.routing_params {
            ser.append_pair(key, val);
        }
        // python-spanner (python-api-core) doesn't encode '/':
        // https://github.com/googleapis/python-api-core/blob/6251eab/google/api_core/gapic_v1/routing_header.py#L85
        ser.finish().replace("%2F", "/")
    }
}

#[cfg(test)]
mod tests {
    use std::{collections::HashMap, str};

    use super::{
        MetadataBuilder, LEADER_AWARE_KEY, METRICS_KEY, PREFIX_KEY, ROUTING_KEY, USER_AGENT,
    };

    // Resource paths should not start with a "/"
    pub const DB: &str = "projects/foo/instances/bar/databases/gorp";
    pub const SESSION: &str = "projects/foo/instances/bar/databases/gorp/sessions/f00B4r_quuX";

    #[test]
    fn metadata_basic() {
        let meta = MetadataBuilder::with_prefix(DB)
            .routing_param("session", SESSION)
            .routing_param("foo", "bar baz")
            .build()
            .unwrap();
        let meta: HashMap<_, _> = meta.into_iter().collect();

        assert_eq!(meta.len(), 3);
        assert_eq!(str::from_utf8(meta.get(PREFIX_KEY).unwrap()).unwrap(), DB);
        assert_eq!(
            str::from_utf8(meta.get(METRICS_KEY).unwrap()).unwrap(),
            USER_AGENT
        );
        assert_eq!(
            str::from_utf8(meta.get(ROUTING_KEY).unwrap()).unwrap(),
            format!("session={SESSION}&foo=bar+baz")
        );
    }

    #[test]
    fn leader_aware() {
        let meta = MetadataBuilder::with_prefix(DB)
            .route_to_leader(true)
            .build()
            .unwrap();
        let meta: HashMap<_, _> = meta.into_iter().collect();

        assert_eq!(meta.len(), 3);
        assert_eq!(
            str::from_utf8(meta.get(LEADER_AWARE_KEY).unwrap()).unwrap(),
            "true"
        );
    }
}