<commit_before>//! This module currently provides utility methods and structs for variable
//! byte codes as described in
//! http://nlp.stanford.edu/IR-book/html/htmledition/variable-byte-codes-1.html.
//!
//! Encode unsigned integers by using the `vbyte_encode` method.
//!
//! Decode a bytestream by instatiating a `VByteDecoder` and using its iterator
//! implementation.
//!
//! #Example
//!
//! ```rust,ignore
//!
//! use perlin::utils::compression::{vbyte_encode, VByteDecoder};
//!
//! let bytes = vbyte_encode(3);
//! let three = VByteDecoder::new(bytes.into_iter()).next().unwrap();
//! assert_eq!(3, three);
//! ```

use std::io;
use std::io::{Seek, SeekFrom, Read, Write, Error};

/// Stores the result of a vbyte encode operation without indirection that a `Vec<u8>` would introduce.
/// Can thus be used to `vbyte_encode` on the stack
pub struct VByteEncoded {
    // Memory layout of the result: [ payload | payload | payload | ... | used_bytes ]
    data: [u8; 11],
}

impl VByteEncoded {
    /// Performs a vbyte encode without allocating memory on the heap
    /// Can then be written to a `Write`-implementor
    pub fn new(mut number: usize) -> Self {
        let mut count = 0;
        let mut result = [0u8; 11];
        loop {
            result[9 - count] = (number % 128) as u8;
            count += 1;
            if number < 128 {
                break;
            } else {
                number /= 128;
            }
        }
        result[9] += 128;
        result[10] = count as u8;
        VByteEncoded { data: result }
    }

    /// Return how many bytes are used to encode the number. Min: 1 Max: 10
    pub fn bytes_used(&self) -> u8 {
        self.data[10]
    }

    /// Access to the written data as slice
    /// Currently only used in tests
    pub fn data_buf(&self) -> &[u8] {
        &self.data[(10 - self.bytes_used()) as usize..10]
    }

    /// Writes the given VByteEncoded number to a target.
    /// Returns the number of bytes written (equal to `bytes_used`) or an `io::Error`
    pub fn write_to<W: Write>(&self, target: &mut W) -> Result<u8, Error> {
        target.write_all(&self.data[(10 - self.bytes_used()) as usize..10]).map(|()| self.bytes_used())
    }
}

/// Iterator that decodes a bytestream to unsigned integers
pub struct VByteDecoder<R> {
    source: R,
    buf: [u8; 10],
    filled: u8,
}

impl<R: Read> VByteDecoder<R> {
    /// Create a new VByteDecoder by passing a bytestream
    pub fn new(source: R) -> Self {
        VByteDecoder {
            source: source,
            buf: [0; 10],
            filled: 0,
        }
    }

    /// Sometimes it is convenient to look at the original bytestream itself
    /// (e.g. when not only vbyte encoded integers are in the bytestream)
    /// This method provides access to the underlying bytestream in form of
    /// a mutable borrow
    pub fn underlying_mut(&mut self) -> &mut R {
        &mut self.source
    }

    /// Sometimes it is convenient to look at the original bytestream itself
    /// (e.g. when not only vbyte encoded integers are in the bytestream)
    /// This method provides access to the underlying bytestream in form of
    /// a borrow
    pub fn underlying(&self) -> &R {
        &self.source
    }
}

impl<R: Read> Read for VByteDecoder<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let mut bytes_read = 0;
        if self.filled > 0 {
            // We have some bytes stored. Give them back
            // Read some bytes from self.buf.
            bytes_read += try!((&self.buf[..self.filled as usize]).read(buf));
            // These lines "shift" self.buf to the left. Eliminating the read bytes
            let mut tmp: [u8; 10] = [0; 10];
            self.filled -= bytes_read as u8;
            tmp[..10 - bytes_read].copy_from_slice(&self.buf[bytes_read..]);
            self.buf = tmp;
            // If buf is full. Return
            if bytes_read >= buf.len() {
                return Ok(bytes_read);
            }
        }
        bytes_read += try!(self.source.read(&mut buf[bytes_read..]));
        Ok(bytes_read)
    }
}

impl<R: Seek + Read> Seek for VByteDecoder<R> {
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        self.buf = [0; 10];
        match pos {
            SeekFrom::Start(_) => {
                self.filled = 0;
                self.source.seek(pos)
            }
            SeekFrom::Current(offset) => {
                let f = self.filled as i64;
                self.filled = 0;
                self.source.seek(SeekFrom::Current(offset - f))
            }
            SeekFrom::End(_) => {
                self.filled = 0;
                self.source.seek(pos)
            }
        }
    }
}

impl<R: Read> Iterator for VByteDecoder<R> {
    type Item = usize;

    /// Returns the next unsigned integer which is encoded in the underlying
    /// bytestream
    /// May iterate the underlying bytestream an arbitrary number of times
    /// Returns None when the underlying bytream returns None or delivers corrupt data
    fn next(&mut self) -> Option<Self::Item> {
        let mut result: usize = 0;
        // Read bytes into buffer
        let read = self.source.read(&mut self.buf[self.filled as usize..]).unwrap();
        let mut ptr = 0;
        // Find the last byte of this number (The first that has the 128bit set)
        while ptr < self.buf.len() && self.buf[ptr] < 128 {
            ptr += 1;
        }
        // If we are not filled and we didnt read anything (Source is empty)
        // Or if we have corrupted data (No 128bit flag set)
        // Return none
        if (self.filled == 0 && read == 0) || ptr == 10 {
            return None;
        }
        // Specialcase where 10 bytes are read
        // In this case the first byte can be maximally 1
        // Else its corrupt data
        if self.buf[0] > 1 && ptr == 9 {
            return None;
        }
        let mut tmp: [u8; 10] = [0; 10];
        {
            // Split buffer so that the lhs is the number we want to decode
            let (a, buf) = self.buf.split_at(ptr + 1);
            // Copy the rest into a temporary
            tmp[..buf.len()].copy_from_slice(buf);
            self.filled = buf.len() as u8;
            // Now decode
            for byte in a {
                result *= 128;
                result += (*byte & 127) as usize;
            }
        }
        // Set self.buff and subtract the 128 from the 128 bit flag
        self.buf = tmp;
        Some(result)
    }
}

#[cfg(test)]
mod tests {

    use super::*;
    use std;
    use std::io::{Cursor, Seek, SeekFrom};

    #[test]
    fn test_heapless_vbyte_encode() {
        assert_eq!(VByteEncoded::new(0).data[9], 0x80);
        assert_eq!(VByteEncoded::new(0).bytes_used(), 1);
        assert_eq!(VByteEncoded::new(128).data[8..10], [0x01, 0x80]);
        assert_eq!(VByteEncoded::new(128).bytes_used(), 0x02);

        assert_eq!(VByteEncoded::new(0xFFFF).data[7..10], [0x03, 0x7F, 0xFF]);
        assert_eq!(VByteEncoded::new(0xFFFF).bytes_used(), 3);
        assert_eq!(VByteEncoded::new(std::u64::MAX as usize).data,
                   [1, 127, 127, 127, 127, 127, 127, 127, 127, 255, 10]);
    }

    #[test]
    fn test_vbyte_decode() {
        assert_eq!(VByteDecoder::new(vec![0x80].as_slice()).collect::<Vec<_>>(),
                   vec![0]);
        assert_eq!(VByteDecoder::new(vec![0x85].as_slice()).collect::<Vec<_>>(),
                   vec![5]);
        assert_eq!(VByteDecoder::new(vec![0xFF].as_slice()).collect::<Vec<_>>(),
                   vec![127]);
        assert_eq!(VByteDecoder::new(vec![0x80, 0x81].as_slice()).collect::<Vec<_>>(),
                   vec![0, 1]);
        assert_eq!(VByteDecoder::new(vec![0x03, 0x7F, 0xFF, 0x01, 0x82, 0x85].as_slice()).collect::<Vec<_>>(),
                   vec![0xFFFF, 130, 5]);
        assert_eq!(VByteDecoder::new(vec![0x80].as_slice()).collect::<Vec<_>>(),
                   vec![0]);
    }

    #[test]
    fn overflowing() {
        assert_eq!(VByteDecoder::new(vec![0x81; 255].as_slice()).collect::<Vec<_>>(),
                   vec![1; 255]);
    }

    #[test]
    fn more_data() {
        let data = vec![0x80, 0x01, 0x82, 0x85, 0x03, 0x7F, 0xFF, 0x80, 0x86, 0x82, 0x85, 0x84, 0x01, 0x83];
        let decoder = VByteDecoder::new(data.as_slice());
        assert_eq!(decoder.collect::<Vec<_>>(),
                   vec![0, 130, 5, 65535, 0, 6, 2, 5, 4, 131]);
    }

    #[test]
    fn seek_basic() {
        let data = vec![0x80];
        let mut decoder = VByteDecoder::new(Cursor::new(&data));
        assert_eq!(decoder.next().unwrap(), 0);
        assert_eq!(decoder.next(), None);
        decoder.seek(SeekFrom::Start(0)).unwrap();
        assert_eq!(decoder.next().unwrap(), 0);
    }

    #[test]
    fn seek_extended() {
        let data = vec![0x80, 0x01, 0x82, 0x85, 0x03, 0x7F, 0xFF, 0x80, 0x86, 0x82, 0x85, 0x84, 0x01, 0x83];
        let mut decoder = VByteDecoder::new(Cursor::new(&data));
        assert_eq!(decoder.next().unwrap(), 0);
        assert_eq!(decoder.seek(SeekFrom::Start(0)).unwrap(), 0);
        assert_eq!(decoder.next().unwrap(), 0);
        assert_eq!(decoder.seek(SeekFrom::Current(2)).unwrap(), 3);
        assert_eq!(decoder.next().unwrap(), 5);
        assert_eq!(decoder.seek(SeekFrom::End(-2)).unwrap(), 12);
        assert_eq!(decoder.next().unwrap(), 131);
        assert_eq!(decoder.next(), None);
    }

    #[test]
    fn seek_edge_case() {
        let data = vec![0x80, 0x01, 0x82, 0x85, 0x03, 0x7F, 0xFF, 0x80, 0x86, 0x82, 0x85, 0x84, 0x01, 0x83];
        let mut decoder = VByteDecoder::new(Cursor::new(&data));
        assert_eq!(decoder.next(), Some(0));
        decoder.seek(SeekFrom::Start(50)).unwrap();
        assert_eq!(decoder.next(), None);
    }

    #[test]
    fn edge_cases() {
        // 0
        assert_eq!(VByteDecoder::new(VByteEncoded::new(0).data_buf()).collect::<Vec<_>>(),
                   vec![0]);

        // MAX
        assert_eq!(VByteDecoder::new(VByteEncoded::new(usize::max_value()).data_buf()).collect::<Vec<_>>(),
                   vec![usize::max_value()]);

        // too many bytes = corrupted data
        assert_eq!(VByteDecoder::new(vec![127, 127, 127, 127, 127, 127, 127, 127, 127, 127, 255].as_slice())
                       .collect::<Vec<_>>(),
                   vec![]);
        // MAX + n = corrupted data
        assert_eq!(VByteDecoder::new(vec![2, 127, 127, 127, 127, 127, 127, 127, 127, 255].as_slice())
                       .collect::<Vec<_>>(),
                   vec![]);
        // zero-bytes
        assert_eq!(VByteDecoder::new(vec![0; 100].as_slice()).collect::<Vec<_>>(),
                   vec![]);
    }

}
<commit_msg>Improved VByteDecoder<commit_after>//! This module currently provides utility methods and structs for variable
//! byte codes as described in
//! http://nlp.stanford.edu/IR-book/html/htmledition/variable-byte-codes-1.html.
//!
//! Encode unsigned integers by using the `vbyte_encode` method.
//!
//! Decode a bytestream by instatiating a `VByteDecoder` and using its iterator
//! implementation.
//!
//! #Example
//!
//! ```rust,ignore
//!
//! use perlin::utils::compression::{vbyte_encode, VByteDecoder};
//!
//! let bytes = vbyte_encode(3);
//! let three = VByteDecoder::new(bytes.into_iter()).next().unwrap();
//! assert_eq!(3, three);
//! ```

use std::mem;
use std::io;
use std::io::{Seek, SeekFrom, Read, Write, Error};

const BUFSIZE: usize = 32;

/// Stores the result of a vbyte encode operation without indirection that a `Vec<u8>` would introduce.
/// Can thus be used to `vbyte_encode` on the stack
pub struct VByteEncoded {
    // Memory layout of the result: [ payload | payload | payload | ... | used_bytes ]
    data: [u8; 11],
}

impl VByteEncoded {
    /// Performs a vbyte encode without allocating memory on the heap
    /// Can then be written to a `Write`-implementor
    pub fn new(mut number: usize) -> Self {
        let mut count = 0;
        let mut result = [0u8; 11];
        loop {
            result[9 - count] = (number % 128) as u8;
            count += 1;
            if number < 128 {
                break;
            } else {
                number /= 128;
            }
        }
        result[9] += 128;
        result[10] = count as u8;
        VByteEncoded { data: result }
    }

    /// Return how many bytes are used to encode the number. Min: 1 Max: 10
    pub fn bytes_used(&self) -> u8 {
        self.data[10]
    }

    /// Access to the written data as slice
    /// Currently only used in tests
    pub fn data_buf(&self) -> &[u8] {
        &self.data[(10 - self.bytes_used()) as usize..10]
    }

    /// Writes the given VByteEncoded number to a target.
    /// Returns the number of bytes written (equal to `bytes_used`) or an `io::Error`
    pub fn write_to<W: Write>(&self, target: &mut W) -> Result<u8, Error> {
        target.write_all(&self.data[(10 - self.bytes_used()) as usize..10]).map(|()| self.bytes_used())
    }
}

/// Iterator that decodes a bytestream to unsigned integers
pub struct VByteDecoder<R> {
    ptr: usize,
    filled: usize,
    buf: [u8; BUFSIZE],
    source: R,
}

impl<R: Read> VByteDecoder<R> {
    /// Create a new VByteDecoder by passing a bytestream
    pub fn new(mut source: R) -> Self {
        let mut buf: [u8; BUFSIZE] = unsafe { mem::uninitialized() };
        let filled = source.read(&mut buf).unwrap();
        VByteDecoder {
            source: source,
            buf: buf,
            ptr: 0,
            filled: filled,
        }
    }

    /// Sometimes it is convenient to look at the original bytestream itself
    /// (e.g. when not only vbyte encoded integers are in the bytestream)
    /// This method provides access to the underlying bytestream in form of
    /// a mutable borrow
    pub fn underlying_mut(&mut self) -> &mut R {
        &mut self.source
    }

    /// Sometimes it is convenient to look at the original bytestream itself
    /// (e.g. when not only vbyte encoded integers are in the bytestream)
    /// This method provides access to the underlying bytestream in form of
    /// a borrow
    pub fn underlying(&self) -> &R {
        &self.source
    }
}

impl<R: Read> Read for VByteDecoder<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let mut bytes_read = 0;
        if self.filled > 0 {
            // We have some bytes stored. Give them back
            // Read some bytes from self.buf.
            bytes_read += (&self.buf[self.ptr..self.ptr+self.filled]).read(buf)?;
            self.filled -= bytes_read;
            self.ptr += bytes_read;
            // If buf is full. Return
            if bytes_read >= buf.len() {
                return Ok(bytes_read);
            }
        }
        bytes_read += self.source.read(&mut buf[bytes_read..])?;
        Ok(bytes_read)
    }
}

impl<R: Seek + Read> Seek for VByteDecoder<R> {

    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        //after a seek we can throw our pointer away and reread the buffer
        //So seek to the correct position in source and then reinitialize
        let result = match pos {
            SeekFrom::Start(_) => {
                self.source.seek(pos)
            }
            SeekFrom::Current(offset) => {
                self.source.seek(SeekFrom::Current(offset - self.filled as i64))
            }
            SeekFrom::End(_) => {
                self.source.seek(pos)
            }
        };
        self.ptr = 0;
        self.filled = self.source.read(&mut self.buf[self.ptr..]).unwrap();
        result
    }
}

impl<R: Read> Iterator for VByteDecoder<R> {
    type Item = usize;

    /// Returns the next unsigned integer which is encoded in the underlying
    /// bytestream
    /// May iterate the underlying bytestream an arbitrary number of times
    /// Returns None when the underlying bytream returns None or delivers corrupt data
    fn next(&mut self) -> Option<Self::Item> {
        let mut result: usize = 0;
        loop {
            // Process as many bytes as possible
            while self.filled > 0 {
                // Read byte, decrease filled, inc ptr
                self.filled -= 1;
                let byte = unsafe {self.buf.get_unchecked(self.ptr)};
                self.ptr = (self.ptr + 1) % BUFSIZE;
                // Shift left by 7 bytes
                result <<= 7;
                // Add the lower 7 bytes to result
                result += (byte & 127) as usize;
                // Check for 128bit flag
                if *byte & 128 == 0 {
                    // Check for overflow
                    if result.leading_zeros() < 7 {
                        return None;
                    }
                    continue;
                }
                return Some(result);
            }
            // No 128bit found. No overflow
            // we have to read further
            self.filled = self.source.read(&mut self.buf[self.ptr..]).unwrap();

            // Source is empty. We are done!
            if self.filled == 0 {
                return None;
            }
        }
    }
}

#[cfg(test)]
mod tests {

    use super::*;
    use std;
    use std::io::{Cursor, Seek, SeekFrom};

    #[test]
    fn test_heapless_vbyte_encode() {
        assert_eq!(VByteEncoded::new(0).data[9], 0x80);
        assert_eq!(VByteEncoded::new(0).bytes_used(), 1);
        assert_eq!(VByteEncoded::new(128).data[8..10], [0x01, 0x80]);
        assert_eq!(VByteEncoded::new(128).bytes_used(), 0x02);

        assert_eq!(VByteEncoded::new(0xFFFF).data[7..10], [0x03, 0x7F, 0xFF]);
        assert_eq!(VByteEncoded::new(0xFFFF).bytes_used(), 3);
        assert_eq!(VByteEncoded::new(std::u64::MAX as usize).data,
                   [1, 127, 127, 127, 127, 127, 127, 127, 127, 255, 10]);
    }

    #[test]
    fn test_vbyte_decode() {
        assert_eq!(VByteDecoder::new(vec![0x80].as_slice()).collect::<Vec<_>>(),
                   vec![0]);
        assert_eq!(VByteDecoder::new(vec![0x85].as_slice()).collect::<Vec<_>>(),
                   vec![5]);
        assert_eq!(VByteDecoder::new(vec![0xFF].as_slice()).collect::<Vec<_>>(),
                   vec![127]);
        assert_eq!(VByteDecoder::new(vec![0x80, 0x81].as_slice()).collect::<Vec<_>>(),
                   vec![0, 1]);
        assert_eq!(VByteDecoder::new(vec![0x03, 0x7F, 0xFF, 0x01, 0x82, 0x85].as_slice()).collect::<Vec<_>>(),
                   vec![0xFFFF, 130, 5]);
        assert_eq!(VByteDecoder::new(vec![0x80].as_slice()).collect::<Vec<_>>(),
                   vec![0]);
    }

    #[test]
    fn overflowing() {
        assert_eq!(VByteDecoder::new(vec![0x81; 255].as_slice()).collect::<Vec<_>>(),
                   vec![1; 255]);
    }

    #[test]
    fn more_data() {
        let data = vec![0x80, 0x01, 0x82, 0x85, 0x03, 0x7F, 0xFF, 0x80, 0x86, 0x82, 0x85, 0x84, 0x01, 0x83];
        let decoder = VByteDecoder::new(data.as_slice());
        assert_eq!(decoder.collect::<Vec<_>>(),
                   vec![0, 130, 5, 65535, 0, 6, 2, 5, 4, 131]);
    }

    #[test]
    fn seek_basic() {
        let data = vec![0x80];
        let mut decoder = VByteDecoder::new(Cursor::new(&data));
        assert_eq!(decoder.next().unwrap(), 0);
        assert_eq!(decoder.next(), None);
        decoder.seek(SeekFrom::Start(0)).unwrap();
        assert_eq!(decoder.next().unwrap(), 0);
    }

    #[test]
    fn seek_extended() {
        let data = vec![0x80, 0x01, 0x82, 0x85, 0x03, 0x7F, 0xFF, 0x80, 0x86, 0x82, 0x85, 0x84, 0x01, 0x83];
        let mut decoder = VByteDecoder::new(Cursor::new(&data));
        assert_eq!(decoder.next().unwrap(), 0);
        assert_eq!(decoder.seek(SeekFrom::Start(0)).unwrap(), 0);
        assert_eq!(decoder.next().unwrap(), 0);
        assert_eq!(decoder.seek(SeekFrom::Current(2)).unwrap(), 3);
        assert_eq!(decoder.next().unwrap(), 5);
        assert_eq!(decoder.seek(SeekFrom::End(-2)).unwrap(), 12);
        assert_eq!(decoder.next().unwrap(), 131);
        assert_eq!(decoder.next(), None);
    }

    #[test]
    fn seek_edge_case() {
        let data = vec![0x80, 0x01, 0x82, 0x85, 0x03, 0x7F, 0xFF, 0x80, 0x86, 0x82, 0x85, 0x84, 0x01, 0x83];
        let mut decoder = VByteDecoder::new(Cursor::new(&data));
        assert_eq!(decoder.next(), Some(0));
        decoder.seek(SeekFrom::Start(50)).unwrap();
        assert_eq!(decoder.next(), None);
    }

    #[test]
    fn edge_cases() {
        // // 0
        // assert_eq!(VByteDecoder::new(VByteEncoded::new(0).data_buf()).collect::<Vec<_>>(),
        //            vec![0]);

        // // MAX
        // assert_eq!(VByteDecoder::new(VByteEncoded::new(usize::max_value()).data_buf()).collect::<Vec<_>>(),
        //            vec![usize::max_value()]);

        // too many bytes = corrupted data
        assert_eq!(VByteDecoder::new(vec![127, 127, 127, 127, 127, 127, 127, 127, 127, 127, 127, 127, 255].as_slice())
                       .collect::<Vec<_>>(),
                   vec![]);
        // MAX + n = corrupted data
        assert_eq!(VByteDecoder::new(vec![2, 127, 127, 127, 127, 127, 127, 127, 127, 255].as_slice())
                       .collect::<Vec<_>>(),
                   vec![]);
        // zero-bytes
        assert_eq!(VByteDecoder::new(vec![0; 100].as_slice()).collect::<Vec<_>>(),
                   vec![]);
    }

}
<|endoftext|>


<commit_before>extern crate serde_json;
extern crate tempdir;

use super::*;

use std::path::{Path};
use std::sync::{
    mpsc::{channel, Receiver, Sender},
};

use task_maker_format::ui::{UIExecutionStatus, UIMessage};
use task_maker_format::ioi::Task;

use award::{AwardName, Score};
use evaluation::*;
use submission::Submission;

use crate::award::{AwardValue, BadgeAwardValue, ScoreAwardValue};
use crate::evaluation::AwardEvent;
use content::TextVariant;
use evaluation::record::ValenceValue;
use evaluation::Event;
use feedback::valence::Valence;
use rusage::{MemoryUsage, TimeUsage};
use std::process::{Command, Stdio};
use std::io::{BufReader, BufRead};

pub fn run_evaluation<P: AsRef<Path>>(task_path: P, submission: Submission) -> Result<Receiver<Event>, failure::Error> {
    let task_path = task_path.as_ref().to_owned();
    let (event_tx, event_rx) = channel();

    // write solution to file
    let dir = tempdir::TempDir::new("evaluation")?;
    let solution_file = &submission.field_values[0].file;
    let solution_path = dir.path().join(&solution_file.name.0);
    std::fs::write(&solution_path, &solution_file.content)?;

    let mut ioi_task: Option<Task> = None;

    let mut task_maker = Command::new("task-maker-rust")
        .arg("--dry-run")
        .arg("--no-statement")
        .arg("--ui=json")
        .arg("--task-dir").arg(task_path)
        .arg("--solution").arg(solution_path)
        .stdin(Stdio::null())
        .stdout(Stdio::piped())
        .spawn()?;

    let stdout_reader = BufReader::new(task_maker.stdout.as_mut().unwrap());
    for line in stdout_reader.lines() {
        let message = serde_json::from_str::<UIMessage>(&line?)?;

        // As the first message task-maker-rust sends us the task... nice!
        if let UIMessage::IOITask { task} = message {
            ioi_task = Some(task);
        } else {
            ui_message_to_events(&ioi_task.as_ref().unwrap(), message, &event_tx);
        }
    }

    task_maker.wait()?;

    Ok(event_rx)
}

fn ui_message_to_events(
    task: &Task,
    ui_message: UIMessage,
    tx: &Sender<Event>,
) -> Result<(), failure::Error> {
    match ui_message {
        UIMessage::IOITestcaseScore {
            testcase,
            score,
            message,
            ..
        } => {
            tx.send(Event::Value(ValueEvent {
                key: record::Key(format!("testcase.{}.score", testcase)),
                value: record::Value::Score(record::ScoreValue {
                    score: Score(score as f64),
                }),
            }))?;
            tx.send(Event::Value(ValueEvent {
                key: record::Key(format!("testcase.{}.message", testcase)),
                value: record::Value::Message(record::TextValue {
                    text: vec![TextVariant {
                        value: message,
                        attributes: vec![],
                    }],
                }),
            }))?;
            tx.send(Event::Value(ValueEvent {
                key: record::Key(format!("testcase.{}.valence", testcase)),
                value: record::Value::Valence(ValenceValue {
                    valence: if score <= 0.0 {
                        Valence::Failure
                    } else if score >= 1.0 {
                        Valence::Success
                    } else {
                        Valence::Partial
                    },
                }),
            }))?;
        }
        UIMessage::IOISubtaskScore {
            subtask,
            score,
            normalized_score,
            ..
        } => {
            tx.send(Event::Value(ValueEvent {
                key: record::Key(format!("subtask.{}.score", subtask)),
                value: record::Value::Score(record::ScoreValue {
                    score: Score(score),
                }),
            }))?;
            tx.send(Event::Award(AwardEvent {
                award_name: AwardName(format!("subtask.{}.score", subtask)),
                value: AwardValue::Score(ScoreAwardValue {
                    score: Score(score),
                }),
            }))?;
            tx.send(Event::Award(AwardEvent {
                award_name: AwardName(format!("subtask.{}.badge", subtask)),
                value: AwardValue::Badge(BadgeAwardValue {
                    badge: normalized_score > 0.0,
                }),
            }))?;
        }
        UIMessage::IOIEvaluation {
            testcase, status, ..
        } => {
            if let UIExecutionStatus::Done { result } = status {
                let time_usage = result.resources.cpu_time;
                tx.send(Event::Value(ValueEvent {
                    key: record::Key(format!("testcase.{}.time_usage", testcase)),
                    value: record::Value::TimeUsage(record::TimeUsageValue {
                        time_usage: TimeUsage(time_usage),
                    }),
                }))?;
                let warning_watermark = 0.25;
                tx.send(Event::Value(ValueEvent {
                    key: record::Key(format!("testcase.{}.time_usage_valence", testcase)),
                    value: record::Value::Valence(ValenceValue {
                        valence: if task
                            .time_limit
                            .map(|limit| time_usage < limit * warning_watermark)
                            .unwrap_or(true)
                        {
                            Valence::Nominal
                        } else if task
                            .time_limit
                            .map(|limit| time_usage < limit)
                            .unwrap_or(true)
                        {
                            Valence::Warning
                        } else {
                            Valence::Failure
                        },
                    }),
                }))?;
                let memory_usage_bytes = (result.resources.memory * 1024) as f64;
                let memory_limit_bytes =
                    task.memory_limit.map(|limit| (limit * 1024 * 1024) as f64);
                tx.send(Event::Value(ValueEvent {
                    key: record::Key(format!("testcase.{}.memory_usage", testcase)),
                    value: record::Value::MemoryUsage(record::MemoryUsageValue {
                        memory_usage: MemoryUsage(memory_usage_bytes as i32),
                    }),
                }))?;
                tx.send(Event::Value(ValueEvent {
                    key: record::Key(format!("testcase.{}.memory_usage_valence", testcase)),
                    value: record::Value::Valence(ValenceValue {
                        valence: if memory_limit_bytes
                            .map(|limit| memory_usage_bytes < limit * warning_watermark)
                            .unwrap_or(true)
                        {
                            Valence::Nominal
                        } else if memory_limit_bytes
                            .map(|limit| memory_usage_bytes < limit)
                            .unwrap_or(true)
                        {
                            Valence::Warning
                        } else {
                            Valence::Failure
                        },
                    }),
                }))?;
            }
        }
        _ => (),
    };
    Ok(())
}
<commit_msg>Search task-maker in path and add option to pass args to it<commit_after>extern crate serde_json;
extern crate tempdir;

use super::*;

use std::path::{Path, PathBuf};
use std::sync::{
    mpsc::{channel, Receiver, Sender},
};

use task_maker_format::ui::{UIExecutionStatus, UIMessage};
use task_maker_format::ioi::Task;

use award::{AwardName, Score};
use evaluation::*;
use submission::Submission;

use crate::award::{AwardValue, BadgeAwardValue, ScoreAwardValue};
use crate::evaluation::AwardEvent;
use content::TextVariant;
use evaluation::record::ValenceValue;
use evaluation::Event;
use feedback::valence::Valence;
use rusage::{MemoryUsage, TimeUsage};
use std::process::{Command, Stdio};
use std::io::{BufReader, BufRead};
use std::env;

/// Searches for an executable in `PATH`
fn find_in_path(exe: &str) -> Option<PathBuf> {
    let path = env::var("PATH").unwrap();

    for dir in env::split_paths(&path) {
        let exe = dir.join(exe);
        if exe.exists() {
            return Some(exe);
        }
    }
    None
}

/// Searches the task-maker executable in the following locations:
/// - env variable `TASK_MAKER_EXE`
/// - executable `task-maker-rust` in `PATH`
/// - executable `task-maker` in `PATH`
fn find_task_maker() -> Result<PathBuf, failure::Error> {
    if let Ok(task_maker) = env::var("TASK_MAKER_EXE") {
        let path = PathBuf::from(task_maker);
        if path.exists() {
            return Ok(path);
        }
    }

    for name in &["task-maker-rust", "task-maker"] {
        if let Some(exe) = find_in_path(name) {
            return Ok(exe)
        }
    }

    Err(failure::err_msg("task-maker executable not found"))
}

pub fn run_evaluation<P: AsRef<Path>>(task_path: P, submission: Submission) -> Result<Receiver<Event>, failure::Error> {
    let task_path = task_path.as_ref().to_owned();
    let (event_tx, event_rx) = channel();

    // write solution to file
    let dir = tempdir::TempDir::new("evaluation")?;
    let solution_file = &submission.field_values[0].file;
    let solution_path = dir.path().join(&solution_file.name.0);
    std::fs::write(&solution_path, &solution_file.content)?;

    let mut ioi_task: Option<Task> = None;

    let mut task_maker = Command::new(find_task_maker()?);
    task_maker
        .arg("--dry-run")
        .arg("--no-statement")
        .arg("--ui=json")
        .arg("--task-dir").arg(task_path)
        .arg("--solution").arg(solution_path);

    if let Ok(args) = env::var("TASK_MAKER_ARGS") {
        task_maker.args(args.split(' '));
    }

    let mut child = task_maker
        .stdin(Stdio::null())
        .stdout(Stdio::piped())
        .spawn()?;

    let stdout_reader = BufReader::new(child.stdout.as_mut().unwrap());
    for line in stdout_reader.lines() {
        let message = serde_json::from_str::<UIMessage>(&line?)?;

        // As the first message task-maker-rust sends us the task... nice!
        if let UIMessage::IOITask { task} = message {
            ioi_task = Some(task);
        } else {
            ui_message_to_events(&ioi_task.as_ref().unwrap(), message, &event_tx);
        }
    }

    child.wait()?;

    Ok(event_rx)
}

fn ui_message_to_events(
    task: &Task,
    ui_message: UIMessage,
    tx: &Sender<Event>,
) -> Result<(), failure::Error> {
    match ui_message {
        UIMessage::IOITestcaseScore {
            testcase,
            score,
            message,
            ..
        } => {
            tx.send(Event::Value(ValueEvent {
                key: record::Key(format!("testcase.{}.score", testcase)),
                value: record::Value::Score(record::ScoreValue {
                    score: Score(score as f64),
                }),
            }))?;
            tx.send(Event::Value(ValueEvent {
                key: record::Key(format!("testcase.{}.message", testcase)),
                value: record::Value::Message(record::TextValue {
                    text: vec![TextVariant {
                        value: message,
                        attributes: vec![],
                    }],
                }),
            }))?;
            tx.send(Event::Value(ValueEvent {
                key: record::Key(format!("testcase.{}.valence", testcase)),
                value: record::Value::Valence(ValenceValue {
                    valence: if score <= 0.0 {
                        Valence::Failure
                    } else if score >= 1.0 {
                        Valence::Success
                    } else {
                        Valence::Partial
                    },
                }),
            }))?;
        }
        UIMessage::IOISubtaskScore {
            subtask,
            score,
            normalized_score,
            ..
        } => {
            tx.send(Event::Value(ValueEvent {
                key: record::Key(format!("subtask.{}.score", subtask)),
                value: record::Value::Score(record::ScoreValue {
                    score: Score(score),
                }),
            }))?;
            tx.send(Event::Award(AwardEvent {
                award_name: AwardName(format!("subtask.{}.score", subtask)),
                value: AwardValue::Score(ScoreAwardValue {
                    score: Score(score),
                }),
            }))?;
            tx.send(Event::Award(AwardEvent {
                award_name: AwardName(format!("subtask.{}.badge", subtask)),
                value: AwardValue::Badge(BadgeAwardValue {
                    badge: normalized_score > 0.0,
                }),
            }))?;
        }
        UIMessage::IOIEvaluation {
            testcase, status, ..
        } => {
            if let UIExecutionStatus::Done { result } = status {
                let time_usage = result.resources.cpu_time;
                tx.send(Event::Value(ValueEvent {
                    key: record::Key(format!("testcase.{}.time_usage", testcase)),
                    value: record::Value::TimeUsage(record::TimeUsageValue {
                        time_usage: TimeUsage(time_usage),
                    }),
                }))?;
                let warning_watermark = 0.25;
                tx.send(Event::Value(ValueEvent {
                    key: record::Key(format!("testcase.{}.time_usage_valence", testcase)),
                    value: record::Value::Valence(ValenceValue {
                        valence: if task
                            .time_limit
                            .map(|limit| time_usage < limit * warning_watermark)
                            .unwrap_or(true)
                        {
                            Valence::Nominal
                        } else if task
                            .time_limit
                            .map(|limit| time_usage < limit)
                            .unwrap_or(true)
                        {
                            Valence::Warning
                        } else {
                            Valence::Failure
                        },
                    }),
                }))?;
                let memory_usage_bytes = (result.resources.memory * 1024) as f64;
                let memory_limit_bytes =
                    task.memory_limit.map(|limit| (limit * 1024 * 1024) as f64);
                tx.send(Event::Value(ValueEvent {
                    key: record::Key(format!("testcase.{}.memory_usage", testcase)),
                    value: record::Value::MemoryUsage(record::MemoryUsageValue {
                        memory_usage: MemoryUsage(memory_usage_bytes as i32),
                    }),
                }))?;
                tx.send(Event::Value(ValueEvent {
                    key: record::Key(format!("testcase.{}.memory_usage_valence", testcase)),
                    value: record::Value::Valence(ValenceValue {
                        valence: if memory_limit_bytes
                            .map(|limit| memory_usage_bytes < limit * warning_watermark)
                            .unwrap_or(true)
                        {
                            Valence::Nominal
                        } else if memory_limit_bytes
                            .map(|limit| memory_usage_bytes < limit)
                            .unwrap_or(true)
                        {
                            Valence::Warning
                        } else {
                            Valence::Failure
                        },
                    }),
                }))?;
            }
        }
        _ => (),
    };
    Ok(())
}
<|endoftext|>


<commit_before>use crate::builders::CrateBuilder;
use crate::util::{ChaosProxy, FreshSchema};
use anyhow::Error;
use cargo_registry::models::{NewUser, User};
use diesel::prelude::*;
use reqwest::blocking::{Client, Response};
use std::collections::HashMap;
use std::io::{BufRead, BufReader, Read};
use std::process::{Child, Command, Stdio};
use std::sync::{mpsc::Sender, Arc};
use std::time::Duration;

const SERVER_BOOT_TIMEOUT_SECONDS: u64 = 30;

#[test]
fn normal_startup() -> Result<(), Error> {
    let server_bin = ServerBin::prepare()?;
    initialize_dummy_crate(&server_bin.db()?);

    let running_server = server_bin.start()?;

    // Ensure the application correctly responds to download requests
    let resp = running_server.get("api/v1/crates/FOO/1.0.0/download")?;
    assert!(resp.status().is_redirection());
    assert!(resp
        .headers()
        .get("location")
        .unwrap()
        .to_str()?
        .ends_with("/crates/foo/foo-1.0.0.crate"));

    Ok(())
}

#[test]
fn startup_without_database() -> Result<(), Error> {
    let server_bin = ServerBin::prepare()?;
    initialize_dummy_crate(&server_bin.db()?);

    // Break the networking *before* starting the binary, to ensure the binary can fully startup
    // without a database connection. Most of crates.io should not work when started without a
    // database, but unconditional redirects will work.
    server_bin.chaosproxy.break_networking();

    let running_server = server_bin.start()?;

    // Ensure unconditional redirects work.
    let resp = running_server.get("api/v1/crates/FOO/1.0.0/download")?;
    assert!(resp.status().is_redirection());
    assert!(resp
        .headers()
        .get("location")
        .unwrap()
        .to_str()?
        .ends_with("/crates/FOO/FOO-1.0.0.crate"));

    Ok(())
}

fn initialize_dummy_crate(conn: &PgConnection) {
    use cargo_registry::schema::users;

    let user: User = diesel::insert_into(users::table)
        .values(NewUser {
            gh_id: 0,
            gh_login: "user",
            ..NewUser::default()
        })
        .get_result(conn)
        .expect("failed to create dummy user");

    CrateBuilder::new("foo", user.id)
        .version("1.0.0")
        .build(conn)
        .expect("failed to create dummy crate");
}

struct ServerBin {
    chaosproxy: Arc<ChaosProxy>,
    db_url: String,
    env: HashMap<String, String>,
    fresh_schema: FreshSchema,
}

impl ServerBin {
    fn prepare() -> Result<Self, Error> {
        let mut env = dotenv::vars().collect::<HashMap<_, _>>();
        // Bind a random port every time the server is started.
        env.insert("PORT".into(), "0".into());
        // Avoid creating too many database connections.
        env.insert("DB_POOL_SIZE".into(), "2".into());
        env.remove("DB_MIN_SIZE");
        // Other configuration variables needed for the application to boot.
        env.insert("WEB_ALLOWED_ORIGINS".into(), "http://localhost:8888".into());
        env.insert(
            "SESSION_KEY".into(),
            std::iter::repeat('a').take(32).collect(),
        );
        env.insert("GH_CLIENT_ID".into(), String::new());
        env.insert("GH_CLIENT_SECRET".into(), String::new());

        // Use a proxied fresh schema as the database url.
        let fresh_schema = FreshSchema::new(env.get("TEST_DATABASE_URL").unwrap());
        let (chaosproxy, db_url) = ChaosProxy::proxy_database_url(fresh_schema.database_url())?;
        env.remove("TEST_DATABASE_URL");
        env.insert("DATABASE_URL".into(), db_url.clone());
        env.insert("READ_ONLY_REPLICA_URL".into(), db_url.clone());

        Ok(ServerBin {
            chaosproxy,
            db_url,
            env,
            fresh_schema,
        })
    }

    fn db(&self) -> Result<PgConnection, Error> {
        Ok(PgConnection::establish(&self.db_url)?)
    }

    fn start(self) -> Result<RunningServer, Error> {
        let mut process = Command::new(env!("CARGO_BIN_EXE_server"))
            .env_clear()
            .envs(self.env.into_iter())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()?;

        let (port_send, port_recv) = std::sync::mpsc::channel();
        stream_processor(process.stdout.take().unwrap(), "stdout", Some(port_send));
        stream_processor(process.stderr.take().unwrap(), "stderr", None);

        // Possible causes for this to fail:
        // - the server binary failed to start
        // - the server binary requires a database connection now
        // - the server binary doesn't print "listening on port {port}" anymore
        let port: u16 = port_recv
            .recv_timeout(Duration::from_secs(SERVER_BOOT_TIMEOUT_SECONDS))
            .map_err(|_| anyhow::anyhow!("the server took too much time to initialize"))?
            .parse()?;

        let http = Client::builder()
            .redirect(reqwest::redirect::Policy::none())
            .build()?;

        Ok(RunningServer {
            process,
            port,
            http,
            _chaosproxy: self.chaosproxy,
            _fresh_schema: self.fresh_schema,
        })
    }
}

struct RunningServer {
    process: Child,
    port: u16,
    http: Client,

    // Keep these two items at the bottom in this order to drop everything in the correct order.
    _chaosproxy: Arc<ChaosProxy>,
    _fresh_schema: FreshSchema,
}

impl RunningServer {
    fn get(&self, url: &str) -> Result<Response, Error> {
        Ok(self
            .http
            .get(format!("http://127.0.0.1:{}/{}", self.port, url))
            .header("User-Agent", "crates.io test suite")
            .send()?)
    }
}

impl Drop for RunningServer {
    fn drop(&mut self) {
        self.process
            .kill()
            .expect("failed to kill the server binary");
    }
}

fn stream_processor<R>(stream: R, kind: &'static str, port_send: Option<Sender<String>>)
where
    R: Read + Send + 'static,
{
    std::thread::spawn(move || {
        let stream = BufReader::new(stream);
        for line in stream.lines() {
            let line = match line {
                Ok(line) => line,
                // We receive an EOF when the process terminates
                Err(err) if err.kind() == std::io::ErrorKind::UnexpectedEof => break,
                Err(err) => panic!("unexpected error while reading process {}: {}", kind, err),
            };

            // If we expect the port number to be logged into this stream, look for it and send it
            // over the channel as soon as it's found.
            if let Some(port) = &port_send {
                if let Some(port_str) = line.strip_prefix("listening on port ") {
                    port.send(port_str.into())
                        .expect("failed to send the port to the test thread")
                }
            }

            println!("[server {}] {}", kind, line);
        }
    });
}
<commit_msg>Address clippy::manual_str_repeat<commit_after>use crate::builders::CrateBuilder;
use crate::util::{ChaosProxy, FreshSchema};
use anyhow::Error;
use cargo_registry::models::{NewUser, User};
use diesel::prelude::*;
use reqwest::blocking::{Client, Response};
use std::collections::HashMap;
use std::io::{BufRead, BufReader, Read};
use std::process::{Child, Command, Stdio};
use std::sync::{mpsc::Sender, Arc};
use std::time::Duration;

const SERVER_BOOT_TIMEOUT_SECONDS: u64 = 30;

#[test]
fn normal_startup() -> Result<(), Error> {
    let server_bin = ServerBin::prepare()?;
    initialize_dummy_crate(&server_bin.db()?);

    let running_server = server_bin.start()?;

    // Ensure the application correctly responds to download requests
    let resp = running_server.get("api/v1/crates/FOO/1.0.0/download")?;
    assert!(resp.status().is_redirection());
    assert!(resp
        .headers()
        .get("location")
        .unwrap()
        .to_str()?
        .ends_with("/crates/foo/foo-1.0.0.crate"));

    Ok(())
}

#[test]
fn startup_without_database() -> Result<(), Error> {
    let server_bin = ServerBin::prepare()?;
    initialize_dummy_crate(&server_bin.db()?);

    // Break the networking *before* starting the binary, to ensure the binary can fully startup
    // without a database connection. Most of crates.io should not work when started without a
    // database, but unconditional redirects will work.
    server_bin.chaosproxy.break_networking();

    let running_server = server_bin.start()?;

    // Ensure unconditional redirects work.
    let resp = running_server.get("api/v1/crates/FOO/1.0.0/download")?;
    assert!(resp.status().is_redirection());
    assert!(resp
        .headers()
        .get("location")
        .unwrap()
        .to_str()?
        .ends_with("/crates/FOO/FOO-1.0.0.crate"));

    Ok(())
}

fn initialize_dummy_crate(conn: &PgConnection) {
    use cargo_registry::schema::users;

    let user: User = diesel::insert_into(users::table)
        .values(NewUser {
            gh_id: 0,
            gh_login: "user",
            ..NewUser::default()
        })
        .get_result(conn)
        .expect("failed to create dummy user");

    CrateBuilder::new("foo", user.id)
        .version("1.0.0")
        .build(conn)
        .expect("failed to create dummy crate");
}

struct ServerBin {
    chaosproxy: Arc<ChaosProxy>,
    db_url: String,
    env: HashMap<String, String>,
    fresh_schema: FreshSchema,
}

impl ServerBin {
    fn prepare() -> Result<Self, Error> {
        let mut env = dotenv::vars().collect::<HashMap<_, _>>();
        // Bind a random port every time the server is started.
        env.insert("PORT".into(), "0".into());
        // Avoid creating too many database connections.
        env.insert("DB_POOL_SIZE".into(), "2".into());
        env.remove("DB_MIN_SIZE");
        // Other configuration variables needed for the application to boot.
        env.insert("WEB_ALLOWED_ORIGINS".into(), "http://localhost:8888".into());
        env.insert("SESSION_KEY".into(), "a".repeat(32));
        env.insert("GH_CLIENT_ID".into(), String::new());
        env.insert("GH_CLIENT_SECRET".into(), String::new());

        // Use a proxied fresh schema as the database url.
        let fresh_schema = FreshSchema::new(env.get("TEST_DATABASE_URL").unwrap());
        let (chaosproxy, db_url) = ChaosProxy::proxy_database_url(fresh_schema.database_url())?;
        env.remove("TEST_DATABASE_URL");
        env.insert("DATABASE_URL".into(), db_url.clone());
        env.insert("READ_ONLY_REPLICA_URL".into(), db_url.clone());

        Ok(ServerBin {
            chaosproxy,
            db_url,
            env,
            fresh_schema,
        })
    }

    fn db(&self) -> Result<PgConnection, Error> {
        Ok(PgConnection::establish(&self.db_url)?)
    }

    fn start(self) -> Result<RunningServer, Error> {
        let mut process = Command::new(env!("CARGO_BIN_EXE_server"))
            .env_clear()
            .envs(self.env.into_iter())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()?;

        let (port_send, port_recv) = std::sync::mpsc::channel();
        stream_processor(process.stdout.take().unwrap(), "stdout", Some(port_send));
        stream_processor(process.stderr.take().unwrap(), "stderr", None);

        // Possible causes for this to fail:
        // - the server binary failed to start
        // - the server binary requires a database connection now
        // - the server binary doesn't print "listening on port {port}" anymore
        let port: u16 = port_recv
            .recv_timeout(Duration::from_secs(SERVER_BOOT_TIMEOUT_SECONDS))
            .map_err(|_| anyhow::anyhow!("the server took too much time to initialize"))?
            .parse()?;

        let http = Client::builder()
            .redirect(reqwest::redirect::Policy::none())
            .build()?;

        Ok(RunningServer {
            process,
            port,
            http,
            _chaosproxy: self.chaosproxy,
            _fresh_schema: self.fresh_schema,
        })
    }
}

struct RunningServer {
    process: Child,
    port: u16,
    http: Client,

    // Keep these two items at the bottom in this order to drop everything in the correct order.
    _chaosproxy: Arc<ChaosProxy>,
    _fresh_schema: FreshSchema,
}

impl RunningServer {
    fn get(&self, url: &str) -> Result<Response, Error> {
        Ok(self
            .http
            .get(format!("http://127.0.0.1:{}/{}", self.port, url))
            .header("User-Agent", "crates.io test suite")
            .send()?)
    }
}

impl Drop for RunningServer {
    fn drop(&mut self) {
        self.process
            .kill()
            .expect("failed to kill the server binary");
    }
}

fn stream_processor<R>(stream: R, kind: &'static str, port_send: Option<Sender<String>>)
where
    R: Read + Send + 'static,
{
    std::thread::spawn(move || {
        let stream = BufReader::new(stream);
        for line in stream.lines() {
            let line = match line {
                Ok(line) => line,
                // We receive an EOF when the process terminates
                Err(err) if err.kind() == std::io::ErrorKind::UnexpectedEof => break,
                Err(err) => panic!("unexpected error while reading process {}: {}", kind, err),
            };

            // If we expect the port number to be logged into this stream, look for it and send it
            // over the channel as soon as it's found.
            if let Some(port) = &port_send {
                if let Some(port_str) = line.strip_prefix("listening on port ") {
                    port.send(port_str.into())
                        .expect("failed to send the port to the test thread")
                }
            }

            println!("[server {}] {}", kind, line);
        }
    });
}
<|endoftext|>


<commit_before>#[cfg(any(feature = "postgres", feature = "mysql"))]
use super::query_helper;
use clap::ArgMatches;
use diesel::dsl::sql;
use diesel::sql_types::Bool;
use diesel::*;
use diesel_migrations::FileBasedMigrations;

use crate::database_error::{DatabaseError, DatabaseResult};

use std::env;
use std::error::Error;
#[cfg(feature = "postgres")]
use std::fs::{self, File};
#[cfg(feature = "postgres")]
use std::io::Write;
use std::path::Path;

enum Backend {
    #[cfg(feature = "postgres")]
    Pg,
    #[cfg(feature = "sqlite")]
    Sqlite,
    #[cfg(feature = "mysql")]
    Mysql,
}

impl Backend {
    fn for_url(database_url: &str) -> Self {
        match database_url {
            _ if database_url.starts_with("postgres://")
                || database_url.starts_with("postgresql://") =>
            {
                #[cfg(feature = "postgres")]
                {
                    Backend::Pg
                }
                #[cfg(not(feature = "postgres"))]
                {
                    panic!(
                        "Database url `{}` requires the `postgres` feature but it's not enabled.",
                        database_url
                    );
                }
            }
            _ if database_url.starts_with("mysql://") =>
            {
                #[cfg(feature = "mysql")]
                {
                    Backend::Mysql
                }
                #[cfg(not(feature = "mysql"))]
                {
                    panic!(
                        "Database url `{}` requires the `mysql` feature but it's not enabled.",
                        database_url
                    );
                }
            }
            #[cfg(feature = "sqlite")]
            _ => Backend::Sqlite,
            #[cfg(not(feature = "sqlite"))]
            _ => {
                if database_url.starts_with("sqlite://") {
                    panic!(
                        "Database url `{}` requires the `sqlite` feature but it's not enabled.",
                        database_url
                    );
                }

                let mut available_schemes: Vec<&str> = Vec::new();

                // One of these will always be true, or you are compiling
                // diesel_cli without a backend. And why would you ever want to
                // do that?
                if cfg!(feature = "postgres") {
                    available_schemes.push("`postgres://`");
                }
                if cfg!(feature = "mysql") {
                    available_schemes.push("`mysql://`");
                }

                panic!(
                    "`{}` is not a valid database URL. It should start with {}, or maybe you meant to use the `sqlite` feature which is not enabled.",
                    database_url,
                    available_schemes.join(" or ")
                );
            }
            #[cfg(not(any(feature = "mysql", feature = "sqlite", feature = "postgres")))]
            _ => compile_error!(
                "At least one backend must be specified for use with this crate. \
                 You may omit the unneeded dependencies in the following command. \n\n \
                 ex. `cargo install diesel_cli --no-default-features --features mysql postgres sqlite` \n"
            ),
        }
    }
}

pub enum InferConnection {
    #[cfg(feature = "postgres")]
    Pg(PgConnection),
    #[cfg(feature = "sqlite")]
    Sqlite(SqliteConnection),
    #[cfg(feature = "mysql")]
    Mysql(MysqlConnection),
}

impl InferConnection {
    pub fn establish(database_url: &str) -> DatabaseResult<Self> {
        match Backend::for_url(database_url) {
            #[cfg(feature = "postgres")]
            Backend::Pg => PgConnection::establish(database_url).map(InferConnection::Pg),
            #[cfg(feature = "sqlite")]
            Backend::Sqlite => {
                SqliteConnection::establish(database_url).map(InferConnection::Sqlite)
            }
            #[cfg(feature = "mysql")]
            Backend::Mysql => MysqlConnection::establish(database_url).map(InferConnection::Mysql),
        }
        .map_err(Into::into)
    }
}

macro_rules! call_with_conn {
    (
        $database_url:expr,
        $($func:ident)::+
    ) => {
        call_with_conn!($database_url, $($func)::+ ())
    };

    (
        $database_url:expr,
        $($func:ident)::+ ($($args:expr),*)
    ) => {
        match crate::database::InferConnection::establish(&$database_url)
            .unwrap_or_else(handle_error)
        {
            #[cfg(feature="postgres")]
            crate::database::InferConnection::Pg(ref conn) => $($func)::+ (conn, $($args),*),
            #[cfg(feature="sqlite")]
            crate::database::InferConnection::Sqlite(ref conn) => $($func)::+ (conn, $($args),*),
            #[cfg(feature="mysql")]
            crate::database::InferConnection::Mysql(ref conn) => $($func)::+ (conn, $($args),*),
        }
    };
}

pub fn reset_database(args: &ArgMatches, migrations_dir: &Path) -> DatabaseResult<()> {
    drop_database(&database_url(args))?;
    setup_database(args, migrations_dir)
}

pub fn setup_database(args: &ArgMatches, migrations_dir: &Path) -> DatabaseResult<()> {
    let database_url = database_url(args);

    create_database_if_needed(&database_url)?;
    create_default_migration_if_needed(&database_url, migrations_dir)?;
    create_schema_table_and_run_migrations_if_needed(&database_url, migrations_dir)?;
    Ok(())
}

pub fn drop_database_command(args: &ArgMatches) -> DatabaseResult<()> {
    drop_database(&database_url(args))
}

/// Creates the database specified in the connection url. It returns an error
/// it was unable to create the database.
fn create_database_if_needed(database_url: &str) -> DatabaseResult<()> {
    match Backend::for_url(database_url) {
        #[cfg(feature = "postgres")]
        Backend::Pg => {
            if PgConnection::establish(database_url).is_err() {
                let (database, postgres_url) = change_database_of_url(database_url, "postgres");
                println!("Creating database: {}", database);
                let conn = PgConnection::establish(&postgres_url)?;
                query_helper::create_database(&database).execute(&conn)?;
            }
        }
        #[cfg(feature = "sqlite")]
        Backend::Sqlite => {
            if !::std::path::Path::new(database_url).exists() {
                println!("Creating database: {}", database_url);
                SqliteConnection::establish(database_url)?;
            }
        }
        #[cfg(feature = "mysql")]
        Backend::Mysql => {
            if MysqlConnection::establish(database_url).is_err() {
                let (database, mysql_url) =
                    change_database_of_url(database_url, "information_schema");
                println!("Creating database: {}", database);
                let conn = MysqlConnection::establish(&mysql_url)?;
                query_helper::create_database(&database).execute(&conn)?;
            }
        }
    }

    Ok(())
}

fn create_default_migration_if_needed(
    database_url: &str,
    migrations_dir: &Path,
) -> DatabaseResult<()> {
    let initial_migration_path = migrations_dir.join("00000000000000_diesel_initial_setup");
    if initial_migration_path.exists() {
        return Ok(());
    }

    #[allow(unreachable_patterns, clippy::single_match)]
    match Backend::for_url(database_url) {
        #[cfg(feature = "postgres")]
        Backend::Pg => {
            fs::create_dir_all(&initial_migration_path)?;
            let mut up_sql = File::create(initial_migration_path.join("up.sql"))?;
            up_sql.write_all(include_bytes!("setup_sql/postgres/initial_setup/up.sql"))?;
            let mut down_sql = File::create(initial_migration_path.join("down.sql"))?;
            down_sql.write_all(include_bytes!("setup_sql/postgres/initial_setup/down.sql"))?;
        }
        _ => {} // No default migration for this backend
    }

    Ok(())
}

/// Creates the `__diesel_schema_migrations` table if it doesn't exist. If the
/// table didn't exist, it also runs any pending migrations. Returns a
/// `DatabaseError::ConnectionError` if it can't create the table, and exits
/// with a migration error if it can't run migrations.
fn create_schema_table_and_run_migrations_if_needed(
    database_url: &str,
    migrations_dir: &Path,
) -> DatabaseResult<()> {
    if !schema_table_exists(database_url).unwrap_or_else(handle_error) {
        let migrations =
            FileBasedMigrations::from_path(migrations_dir).unwrap_or_else(handle_error);
        call_with_conn!(database_url, super::run_migrations_with_output(migrations))?;
    };
    Ok(())
}

/// Drops the database specified in the connection url. It returns an error
/// if it was unable to drop the database.
fn drop_database(database_url: &str) -> DatabaseResult<()> {
    match Backend::for_url(database_url) {
        #[cfg(feature = "postgres")]
        Backend::Pg => {
            let (database, postgres_url) = change_database_of_url(database_url, "postgres");
            let conn = PgConnection::establish(&postgres_url)?;
            if pg_database_exists(&conn, &database)? {
                println!("Dropping database: {}", database);
                query_helper::drop_database(&database)
                    .if_exists()
                    .execute(&conn)?;
            }
        }
        #[cfg(feature = "sqlite")]
        Backend::Sqlite => {
            if Path::new(database_url).exists() {
                println!("Dropping database: {}", database_url);
                std::fs::remove_file(&database_url)?;
            }
        }
        #[cfg(feature = "mysql")]
        Backend::Mysql => {
            let (database, mysql_url) = change_database_of_url(database_url, "information_schema");
            let conn = MysqlConnection::establish(&mysql_url)?;
            if mysql_database_exists(&conn, &database)? {
                println!("Dropping database: {}", database);
                query_helper::drop_database(&database)
                    .if_exists()
                    .execute(&conn)?;
            }
        }
    }
    Ok(())
}

#[cfg(feature = "postgres")]
table! {
    pg_database (datname) {
        datname -> Text,
        datistemplate -> Bool,
    }
}

#[cfg(feature = "postgres")]
fn pg_database_exists(conn: &PgConnection, database_name: &str) -> QueryResult<bool> {
    use self::pg_database::dsl::*;

    pg_database
        .select(datname)
        .filter(datname.eq(database_name))
        .filter(datistemplate.eq(false))
        .get_result::<String>(conn)
        .optional()
        .map(|x| x.is_some())
}

#[cfg(feature = "mysql")]
table! {
    information_schema.schemata (schema_name) {
        schema_name -> Text,
    }
}

#[cfg(feature = "mysql")]
fn mysql_database_exists(conn: &MysqlConnection, database_name: &str) -> QueryResult<bool> {
    use self::schemata::dsl::*;

    schemata
        .select(schema_name)
        .filter(schema_name.eq(database_name))
        .get_result::<String>(conn)
        .optional()
        .map(|x| x.is_some())
}

/// Returns true if the `__diesel_schema_migrations` table exists in the
/// database we connect to, returns false if it does not.
pub fn schema_table_exists(database_url: &str) -> DatabaseResult<bool> {
    match InferConnection::establish(database_url).unwrap() {
        #[cfg(feature = "postgres")]
        InferConnection::Pg(conn) => select(sql::<Bool>(
            "EXISTS \
             (SELECT 1 \
             FROM information_schema.tables \
             WHERE table_name = '__diesel_schema_migrations')",
        ))
        .get_result(&conn),
        #[cfg(feature = "sqlite")]
        InferConnection::Sqlite(conn) => select(sql::<Bool>(
            "EXISTS \
             (SELECT 1 \
             FROM sqlite_master \
             WHERE type = 'table' \
             AND name = '__diesel_schema_migrations')",
        ))
        .get_result(&conn),
        #[cfg(feature = "mysql")]
        InferConnection::Mysql(conn) => select(sql::<Bool>(
            "EXISTS \
                    (SELECT 1 \
                     FROM information_schema.tables \
                     WHERE table_name = '__diesel_schema_migrations'
                     AND table_schema = DATABASE())",
        ))
        .get_result(&conn),
    }
    .map_err(Into::into)
}

pub fn database_url(matches: &ArgMatches) -> String {
    matches
        .value_of("DATABASE_URL")
        .map(|s| s.into())
        .or_else(|| env::var("DATABASE_URL").ok())
        .unwrap_or_else(|| handle_error(DatabaseError::DatabaseUrlMissing))
}

#[cfg(any(feature = "postgres", feature = "mysql"))]
fn change_database_of_url(database_url: &str, default_database: &str) -> (String, String) {
    let base = ::url::Url::parse(database_url).unwrap();
    let database = base.path_segments().unwrap().last().unwrap().to_owned();
    let mut new_url = base.join(default_database).unwrap();
    new_url.set_query(base.query());
    (database, new_url.into_string())
}

#[allow(clippy::needless_pass_by_value)]
fn handle_error<E: Error, T>(error: E) -> T {
    println!("{}", error);
    ::std::process::exit(1);
}

#[cfg(all(test, any(feature = "postgres", feature = "mysql")))]
mod tests {
    use super::change_database_of_url;

    #[test]
    fn split_pg_connection_string_returns_postgres_url_and_database() {
        let database = "database".to_owned();
        let base_url = "postgresql://localhost:5432".to_owned();
        let database_url = format!("{}/{}", base_url, database);
        let postgres_url = format!("{}/{}", base_url, "postgres");
        assert_eq!(
            (database, postgres_url),
            change_database_of_url(&database_url, "postgres")
        );
    }

    #[test]
    fn split_pg_connection_string_handles_user_and_password() {
        let database = "database".to_owned();
        let base_url = "postgresql://user:password@localhost:5432".to_owned();
        let database_url = format!("{}/{}", base_url, database);
        let postgres_url = format!("{}/{}", base_url, "postgres");
        assert_eq!(
            (database, postgres_url),
            change_database_of_url(&database_url, "postgres")
        );
    }

    #[test]
    fn split_pg_connection_string_handles_query_string() {
        let database = "database".to_owned();
        let query = "?sslmode=true".to_owned();
        let base_url = "postgresql://user:password@localhost:5432".to_owned();
        let database_url = format!("{}/{}{}", base_url, database, query);
        let postgres_url = format!("{}/{}{}", base_url, "postgres", query);
        assert_eq!(
            (database, postgres_url),
            change_database_of_url(&database_url, "postgres")
        );
    }
}
<commit_msg>add more details when error on call_with_conn<commit_after>#[cfg(any(feature = "postgres", feature = "mysql"))]
use super::query_helper;
use clap::ArgMatches;
use diesel::dsl::sql;
use diesel::sql_types::Bool;
use diesel::*;
use diesel_migrations::FileBasedMigrations;

use crate::database_error::{DatabaseError, DatabaseResult};

use std::env;
use std::error::Error;
#[cfg(feature = "postgres")]
use std::fs::{self, File};
#[cfg(feature = "postgres")]
use std::io::Write;
use std::path::Path;

enum Backend {
    #[cfg(feature = "postgres")]
    Pg,
    #[cfg(feature = "sqlite")]
    Sqlite,
    #[cfg(feature = "mysql")]
    Mysql,
}

impl Backend {
    fn for_url(database_url: &str) -> Self {
        match database_url {
            _ if database_url.starts_with("postgres://")
                || database_url.starts_with("postgresql://") =>
            {
                #[cfg(feature = "postgres")]
                {
                    Backend::Pg
                }
                #[cfg(not(feature = "postgres"))]
                {
                    panic!(
                        "Database url `{}` requires the `postgres` feature but it's not enabled.",
                        database_url
                    );
                }
            }
            _ if database_url.starts_with("mysql://") =>
            {
                #[cfg(feature = "mysql")]
                {
                    Backend::Mysql
                }
                #[cfg(not(feature = "mysql"))]
                {
                    panic!(
                        "Database url `{}` requires the `mysql` feature but it's not enabled.",
                        database_url
                    );
                }
            }
            #[cfg(feature = "sqlite")]
            _ => Backend::Sqlite,
            #[cfg(not(feature = "sqlite"))]
            _ => {
                if database_url.starts_with("sqlite://") {
                    panic!(
                        "Database url `{}` requires the `sqlite` feature but it's not enabled.",
                        database_url
                    );
                }

                let mut available_schemes: Vec<&str> = Vec::new();

                // One of these will always be true, or you are compiling
                // diesel_cli without a backend. And why would you ever want to
                // do that?
                if cfg!(feature = "postgres") {
                    available_schemes.push("`postgres://`");
                }
                if cfg!(feature = "mysql") {
                    available_schemes.push("`mysql://`");
                }

                panic!(
                    "`{}` is not a valid database URL. It should start with {}, or maybe you meant to use the `sqlite` feature which is not enabled.",
                    database_url,
                    available_schemes.join(" or ")
                );
            }
            #[cfg(not(any(feature = "mysql", feature = "sqlite", feature = "postgres")))]
            _ => compile_error!(
                "At least one backend must be specified for use with this crate. \
                 You may omit the unneeded dependencies in the following command. \n\n \
                 ex. `cargo install diesel_cli --no-default-features --features mysql postgres sqlite` \n"
            ),
        }
    }
}

pub enum InferConnection {
    #[cfg(feature = "postgres")]
    Pg(PgConnection),
    #[cfg(feature = "sqlite")]
    Sqlite(SqliteConnection),
    #[cfg(feature = "mysql")]
    Mysql(MysqlConnection),
}

impl InferConnection {
    pub fn establish(database_url: &str) -> DatabaseResult<Self> {
        match Backend::for_url(database_url) {
            #[cfg(feature = "postgres")]
            Backend::Pg => PgConnection::establish(database_url).map(InferConnection::Pg),
            #[cfg(feature = "sqlite")]
            Backend::Sqlite => {
                SqliteConnection::establish(database_url).map(InferConnection::Sqlite)
            }
            #[cfg(feature = "mysql")]
            Backend::Mysql => MysqlConnection::establish(database_url).map(InferConnection::Mysql),
        }
        .map_err(Into::into)
    }
}

macro_rules! call_with_conn {
    (
        $database_url:expr,
        $($func:ident)::+
    ) => {
        call_with_conn!($database_url, $($func)::+ ())
    };

    (
        $database_url:expr,
        $($func:ident)::+ ($($args:expr),*)
    ) => {
        match crate::database::InferConnection::establish(&$database_url)
            .unwrap_or_else(|err| {crate::database::handle_error_with_database_url(&$database_url, err)})
        {
            #[cfg(feature="postgres")]
            crate::database::InferConnection::Pg(ref conn) => $($func)::+ (conn, $($args),*),
            #[cfg(feature="sqlite")]
            crate::database::InferConnection::Sqlite(ref conn) => $($func)::+ (conn, $($args),*),
            #[cfg(feature="mysql")]
            crate::database::InferConnection::Mysql(ref conn) => $($func)::+ (conn, $($args),*),
        }
    };
}

pub fn reset_database(args: &ArgMatches, migrations_dir: &Path) -> DatabaseResult<()> {
    drop_database(&database_url(args))?;
    setup_database(args, migrations_dir)
}

pub fn setup_database(args: &ArgMatches, migrations_dir: &Path) -> DatabaseResult<()> {
    let database_url = database_url(args);

    create_database_if_needed(&database_url)?;
    create_default_migration_if_needed(&database_url, migrations_dir)?;
    create_schema_table_and_run_migrations_if_needed(&database_url, migrations_dir)?;
    Ok(())
}

pub fn drop_database_command(args: &ArgMatches) -> DatabaseResult<()> {
    drop_database(&database_url(args))
}

/// Creates the database specified in the connection url. It returns an error
/// it was unable to create the database.
fn create_database_if_needed(database_url: &str) -> DatabaseResult<()> {
    match Backend::for_url(database_url) {
        #[cfg(feature = "postgres")]
        Backend::Pg => {
            if PgConnection::establish(database_url).is_err() {
                let (database, postgres_url) = change_database_of_url(database_url, "postgres");
                println!("Creating database: {}", database);
                let conn = PgConnection::establish(&postgres_url)?;
                query_helper::create_database(&database).execute(&conn)?;
            }
        }
        #[cfg(feature = "sqlite")]
        Backend::Sqlite => {
            if !::std::path::Path::new(database_url).exists() {
                println!("Creating database: {}", database_url);
                SqliteConnection::establish(database_url)?;
            }
        }
        #[cfg(feature = "mysql")]
        Backend::Mysql => {
            if MysqlConnection::establish(database_url).is_err() {
                let (database, mysql_url) =
                    change_database_of_url(database_url, "information_schema");
                println!("Creating database: {}", database);
                let conn = MysqlConnection::establish(&mysql_url)?;
                query_helper::create_database(&database).execute(&conn)?;
            }
        }
    }

    Ok(())
}

fn create_default_migration_if_needed(
    database_url: &str,
    migrations_dir: &Path,
) -> DatabaseResult<()> {
    let initial_migration_path = migrations_dir.join("00000000000000_diesel_initial_setup");
    if initial_migration_path.exists() {
        return Ok(());
    }

    #[allow(unreachable_patterns, clippy::single_match)]
    match Backend::for_url(database_url) {
        #[cfg(feature = "postgres")]
        Backend::Pg => {
            fs::create_dir_all(&initial_migration_path)?;
            let mut up_sql = File::create(initial_migration_path.join("up.sql"))?;
            up_sql.write_all(include_bytes!("setup_sql/postgres/initial_setup/up.sql"))?;
            let mut down_sql = File::create(initial_migration_path.join("down.sql"))?;
            down_sql.write_all(include_bytes!("setup_sql/postgres/initial_setup/down.sql"))?;
        }
        _ => {} // No default migration for this backend
    }

    Ok(())
}

/// Creates the `__diesel_schema_migrations` table if it doesn't exist. If the
/// table didn't exist, it also runs any pending migrations. Returns a
/// `DatabaseError::ConnectionError` if it can't create the table, and exits
/// with a migration error if it can't run migrations.
fn create_schema_table_and_run_migrations_if_needed(
    database_url: &str,
    migrations_dir: &Path,
) -> DatabaseResult<()> {
    if !schema_table_exists(database_url).unwrap_or_else(handle_error) {
        let migrations =
            FileBasedMigrations::from_path(migrations_dir).unwrap_or_else(handle_error);
        call_with_conn!(database_url, super::run_migrations_with_output(migrations))?;
    };
    Ok(())
}

/// Drops the database specified in the connection url. It returns an error
/// if it was unable to drop the database.
fn drop_database(database_url: &str) -> DatabaseResult<()> {
    match Backend::for_url(database_url) {
        #[cfg(feature = "postgres")]
        Backend::Pg => {
            let (database, postgres_url) = change_database_of_url(database_url, "postgres");
            let conn = PgConnection::establish(&postgres_url)?;
            if pg_database_exists(&conn, &database)? {
                println!("Dropping database: {}", database);
                query_helper::drop_database(&database)
                    .if_exists()
                    .execute(&conn)?;
            }
        }
        #[cfg(feature = "sqlite")]
        Backend::Sqlite => {
            if Path::new(database_url).exists() {
                println!("Dropping database: {}", database_url);
                std::fs::remove_file(&database_url)?;
            }
        }
        #[cfg(feature = "mysql")]
        Backend::Mysql => {
            let (database, mysql_url) = change_database_of_url(database_url, "information_schema");
            let conn = MysqlConnection::establish(&mysql_url)?;
            if mysql_database_exists(&conn, &database)? {
                println!("Dropping database: {}", database);
                query_helper::drop_database(&database)
                    .if_exists()
                    .execute(&conn)?;
            }
        }
    }
    Ok(())
}

#[cfg(feature = "postgres")]
table! {
    pg_database (datname) {
        datname -> Text,
        datistemplate -> Bool,
    }
}

#[cfg(feature = "postgres")]
fn pg_database_exists(conn: &PgConnection, database_name: &str) -> QueryResult<bool> {
    use self::pg_database::dsl::*;

    pg_database
        .select(datname)
        .filter(datname.eq(database_name))
        .filter(datistemplate.eq(false))
        .get_result::<String>(conn)
        .optional()
        .map(|x| x.is_some())
}

#[cfg(feature = "mysql")]
table! {
    information_schema.schemata (schema_name) {
        schema_name -> Text,
    }
}

#[cfg(feature = "mysql")]
fn mysql_database_exists(conn: &MysqlConnection, database_name: &str) -> QueryResult<bool> {
    use self::schemata::dsl::*;

    schemata
        .select(schema_name)
        .filter(schema_name.eq(database_name))
        .get_result::<String>(conn)
        .optional()
        .map(|x| x.is_some())
}

/// Returns true if the `__diesel_schema_migrations` table exists in the
/// database we connect to, returns false if it does not.
pub fn schema_table_exists(database_url: &str) -> DatabaseResult<bool> {
    match InferConnection::establish(database_url).unwrap() {
        #[cfg(feature = "postgres")]
        InferConnection::Pg(conn) => select(sql::<Bool>(
            "EXISTS \
             (SELECT 1 \
             FROM information_schema.tables \
             WHERE table_name = '__diesel_schema_migrations')",
        ))
        .get_result(&conn),
        #[cfg(feature = "sqlite")]
        InferConnection::Sqlite(conn) => select(sql::<Bool>(
            "EXISTS \
             (SELECT 1 \
             FROM sqlite_master \
             WHERE type = 'table' \
             AND name = '__diesel_schema_migrations')",
        ))
        .get_result(&conn),
        #[cfg(feature = "mysql")]
        InferConnection::Mysql(conn) => select(sql::<Bool>(
            "EXISTS \
                    (SELECT 1 \
                     FROM information_schema.tables \
                     WHERE table_name = '__diesel_schema_migrations'
                     AND table_schema = DATABASE())",
        ))
        .get_result(&conn),
    }
    .map_err(Into::into)
}

pub fn database_url(matches: &ArgMatches) -> String {
    matches
        .value_of("DATABASE_URL")
        .map(|s| s.into())
        .or_else(|| env::var("DATABASE_URL").ok())
        .unwrap_or_else(|| handle_error(DatabaseError::DatabaseUrlMissing))
}

#[cfg(any(feature = "postgres", feature = "mysql"))]
fn change_database_of_url(database_url: &str, default_database: &str) -> (String, String) {
    let base = ::url::Url::parse(database_url).unwrap();
    let database = base.path_segments().unwrap().last().unwrap().to_owned();
    let mut new_url = base.join(default_database).unwrap();
    new_url.set_query(base.query());
    (database, new_url.into_string())
}

#[allow(clippy::needless_pass_by_value)]
fn handle_error<E: Error, T>(error: E) -> T {
    println!("{}", error);
    ::std::process::exit(1);
}

#[allow(clippy::needless_pass_by_value)]
pub fn handle_error_with_database_url<E: Error, T>(database_url: &str, error: E) -> T {
    eprintln!(
        "Could not connect to database via `{}`: {}",
        database_url, error
    );
    ::std::process::exit(1);
}

#[cfg(all(test, any(feature = "postgres", feature = "mysql")))]
mod tests {
    use super::change_database_of_url;

    #[test]
    fn split_pg_connection_string_returns_postgres_url_and_database() {
        let database = "database".to_owned();
        let base_url = "postgresql://localhost:5432".to_owned();
        let database_url = format!("{}/{}", base_url, database);
        let postgres_url = format!("{}/{}", base_url, "postgres");
        assert_eq!(
            (database, postgres_url),
            change_database_of_url(&database_url, "postgres")
        );
    }

    #[test]
    fn split_pg_connection_string_handles_user_and_password() {
        let database = "database".to_owned();
        let base_url = "postgresql://user:password@localhost:5432".to_owned();
        let database_url = format!("{}/{}", base_url, database);
        let postgres_url = format!("{}/{}", base_url, "postgres");
        assert_eq!(
            (database, postgres_url),
            change_database_of_url(&database_url, "postgres")
        );
    }

    #[test]
    fn split_pg_connection_string_handles_query_string() {
        let database = "database".to_owned();
        let query = "?sslmode=true".to_owned();
        let base_url = "postgresql://user:password@localhost:5432".to_owned();
        let database_url = format!("{}/{}{}", base_url, database, query);
        let postgres_url = format!("{}/{}{}", base_url, "postgres", query);
        assert_eq!(
            (database, postgres_url),
            change_database_of_url(&database_url, "postgres")
        );
    }
}
<|endoftext|>