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

use std::error::Error;
use std::io::Write;

/// Structure to define column name and and column value
///
/// # Example
/// ```
/// extern crate redshift;
/// use std::str;
///
/// struct TestItem<'a> {
///     pub A: &'a str,
///     pub B: &'a str,
///     pub C: &'a str,
/// }
///
/// fn main() {
///     let test_column_definitions = vec![
///         redshift::writer::ColumnDefinition::<TestItem> {
///             name:  "Acolumn",
///             extract_column: Box::new(move |i: &TestItem| i.A.to_string().clone()),
///         },
///         redshift::writer::ColumnDefinition::<TestItem> {
///             name: "Bcolumn",
///             extract_column: Box::new(move |i: &TestItem| i.B.to_string().clone()),
///         },
///         redshift::writer::ColumnDefinition::<TestItem> {
///             name: "Ccolumn",
///             extract_column: Box::new(move |i: &TestItem| i.C.to_string().clone()),
///         },
///     ];
/// }
/// ```
pub struct ColumnDefinition<'a, T> {

    /// Column name
    pub name: &'a str,

    /// Method to extract the column value from item `T`
    pub extract_column: Box<(Fn (&T) -> String)>,
}

/// Writer for redshift files
///
/// # Examples
/// ```
/// extern crate redshift;
/// 
/// use std::str;
/// 
/// struct TestItem<'a> {
///     pub A: &'a str,
///     pub B: &'a str,
///     pub C: &'a str,
/// }
/// 
/// fn main() {
/// 
///     // Arrange
///     let test_column_definitions = vec![
///         redshift::writer::ColumnDefinition::<TestItem> {
///             name:  "Acolumn",
///             extract_column: Box::new(move |i: &TestItem| i.A.to_string().clone()),
///         },
///         redshift::writer::ColumnDefinition::<TestItem> {
/// 
///             name: "Bcolumn",
///             extract_column: Box::new(move |i: &TestItem| i.B.to_string().clone()),
///         },
///         redshift::writer::ColumnDefinition::<TestItem> {
///             name: "Ccolumn",
///             extract_column: Box::new(move |i: &TestItem| i.C.to_string().clone()),
///         },
///     ];
///     let items = vec![
///         TestItem { A: "a1", B: "b1", C: "c1" },
///         TestItem { A: "a2", B: "b2", C: "c2" },
///         TestItem { A: "a3", B: "b3", C: "c3" },
///     ];
/// 
///     // Act
///     let mut byte_vec: Vec<u8> = Vec::new();
///     let mut writer = redshift::writer::Writer::new();
///     let res = writer.write_columns(&mut byte_vec, test_column_definitions, items);
/// }
/// ```
pub struct Writer {
    output_header: bool,
}

impl Writer {

    /// Construct a new writer
    ///
    /// # Examples
    /// ```
    /// extern crate redshift;
    /// use std::str;
    ///
    /// let mut writer = redshift::writer::Writer::new();
    ///
    /// ```
    pub fn new<'a>() -> Writer {
        Writer {
            output_header: false,
        }
    }


    /// Construct a new writer
    ///
    /// # Examples
    /// ```
    /// extern crate redshift;
    /// use std::str;
    ///
    /// let mut writer = redshift::writer::Writer::new();
    /// let mut writer_with_header = writer.with_header();
    /// ```
    pub fn with_header(&mut self) -> &mut Writer {
        self.output_header = true;
        self
    }

    /// Write out rows to a writer `W` using the provided column definitions and item list
    pub fn write_columns<W: Write, T>(&self, mut writer: W, column_definitions: Vec<ColumnDefinition<T>>, items: Vec<T>) -> Result<(), Box<Error>> {

        // write headers
        if self.output_header {
            for i in 0..column_definitions.len() {
                let escaped = escape_column_value(String::from(column_definitions[i].name));
                try!(writer.write(escaped.as_bytes()));
                if i < column_definitions.len() - 1 {
                    try!(writer.write(b"|"));
                }
            }
            try!(writer.write(b"\n"));
        }

        // encode each item and write it out to the stream
        for item in &items {

            // for each column, extract the column value from the item
            for i in 0..column_definitions.len() {
                let escaped = escape_column_value((column_definitions[i].extract_column)(item));
                try!(writer.write(escaped.as_bytes()));
                if i < column_definitions.len() - 1 {
                    try!(writer.write(b"|"));
                }
            }
            try!(writer.write(b"\n"));
            try!(writer.flush());
        }
        Ok(())
    }
}

fn escape_column_value(value: String) -> String {
    let escaped = value
        .replace("\\", "\\\\")
        .replace("|", "\\|")
        .replace("\n", "\\\n")
        .replace("\r", "\\\r")
        .replace("'", "\'")
        .replace("\"", "\\\"");
    return format!("\"{}\"", escaped);
}