Skip to main content

tuwunel_admin/query/raw/
mod.rs

1mod clear;
2mod compact;
3mod count;
4mod del;
5mod flush;
6mod get;
7mod iter;
8mod keys;
9mod keys_sizes;
10mod keys_total;
11mod maps;
12mod put;
13mod sequence;
14mod vals_sizes;
15mod vals_total;
16
17use std::{fmt::Write, sync::Arc};
18
19use clap::Subcommand;
20use tuwunel_core::{Result, err, expected, itertools::Itertools, utils::math::Expected};
21use tuwunel_database::Map;
22use tuwunel_service::Services;
23
24use crate::admin_command_dispatch;
25
26#[admin_command_dispatch(handler_prefix = "raw")]
27#[derive(Debug, Subcommand)]
28/// Query tables from database
29pub(crate) enum RawCommand {
30	/// - List database maps
31	Maps,
32
33	/// - Current rocksdb sequence number.
34	Sequence,
35
36	/// - Raw database query
37	Get {
38		/// Map name
39		map: String,
40
41		/// Key
42		key: String,
43
44		/// Encode as base64
45		#[arg(long, short)]
46		base64: bool,
47	},
48
49	/// - Raw database keys iteration
50	Keys {
51		/// Map name
52		map: String,
53
54		/// Key prefix
55		prefix: Option<String>,
56
57		/// Limit
58		#[arg(short, long)]
59		limit: Option<usize>,
60
61		/// Lower bound
62		#[arg(short, long)]
63		from: Option<String>,
64
65		/// Reverse iteration order
66		#[arg(short, long, default_value("false"))]
67		backwards: bool,
68	},
69
70	/// - Raw database items iteration
71	Iter {
72		/// Map name
73		map: String,
74
75		/// Key prefix
76		prefix: Option<String>,
77
78		/// Limit
79		#[arg(short, long)]
80		limit: Option<usize>,
81
82		/// Lower bound
83		#[arg(short, long)]
84		from: Option<String>,
85
86		/// Reverse iteration order
87		#[arg(short, long, default_value("false"))]
88		backwards: bool,
89	},
90
91	/// - Raw database key size breakdown
92	KeysSizes {
93		/// Map name
94		map: Option<String>,
95
96		/// Key prefix
97		prefix: Option<String>,
98	},
99
100	/// - Raw database keys total bytes
101	KeysTotal {
102		/// Map name
103		map: Option<String>,
104
105		/// Key prefix
106		prefix: Option<String>,
107	},
108
109	/// - Raw database values size breakdown
110	ValsSizes {
111		/// Map name
112		map: Option<String>,
113
114		/// Key prefix
115		prefix: Option<String>,
116	},
117
118	/// - Raw database values total bytes
119	ValsTotal {
120		/// Map name
121		map: Option<String>,
122
123		/// Key prefix
124		prefix: Option<String>,
125	},
126
127	/// - Raw database record count
128	Count {
129		/// Map name
130		map: Option<String>,
131
132		/// Key prefix
133		prefix: Option<String>,
134	},
135
136	/// - Raw database put
137	Put {
138		/// Map name
139		map: String,
140
141		/// Key
142		key: String,
143
144		/// Value
145		value: String,
146	},
147
148	/// - Raw database delete (for string keys) DANGER!!!
149	Del {
150		/// Map name
151		map: String,
152
153		/// Key
154		key: String,
155	},
156
157	/// - Clear database table DANGER!!!
158	Clear {
159		/// Map name
160		map: String,
161
162		/// Confirm
163		#[arg(long)]
164		confirm: bool,
165	},
166
167	/// - Compact database DANGER!!!
168	Compact {
169		#[arg(short, long, alias("column"))]
170		maps: Option<Vec<String>>,
171
172		#[arg(long)]
173		start: Option<String>,
174
175		#[arg(long)]
176		stop: Option<String>,
177
178		#[arg(long)]
179		from: Option<usize>,
180
181		#[arg(long)]
182		into: Option<usize>,
183
184		/// There is one compaction job per column; then this controls how many
185		/// columns are compacted in parallel. If zero, one compaction job is
186		/// still run at a time here, but in exclusive-mode blocking any other
187		/// automatic compaction jobs until complete.
188		#[arg(long)]
189		parallelism: Option<usize>,
190
191		#[arg(long, default_value("false"))]
192		exhaustive: bool,
193	},
194
195	/// - Flush RocksDB memtables to SST files (LSM flush, not fsync/fflush)
196	Flush,
197}
198
199fn with_map_or(map: Option<&str>, services: &Services) -> Result<Vec<Arc<Map>>> {
200	with_maps_or(
201		map.map(|map| [map])
202			.as_ref()
203			.map(<[&str; 1]>::as_slice),
204		services,
205	)
206}
207
208fn with_maps_or<S: AsRef<str>>(maps: Option<&[S]>, services: &Services) -> Result<Vec<Arc<Map>>> {
209	Ok(if let Some(maps) = maps {
210		maps.iter()
211			.map(|map| {
212				let map = map.as_ref();
213				services
214					.db
215					.get(map)
216					.cloned()
217					.map_err(|_| err!("map {map} not found"))
218			})
219			.try_collect()?
220	} else {
221		services.db.iter().map(|x| x.1.clone()).collect()
222	})
223}
224
225fn from_hex(byte: u8) -> Option<u8> {
226	match byte {
227		| 0x30..=0x39 => Some(expected!(byte - 0x30)),
228		| 0x41..=0x46 => Some(expected!((byte - 0x41) + 10)),
229		| 0x61..=0x66 => Some(expected!((byte - 0x61) + 10)),
230		| _ => None,
231	}
232}
233
234fn decode(data: &str) -> Vec<u8> {
235	let mut res = Vec::with_capacity(data.len());
236
237	for byte in data.bytes() {
238		res.push(byte);
239
240		let length = res.len();
241
242		if length >= 4
243			&& let Some(slice) = res.get(expected!(length - 4)..length)
244			&& slice.starts_with(b"\\x")
245			&& let Some(a) = from_hex(slice[2])
246			&& let Some(b) = from_hex(slice[3])
247		{
248			res.truncate(expected!(length - 4));
249
250			let byte = (a << 4) | b;
251			res.push(byte);
252		}
253	}
254
255	res
256}
257
258#[expect(clippy::as_conversions)]
259fn encode(data: &[u8]) -> String {
260	let mut res = String::with_capacity(data.len().expected_mul(4));
261
262	for byte in data {
263		if *byte < 0x20 || *byte > 0x7E {
264			_ = write!(res, "\\x{byte:02x}");
265		} else {
266			res.push(*byte as char);
267		}
268	}
269
270	res
271}