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
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
/*
 * Copyright 2023, Sayan Nandan <nandansayan@outlook.com>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
*/

//! # Queries
//!
//! This module provides the basic tools needed to create and run queries.
//!
//! ## Example
//! ```no_run
//! use skytable::{Config, Query};
//!
//! let mut db = Config::new_default("username", "password").connect().unwrap();
//! let mut query = Query::new("select * from myspace where username = ?");
//! query.push_param(".... get user name from somewhere ...");
//!
//! let ret = db.query(&query).unwrap();
//! ```
//!

use std::{
    io::{self, Write},
    iter::FromIterator,
    num::{
        NonZeroI16, NonZeroI32, NonZeroI64, NonZeroI8, NonZeroIsize, NonZeroU16, NonZeroU32,
        NonZeroU64, NonZeroU8, NonZeroUsize,
    },
};

/*
    query impl
*/

#[derive(Debug, PartialEq, Clone)]
/// A [`Query`] represents a Skyhash query. This is the "standard query" that you will normally use for almost all operations.
///
/// Specification: `QTDEX-A/BQL-S1`
pub struct Query {
    buf: Vec<u8>,
    param_cnt: usize,
    q_window: usize,
}

impl From<String> for Query {
    fn from(q: String) -> Self {
        Self::new_string(q)
    }
}

impl<'a> From<&'a str> for Query {
    fn from(q: &'a str) -> Self {
        Self::new(q)
    }
}

impl Query {
    /// Create a new query from a [`str`]
    pub fn new(query: &str) -> Self {
        Self::_new(query.to_owned())
    }
    /// Create a new query from a [`String`]
    pub fn new_string(query: String) -> Self {
        Self::_new(query)
    }
    fn _new(query: String) -> Self {
        let l = query.len();
        Self {
            buf: query.into_bytes(),
            param_cnt: 0,
            q_window: l,
        }
    }
    /// Returns a reference to the query string
    pub fn query_str(&self) -> &str {
        unsafe { core::str::from_utf8_unchecked(&self.buf[..self.q_window]) }
    }
    /// Add a new parameter to the query
    pub fn push_param(&mut self, param: impl SQParam) -> &mut Self {
        self.param_cnt += param.append_param(&mut self.buf);
        self
    }
    /// Get the number of parameters
    pub fn param_cnt(&self) -> usize {
        self.param_cnt
    }
    #[inline(always)]
    pub(crate) fn write_packet(&self, buf: &mut impl Write) -> io::Result<()> {
        /*
            [[total packet size][query window]][[dataframe][qframe]]
            ^meta1            ^meta2           ^payload
        */
        // compute the total packet size
        // q window
        let mut query_window_buffer = itoa::Buffer::new();
        let query_window_str = query_window_buffer.format(self.q_window);
        // full packet
        let total_packet_size = query_window_str.len() + 1 + self.buf.len();
        let mut total_packet_size_buffer = itoa::Buffer::new();
        let total_packet_size_str = total_packet_size_buffer.format(total_packet_size);
        // segment 1: meta
        buf.write_all(b"S")?;
        buf.write_all(total_packet_size_str.as_bytes())?;
        buf.write_all(b"\n")?;
        // segment 2: variable meta
        buf.write_all(query_window_str.as_bytes())?;
        buf.write_all(b"\n")?;
        // segment 3: payload
        buf.write_all(&self.buf)?;
        Ok(())
    }
    #[inline(always)]
    /// Encodes the packet using Skyhash and returns a raw packet for debugging purposes
    pub fn debug_encode_packet(&self) -> Vec<u8> {
        let mut v = vec![];
        self.write_packet(&mut v).unwrap();
        v
    }
}

/// # Pipeline
///
/// A pipeline can be used to send multiple queries at once to the server. Queries in a pipeline are executed independently
/// of one another, but they are executed serially unless otherwise configured
pub struct Pipeline {
    cnt: usize,
    buf: Vec<u8>,
}

impl Pipeline {
    /// Create a new pipeline
    pub const fn new() -> Self {
        Self {
            cnt: 0,
            buf: Vec::new(),
        }
    }
    pub(crate) fn buf(&self) -> &[u8] {
        &self.buf
    }
    /// Returns the number of queries that were appended to this pipeline
    pub fn query_count(&self) -> usize {
        self.cnt
    }
    /// Same as [`Self::push`], but passes ownership to the [`Pipeline`]
    pub fn push_owned(&mut self, q: Query) {
        self.push(&q);
    }
    /// Add a query to this pipeline
    ///
    /// Note: It's not possible to get the query back from the pipeline since it's not indexed (and doing so would be an unnecessary
    /// waste of space and time). That's why we take a reference which allows the caller to continue owning the [`Query`] item
    pub fn push(&mut self, q: &Query) {
        // qlen
        self.buf
            .extend(itoa::Buffer::new().format(q.q_window).as_bytes());
        self.buf.push(b'\n');
        // plen
        self.buf.extend(
            itoa::Buffer::new()
                .format(q.buf.len() - q.q_window)
                .as_bytes(),
        );
        self.buf.push(b'\n');
        // body
        self.buf.extend(&q.buf);
        self.cnt += 1;
    }
    /// Add a query to this pipeline (builder pattern)
    ///
    /// This is intended to be used with the
    /// ["builder pattern"](https://rust-unofficial.github.io/patterns/patterns/creational/builder.html). For example:
    /// ```
    /// use skytable::{query, Pipeline};
    ///
    /// let pipeline = Pipeline::new()
    ///     .add(&query!("create space myspace"))
    ///     .add(&query!("drop space myspace"));
    /// assert_eq!(pipeline.query_count(), 2);
    /// ```
    pub fn add(mut self, q: &Query) -> Self {
        self.push(q);
        self
    }
}

impl<Q: AsRef<Query>, I> From<I> for Pipeline
where
    I: Iterator<Item = Q>,
{
    fn from(iter: I) -> Self {
        let mut pipeline = Pipeline::new();
        iter.into_iter().for_each(|q| pipeline.push(q.as_ref()));
        pipeline
    }
}

impl<Q: AsRef<Query>> Extend<Q> for Pipeline {
    fn extend<T: IntoIterator<Item = Q>>(&mut self, iter: T) {
        iter.into_iter().for_each(|q| self.push(q.as_ref()))
    }
}

impl<Q: AsRef<Query>> FromIterator<Q> for Pipeline {
    fn from_iter<T: IntoIterator<Item = Q>>(iter: T) -> Self {
        let mut pipe = Pipeline::new();
        iter.into_iter().for_each(|q| pipe.push(q.as_ref()));
        pipe
    }
}

impl AsRef<Query> for Query {
    fn as_ref(&self) -> &Query {
        self
    }
}

/*
    Query parameters
*/

/// An [`SQParam`] should be implemented by any type that is expected to be used as a parameter
///
/// ## Example implementation
///
/// Say you have a custom type which has to store `<username>-<id>` in your database and your database schema looks like
/// `create model myspace.mymodel(username: string, id: uint64, password: string)`. You can do that directly:
///
/// ```
/// use skytable::{query, query::SQParam};
///
/// struct MyType {
///     username: String,
///     id: u64,
///     password: String,
/// }
///
/// impl MyType {
///     fn new(username: String, id: u64, password: String) -> Self {
///         Self { username, id, password }
///     }
/// }
///
/// impl SQParam for MyType {
///     fn append_param(&self, buf: &mut Vec<u8>) -> usize {
///         self.username.append_param(buf) +
///         self.id.append_param(buf) +
///         self.password.append_param(buf)
///     }
/// }
///
/// // You can now directly do this!
/// let query = query!("insert into myspace.mymodel(?, ?, ?)", MyType::new("sayan".to_owned(), 0, "pass123".to_owned()));
/// assert_eq!(query.param_cnt(), 3);
/// // You can also used it beside normal params
/// // assume schema is `create model mymomdel2(uname: string, id: uint64, pass: string, age: uint8)`
/// let query = query!("insert into myspace.mymodel(?, ?, ?, ?)", MyType::new("sayan".to_owned(), 0, "pass123".to_owned()), 101);
/// assert_eq!(query.param_cnt(), 4);
/// ```
pub trait SQParam {
    /// Append this element to the raw parameter buffer
    ///
    /// Return the number of parameters appended (see example above)
    fn append_param(&self, q: &mut Vec<u8>) -> usize;
}
// null
impl<T> SQParam for Option<T>
where
    T: SQParam,
{
    fn append_param(&self, buf: &mut Vec<u8>) -> usize {
        match self {
            None => {
                buf.push(0);
                1
            }
            Some(e) => e.append_param(buf),
        }
    }
}

/// Use this when you need to use `null`
pub struct Null;
impl SQParam for Null {
    fn append_param(&self, buf: &mut Vec<u8>) -> usize {
        buf.push(0);
        1
    }
}
// bool
impl SQParam for bool {
    fn append_param(&self, buf: &mut Vec<u8>) -> usize {
        let a = [1, *self as u8];
        buf.extend(a);
        1
    }
}
macro_rules! imp_number {
    ($($code:literal => $($ty:ty as $base:ty),*),* $(,)?) => {
        $($(impl SQParam for $ty { fn append_param(&self, b: &mut Vec<u8>) -> usize {
            let mut buf = ::itoa::Buffer::new();
            let str = buf.format(<$base>::from(*self));
            b.push($code); b.extend(str.as_bytes()); b.push(b'\n');
            1
        } })*)*
    }
}

macro_rules! imp_terminated_str_type {
    ($($code:literal => $($ty:ty),*),* $(,)?) => {
        $($(impl SQParam for $ty { fn append_param(&self, buf: &mut Vec<u8>) -> usize { buf.push($code); buf.extend(self.to_string().as_bytes()); buf.push(b'\n'); 1} })*)*
    }
}

// uint, sint, float
imp_number!(
    2 => u8 as u8, NonZeroU8 as u8, u16 as u16, NonZeroU16 as u16, u32 as u32, NonZeroU32 as u32, u64 as u64, NonZeroU64 as u64, usize as usize, NonZeroUsize as usize,
    3 => i8 as i8, NonZeroI8 as i8, i16 as i16, NonZeroI16 as i16, i32 as i32, NonZeroI32 as i32, i64 as i64, NonZeroI64 as i64, isize as isize, NonZeroIsize as isize,
);

imp_terminated_str_type!(
    4 => f32, f64
);

// bin
impl<'a> SQParam for &'a [u8] {
    fn append_param(&self, buf: &mut Vec<u8>) -> usize {
        buf.push(5);
        pushlen!(buf, self.len());
        buf.extend(*self);
        1
    }
}
impl<const N: usize> SQParam for [u8; N] {
    fn append_param(&self, buf: &mut Vec<u8>) -> usize {
        buf.push(5);
        pushlen!(buf, self.len());
        buf.extend(self);
        1
    }
}
impl<'a, const N: usize> SQParam for &'a [u8; N] {
    fn append_param(&self, buf: &mut Vec<u8>) -> usize {
        buf.push(5);
        pushlen!(buf, self.len());
        buf.extend(*self);
        1
    }
}
impl SQParam for Vec<u8> {
    fn append_param(&self, buf: &mut Vec<u8>) -> usize {
        buf.push(5);
        pushlen!(buf, self.len());
        buf.extend(self);
        1
    }
}
// str
impl<'a> SQParam for &'a str {
    fn append_param(&self, buf: &mut Vec<u8>) -> usize {
        buf.push(6);
        pushlen!(buf, self.len());
        buf.extend(self.as_bytes());
        1
    }
}
impl<'a> SQParam for &'a String {
    fn append_param(&self, q: &mut Vec<u8>) -> usize {
        self.as_str().append_param(q)
    }
}
impl SQParam for String {
    fn append_param(&self, buf: &mut Vec<u8>) -> usize {
        self.as_str().append_param(buf)
    }
}