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
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
//! High-performance extractors for [tokio_postgres].
//!
//! This crate contains traits and proc macros for creating high-performance extractors
//! for Rust types from [`Row`]s.
//!
//! # Examples
//!
//! ## Extracting a single row
//!
//! ```
//! # use tokio_postgres::{Client, Error};
//! # use tokio_postgres_extractor::RowExtractExt;
//! # use tokio_postgres_extractor::Columns;
//! # use tokio_postgres_extractor::Extract;
//! #[derive(Columns, Extract)]
//! struct User {
//!     id: i32,
//!     name: String,
//! }
//!
//! async fn get_user(client: &Client, id: i32) -> Result<User, Error> {
//!     client
//!         .query_one("select * from user where id = $1", &[&id])
//!         .await
//!         .map(|r| r.extract_once())
//! }
//! ```
//!
//! ## Extracting a stream of rows
//!
//! ```
//! # use futures_util::TryStreamExt;
//! # use tokio_postgres::{Client, Error};
//! # use tokio_postgres_extractor::stream::RowStreamExtractExt;
//! # use tokio_postgres_extractor::Columns;
//! # use tokio_postgres_extractor::Extract;
//! #[derive(Columns, Extract)]
//! struct User {
//!     id: i32,
//!     name: String,
//! }
//!
//! async fn get_users(client: &Client) -> Result<Vec<User>, Error> {
//!     client
//!         .query_raw("select * from user", None::<i32>)
//!         .await?
//!         .extract()
//!         .try_collect()
//!         .await
//! }
//! ```
//!
//! # Generic types
//!
//! Generic types are fully supported.
//!
//! ```
//! # use tokio_postgres_extractor::Columns;
//! # use tokio_postgres_extractor::Extract;
//! #[derive(Columns, Extract)]
//! struct User<'a, T: ?Sized> {
//!     id: i32,
//!     name: &'a T,
//! }
//!
//! fn assert_is_extract<'a, T: Extract<'a>>() { }
//!
//! assert_is_extract::<User<str>>();
//! ```
//!
//! # Custom column names
//!
//! You can specify column names that are different from field names. See the
//! documentation of the [`Columns`][macro@Columns] proc macro.
//!
//! # Design
//!
//! A naive mapping function such as
//!
//! ```rust,ignore
//! User {
//!     id: row.get("id"),
//!     name: row.get("name"),
//! }
//! ```
//!
//! has `O(N^2)` runtime where `N` is the number of columns. Each invocation of `row.get`
//! must walk all columns and compare their name to the name of the requested column. This
//! crate solves this by
//!
//! 1. Constructing an efficient data structure at compile time that reduces lookup time
//!    to `O(N)`.
//!
//!    This data structure is similar to a perfect hash function but more efficient.
//!
//! 2. Memorizing the mapping from fields to columns whenever possible.

extern crate self as tokio_postgres_extractor;

/// Proc macro for deriving the [`Columns`] trait.
///
/// # Custom column names
///
/// If the column name is different from the name of the field, you can use
///
/// ```rust,ignore
/// #[column(name = "Type")]`
/// ty: i32,
/// ```
///
/// to explicitly specify a name. The name must be a string literal.
///
/// # Explicit indices
///
/// If you already know the index a field maps to, you can use
///
/// ```rust,ignore
/// #[column(idx = 123)]`
/// ty: i32,
/// ```
///
/// to specify it.
///
/// # Implementation
///
/// The critical section in the expansion of
///
/// ```
/// # use tokio_postgres_extractor_macros::Columns;
/// #[derive(Columns)]
/// struct Account {
///     account_id: i32,
///     account_name: String,
///     account_role: i64,
/// }
/// ```
///
/// is
///
/// ```rust,ignore
/// for column in row.columns().iter() {
///     let name = column.name();
///     let idx = match name.len() {
///         10 => match name {
///             "account_id" => 0,
///             _ => continue,
///         },
///         12 => {
///             let b = name.as_bytes();
///             let disc = b[8];
///             match disc {
///                 110 => match name {
///                     "account_name" => 1,
///                     _ => continue,
///                 },
///                 114 => match name {
///                     "account_role" => 2,
///                     _ => continue,
///                 },
///                 _ => continue,
///             }
///         }
///         _ => continue,
///     };
///     // ...
/// }
/// ```
///
/// meaning that for each column the code
///
/// 1. uses a jump table indexed with the length of the column name,
/// 2. extracts an integer of size 1, 2, 4, or 8 from the column name,
/// 3. compares this integer to a number of candidates,
/// 4. compares the column name to the candidate.
///
/// This is very fast on current hardware. If the name of the candidate is no more than 16
/// bytes, then the final comparison will compile to a number of assembly instructions on
/// x86_64. Otherwise it compiles to a call to `bcmp`. If you compile with `-C
/// target-cpu=native`, then even these calls to `bcmp` will be inlined.
///
/// In practice, this means that the implementation of `Columns` is usually `O(N)` in the
/// number of columns with a "small" constant. In some degenerate cases the construction
/// above is not possible and the implementation will have to perform multiple string
/// comparisons. Even this is still much faster than using the phf crate or similar.
pub use tokio_postgres_extractor_macros::Columns;
/// Proc macro for deriving the [`Extract`] trait.
pub use tokio_postgres_extractor_macros::Extract;
use {crate::sealed::Sealed, std::ops::Index, tokio_postgres::Row};

pub mod iter;
pub mod stream;

#[cfg(test)]
mod tests;

#[doc(hidden)]
pub mod private {
    pub use tokio_postgres;
}

/// A type whose fields map to Postgres columns.
///
/// This trait is almost always derived with the [Columns](macro@Columns) proc macro:
///
/// ```
/// # use tokio_postgres_extractor::Columns;
/// #[derive(Columns)]
/// struct User<'a> {
///     id: i32,
///     name: &'a str,
/// }
/// ```
///
/// In this case the associated `Columns` type is `[usize; N]` where `N` is the number of
/// fields.
///
/// Assume that a [`Row`] was created from the following query:
///
/// ```sql
/// select 'john' user, 123 balance, 1 id;
/// ```
///
/// Then the implementation of `Columns` for the type `User` above should return `[2, 0]`
/// because the `id` field maps to the third column and the `name` field maps to the
/// first column.
pub trait Columns {
    /// The type identifying the columns.
    ///
    /// This should always be `[usize; N]` where `N` is the number of fields. In a future
    /// version of this crate, this field will be replaced by
    ///
    /// ```
    /// pub trait Columns {
    ///     const NUM_COLUMNS: usize;
    /// }
    /// ```
    type Columns: Unpin + Index<usize, Output = usize>;

    /// Returns the mapping from the type's fields to the columns in a [`Row`].
    fn columns(row: &Row) -> Self::Columns;
}

/// A type that can be extracted from a [`Row`].
///
/// This trait is usually derived:
///
/// ```
/// # use tokio_postgres_extractor_macros::{Columns, Extract};
/// #[derive(Columns, Extract)]
/// struct User<'a> {
///     id: i32,
///     name: &'a str,
/// }
/// ```
///
/// Most of the time you will not use this trait directly but one of the utility functions:
///
/// - [`IterExtractExt::extract`][iter::IterExtractExt::extract]
/// - [`IterExtractRefExt::extract_ref`][iter::IterExtractRefExt::extract_ref]
/// - [`RowStreamExtractExt::extract`][stream::RowStreamExtractExt::extract]
/// - [`RowStreamExtractExt::extract_mut`][stream::RowStreamExtractExt::extract_mut]
///
/// # Examples
///
/// ```
/// # use tokio_postgres::{Client, Error};
/// # use tokio_postgres_extractor::{Columns, Extract};
/// #[derive(Columns, Extract)]
/// struct User {
///     id: i32,
///     name: String,
/// }
///
/// async fn get_user(client: &Client, id: u32) -> Result<User, Error> {
///     client
///         .query_one("select * from user where id = $1", &[&id])
///         .await
///         .map(|r| User::extract_once(&r))
/// }
/// ```
pub trait Extract<'row>: Columns + Sized {
    /// Extracts an instance of the type from a [`Row`].
    ///
    /// If the implementation of [`Columns`] was derived, then this function is almost
    /// always faster than manually calling [`Row::get`] with the column names. Usually it
    /// is much faster.
    ///
    /// However, if you are extracting more than one row at at time, then you will want to
    /// instead call
    ///
    /// - [`Extract::extract`],
    /// - [`IterExtractExt::extract`][iter::IterExtractExt::extract],
    /// - [`IterExtractRefExt::extract_ref`][iter::IterExtractRefExt::extract_ref],
    /// - [`RowStreamExtractExt::extract`][stream::RowStreamExtractExt::extract], or
    /// - [`RowStreamExtractExt::extract_mut`][stream::RowStreamExtractExt::extract_mut].
    ///
    /// These function memorize the output of [`Columns::columns`] and are often 2x as fast
    /// per row as this function.
    ///
    /// # Panics
    ///
    /// Panics if [`Columns::columns`] or [`Row::get`] panics.
    ///
    /// # Examples
    ///
    /// ```
    /// # use tokio_postgres::Row;
    /// # use tokio_postgres_extractor::{Columns, Extract};
    /// #[derive(Columns, Extract)]
    /// struct User<'a> {
    ///     id: i32,
    ///     name: &'a str,
    /// }
    ///
    /// fn map_user(row: &Row) -> User<'_> {
    ///     User::extract_once(row)
    /// }
    /// ```
    fn extract_once(row: &'row Row) -> Self {
        Self::extract(&mut None, row)
    }

    /// Extracts an instance of the type from a [`Row`], memorizing the mapping between
    /// fields and columns.
    ///
    /// If you call this function multiple times with a memorized mapping, then you should
    /// make sure that the columns in the rows are in the same order. This is always the
    /// case if the rows were produced by a single SQL statement. Otherwise this function
    /// might panic or the mapping might be incorrect.
    ///
    /// Often you will want to call
    ///
    /// - [`IterExtractExt::extract`][iter::IterExtractExt::extract],
    /// - [`IterExtractRefExt::extract_ref`][iter::IterExtractRefExt::extract_ref],
    /// - [`RowStreamExtractExt::extract`][stream::RowStreamExtractExt::extract], or
    /// - [`RowStreamExtractExt::extract_mut`][stream::RowStreamExtractExt::extract_mut].
    ///
    /// instead, which hide the memorization behind and abstraction.
    ///
    /// # Panics
    ///
    /// Panics if [`Columns::columns`] or [`Row::get`] panics.
    ///
    /// # Examples
    ///
    /// ```
    /// # use futures_util::StreamExt;
    /// # use tokio_postgres::Row;
    /// # use tokio_postgres_extractor::{Columns, Extract};
    /// #[derive(Columns, Extract)]
    /// struct User<'a> {
    ///     id: i32,
    ///     name: &'a str,
    /// }
    ///
    /// fn map_users(rows: &[Row]) -> Vec<User<'_>> {
    ///     let mut columns = None;
    ///     rows.iter().map(|row| User::extract(&mut columns, row)).collect()
    /// }
    /// ```
    fn extract(columns: &mut Option<<Self as Columns>::Columns>, row: &'row Row) -> Self {
        Self::extract_with_columns(
            columns.get_or_insert_with(|| <Self as Columns>::columns(row)),
            row,
        )
    }

    /// Extracts an instance of the type from a [`Row`] and a mapping between the
    /// fields and columns.
    ///
    /// This function is usually derived with the [`Extract`](macro@Extract) proc macro.
    /// In this case the implementation looks as follows:
    ///
    /// ```
    /// # use futures_util::StreamExt;
    /// # use tokio_postgres::Row;
    /// # use tokio_postgres::types::FromSql;
    /// # use tokio_postgres_extractor::{Columns, Extract};
    /// #[derive(Columns)]
    /// struct User<'a> {
    ///     id: i32,
    ///     name: &'a str,
    /// }
    ///
    /// impl<'a, 'row> Extract<'row> for User<'a>
    /// where
    ///     i32: FromSql<'row>,
    ///     &'a str: FromSql<'row>,
    /// {
    ///     fn extract_with_columns(columns: &Self::Columns, row: &'row Row) -> Self {
    ///         Self {
    ///             id: row.get(columns[0]),
    ///             name: row.get(columns[1]),
    ///         }
    ///     }
    /// }
    /// ```
    ///
    /// # Panics
    ///
    /// Panics if [`Row::get`] panics.
    fn extract_with_columns(columns: &<Self as Columns>::Columns, row: &'row Row) -> Self;
}

/// A type that can be extracted from a [`Row`] without borrowing the [`Row`].
///
/// This is primarily useful for trait bounds on functions. For example
///
/// ```
/// # use tokio_postgres::Row;
/// # use tokio_postgres_extractor::ExtractOwned;
/// # fn get_row() -> Row {
/// #     unimplemented!()
/// # }
/// fn extract<T: ExtractOwned>() -> T {
///     let row: Row = get_row();
///     T::extract_once(&row)
/// }
/// ```
pub trait ExtractOwned: for<'a> Extract<'a> {}

impl<T> ExtractOwned for T where T: for<'a> Extract<'a> {}

/// Extension trait for extracting from a [`Row`].
pub trait RowExtractExt: Sealed {
    /// Extracts an instance of `T` from this [`Row`].
    ///
    /// This is equivalent to [`T::extract_once(self)`][Extract::extract_once].
    ///
    /// # Panics
    ///
    /// Panics if [`Extract::extract_once`] panics.
    ///
    /// # Examples
    ///
    /// ```
    /// # use tokio_postgres::Row;
    /// # use tokio_postgres_extractor::RowExtractExt;
    /// # use tokio_postgres_extractor_macros::{Columns, Extract};
    /// #[derive(Columns, Extract)]
    /// struct User<'a> {
    ///     id: i32,
    ///     name: &'a str,
    /// }
    ///
    /// fn map_user(row: &Row) -> User<'_> {
    ///     row.extract_once()
    /// }
    /// ```
    fn extract_once<'row, T>(&'row self) -> T
    where
        T: Extract<'row>;

    /// Extracts an instance of `T` from this [`Row`], memorizing the mapping between
    /// fields and columns.
    ///
    /// This is equivalent to [`T::extract(columns, self)`][Extract::extract].
    ///
    /// # Panics
    ///
    /// Panics if [`Extract::extract`] panics.
    ///
    /// # Examples
    ///
    /// ```
    /// # use futures_util::StreamExt;
    /// # use tokio_postgres::Row;
    /// # use tokio_postgres_extractor::{Columns, Extract, RowExtractExt};
    /// #[derive(Columns, Extract)]
    /// struct User<'a> {
    ///     id: i32,
    ///     name: &'a str,
    /// }
    ///
    /// fn map_users(rows: &[Row]) -> Vec<User<'_>> {
    ///     let mut columns = None;
    ///     rows.iter().map(|row| row.extract(&mut columns)).collect()
    /// }
    /// ```
    fn extract<'row, T>(&'row self, columns: &mut Option<<T as Columns>::Columns>) -> T
    where
        T: Extract<'row>;

    /// Extracts an instance of `T` from this [`Row`] and a mapping between the
    /// fields and columns.
    ///
    /// This is equivalent to [`T::extract_with_columns(columns, self)`][Extract::extract_with_columns].
    ///
    /// # Panics
    ///
    /// Panics if [`Extract::extract_with_columns`] panics.
    fn extract_with_columns<'row, T>(&'row self, columns: &<T as Columns>::Columns) -> T
    where
        T: Extract<'row>;
}

impl Sealed for Row {}

impl RowExtractExt for Row {
    fn extract_once<'row, T: Extract<'row>>(&'row self) -> T {
        T::extract_once(self)
    }

    fn extract<'row, T>(&'row self, columns: &mut Option<<T as Columns>::Columns>) -> T
    where
        T: Extract<'row>,
    {
        T::extract(columns, self)
    }
    fn extract_with_columns<'row, T>(&'row self, columns: &<T as Columns>::Columns) -> T
    where
        T: Extract<'row>,
    {
        T::extract_with_columns(columns, self)
    }
}

mod sealed {
    pub trait Sealed {}
}