Skip to main content

tank_tests/
service.rs

1#![allow(unused_imports)]
2use std::sync::LazyLock;
3use tank::{
4    AsValue, Entity, Error, Executor, QueryBuilder, Result, Value, cols, expr,
5    stream::{StreamExt, TryStreamExt},
6};
7use tokio::sync::Mutex;
8
9static MUTEX: LazyLock<Mutex<()>> = LazyLock::new(|| Mutex::new(()));
10
11#[derive(Clone, Debug, PartialEq, Eq)]
12pub struct HostPort {
13    pub host: String,
14    pub port: u16,
15}
16
17impl HostPort {
18    pub fn new(host: impl Into<String>, port: u16) -> Self {
19        Self {
20            host: host.into(),
21            port,
22        }
23    }
24}
25
26impl AsValue for HostPort {
27    fn as_empty_value() -> Value {
28        Value::Varchar(None)
29    }
30
31    fn as_value(self) -> Value {
32        Value::Varchar(Some(format!("{}:{}", self.host, self.port).into()))
33    }
34
35    fn try_from_value(value: Value) -> Result<Self>
36    where
37        Self: Sized,
38    {
39        // Always call try_as before checking the received value
40        match value.try_as(&Value::Varchar(None)) {
41            Ok(Value::Varchar(Some(v))) => {
42                let context = || Error::msg(format!("Failed to parse HostPort from value `{v}`"));
43                let (host, port) = v.split_once(':').ok_or_else(context)?;
44                Ok(Self {
45                    host: host.to_string(),
46                    port: port
47                        .parse::<u16>()
48                        .map_err(|e| Error::new(e).context(context()))?,
49                })
50            }
51            _ => Err(Error::msg(
52                "Could not convert value into HostPort (expected Value::Varchar)",
53            )),
54        }
55    }
56}
57
58#[derive(Entity, Clone, Debug, PartialEq, Eq)]
59#[tank(schema = "ops")]
60#[tank(primary_key = (name, addr))]
61struct Service {
62    name: String,
63    #[tank(clustering_key)]
64    addr: HostPort,
65    backup_addr: Option<HostPort>,
66}
67
68pub async fn service(executor: &mut impl Executor) {
69    let _lock = MUTEX.lock().await;
70
71    // Setup
72    Service::drop_table(executor, true, false)
73        .await
74        .expect("Failed to drop Service");
75    Service::create_table(executor, false, true)
76        .await
77        .expect("Failed to create Service");
78
79    // Query
80    let mut api = Service {
81        addr: HostPort::new("api.internal", 443),
82        name: "api".into(),
83        backup_addr: None,
84    };
85    api.save(executor).await.expect("Failed to save api");
86    let loaded = Service::find_one(executor, api.primary_key_expr())
87        .await
88        .expect("Failed to load api")
89        .expect("Missing api");
90    assert_eq!(
91        loaded,
92        Service {
93            addr: HostPort::new("api.internal", 443),
94            name: "api".into(),
95            backup_addr: None,
96        }
97    );
98
99    api.backup_addr = Some(HostPort::new("api.internal", 8443));
100    api.save(executor)
101        .await
102        .expect("Failed to update api backup");
103    let loaded = Service::find_one(executor, api.primary_key_expr())
104        .await
105        .expect("Failed to reload api")
106        .expect("Missing api after update");
107    assert_eq!(
108        loaded,
109        Service {
110            addr: HostPort::new("api.internal", 443),
111            name: "api".into(),
112            backup_addr: HostPort::new("api.internal", 8443).into(),
113        }
114    );
115
116    let mut query = Service::prepare_find(executor, expr!(Service::name == ?), None)
117        .await
118        .expect("Failed to prepare query by name");
119    query.bind("api").expect("Failed to bind name parameter");
120
121    let api_by_addr = executor
122        .fetch(query)
123        .map_ok(Service::from_row)
124        .map(Result::flatten)
125        .try_collect::<Vec<_>>()
126        .await
127        .expect("Failed to query by addr");
128    assert_eq!(
129        api_by_addr,
130        [Service {
131            addr: HostPort::new("api.internal", 443),
132            name: "api".into(),
133            backup_addr: HostPort::new("api.internal", 8443).into(),
134        }]
135    );
136
137    let api_canary = Service {
138        addr: HostPort::new("api.internal", 8444),
139        name: "api".into(),
140        backup_addr: None,
141    };
142    Service::insert_one(executor, &api_canary)
143        .await
144        .expect("Failed to insert api canary");
145
146    let web = Service {
147        addr: HostPort::new("web.internal", 80),
148        name: "web".into(),
149        backup_addr: Some(HostPort::new("web.internal", 8080)),
150    };
151    Service::insert_one(executor, &web)
152        .await
153        .expect("Failed to insert web");
154    #[cfg(not(feature = "disable-ordering"))]
155    {
156        let rows = executor
157            .fetch(
158                QueryBuilder::new()
159                    .select(Service::columns())
160                    .from(Service::table())
161                    .where_expr(expr!(Service::name == "api"))
162                    .order_by(cols!(Service::addr ASC))
163                    .build(&executor.driver()),
164            )
165            .map_ok(Service::from_row)
166            .map(Result::flatten)
167            .try_collect::<Vec<_>>()
168            .await
169            .expect("Failed to select ordered services");
170
171        assert_eq!(
172            rows,
173            [
174                Service {
175                    addr: HostPort::new("api.internal", 443),
176                    name: "api".into(),
177                    backup_addr: HostPort::new("api.internal", 8443).into(),
178                },
179                Service {
180                    addr: HostPort::new("api.internal", 8444),
181                    name: "api".into(),
182                    backup_addr: None,
183                },
184            ]
185        );
186    }
187}