1#![allow(unused_imports)]
2use std::collections::BTreeMap;
3use std::str::FromStr;
4use std::sync::{Arc, LazyLock};
5use tank::{Entity, Executor, expr};
6use tokio::sync::Mutex;
7use uuid::Uuid;
8
9use crate::silent_logs;
10
11static MUTEX: LazyLock<Mutex<()>> = LazyLock::new(|| Mutex::new(()));
12
13#[derive(Entity, PartialEq, Debug, Clone)]
14#[tank(schema = "testing", primary_key = (Self::key2, Self::key1))]
15pub struct KV {
16 pub key1: String,
17 pub key2: String,
18 pub string_val: String,
19 pub int_val: i64,
20 pub small_int: i16,
21 pub unsigned_val: u32,
22 pub float_val: f32,
23 pub double_val: f64,
24 pub boolean_val: bool,
25 pub opt_string: Option<String>,
26 pub opt_int: Option<i64>,
27 pub opt_float: Option<f64>,
28 pub opt_bool: Option<bool>,
29 pub uuid: Uuid,
30 pub opt_uuid: Option<Uuid>,
31 #[cfg(not(feature = "disable-arrays"))]
32 pub fixed_bytes: [u8; 16],
33 #[cfg(not(feature = "disable-arrays"))]
34 pub numbers: [i32; 4],
35 #[cfg(not(feature = "disable-lists"))]
36 pub tags: Vec<String>,
37 #[cfg(not(feature = "disable-lists"))]
38 pub scores: Vec<i32>,
39 #[cfg(not(feature = "disable-lists"))]
40 pub floats: Option<Vec<f64>>,
41 #[cfg(not(feature = "disable-lists"))]
42 pub shared_strings: Arc<Vec<String>>,
43 #[cfg(not(feature = "disable-maps"))]
44 pub metadata: BTreeMap<String, String>,
45 #[cfg(not(feature = "disable-maps"))]
46 pub counters: Option<BTreeMap<String, i64>>,
47}
48
49pub async fn kv_storage(executor: &mut impl Executor) {
50 let _lock = MUTEX.lock().await;
51
52 silent_logs! {
54 KV::drop_table(executor, true, false)
56 .await
57 .expect("Failed to drop KV table");
58 }
59 KV::create_table(executor, false, true)
60 .await
61 .expect("Failed to create KV table");
62
63 let first = KV {
65 key1: "first".into(),
66 key2: "aa".into(),
67 string_val: "hello".into(),
68 int_val: 42,
69 small_int: -7,
70 unsigned_val: 99,
71 float_val: 3.14,
72 double_val: 9.87654321,
73 boolean_val: true,
74 opt_string: Some("optional".into()),
75 opt_int: None,
76 opt_float: Some(1.5),
77 opt_bool: None,
78 uuid: Uuid::from_str("c41af414-54cf-49cb-96c6-364e9c42f294").unwrap(),
79 opt_uuid: Some(Uuid::from_str("b1cacbcb-b3e0-4502-bbcd-d4ef014d189b").unwrap()),
80 #[cfg(not(feature = "disable-arrays"))]
81 fixed_bytes: *b"abcdefghijklmnop",
82 #[cfg(not(feature = "disable-arrays"))]
83 numbers: [1, 2, 3, 4],
84 #[cfg(not(feature = "disable-lists"))]
85 tags: vec!["kv".into(), "test".into(), "valkey".into()],
86 #[cfg(not(feature = "disable-lists"))]
87 scores: vec![10, 20, 30],
88 #[cfg(not(feature = "disable-lists"))]
89 floats: Some(vec![1.1, 2.2, 3.3]),
90 #[cfg(not(feature = "disable-lists"))]
91 shared_strings: Arc::new(vec!["a".into(), "b".into(), "c".into()]),
92 #[cfg(not(feature = "disable-maps"))]
93 metadata: BTreeMap::from_iter([
94 ("env".into(), "test".into()),
95 ("region".into(), "eu".into()),
96 ]),
97 #[cfg(not(feature = "disable-maps"))]
98 counters: Some(BTreeMap::from_iter([
99 ("views".into(), 100),
100 ("likes".into(), 5),
101 ])),
102 };
103 let second = KV {
104 key1: "second".into(),
105 key2: "bb".into(),
106 string_val: "world".into(),
107 int_val: 99,
108 small_int: 16,
109 unsigned_val: 0,
110 float_val: 0.005,
111 double_val: 0.00001,
112 boolean_val: false,
113 opt_string: None,
114 opt_int: 19.into(),
115 opt_float: None,
116 opt_bool: Some(true),
117 uuid: Uuid::from_str("4ba01a6a-e1b5-47c8-8d1b-e38b64a73989").unwrap(),
118 opt_uuid: None,
119 #[cfg(not(feature = "disable-arrays"))]
120 fixed_bytes: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
121 #[cfg(not(feature = "disable-arrays"))]
122 numbers: [-10, -20, -30, -40],
123 #[cfg(not(feature = "disable-lists"))]
124 tags: vec![],
125 #[cfg(not(feature = "disable-lists"))]
126 scores: vec![],
127 #[cfg(not(feature = "disable-lists"))]
128 floats: None,
129 #[cfg(not(feature = "disable-lists"))]
130 shared_strings: Arc::new(vec![]),
131 #[cfg(not(feature = "disable-maps"))]
132 metadata: BTreeMap::new(),
133 #[cfg(not(feature = "disable-maps"))]
134 counters: None,
135 };
136 let result = KV::insert_many(executor, [&first, &second])
137 .await
138 .expect("Failed to insert KV entity");
139
140 if let Some(rows) = result.rows_affected {
141 assert_eq!(rows, 2);
142 }
143
144 let value = KV::find_one(executor, expr!(KV::key1 == "first" && KV::key2 == "aa"))
146 .await
147 .expect("Failed to query KV");
148 assert_eq!(value, Some(first));
149
150 let value = KV::find_one(executor, expr!(KV::key1 == "second" && KV::key2 == "bb"))
152 .await
153 .expect("Failed to query KV");
154 assert_eq!(value, Some(second));
155}