<commit_before>// Copyright (c) 2016 The vulkano developers
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>,
// at your option. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

use std::error;
use std::fmt;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;

use SafeDeref;
use VulkanObject;
use buffer::BufferAccess;
use command_buffer::submit::SubmitAnyBuilder;
use command_buffer::submit::SubmitCommandBufferBuilder;
use command_buffer::sys::UnsafeCommandBuffer;
use device::Device;
use device::DeviceOwned;
use device::Queue;
use image::ImageAccess;
use image::ImageLayout;
use sync::AccessCheckError;
use sync::AccessError;
use sync::AccessFlagBits;
use sync::FlushError;
use sync::GpuFuture;
use sync::NowFuture;
use sync::PipelineStages;
use sync::now;

pub unsafe trait CommandBuffer: DeviceOwned {
    /// The command pool of the command buffer.
    type PoolAlloc;

    /// Returns the underlying `UnsafeCommandBuffer` of this command buffer.
    fn inner(&self) -> &UnsafeCommandBuffer<Self::PoolAlloc>;

    /*/// Returns the queue family of the command buffer.
    #[inline]
    fn queue_family(&self) -> QueueFamily
        where Self::PoolAlloc: CommandPoolAlloc
    {
        self.inner().queue_family()
    }*/

    #[deprecated(note = "Renamed to `lock_submit`")]
    fn prepare_submit(&self, future: &GpuFuture, queue: &Queue)
                      -> Result<(), CommandBufferExecError>
    {
        self.lock_submit(future, queue)
    }

    /// Checks whether this command buffer is allowed to be submitted after the `future` and on
    /// the given queue, and if so locks it.
    ///
    /// If you call this function, then you should call `unlock` afterwards.
    fn lock_submit(&self, future: &GpuFuture, queue: &Queue)
                   -> Result<(), CommandBufferExecError>;

    /// Unlocks the command buffer. Should be called once for each call to `lock_submit`.
    ///
    /// # Safety
    ///
    /// Must not be called if you haven't called `lock_submit` before.
    unsafe fn unlock(&self);

    /// Executes this command buffer on a queue.
    ///
    /// This function returns an object that implements the `GpuFuture` trait. See the
    /// documentation of the `sync` module for more information.
    ///
    /// The command buffer is not actually executed until you call `flush()` on the object.
    /// You are encouraged to chain together as many futures as possible before calling `flush()`,
    /// and call `.then_signal_future()` before doing so. Note however that once you called
    /// `execute()` there is no way to cancel the execution, even if you didn't flush yet.
    ///
    /// > **Note**: In the future this function may return `-> impl GpuFuture` instead of a
    /// > concrete type.
    ///
    /// > **Note**: This is just a shortcut for `execute_after(vulkano::sync::now(), queue)`.
    ///
    /// # Panic
    ///
    /// Panics if the device of the command buffer is not the same as the device of the future.
    #[inline]
    fn execute(self, queue: Arc<Queue>)
               -> Result<CommandBufferExecFuture<NowFuture, Self>, CommandBufferExecError>
        where Self: Sized + 'static
    {
        let device = queue.device().clone();
        self.execute_after(now(device), queue)
    }

    /// Executes the command buffer after an existing future.
    ///
    /// This function returns an object that implements the `GpuFuture` trait. See the
    /// documentation of the `sync` module for more information.
    ///
    /// The command buffer is not actually executed until you call `flush()` on the object.
    /// You are encouraged to chain together as many futures as possible before calling `flush()`,
    /// and call `.then_signal_future()` before doing so. Note however that once you called
    /// `execute()` there is no way to cancel the execution, even if you didn't flush yet.
    ///
    /// > **Note**: In the future this function may return `-> impl GpuFuture` instead of a
    /// > concrete type.
    ///
    /// This function requires the `'static` lifetime to be on the command buffer. This is because
    /// this function returns a `CommandBufferExecFuture` whose job is to lock resources and keep
    /// them alive while they are in use by the GPU. If `'static` wasn't required, you could call
    /// `std::mem::forget` on that object and "unlock" these resources. For more information about
    /// this problem, search the web for "rust thread scoped leakpocalypse".
    ///
    /// # Panic
    ///
    /// Panics if the device of the command buffer is not the same as the device of the future.
    #[inline]
    fn execute_after<F>(self, future: F, queue: Arc<Queue>)
                        -> Result<CommandBufferExecFuture<F, Self>, CommandBufferExecError>
        where Self: Sized + 'static,
              F: GpuFuture
    {
        assert_eq!(self.device().internal_object(),
                   future.device().internal_object());

        if !future.queue_change_allowed() {
            assert!(future.queue().unwrap().is_same(&queue));
        }

        self.lock_submit(&future, &queue)?;

        Ok(CommandBufferExecFuture {
               previous: future,
               command_buffer: self,
               queue: queue,
               submitted: Mutex::new(false),
               finished: AtomicBool::new(false),
           })
    }

    fn check_buffer_access(&self, buffer: &BufferAccess, exclusive: bool, queue: &Queue)
                           -> Result<Option<(PipelineStages, AccessFlagBits)>, AccessCheckError>;

    fn check_image_access(&self, image: &ImageAccess, layout: ImageLayout, exclusive: bool,
                          queue: &Queue)
                          -> Result<Option<(PipelineStages, AccessFlagBits)>, AccessCheckError>;

    // FIXME: lots of other methods
}

/// Turns a command buffer builder into a real command buffer.
pub unsafe trait CommandBufferBuild {
    /// The type of the built command buffer.
    type Out;
    /// Error that can be returned when building.
    type Err;

    /// Builds the command buffer.
    fn build(self) -> Result<Self::Out, Self::Err>;
}

unsafe impl<T> CommandBuffer for T
    where T: SafeDeref,
          T::Target: CommandBuffer
{
    type PoolAlloc = <T::Target as CommandBuffer>::PoolAlloc;

    #[inline]
    fn inner(&self) -> &UnsafeCommandBuffer<Self::PoolAlloc> {
        (**self).inner()
    }

    #[inline]
    fn lock_submit(&self, future: &GpuFuture, queue: &Queue)
                   -> Result<(), CommandBufferExecError> {
        (**self).lock_submit(future, queue)
    }

    #[inline]
    unsafe fn unlock(&self) {
        (**self).unlock();
    }

    #[inline]
    fn check_buffer_access(
        &self, buffer: &BufferAccess, exclusive: bool, queue: &Queue)
        -> Result<Option<(PipelineStages, AccessFlagBits)>, AccessCheckError> {
        (**self).check_buffer_access(buffer, exclusive, queue)
    }

    #[inline]
    fn check_image_access(&self, image: &ImageAccess, layout: ImageLayout, exclusive: bool,
                          queue: &Queue)
                          -> Result<Option<(PipelineStages, AccessFlagBits)>, AccessCheckError> {
        (**self).check_image_access(image, layout, exclusive, queue)
    }
}

/// Represents a command buffer being executed by the GPU and the moment when the execution
/// finishes.
#[must_use = "Dropping this object will immediately block the thread until the GPU has finished processing the submission"]
pub struct CommandBufferExecFuture<F, Cb>
    where F: GpuFuture,
          Cb: CommandBuffer
{
    previous: F,
    command_buffer: Cb,
    queue: Arc<Queue>,
    // True if the command buffer has already been submitted.
    // If flush is called multiple times, we want to block so that only one flushing is executed.
    // Therefore we use a `Mutex<bool>` and not an `AtomicBool`.
    submitted: Mutex<bool>,
    finished: AtomicBool,
}

unsafe impl<F, Cb> GpuFuture for CommandBufferExecFuture<F, Cb>
    where F: GpuFuture,
          Cb: CommandBuffer
{
    #[inline]
    fn cleanup_finished(&mut self) {
        self.previous.cleanup_finished();
    }

    unsafe fn build_submission(&self) -> Result<SubmitAnyBuilder, FlushError> {
        Ok(match self.previous.build_submission()? {
               SubmitAnyBuilder::Empty => {
                   let mut builder = SubmitCommandBufferBuilder::new();
                   builder.add_command_buffer(self.command_buffer.inner());
                   SubmitAnyBuilder::CommandBuffer(builder)
               },
               SubmitAnyBuilder::SemaphoresWait(sem) => {
                   let mut builder: SubmitCommandBufferBuilder = sem.into();
                   builder.add_command_buffer(self.command_buffer.inner());
                   SubmitAnyBuilder::CommandBuffer(builder)
               },
               SubmitAnyBuilder::CommandBuffer(mut builder) => {
                   // FIXME: add pipeline barrier
                   builder.add_command_buffer(self.command_buffer.inner());
                   SubmitAnyBuilder::CommandBuffer(builder)
               },
               SubmitAnyBuilder::QueuePresent(_) |
               SubmitAnyBuilder::BindSparse(_) => {
                   unimplemented!() // TODO:
                /*present.submit();     // TODO: wrong
                let mut builder = SubmitCommandBufferBuilder::new();
                builder.add_command_buffer(self.command_buffer.inner());
                SubmitAnyBuilder::CommandBuffer(builder)*/
               },
           })
    }

    #[inline]
    fn flush(&self) -> Result<(), FlushError> {
        unsafe {
            let mut submitted = self.submitted.lock().unwrap();
            if *submitted {
                return Ok(());
            }

            let queue = self.queue.clone();

            match self.build_submission()? {
                SubmitAnyBuilder::Empty => {},
                SubmitAnyBuilder::CommandBuffer(builder) => {
                    builder.submit(&queue)?;
                },
                _ => unreachable!(),
            };

            // Only write `true` here in order to try again next time if we failed to submit.
            *submitted = true;
            Ok(())
        }
    }

    #[inline]
    unsafe fn signal_finished(&self) {
        self.finished.store(true, Ordering::SeqCst);
        self.command_buffer.unlock();
        self.previous.signal_finished();
    }

    #[inline]
    fn queue_change_allowed(&self) -> bool {
        false
    }

    #[inline]
    fn queue(&self) -> Option<Arc<Queue>> {
        Some(self.queue.clone())
    }

    #[inline]
    fn check_buffer_access(
        &self, buffer: &BufferAccess, exclusive: bool, queue: &Queue)
        -> Result<Option<(PipelineStages, AccessFlagBits)>, AccessCheckError> {
        match self.command_buffer
            .check_buffer_access(buffer, exclusive, queue) {
            Ok(v) => Ok(v),
            Err(AccessCheckError::Denied(err)) => Err(AccessCheckError::Denied(err)),
            Err(AccessCheckError::Unknown) => {
                self.previous.check_buffer_access(buffer, exclusive, queue)
            },
        }
    }

    #[inline]
    fn check_image_access(&self, image: &ImageAccess, layout: ImageLayout, exclusive: bool,
                          queue: &Queue)
                          -> Result<Option<(PipelineStages, AccessFlagBits)>, AccessCheckError> {
        match self.command_buffer
            .check_image_access(image, layout, exclusive, queue) {
            Ok(v) => Ok(v),
            Err(AccessCheckError::Denied(err)) => Err(AccessCheckError::Denied(err)),
            Err(AccessCheckError::Unknown) => {
                self.previous
                    .check_image_access(image, layout, exclusive, queue)
            },
        }
    }
}

unsafe impl<F, Cb> DeviceOwned for CommandBufferExecFuture<F, Cb>
    where F: GpuFuture,
          Cb: CommandBuffer
{
    #[inline]
    fn device(&self) -> &Arc<Device> {
        self.command_buffer.device()
    }
}

impl<F, Cb> Drop for CommandBufferExecFuture<F, Cb>
    where F: GpuFuture,
          Cb: CommandBuffer
{
    fn drop(&mut self) {
        unsafe {
            if !*self.finished.get_mut() {
                // TODO: handle errors?
                self.flush().unwrap();
                // Block until the queue finished.
                self.queue.wait().unwrap();
                self.command_buffer.unlock();
                self.previous.signal_finished();
            }
        }
    }
}

/// Error that can happen when attempting to execute a command buffer.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum CommandBufferExecError {
    /// Access to a resource has been denied.
    AccessError(AccessError),

    // TODO: missing entries (eg. wrong queue family, secondary command buffer)
}

impl error::Error for CommandBufferExecError {
    #[inline]
    fn description(&self) -> &str {
        match *self {
            CommandBufferExecError::AccessError(_) => "access to a resource has been denied",
        }
    }

    #[inline]
    fn cause(&self) -> Option<&error::Error> {
        match *self {
            CommandBufferExecError::AccessError(ref err) => Some(err),
        }
    }
}

impl fmt::Display for CommandBufferExecError {
    #[inline]
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "{}", error::Error::description(self))
    }
}

impl From<AccessError> for CommandBufferExecError {
    #[inline]
    fn from(err: AccessError) -> CommandBufferExecError {
        CommandBufferExecError::AccessError(err)
    }
}
<commit_msg>Add a small TODO<commit_after>// Copyright (c) 2016 The vulkano developers
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>,
// at your option. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

use std::error;
use std::fmt;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;

use SafeDeref;
use VulkanObject;
use buffer::BufferAccess;
use command_buffer::submit::SubmitAnyBuilder;
use command_buffer::submit::SubmitCommandBufferBuilder;
use command_buffer::sys::UnsafeCommandBuffer;
use device::Device;
use device::DeviceOwned;
use device::Queue;
use image::ImageAccess;
use image::ImageLayout;
use sync::AccessCheckError;
use sync::AccessError;
use sync::AccessFlagBits;
use sync::FlushError;
use sync::GpuFuture;
use sync::NowFuture;
use sync::PipelineStages;
use sync::now;

pub unsafe trait CommandBuffer: DeviceOwned {
    /// The command pool of the command buffer.
    type PoolAlloc;

    /// Returns the underlying `UnsafeCommandBuffer` of this command buffer.
    fn inner(&self) -> &UnsafeCommandBuffer<Self::PoolAlloc>;

    /*/// Returns the queue family of the command buffer.
    #[inline]
    fn queue_family(&self) -> QueueFamily
        where Self::PoolAlloc: CommandPoolAlloc
    {
        self.inner().queue_family()
    }*/

    #[deprecated(note = "Renamed to `lock_submit`")]
    fn prepare_submit(&self, future: &GpuFuture, queue: &Queue)
                      -> Result<(), CommandBufferExecError>
    {
        self.lock_submit(future, queue)
    }

    /// Checks whether this command buffer is allowed to be submitted after the `future` and on
    /// the given queue, and if so locks it.
    ///
    /// If you call this function, then you should call `unlock` afterwards.
    // TODO: require `&mut self` instead, but this has some consequences on other parts of the lib
    fn lock_submit(&self, future: &GpuFuture, queue: &Queue)
                   -> Result<(), CommandBufferExecError>;

    /// Unlocks the command buffer. Should be called once for each call to `lock_submit`.
    ///
    /// # Safety
    ///
    /// Must not be called if you haven't called `lock_submit` before.
    // TODO: require `&mut self` instead, but this has some consequences on other parts of the lib
    unsafe fn unlock(&self);

    /// Executes this command buffer on a queue.
    ///
    /// This function returns an object that implements the `GpuFuture` trait. See the
    /// documentation of the `sync` module for more information.
    ///
    /// The command buffer is not actually executed until you call `flush()` on the object.
    /// You are encouraged to chain together as many futures as possible before calling `flush()`,
    /// and call `.then_signal_future()` before doing so. Note however that once you called
    /// `execute()` there is no way to cancel the execution, even if you didn't flush yet.
    ///
    /// > **Note**: In the future this function may return `-> impl GpuFuture` instead of a
    /// > concrete type.
    ///
    /// > **Note**: This is just a shortcut for `execute_after(vulkano::sync::now(), queue)`.
    ///
    /// # Panic
    ///
    /// Panics if the device of the command buffer is not the same as the device of the future.
    #[inline]
    fn execute(self, queue: Arc<Queue>)
               -> Result<CommandBufferExecFuture<NowFuture, Self>, CommandBufferExecError>
        where Self: Sized + 'static
    {
        let device = queue.device().clone();
        self.execute_after(now(device), queue)
    }

    /// Executes the command buffer after an existing future.
    ///
    /// This function returns an object that implements the `GpuFuture` trait. See the
    /// documentation of the `sync` module for more information.
    ///
    /// The command buffer is not actually executed until you call `flush()` on the object.
    /// You are encouraged to chain together as many futures as possible before calling `flush()`,
    /// and call `.then_signal_future()` before doing so. Note however that once you called
    /// `execute()` there is no way to cancel the execution, even if you didn't flush yet.
    ///
    /// > **Note**: In the future this function may return `-> impl GpuFuture` instead of a
    /// > concrete type.
    ///
    /// This function requires the `'static` lifetime to be on the command buffer. This is because
    /// this function returns a `CommandBufferExecFuture` whose job is to lock resources and keep
    /// them alive while they are in use by the GPU. If `'static` wasn't required, you could call
    /// `std::mem::forget` on that object and "unlock" these resources. For more information about
    /// this problem, search the web for "rust thread scoped leakpocalypse".
    ///
    /// # Panic
    ///
    /// Panics if the device of the command buffer is not the same as the device of the future.
    #[inline]
    fn execute_after<F>(self, future: F, queue: Arc<Queue>)
                        -> Result<CommandBufferExecFuture<F, Self>, CommandBufferExecError>
        where Self: Sized + 'static,
              F: GpuFuture
    {
        assert_eq!(self.device().internal_object(),
                   future.device().internal_object());

        if !future.queue_change_allowed() {
            assert!(future.queue().unwrap().is_same(&queue));
        }

        self.lock_submit(&future, &queue)?;

        Ok(CommandBufferExecFuture {
               previous: future,
               command_buffer: self,
               queue: queue,
               submitted: Mutex::new(false),
               finished: AtomicBool::new(false),
           })
    }

    fn check_buffer_access(&self, buffer: &BufferAccess, exclusive: bool, queue: &Queue)
                           -> Result<Option<(PipelineStages, AccessFlagBits)>, AccessCheckError>;

    fn check_image_access(&self, image: &ImageAccess, layout: ImageLayout, exclusive: bool,
                          queue: &Queue)
                          -> Result<Option<(PipelineStages, AccessFlagBits)>, AccessCheckError>;

    // FIXME: lots of other methods
}

/// Turns a command buffer builder into a real command buffer.
pub unsafe trait CommandBufferBuild {
    /// The type of the built command buffer.
    type Out;
    /// Error that can be returned when building.
    type Err;

    /// Builds the command buffer.
    fn build(self) -> Result<Self::Out, Self::Err>;
}

unsafe impl<T> CommandBuffer for T
    where T: SafeDeref,
          T::Target: CommandBuffer
{
    type PoolAlloc = <T::Target as CommandBuffer>::PoolAlloc;

    #[inline]
    fn inner(&self) -> &UnsafeCommandBuffer<Self::PoolAlloc> {
        (**self).inner()
    }

    #[inline]
    fn lock_submit(&self, future: &GpuFuture, queue: &Queue)
                   -> Result<(), CommandBufferExecError> {
        (**self).lock_submit(future, queue)
    }

    #[inline]
    unsafe fn unlock(&self) {
        (**self).unlock();
    }

    #[inline]
    fn check_buffer_access(
        &self, buffer: &BufferAccess, exclusive: bool, queue: &Queue)
        -> Result<Option<(PipelineStages, AccessFlagBits)>, AccessCheckError> {
        (**self).check_buffer_access(buffer, exclusive, queue)
    }

    #[inline]
    fn check_image_access(&self, image: &ImageAccess, layout: ImageLayout, exclusive: bool,
                          queue: &Queue)
                          -> Result<Option<(PipelineStages, AccessFlagBits)>, AccessCheckError> {
        (**self).check_image_access(image, layout, exclusive, queue)
    }
}

/// Represents a command buffer being executed by the GPU and the moment when the execution
/// finishes.
#[must_use = "Dropping this object will immediately block the thread until the GPU has finished processing the submission"]
pub struct CommandBufferExecFuture<F, Cb>
    where F: GpuFuture,
          Cb: CommandBuffer
{
    previous: F,
    command_buffer: Cb,
    queue: Arc<Queue>,
    // True if the command buffer has already been submitted.
    // If flush is called multiple times, we want to block so that only one flushing is executed.
    // Therefore we use a `Mutex<bool>` and not an `AtomicBool`.
    submitted: Mutex<bool>,
    finished: AtomicBool,
}

unsafe impl<F, Cb> GpuFuture for CommandBufferExecFuture<F, Cb>
    where F: GpuFuture,
          Cb: CommandBuffer
{
    #[inline]
    fn cleanup_finished(&mut self) {
        self.previous.cleanup_finished();
    }

    unsafe fn build_submission(&self) -> Result<SubmitAnyBuilder, FlushError> {
        Ok(match self.previous.build_submission()? {
               SubmitAnyBuilder::Empty => {
                   let mut builder = SubmitCommandBufferBuilder::new();
                   builder.add_command_buffer(self.command_buffer.inner());
                   SubmitAnyBuilder::CommandBuffer(builder)
               },
               SubmitAnyBuilder::SemaphoresWait(sem) => {
                   let mut builder: SubmitCommandBufferBuilder = sem.into();
                   builder.add_command_buffer(self.command_buffer.inner());
                   SubmitAnyBuilder::CommandBuffer(builder)
               },
               SubmitAnyBuilder::CommandBuffer(mut builder) => {
                   // FIXME: add pipeline barrier
                   builder.add_command_buffer(self.command_buffer.inner());
                   SubmitAnyBuilder::CommandBuffer(builder)
               },
               SubmitAnyBuilder::QueuePresent(_) |
               SubmitAnyBuilder::BindSparse(_) => {
                   unimplemented!() // TODO:
                /*present.submit();     // TODO: wrong
                let mut builder = SubmitCommandBufferBuilder::new();
                builder.add_command_buffer(self.command_buffer.inner());
                SubmitAnyBuilder::CommandBuffer(builder)*/
               },
           })
    }

    #[inline]
    fn flush(&self) -> Result<(), FlushError> {
        unsafe {
            let mut submitted = self.submitted.lock().unwrap();
            if *submitted {
                return Ok(());
            }

            let queue = self.queue.clone();

            match self.build_submission()? {
                SubmitAnyBuilder::Empty => {},
                SubmitAnyBuilder::CommandBuffer(builder) => {
                    builder.submit(&queue)?;
                },
                _ => unreachable!(),
            };

            // Only write `true` here in order to try again next time if we failed to submit.
            *submitted = true;
            Ok(())
        }
    }

    #[inline]
    unsafe fn signal_finished(&self) {
        self.finished.store(true, Ordering::SeqCst);
        self.command_buffer.unlock();
        self.previous.signal_finished();
    }

    #[inline]
    fn queue_change_allowed(&self) -> bool {
        false
    }

    #[inline]
    fn queue(&self) -> Option<Arc<Queue>> {
        Some(self.queue.clone())
    }

    #[inline]
    fn check_buffer_access(
        &self, buffer: &BufferAccess, exclusive: bool, queue: &Queue)
        -> Result<Option<(PipelineStages, AccessFlagBits)>, AccessCheckError> {
        match self.command_buffer
            .check_buffer_access(buffer, exclusive, queue) {
            Ok(v) => Ok(v),
            Err(AccessCheckError::Denied(err)) => Err(AccessCheckError::Denied(err)),
            Err(AccessCheckError::Unknown) => {
                self.previous.check_buffer_access(buffer, exclusive, queue)
            },
        }
    }

    #[inline]
    fn check_image_access(&self, image: &ImageAccess, layout: ImageLayout, exclusive: bool,
                          queue: &Queue)
                          -> Result<Option<(PipelineStages, AccessFlagBits)>, AccessCheckError> {
        match self.command_buffer
            .check_image_access(image, layout, exclusive, queue) {
            Ok(v) => Ok(v),
            Err(AccessCheckError::Denied(err)) => Err(AccessCheckError::Denied(err)),
            Err(AccessCheckError::Unknown) => {
                self.previous
                    .check_image_access(image, layout, exclusive, queue)
            },
        }
    }
}

unsafe impl<F, Cb> DeviceOwned for CommandBufferExecFuture<F, Cb>
    where F: GpuFuture,
          Cb: CommandBuffer
{
    #[inline]
    fn device(&self) -> &Arc<Device> {
        self.command_buffer.device()
    }
}

impl<F, Cb> Drop for CommandBufferExecFuture<F, Cb>
    where F: GpuFuture,
          Cb: CommandBuffer
{
    fn drop(&mut self) {
        unsafe {
            if !*self.finished.get_mut() {
                // TODO: handle errors?
                self.flush().unwrap();
                // Block until the queue finished.
                self.queue.wait().unwrap();
                self.command_buffer.unlock();
                self.previous.signal_finished();
            }
        }
    }
}

/// Error that can happen when attempting to execute a command buffer.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum CommandBufferExecError {
    /// Access to a resource has been denied.
    AccessError(AccessError),

    // TODO: missing entries (eg. wrong queue family, secondary command buffer)
}

impl error::Error for CommandBufferExecError {
    #[inline]
    fn description(&self) -> &str {
        match *self {
            CommandBufferExecError::AccessError(_) => "access to a resource has been denied",
        }
    }

    #[inline]
    fn cause(&self) -> Option<&error::Error> {
        match *self {
            CommandBufferExecError::AccessError(ref err) => Some(err),
        }
    }
}

impl fmt::Display for CommandBufferExecError {
    #[inline]
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "{}", error::Error::description(self))
    }
}

impl From<AccessError> for CommandBufferExecError {
    #[inline]
    fn from(err: AccessError) -> CommandBufferExecError {
        CommandBufferExecError::AccessError(err)
    }
}
<|endoftext|>


<commit_before>extern crate "cargo-registry" as cargo_registry;
extern crate migrate;
extern crate postgres;
extern crate r2d2;

use std::os;
use std::collections::HashSet;
use migrate::Migration;
use postgres::{PostgresTransaction, PostgresResult};

use cargo_registry::krate::Crate;
use cargo_registry::model::Model;

fn main() {
    let db_config = r2d2::Config {
        pool_size: 1,
        helper_tasks: 1,
        test_on_check_out: false,
    };
    let database = cargo_registry::db::pool(env("DATABASE_URL").as_slice(),
                                            db_config);
    let conn = database.get().unwrap();
    let migrations = migrations();

    if os::args().as_slice().get(1).map(|s| s.as_slice()) == Some("rollback") {
        rollback(conn.transaction().unwrap(), migrations).unwrap();
    } else {
        apply(conn.transaction().unwrap(), migrations).unwrap();
    }

    fn env(s: &str) -> String {
        match os::getenv(s) {
            Some(s) => s,
            None => fail!("must have `{}` defined", s),
        }
    }
}

fn apply(tx: PostgresTransaction,
         migrations: Vec<Migration>) -> PostgresResult<()> {
    let mut mgr = try!(migrate::Manager::new(tx));
    for m in migrations.into_iter() {
        try!(mgr.apply(m));
    }
    mgr.set_commit();
    mgr.finish()
}

fn rollback(tx: PostgresTransaction,
            migrations: Vec<Migration>) -> PostgresResult<()> {
    let mut mgr = try!(migrate::Manager::new(tx));
    for m in migrations.into_iter().rev() {
        if mgr.contains(m.version()) {
            try!(mgr.rollback(m));
            break
        }
    }
    mgr.set_commit();
    mgr.finish()
}

fn migrations() -> Vec<Migration> {
    let migrations = vec![
        Migration::add_table(20140924113530, "users", "
            id              SERIAL PRIMARY KEY,
            email           VARCHAR NOT NULL UNIQUE,
            gh_access_token VARCHAR NOT NULL,
            api_token       VARCHAR NOT NULL
        "),
        Migration::add_table(20140924114003, "packages", "
            id              SERIAL PRIMARY KEY,
            name            VARCHAR NOT NULL UNIQUE,
            user_id         INTEGER NOT NULL
        "),
        Migration::add_table(20140924114059, "versions", "
            id              SERIAL PRIMARY KEY,
            package_id      INTEGER NOT NULL,
            num             VARCHAR NOT NULL
        "),
        Migration::run(20140924115329,
                       format!("ALTER TABLE versions ADD CONSTRAINT \
                                unique_num UNIQUE (package_id, num)"),
                       format!("ALTER TABLE versions DROP CONSTRAINT \
                                unique_num")),
        Migration::add_table(20140924120803, "version_dependencies", "
            version_id      INTEGER NOT NULL,
            depends_on_id   INTEGER NOT NULL
        "),
        Migration::add_column(20140925132248, "packages", "updated_at",
                              "TIMESTAMP NOT NULL DEFAULT now()"),
        Migration::add_column(20140925132249, "packages", "created_at",
                              "TIMESTAMP NOT NULL DEFAULT now()"),
        Migration::new(20140925132250, proc(tx) {
            try!(tx.execute("UPDATE packages SET updated_at = now() \
                             WHERE updated_at IS NULL", []));
            try!(tx.execute("UPDATE packages SET created_at = now() \
                             WHERE created_at IS NULL", []));
            Ok(())
        }, proc(_) Ok(())),
        Migration::add_column(20140925132251, "versions", "updated_at",
                              "TIMESTAMP NOT NULL DEFAULT now()"),
        Migration::add_column(20140925132252, "versions", "created_at",
                              "TIMESTAMP NOT NULL DEFAULT now()"),
        Migration::new(20140925132253, proc(tx) {
            try!(tx.execute("UPDATE versions SET updated_at = now() \
                             WHERE updated_at IS NULL", []));
            try!(tx.execute("UPDATE versions SET created_at = now() \
                             WHERE created_at IS NULL", []));
            Ok(())
        }, proc(_) Ok(())),
        Migration::new(20140925132254, proc(tx) {
            try!(tx.execute("ALTER TABLE versions ALTER COLUMN updated_at \
                             DROP DEFAULT", []));
            try!(tx.execute("ALTER TABLE versions ALTER COLUMN created_at \
                             DROP DEFAULT", []));
            try!(tx.execute("ALTER TABLE packages ALTER COLUMN updated_at \
                             DROP DEFAULT", []));
            try!(tx.execute("ALTER TABLE packages ALTER COLUMN created_at \
                             DROP DEFAULT", []));
            Ok(())
        }, proc(_) Ok(())),
        Migration::add_table(20140925153704, "metadata", "
            total_downloads        BIGINT NOT NULL
        "),
        Migration::new(20140925153705, proc(tx) {
            try!(tx.execute("INSERT INTO metadata (total_downloads) \
                             VALUES ($1)", &[&0i64]));
            Ok(())
        }, proc(tx) {
            try!(tx.execute("DELETE FROM metadata", [])); Ok(())
        }),
        Migration::add_column(20140925161623, "packages", "downloads",
                              "INTEGER NOT NULL DEFAULT 0"),
        Migration::add_column(20140925161624, "versions", "downloads",
                              "INTEGER NOT NULL DEFAULT 0"),
        Migration::new(20140925161625, proc(tx) {
            try!(tx.execute("ALTER TABLE versions ALTER COLUMN downloads \
                             DROP DEFAULT", []));
            try!(tx.execute("ALTER TABLE packages ALTER COLUMN downloads \
                             DROP DEFAULT", []));
            Ok(())
        }, proc(_) Ok(())),
        Migration::add_column(20140926130044, "packages", "max_version",
                              "VARCHAR"),
        Migration::new(20140926130045, proc(tx) {
            let stmt = try!(tx.prepare("SELECT * FROM packages"));
            for row in try!(stmt.query(&[])) {
                let pkg: Crate = Model::from_row(&row);
                let versions = pkg.versions(tx).unwrap();
                let v = versions.iter().max_by(|v| &v.num).unwrap();
                let max = v.num.to_string();
                try!(tx.execute("UPDATE packages SET max_version = $1 \
                                 WHERE id = $2",
                                &[&max, &pkg.id]));
            }
            Ok(())
        }, proc(_) Ok(())),
        Migration::new(20140926130046, proc(tx) {
            try!(tx.execute("ALTER TABLE versions ALTER COLUMN downloads \
                             SET NOT NULL", []));
            Ok(())
        }, proc(tx) {
            try!(tx.execute("ALTER TABLE versions ALTER COLUMN downloads \
                             DROP NOT NULL", []));
            Ok(())
        }),
        Migration::new(20140926174020, proc(tx) {
            try!(tx.execute("ALTER TABLE packages RENAME TO crates", []));
            try!(tx.execute("ALTER TABLE versions RENAME COLUMN package_id \
                             TO crate_id", []));
            Ok(())
        }, proc(tx) {
            try!(tx.execute("ALTER TABLE crates RENAME TO packages", []));
            try!(tx.execute("ALTER TABLE versions RENAME COLUMN crate_id \
                             TO package_id", []));
            Ok(())
        }),
        Migration::run(20140929103749,
                       "CREATE INDEX index_crate_updated_at ON crates (updated_at)",
                       "DROP INDEX index_crate_updated_at"),
        Migration::run(20140929103750,
                       "CREATE INDEX index_crate_created_at ON crates (created_at)",
                       "DROP INDEX index_crate_created_at"),
        Migration::run(20140929103751,
                       "CREATE INDEX index_crate_downloads ON crates (downloads)",
                       "DROP INDEX index_crate_downloads"),
        Migration::run(20140929103752,
                       "CREATE INDEX index_version_crate_id ON versions (crate_id)",
                       "DROP INDEX index_version_crate_id"),
        Migration::run(20140929103753,
                       "CREATE INDEX index_version_num ON versions (num)",
                       "DROP INDEX index_version_num"),
        Migration::run(20140929103754,
                       "CREATE INDEX index_version_dependencies_version_id \
                        ON version_dependencies (version_id)",
                       "DROP INDEX index_version_dependencies_version_id"),
        Migration::run(20140929103755,
                       "CREATE INDEX index_version_dependencies_depends_on_id \
                        ON version_dependencies (depends_on_id)",
                       "DROP INDEX index_version_dependencies_depends_on_id"),
        Migration::add_table(20140929103756, "crate_downloads", "
            id              SERIAL PRIMARY KEY,
            crate_id        INTEGER NOT NULL,
            downloads       INTEGER NOT NULL,
            date            TIMESTAMP NOT NULL
        "),
        Migration::run(20140929103757,
                       "CREATE INDEX index_crate_downloads_crate_id \
                        ON crate_downloads (crate_id)",
                       "DROP INDEX index_crate_downloads_crate_id"),
        Migration::run(20140929103758,
                       "CREATE INDEX index_crate_downloads_date \
                        ON crate_downloads (date(date))",
                       "DROP INDEX index_crate_downloads_date"),
        Migration::add_table(20140929103759, "version_downloads", "
            id              SERIAL PRIMARY KEY,
            version_id      INTEGER NOT NULL,
            downloads       INTEGER NOT NULL,
            counted         INTEGER NOT NULL,
            date            TIMESTAMP NOT NULL,
            processed       BOOLEAN NOT NULL
        "),
        Migration::run(20140929103760,
                       "CREATE INDEX index_version_downloads_version_id \
                        ON version_downloads (version_id)",
                       "DROP INDEX index_version_downloads_version_id"),
        Migration::run(20140929103761,
                       "CREATE INDEX index_version_downloads_date \
                        ON version_downloads (date(date))",
                       "DROP INDEX index_version_downloads_date"),
        Migration::run(20140929103763,
                       "CREATE INDEX index_version_downloads_processed \
                        ON version_downloads (processed)
                        WHERE processed = FALSE",
                       "DROP INDEX index_version_downloads_processed"),
        Migration::run(20140929185718,
                       "CREATE INDEX index_crates_name_search \
                        ON crates USING gin(to_tsvector('english', name))",
                       "DROP INDEX index_crates_name_search"),
        Migration::run(20140930082104,
                       "DROP TABLE version_dependencies",
                       "CREATE TABLE version_dependencies (
                            version_id INTEGER
                        )"),
        Migration::add_table(20140930082105, "dependencies", "
            id               SERIAL PRIMARY KEY,
            version_id       INTEGER NOT NULL,
            crate_id         INTEGER NOT NULL,
            req              VARCHAR NOT NULL,
            optional         BOOLEAN NOT NULL,
            default_features BOOLEAN NOT NULL,
            features         VARCHAR NOT NULL
        "),
        Migration::add_column(20140930085441, "versions", "features",
                              "VARCHAR"),
        Migration::run(20140930203145,
                       "CREATE INDEX index_dependencies_version_id \
                        ON dependencies (version_id)",
                       "DROP INDEX index_dependencies_version_id"),
        Migration::run(20140930203146,
                       "CREATE INDEX index_dependencies_crate_id \
                        ON dependencies (crate_id)",
                       "DROP INDEX index_dependencies_crate_id"),
        Migration::add_column(20141001190227, "users", "gh_login",
                              "VARCHAR NOT NULL"),
        Migration::add_column(20141001190228, "users", "name", "VARCHAR"),
        Migration::run(20141001190229,
                       "CREATE INDEX index_users_gh_login \
                        ON users (gh_login)",
                       "DROP INDEX index_users_gh_login"),
        Migration::run(20141001190230,
                       "ALTER TABLE users ALTER COLUMN email DROP NOT NULL",
                       "ALTER TABLE users ALTER COLUMN email SET NOT NULL"),
        Migration::add_column(20141001190231, "users", "gh_avatar", "VARCHAR"),
        Migration::run(20141002195939,
                       "CREATE INDEX index_crates_user_id \
                        ON crates (user_id)",
                       "DROP INDEX index_crates_user_id"),
        Migration::add_table(20141002195940, "follows", "
            user_id          INTEGER NOT NULL,
            crate_id         INTEGER NOT NULL
        "),
        Migration::run(20141002195941,
                       "CREATE INDEX index_follows_user_id \
                        ON follows (user_id)",
                       "DROP INDEX index_follows_user_id"),
        foreign_key(20141002222426, "crate_downloads", "crate_id", "crates (id)"),
        foreign_key(20141002222427, "crates", "user_id", "users (id)"),
        foreign_key(20141002222428, "dependencies", "version_id", "versions (id)"),
        foreign_key(20141002222429, "dependencies", "crate_id", "crates (id)"),
        foreign_key(20141002222430, "follows", "crate_id", "crates (id)"),
        foreign_key(20141002222431, "version_downloads", "version_id",
                    "versions (id)"),
        foreign_key(20141002222432, "versions", "crate_id", "crates (id)"),
    ];
    // NOTE: Generate a new id via `date +"%Y%m%d%H%M%S"`

    let mut seen = HashSet::new();
    for m in migrations.iter() {
        if !seen.insert(m.version()) {
            fail!("duplicate id: {}", m.version());
        }
    }
    return migrations;

    fn foreign_key(id: i64, table: &str, column: &str,
                   references: &str) -> Migration {
        let add = format!("ALTER TABLE {table} ADD CONSTRAINT fk_{table}_{col}
                                 FOREIGN KEY ({col}) REFERENCES {reference}",
                          table = table, col = column, reference = references);
        let rm = format!("ALTER TABLE {table} DROP CONSTRAINT fk_{table}_{col}",
                          table = table, col = column);
        Migration::run(id, add.as_slice(), rm.as_slice())
    }
}
<commit_msg>Add one missing foreign key constraint<commit_after>extern crate "cargo-registry" as cargo_registry;
extern crate migrate;
extern crate postgres;
extern crate r2d2;

use std::os;
use std::collections::HashSet;
use migrate::Migration;
use postgres::{PostgresTransaction, PostgresResult};

use cargo_registry::krate::Crate;
use cargo_registry::model::Model;

fn main() {
    let db_config = r2d2::Config {
        pool_size: 1,
        helper_tasks: 1,
        test_on_check_out: false,
    };
    let database = cargo_registry::db::pool(env("DATABASE_URL").as_slice(),
                                            db_config);
    let conn = database.get().unwrap();
    let migrations = migrations();

    if os::args().as_slice().get(1).map(|s| s.as_slice()) == Some("rollback") {
        rollback(conn.transaction().unwrap(), migrations).unwrap();
    } else {
        apply(conn.transaction().unwrap(), migrations).unwrap();
    }

    fn env(s: &str) -> String {
        match os::getenv(s) {
            Some(s) => s,
            None => fail!("must have `{}` defined", s),
        }
    }
}

fn apply(tx: PostgresTransaction,
         migrations: Vec<Migration>) -> PostgresResult<()> {
    let mut mgr = try!(migrate::Manager::new(tx));
    for m in migrations.into_iter() {
        try!(mgr.apply(m));
    }
    mgr.set_commit();
    mgr.finish()
}

fn rollback(tx: PostgresTransaction,
            migrations: Vec<Migration>) -> PostgresResult<()> {
    let mut mgr = try!(migrate::Manager::new(tx));
    for m in migrations.into_iter().rev() {
        if mgr.contains(m.version()) {
            try!(mgr.rollback(m));
            break
        }
    }
    mgr.set_commit();
    mgr.finish()
}

fn migrations() -> Vec<Migration> {
    let migrations = vec![
        Migration::add_table(20140924113530, "users", "
            id              SERIAL PRIMARY KEY,
            email           VARCHAR NOT NULL UNIQUE,
            gh_access_token VARCHAR NOT NULL,
            api_token       VARCHAR NOT NULL
        "),
        Migration::add_table(20140924114003, "packages", "
            id              SERIAL PRIMARY KEY,
            name            VARCHAR NOT NULL UNIQUE,
            user_id         INTEGER NOT NULL
        "),
        Migration::add_table(20140924114059, "versions", "
            id              SERIAL PRIMARY KEY,
            package_id      INTEGER NOT NULL,
            num             VARCHAR NOT NULL
        "),
        Migration::run(20140924115329,
                       format!("ALTER TABLE versions ADD CONSTRAINT \
                                unique_num UNIQUE (package_id, num)"),
                       format!("ALTER TABLE versions DROP CONSTRAINT \
                                unique_num")),
        Migration::add_table(20140924120803, "version_dependencies", "
            version_id      INTEGER NOT NULL,
            depends_on_id   INTEGER NOT NULL
        "),
        Migration::add_column(20140925132248, "packages", "updated_at",
                              "TIMESTAMP NOT NULL DEFAULT now()"),
        Migration::add_column(20140925132249, "packages", "created_at",
                              "TIMESTAMP NOT NULL DEFAULT now()"),
        Migration::new(20140925132250, proc(tx) {
            try!(tx.execute("UPDATE packages SET updated_at = now() \
                             WHERE updated_at IS NULL", []));
            try!(tx.execute("UPDATE packages SET created_at = now() \
                             WHERE created_at IS NULL", []));
            Ok(())
        }, proc(_) Ok(())),
        Migration::add_column(20140925132251, "versions", "updated_at",
                              "TIMESTAMP NOT NULL DEFAULT now()"),
        Migration::add_column(20140925132252, "versions", "created_at",
                              "TIMESTAMP NOT NULL DEFAULT now()"),
        Migration::new(20140925132253, proc(tx) {
            try!(tx.execute("UPDATE versions SET updated_at = now() \
                             WHERE updated_at IS NULL", []));
            try!(tx.execute("UPDATE versions SET created_at = now() \
                             WHERE created_at IS NULL", []));
            Ok(())
        }, proc(_) Ok(())),
        Migration::new(20140925132254, proc(tx) {
            try!(tx.execute("ALTER TABLE versions ALTER COLUMN updated_at \
                             DROP DEFAULT", []));
            try!(tx.execute("ALTER TABLE versions ALTER COLUMN created_at \
                             DROP DEFAULT", []));
            try!(tx.execute("ALTER TABLE packages ALTER COLUMN updated_at \
                             DROP DEFAULT", []));
            try!(tx.execute("ALTER TABLE packages ALTER COLUMN created_at \
                             DROP DEFAULT", []));
            Ok(())
        }, proc(_) Ok(())),
        Migration::add_table(20140925153704, "metadata", "
            total_downloads        BIGINT NOT NULL
        "),
        Migration::new(20140925153705, proc(tx) {
            try!(tx.execute("INSERT INTO metadata (total_downloads) \
                             VALUES ($1)", &[&0i64]));
            Ok(())
        }, proc(tx) {
            try!(tx.execute("DELETE FROM metadata", [])); Ok(())
        }),
        Migration::add_column(20140925161623, "packages", "downloads",
                              "INTEGER NOT NULL DEFAULT 0"),
        Migration::add_column(20140925161624, "versions", "downloads",
                              "INTEGER NOT NULL DEFAULT 0"),
        Migration::new(20140925161625, proc(tx) {
            try!(tx.execute("ALTER TABLE versions ALTER COLUMN downloads \
                             DROP DEFAULT", []));
            try!(tx.execute("ALTER TABLE packages ALTER COLUMN downloads \
                             DROP DEFAULT", []));
            Ok(())
        }, proc(_) Ok(())),
        Migration::add_column(20140926130044, "packages", "max_version",
                              "VARCHAR"),
        Migration::new(20140926130045, proc(tx) {
            let stmt = try!(tx.prepare("SELECT * FROM packages"));
            for row in try!(stmt.query(&[])) {
                let pkg: Crate = Model::from_row(&row);
                let versions = pkg.versions(tx).unwrap();
                let v = versions.iter().max_by(|v| &v.num).unwrap();
                let max = v.num.to_string();
                try!(tx.execute("UPDATE packages SET max_version = $1 \
                                 WHERE id = $2",
                                &[&max, &pkg.id]));
            }
            Ok(())
        }, proc(_) Ok(())),
        Migration::new(20140926130046, proc(tx) {
            try!(tx.execute("ALTER TABLE versions ALTER COLUMN downloads \
                             SET NOT NULL", []));
            Ok(())
        }, proc(tx) {
            try!(tx.execute("ALTER TABLE versions ALTER COLUMN downloads \
                             DROP NOT NULL", []));
            Ok(())
        }),
        Migration::new(20140926174020, proc(tx) {
            try!(tx.execute("ALTER TABLE packages RENAME TO crates", []));
            try!(tx.execute("ALTER TABLE versions RENAME COLUMN package_id \
                             TO crate_id", []));
            Ok(())
        }, proc(tx) {
            try!(tx.execute("ALTER TABLE crates RENAME TO packages", []));
            try!(tx.execute("ALTER TABLE versions RENAME COLUMN crate_id \
                             TO package_id", []));
            Ok(())
        }),
        Migration::run(20140929103749,
                       "CREATE INDEX index_crate_updated_at ON crates (updated_at)",
                       "DROP INDEX index_crate_updated_at"),
        Migration::run(20140929103750,
                       "CREATE INDEX index_crate_created_at ON crates (created_at)",
                       "DROP INDEX index_crate_created_at"),
        Migration::run(20140929103751,
                       "CREATE INDEX index_crate_downloads ON crates (downloads)",
                       "DROP INDEX index_crate_downloads"),
        Migration::run(20140929103752,
                       "CREATE INDEX index_version_crate_id ON versions (crate_id)",
                       "DROP INDEX index_version_crate_id"),
        Migration::run(20140929103753,
                       "CREATE INDEX index_version_num ON versions (num)",
                       "DROP INDEX index_version_num"),
        Migration::run(20140929103754,
                       "CREATE INDEX index_version_dependencies_version_id \
                        ON version_dependencies (version_id)",
                       "DROP INDEX index_version_dependencies_version_id"),
        Migration::run(20140929103755,
                       "CREATE INDEX index_version_dependencies_depends_on_id \
                        ON version_dependencies (depends_on_id)",
                       "DROP INDEX index_version_dependencies_depends_on_id"),
        Migration::add_table(20140929103756, "crate_downloads", "
            id              SERIAL PRIMARY KEY,
            crate_id        INTEGER NOT NULL,
            downloads       INTEGER NOT NULL,
            date            TIMESTAMP NOT NULL
        "),
        Migration::run(20140929103757,
                       "CREATE INDEX index_crate_downloads_crate_id \
                        ON crate_downloads (crate_id)",
                       "DROP INDEX index_crate_downloads_crate_id"),
        Migration::run(20140929103758,
                       "CREATE INDEX index_crate_downloads_date \
                        ON crate_downloads (date(date))",
                       "DROP INDEX index_crate_downloads_date"),
        Migration::add_table(20140929103759, "version_downloads", "
            id              SERIAL PRIMARY KEY,
            version_id      INTEGER NOT NULL,
            downloads       INTEGER NOT NULL,
            counted         INTEGER NOT NULL,
            date            TIMESTAMP NOT NULL,
            processed       BOOLEAN NOT NULL
        "),
        Migration::run(20140929103760,
                       "CREATE INDEX index_version_downloads_version_id \
                        ON version_downloads (version_id)",
                       "DROP INDEX index_version_downloads_version_id"),
        Migration::run(20140929103761,
                       "CREATE INDEX index_version_downloads_date \
                        ON version_downloads (date(date))",
                       "DROP INDEX index_version_downloads_date"),
        Migration::run(20140929103763,
                       "CREATE INDEX index_version_downloads_processed \
                        ON version_downloads (processed)
                        WHERE processed = FALSE",
                       "DROP INDEX index_version_downloads_processed"),
        Migration::run(20140929185718,
                       "CREATE INDEX index_crates_name_search \
                        ON crates USING gin(to_tsvector('english', name))",
                       "DROP INDEX index_crates_name_search"),
        Migration::run(20140930082104,
                       "DROP TABLE version_dependencies",
                       "CREATE TABLE version_dependencies (
                            version_id INTEGER
                        )"),
        Migration::add_table(20140930082105, "dependencies", "
            id               SERIAL PRIMARY KEY,
            version_id       INTEGER NOT NULL,
            crate_id         INTEGER NOT NULL,
            req              VARCHAR NOT NULL,
            optional         BOOLEAN NOT NULL,
            default_features BOOLEAN NOT NULL,
            features         VARCHAR NOT NULL
        "),
        Migration::add_column(20140930085441, "versions", "features",
                              "VARCHAR"),
        Migration::run(20140930203145,
                       "CREATE INDEX index_dependencies_version_id \
                        ON dependencies (version_id)",
                       "DROP INDEX index_dependencies_version_id"),
        Migration::run(20140930203146,
                       "CREATE INDEX index_dependencies_crate_id \
                        ON dependencies (crate_id)",
                       "DROP INDEX index_dependencies_crate_id"),
        Migration::add_column(20141001190227, "users", "gh_login",
                              "VARCHAR NOT NULL"),
        Migration::add_column(20141001190228, "users", "name", "VARCHAR"),
        Migration::run(20141001190229,
                       "CREATE INDEX index_users_gh_login \
                        ON users (gh_login)",
                       "DROP INDEX index_users_gh_login"),
        Migration::run(20141001190230,
                       "ALTER TABLE users ALTER COLUMN email DROP NOT NULL",
                       "ALTER TABLE users ALTER COLUMN email SET NOT NULL"),
        Migration::add_column(20141001190231, "users", "gh_avatar", "VARCHAR"),
        Migration::run(20141002195939,
                       "CREATE INDEX index_crates_user_id \
                        ON crates (user_id)",
                       "DROP INDEX index_crates_user_id"),
        Migration::add_table(20141002195940, "follows", "
            user_id          INTEGER NOT NULL,
            crate_id         INTEGER NOT NULL
        "),
        Migration::run(20141002195941,
                       "CREATE INDEX index_follows_user_id \
                        ON follows (user_id)",
                       "DROP INDEX index_follows_user_id"),
        foreign_key(20141002222426, "crate_downloads", "crate_id", "crates (id)"),
        foreign_key(20141002222427, "crates", "user_id", "users (id)"),
        foreign_key(20141002222428, "dependencies", "version_id", "versions (id)"),
        foreign_key(20141002222429, "dependencies", "crate_id", "crates (id)"),
        foreign_key(20141002222430, "follows", "crate_id", "crates (id)"),
        foreign_key(20141002222431, "version_downloads", "version_id",
                    "versions (id)"),
        foreign_key(20141002222432, "versions", "crate_id", "crates (id)"),
        foreign_key(20141002222433, "follows", "user_id", "users (id)"),
    ];
    // NOTE: Generate a new id via `date +"%Y%m%d%H%M%S"`

    let mut seen = HashSet::new();
    for m in migrations.iter() {
        if !seen.insert(m.version()) {
            fail!("duplicate id: {}", m.version());
        }
    }
    return migrations;

    fn foreign_key(id: i64, table: &str, column: &str,
                   references: &str) -> Migration {
        let add = format!("ALTER TABLE {table} ADD CONSTRAINT fk_{table}_{col}
                                 FOREIGN KEY ({col}) REFERENCES {reference}",
                          table = table, col = column, reference = references);
        let rm = format!("ALTER TABLE {table} DROP CONSTRAINT fk_{table}_{col}",
                          table = table, col = column);
        Migration::run(id, add.as_slice(), rm.as_slice())
    }
}
<|endoftext|>


<commit_before>//! Checks for any invariants we expect to be true, and pages whoever is on call
//! if they are not.
//!
//! Usage:
//!     cargo run --bin monitor

#![warn(clippy::all, rust_2018_idioms)]

mod on_call;

use cargo_registry::{db, schema::*, util::Error};
use diesel::prelude::*;

fn main() -> Result<(), Error> {
    let conn = db::connect_now()?;

    check_stalled_background_jobs(&conn)?;
    check_spam_attack(&conn)?;
    Ok(())
}

fn check_stalled_background_jobs(conn: &PgConnection) -> Result<(), Error> {
    use cargo_registry::schema::background_jobs::dsl::*;
    use diesel::dsl::*;

    const EVENT_KEY: &str = "background_jobs";

    println!("Checking for stalled background jobs");

    let max_job_time = dotenv::var("MAX_JOB_TIME")
        .map(|s| s.parse::<i32>().unwrap())
        .unwrap_or(15);

    let stalled_job_count = background_jobs
        .filter(created_at.lt(now - max_job_time.minutes()))
        .count()
        .get_result::<i64>(conn)?;

    let event = if stalled_job_count > 0 {
        on_call::Event::Trigger {
            incident_key: Some(EVENT_KEY.into()),
            description: format!(
                "{} jobs have been in the queue for more than {} minutes",
                stalled_job_count, max_job_time
            ),
        }
    } else {
        on_call::Event::Resolve {
            incident_key: EVENT_KEY.into(),
            description: Some("No stalled background jobs".into()),
        }
    };

    log_and_trigger_event(event)?;
    Ok(())
}

fn check_spam_attack(conn: &PgConnection) -> Result<(), Error> {
    use cargo_registry::models::krate::canon_crate_name;
    use diesel::dsl::*;
    use diesel::sql_types::Bool;

    const EVENT_KEY: &str = "spam_attack";

    println!("Checking for crates indicating someone is spamming us");

    let bad_crate_names = dotenv::var("SPAM_CRATE_NAMES");
    let bad_crate_names: Vec<_> = bad_crate_names
        .as_ref()
        .map(|s| s.split(',').collect())
        .unwrap_or_default();
    let bad_author_patterns = dotenv::var("SPAM_AUTHOR_PATTERNS");
    let bad_author_patterns: Vec<_> = bad_author_patterns
        .as_ref()
        .map(|s| s.split(',').collect())
        .unwrap_or_default();

    let mut event_description = None;

    let bad_crate = crates::table
        .filter(canon_crate_name(crates::name).eq(any(bad_crate_names)))
        .select(crates::name)
        .first::<String>(conn)
        .optional()?;

    if let Some(bad_crate) = bad_crate {
        event_description = Some(format!("Crate named {} published", bad_crate));
    }

    let mut query = version_authors::table
        .select(version_authors::name)
        .filter(false.into_sql::<Bool>()) // Never return anything if we have no patterns
        .into_boxed();
    for author_pattern in bad_author_patterns {
        query = query.or_filter(version_authors::name.like(author_pattern));
    }
    let bad_author = query.first::<String>(conn).optional()?;

    if let Some(bad_author) = bad_author {
        event_description = Some(format!("Crate with author {} published", bad_author));
    }

    let event = if let Some(event_description) = event_description {
        on_call::Event::Trigger {
            incident_key: Some(EVENT_KEY.into()),
            description: format!("{}, possible spam attack underway", event_description,),
        }
    } else {
        on_call::Event::Resolve {
            incident_key: EVENT_KEY.into(),
            description: Some("No spam crates detected".into()),
        }
    };

    log_and_trigger_event(event)?;
    Ok(())
}

fn log_and_trigger_event(event: on_call::Event) -> Result<(), Error> {
    match event {
        on_call::Event::Trigger {
            ref description, ..
        } => println!("Paging on-call: {}", description),
        on_call::Event::Resolve {
            description: Some(ref description),
            ..
        } => println!("{}", description),
        _ => {} // noop
    }
    event.send()
}
<commit_msg>Auto merge of #2434 - sgrif:sg-dont-page-for-running-jobs, r=jtgeibel<commit_after>//! Checks for any invariants we expect to be true, and pages whoever is on call
//! if they are not.
//!
//! Usage:
//!     cargo run --bin monitor

#![warn(clippy::all, rust_2018_idioms)]

mod on_call;

use cargo_registry::{db, schema::*, util::Error};
use diesel::prelude::*;

fn main() -> Result<(), Error> {
    let conn = db::connect_now()?;

    check_stalled_background_jobs(&conn)?;
    check_spam_attack(&conn)?;
    Ok(())
}

fn check_stalled_background_jobs(conn: &PgConnection) -> Result<(), Error> {
    use cargo_registry::schema::background_jobs::dsl::*;
    use diesel::dsl::*;
    use diesel::sql_types::Integer;

    const EVENT_KEY: &str = "background_jobs";

    println!("Checking for stalled background jobs");

    let max_job_time = dotenv::var("MAX_JOB_TIME")
        .map(|s| s.parse::<i32>().unwrap())
        .unwrap_or(15);

    let stalled_job_count = background_jobs
        .select(1.into_sql::<Integer>())
        .filter(created_at.lt(now - max_job_time.minutes()))
        .for_update()
        .skip_locked()
        .load::<i32>(conn)?
        .len();

    let event = if stalled_job_count > 0 {
        on_call::Event::Trigger {
            incident_key: Some(EVENT_KEY.into()),
            description: format!(
                "{} jobs have been in the queue for more than {} minutes",
                stalled_job_count, max_job_time
            ),
        }
    } else {
        on_call::Event::Resolve {
            incident_key: EVENT_KEY.into(),
            description: Some("No stalled background jobs".into()),
        }
    };

    log_and_trigger_event(event)?;
    Ok(())
}

fn check_spam_attack(conn: &PgConnection) -> Result<(), Error> {
    use cargo_registry::models::krate::canon_crate_name;
    use diesel::dsl::*;
    use diesel::sql_types::Bool;

    const EVENT_KEY: &str = "spam_attack";

    println!("Checking for crates indicating someone is spamming us");

    let bad_crate_names = dotenv::var("SPAM_CRATE_NAMES");
    let bad_crate_names: Vec<_> = bad_crate_names
        .as_ref()
        .map(|s| s.split(',').collect())
        .unwrap_or_default();
    let bad_author_patterns = dotenv::var("SPAM_AUTHOR_PATTERNS");
    let bad_author_patterns: Vec<_> = bad_author_patterns
        .as_ref()
        .map(|s| s.split(',').collect())
        .unwrap_or_default();

    let mut event_description = None;

    let bad_crate = crates::table
        .filter(canon_crate_name(crates::name).eq(any(bad_crate_names)))
        .select(crates::name)
        .first::<String>(conn)
        .optional()?;

    if let Some(bad_crate) = bad_crate {
        event_description = Some(format!("Crate named {} published", bad_crate));
    }

    let mut query = version_authors::table
        .select(version_authors::name)
        .filter(false.into_sql::<Bool>()) // Never return anything if we have no patterns
        .into_boxed();
    for author_pattern in bad_author_patterns {
        query = query.or_filter(version_authors::name.like(author_pattern));
    }
    let bad_author = query.first::<String>(conn).optional()?;

    if let Some(bad_author) = bad_author {
        event_description = Some(format!("Crate with author {} published", bad_author));
    }

    let event = if let Some(event_description) = event_description {
        on_call::Event::Trigger {
            incident_key: Some(EVENT_KEY.into()),
            description: format!("{}, possible spam attack underway", event_description,),
        }
    } else {
        on_call::Event::Resolve {
            incident_key: EVENT_KEY.into(),
            description: Some("No spam crates detected".into()),
        }
    };

    log_and_trigger_event(event)?;
    Ok(())
}

fn log_and_trigger_event(event: on_call::Event) -> Result<(), Error> {
    match event {
        on_call::Event::Trigger {
            ref description, ..
        } => println!("Paging on-call: {}", description),
        on_call::Event::Resolve {
            description: Some(ref description),
            ..
        } => println!("{}", description),
        _ => {} // noop
    }
    event.send()
}
<|endoftext|>


<commit_before>extern crate pirate;
extern crate tomllib;
extern crate csv;
use std::fs::File;
use std::env;
use std::io::{Read, BufReader, Error};
use pirate::{Matches, Match, Vars, matches, usage, vars};
use tomllib::TOMLParser;
use tomllib::types::{ParseResult, Children};
use csv::Reader;

macro_rules! usage(
  ($tval:expr) => (
    usage($tval);
    return;
  );
  ($submac:ident!( $($args:tt)* ), $tval:expr) => (
    $submac!($($args)*);
    usage($tval);
    return;
  );
);

fn main() {
  let options = vec![
    "Required arguments",
    "f/file#The path to the TOML document to parse and manipulate.:",
    "#File-preserving commands",
    "g/get-value#Given a key or comma separated list of keys, retrieve the values for those keys. If multiple keys are \
      specified and any one of them fails to retrieve a value, the whole command will fail with an error message.:",
    "/has-value#Given a key or comma separated list of keys, print \"true\" if the key has a value and \"false\" if \
      the key doesn't have child keys. Optionally use the --set-true and --set-false flags to change what values are \
      printed instead of \"true\" and \"false\".:",
    "c/get-children#Given a key or comma separated list of keys, retrieve each key's set of child keys. If multiple \
      keys are specified and any one of them fails to retrieve a child keys, the whole command will fail with an error \
      message.:",
    "/has-children#Given a key or comma separated list of keys, print \"true\" if the key has child keys and \"false\" \
      if the key doesn't have child keys. Optionally use the --set-true and --set-false flags to change what values \
      are printed instead of \"true\" and \"false\".:",
    "#File-modifying commands",
    "s/set-value#Given a comma separated list of key followed by value followed by type, set the key's value to the \
      specified value. For instance: \"foo.bar,hello,basic-string,baz.qux,82374,int\" will set the key \"foo.bar\" to \
      a basic string value of \"hello\" and the key \"baz.qux\" to an integer value of 82374. If multiple keys are \
      specified and any one of them fails to set a value, the whole command will fail with an error message.:",
    "#Pre-command Options.",
    "h/help#Show this screen.",
    "/set-true#For commands that print \"true\" or \"false\", this will change what value is printed for \"true\", \
      i.e. --set-true=1 will cause tomlkit to print \"1\" for \"true\" instead of \"true\".:",
    "/set-false#For commands that print \"true\" or \"false\", this will change what value is printed for \"false\", \
      i.e. --set-false=0 will cause tomlkit to print \"0\" for \"false\" instead of \"false\".:",
    "#Post-command Options.",
    "/print-doc#Print out the resultant TOML document after all requested changes have been made.",
  ];

  let mut vars: Vars = match vars("tomlkit", &options) {
    Ok(v) => v,
    Err(e) => panic!("There was an error parsing argument definitions: {}", e)
  };

  let args: Vec<String> = env::args().collect();
  let matches: Matches = match matches(&args, &mut vars) {
    Ok(m) => m,
    Err(e) => {
      usage!(println!("Error: {}", e), &vars);
    }
  };

  if matches.has_match("help") {
    usage!(&vars);
  }
  let mut true_val = &"true".to_string();
  let mut false_val = &"false".to_string();
  let mut output = String::new();

  // The file we're operating on
  let mut file = String::new();
  if matches.has_match("file") {
    if let Some(f) = matches.get("file") {
      match get_file(f, &mut file) {
        Err(err) => {
          println!("Error: Unable to read file: \"{}\". Reason: {}", f, err);
          return;
        },
        _ => ()
      }
    } else {
      usage!(println!("Error: A required argument is missing for f/file."), &vars);
    }
  } else {
    println!("Error: -f/--file is a required argument.");
    return;
  }

  // Parse the document
  let parser: TOMLParser = TOMLParser::new();
  let (parser, result) = parser.parse(&file);
  match result {
    ParseResult::Partial(_,_,_) => {
      println!("Error: Document only partially parsed. Please correct any errors before trying again.");
      return;
    },
    ParseResult::PartialError(_,_,_,_) => {
      println!("Error: Document only partially parsed with errors. Please correct any errors before trying again.");
      return;
    },
    ParseResult::Failure(_,_) => {
      println!("Error: Completely failed to parse document. Please correct any error before trying again.");
      return;
    },
    _ => (), // If verbose output Full or FullError
  }

  // Pre-command options
  if matches.has_match("set-true") {
    if let Some(t) = matches.get("set-true") {
      true_val = t;
    } else {
      usage!(println!("Error: A required argument is missing for set-true."), &vars);
    }
  }
  if matches.has_match("set-false") {
    if let Some(f) = matches.get("set-false") {
      false_val = f;
    } else {
      usage!(println!("Error: A required argument is missing for set-false."), &vars);
    }
  }

  // Commands only one command allowed per invocation for this version
  if matches.has_match("get-value") {
    if let Some(k) = matches.get("get-value") {
      if let Some(val) = get_value(k, &parser) {
        output = val;
      }
    } else {
      usage!(println!("Error: A required argument is missing for g/get-value."), &vars);
    }
  } else if matches.has_match("has-value") {
    if let Some(k) = matches.get("has-value") {
      output = has_value(k, &parser, true_val, false_val);
    } else {
      usage!(println!("Error: A required argument is missing for has-value."), &vars);
    }
  } else if matches.has_match("get-children") {
    if let Some(k) = matches.get("get-children") {
      if let Some(c) = get_children(k, &parser) {
        output = c;
      }
    } else {
      usage!(println!("Error: A required argument is missing for c/get-children."), &vars);
    }
  } else if matches.has_match("has-children") {
    if let Some(k) = matches.get("has-children") {
      output = has_children(k, &parser, true_val, false_val);
    } else {
      usage!(println!("Error: A required argument is missing for has-children."), &vars);
    }
  } else if matches.has_match("set-value") {
    if let Some(kv) = matches.get("set-value") {
      if let Some(out) = set_value(kv, &parser) {
        output = out;
      }
    } else {
      usage!(println!("Error: A required argument is missing for s/set-value."), &vars);
    }
  } else {
    // No command specified print usage
    usage!(&vars);
  }

  // ************** Print output here! *******************
  println!("{}", output);

  // Post-command options
  if matches.has_match("print-doc") {
    println!("\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>DOCUMENT<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<\n");
    print_doc(&parser);
  }
}

fn get_file(file_path: &String, out_file: &mut String) -> Result<usize, Error> {
  match File::open(file_path) {
    Ok(file) => {
      let mut contents = BufReader::new(&file);
      match contents.read_to_string(out_file) {
        Ok(bytes) => return Ok(bytes),
        Err(err) => return Err(err),
      }
    },
    Err(err) => return Err(err),
  }
}

fn get_value(key: &str, doc: &TOMLParser) -> Option<String> {
  if let Some(value) = doc.get_value(key) {
    return Some(format!("{}", value));
  }
  None
}

fn has_value(key: &str, doc: &TOMLParser, true_val: &String, false_val: &String) -> String {
  if let Some(_) = doc.get_value(key) {
    return true_val.clone();
  }
  return false_val.clone();
}

fn get_children(key: &str, doc: &TOMLParser) -> Option<String> {
  if let Some(c) = doc.get_children(key) {
    match c {
      &Children::Keys(ref keys) => {
        let mut retval = String::new();
        retval.push('[');
        if keys.borrow().len() > 0 {
          for i in 0..keys.borrow().len() - 1 {
            retval.push_str(&keys.borrow()[i]);
            retval.push_str(", ");
          }
          retval.push_str(&keys.borrow()[keys.borrow().len() - 1]);
        }
        retval.push(']');
        return Some(retval);
      },
      &Children::Count(ref size) => {
        if size.get() == 0 {
          return None;
        }
        return Some(format!("{}", size.get()))
      },
    }
  }
  None
}

fn has_children(key: &str, doc: &TOMLParser, true_val: &String, false_val: &String) -> String {
  if let Some(c) = doc.get_children(key) {
    if let &Children::Count(ref count) = c {
      if count.get() == 0 {
        return false_val.clone();
      }
    }
    return true_val.clone();
  }
  return false_val.clone();
}

fn set_value(keyvals: &str, doc: &TOMLParser) -> Option<String> {
  println!("Called set_value with these keyvals: \"{:?}\"", csv_to_vec(keyvals));
  return None;
}

fn csv_to_vec<'a>(csv: &str) -> Result<Vec<String>, csv::Error> {
  let mut fields = vec![];
  let mut rdr = Reader::from_string(csv).has_headers(false).escape(Some(b'\\'));
  while !rdr.done() {
    while let Some(result) = rdr.next_str().into_iter_result() {
      match result {
        Ok(field) => fields.push(field.to_string()),
        Err(err)  => return Err(err),
      }
    }
  }
  Ok(fields)
}

fn print_doc(doc: &TOMLParser) {
  //unimplemented!();
  println!("{}", doc);
}
<commit_msg>Got set-value working, not for entire arrays or inline tables though and because of a bad design decision the datetime parsing is a mess. That will need to be fixed in 0.2.0.<commit_after>extern crate pirate;
extern crate tomllib;
extern crate csv;
use std::fs::File;
use std::env;
use std::io::{Read, BufReader, Error};
use pirate::{Matches, Match, Vars, matches, usage, vars};
use tomllib::TOMLParser;
use tomllib::types::{ParseResult, Children, Value, TOMLError, TimeOffset, DateTime, Date, Time, TimeOffsetAmount,
                     PosNeg};
use csv::Reader;

macro_rules! usage(
  ($tval:expr) => (
    usage($tval);
    return;
  );
  ($submac:ident!( $($args:tt)* ), $tval:expr) => (
    $submac!($($args)*);
    usage($tval);
    return;
  );
);

fn main() {
  let options = vec![
    "Required arguments",
    "f/file#The path to the TOML document to parse and manipulate.:",
    "#File-preserving commands",
    "g/get-value#Given a key or comma separated list of keys, retrieve the values for those keys. If multiple keys are \
      specified and any one of them fails to retrieve a value, the whole command will fail with an error message.:",
    "/has-value#Given a key or comma separated list of keys, print \"true\" if the key has a value and \"false\" if \
      the key doesn't have child keys. Optionally use the --set-true and --set-false flags to change what values are \
      printed instead of \"true\" and \"false\".:",
    "c/get-children#Given a key or comma separated list of keys, retrieve each key's set of child keys. If multiple \
      keys are specified and any one of them fails to retrieve a child keys, the whole command will fail with an error \
      message.:",
    "/has-children#Given a key or comma separated list of keys, print \"true\" if the key has child keys and \"false\" \
      if the key doesn't have child keys. Optionally use the --set-true and --set-false flags to change what values \
      are printed instead of \"true\" and \"false\".:",
    "#File-modifying commands",
    "s/set-value#Given a comma separated list of key followed by value followed by type, set the key's value to the \
      specified value. For instance: \"foo.bar,hello,basic-string,baz.qux,82374,int\" will set the key \"foo.bar\" to \
      a basic string value of \"hello\" and the key \"baz.qux\" to an integer value of 82374. If multiple keys are \
      specified and any one of them fails to set a value, the whole command will fail with an error message.:",
    "#Pre-command Options.",
    "h/help#Show this screen.",
    "/set-true#For commands that print \"true\" or \"false\", this will change what value is printed for \"true\", \
      i.e. --set-true=1 will cause tomlkit to print \"1\" for \"true\" instead of \"true\".:",
    "/set-false#For commands that print \"true\" or \"false\", this will change what value is printed for \"false\", \
      i.e. --set-false=0 will cause tomlkit to print \"0\" for \"false\" instead of \"false\".:",
    "#Post-command Options.",
    "/print-doc#Print out the resultant TOML document after all requested changes have been made.",
  ];

  let mut vars: Vars = match vars("tomlkit", &options) {
    Ok(v) => v,
    Err(e) => panic!("There was an error parsing argument definitions: {}", e)
  };

  let args: Vec<String> = env::args().collect();
  let matches: Matches = match matches(&args, &mut vars) {
    Ok(m) => m,
    Err(e) => {
      usage!(println!("Error: {}", e), &vars);
    }
  };

  if matches.has_match("help") {
    usage!(&vars);
  }
  let mut true_val = &"true".to_string();
  let mut false_val = &"false".to_string();
  let mut output = String::new();

  // The file we're operating on
  let mut file = String::new();
  if matches.has_match("file") {
    if let Some(f) = matches.get("file") {
      match get_file(f, &mut file) {
        Err(err) => {
          println!("Error: Unable to read file: \"{}\". Reason: {}", f, err);
          return;
        },
        _ => ()
      }
    } else {
      usage!(println!("Error: A required argument is missing for f/file."), &vars);
    }
  } else {
    println!("Error: -f/--file is a required argument.");
    return;
  }

  // Parse the document
  let mut parser: TOMLParser = TOMLParser::new();
  let (mut parser, result) = parser.parse(&file);
  match result {
    ParseResult::Partial(_,_,_) => {
      println!("Error: Document only partially parsed. Please correct any errors before trying again.");
      return;
    },
    ParseResult::PartialError(_,_,_,_) => {
      println!("Error: Document only partially parsed with errors. Please correct any errors before trying again.");
      return;
    },
    ParseResult::Failure(_,_) => {
      println!("Error: Completely failed to parse document. Please correct any error before trying again.");
      return;
    },
    _ => (), // If verbose output Full or FullError
  }

  // Pre-command options
  if matches.has_match("set-true") {
    if let Some(t) = matches.get("set-true") {
      true_val = t;
    } else {
      usage!(println!("Error: A required argument is missing for set-true."), &vars);
    }
  }
  if matches.has_match("set-false") {
    if let Some(f) = matches.get("set-false") {
      false_val = f;
    } else {
      usage!(println!("Error: A required argument is missing for set-false."), &vars);
    }
  }

  // Commands only one command allowed per invocation for this version
  if matches.has_match("get-value") {
    if let Some(k) = matches.get("get-value") {
      if let Some(val) = get_value(k, &parser) {
        output = val;
      }
    } else {
      usage!(println!("Error: A required argument is missing for g/get-value."), &vars);
    }
  } else if matches.has_match("has-value") {
    if let Some(k) = matches.get("has-value") {
      output = has_value(k, &parser, true_val, false_val);
    } else {
      usage!(println!("Error: A required argument is missing for has-value."), &vars);
    }
  } else if matches.has_match("get-children") {
    if let Some(k) = matches.get("get-children") {
      if let Some(c) = get_children(k, &parser) {
        output = c;
      }
    } else {
      usage!(println!("Error: A required argument is missing for c/get-children."), &vars);
    }
  } else if matches.has_match("has-children") {
    if let Some(k) = matches.get("has-children") {
      output = has_children(k, &parser, true_val, false_val);
    } else {
      usage!(println!("Error: A required argument is missing for has-children."), &vars);
    }
  } else if matches.has_match("set-value") {
    if let Some(kv) = matches.get("set-value") {
      if let Some(out) = set_value(kv, &mut parser) {
        output = out;
      }
    } else {
      usage!(println!("Error: A required argument is missing for s/set-value."), &vars);
    }
  } else {
    // No command specified print usage
    usage!(&vars);
  }

  // ************** Print output here! *******************
  println!("{}", output);

  // Post-command options
  if matches.has_match("print-doc") {
    println!("\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>DOCUMENT<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<\n");
    print_doc(&parser);
  }
}

fn get_file(file_path: &String, out_file: &mut String) -> Result<usize, Error> {
  match File::open(file_path) {
    Ok(file) => {
      let mut contents = BufReader::new(&file);
      match contents.read_to_string(out_file) {
        Ok(bytes) => return Ok(bytes),
        Err(err) => return Err(err),
      }
    },
    Err(err) => return Err(err),
  }
}

fn get_value(key: &str, doc: &TOMLParser) -> Option<String> {
  if let Some(value) = doc.get_value(key) {
    return Some(format!("{}", value));
  }
  None
}

fn has_value(key: &str, doc: &TOMLParser, true_val: &String, false_val: &String) -> String {
  if let Some(_) = doc.get_value(key) {
    return true_val.clone();
  }
  return false_val.clone();
}

fn get_children(key: &str, doc: &TOMLParser) -> Option<String> {
  if let Some(c) = doc.get_children(key) {
    match c {
      &Children::Keys(ref keys) => {
        let mut retval = String::new();
        retval.push('[');
        if keys.borrow().len() > 0 {
          for i in 0..keys.borrow().len() - 1 {
            retval.push_str(&keys.borrow()[i]);
            retval.push_str(", ");
          }
          retval.push_str(&keys.borrow()[keys.borrow().len() - 1]);
        }
        retval.push(']');
        return Some(retval);
      },
      &Children::Count(ref size) => {
        if size.get() == 0 {
          return None;
        }
        return Some(format!("{}", size.get()))
      },
    }
  }
  None
}

fn has_children(key: &str, doc: &TOMLParser, true_val: &String, false_val: &String) -> String {
  if let Some(c) = doc.get_children(key) {
    if let &Children::Count(ref count) = c {
      if count.get() == 0 {
        return false_val.clone();
      }
    }
    return true_val.clone();
  }
  return false_val.clone();
}

fn set_value<'a>(kvs: &str, doc: &mut TOMLParser) -> Option<String> {
  let keyval_results = csv_to_vec(kvs);
  if let Ok(keyvals) = keyval_results {
    if keyvals.len() % 3 != 0 || keyvals.len() == 0 {
      return None;
    }
    for i in 0..keyvals.len() / 3 {
      let key: &str = &keyvals[i];
      let val: &str = &keyvals[i+1];
      let typ: &str = &keyvals[i+2];
      let val_result: Result<Value, TOMLError>;
      match typ {
        "basic-string" | "bs" => val_result = Value::basic_string(val),
        "ml-basic-string" | "mbs" => val_result = Value::ml_basic_string(val),
        "literal-string" | "ls" => val_result = Value::literal_string(val),
        "ml-literal-string" | "mls" => val_result = Value::ml_literal_string(val),
        "integer" | "int" => val_result = Value::int_from_str(val),
        "float" | "flt" => val_result = Value::float_from_str(val),
        "boolean" | "bool" => val_result = Value::bool_from_str(val),
        "datetime" | "dt" => {
          let str_val: &str = &val;
          let tmp_result = Value::datetime_parse(str_val);
          let mut new_dt: DateTime = DateTime{date: Date{year: "".into(), month: "".into(), day: "".into()}, time: None};
          let (mut year, mut month, mut day, mut hour, mut minute, mut second, mut fraction) = ("".into(), "".into(),
            "".into(), "".into(), "".into(), "".into(), "".into());
          let (mut off_hour, mut off_minute, mut pos_neg) = ("".into(), "".into(), PosNeg::Pos);
          let (mut has_time, mut has_fraction, mut has_offset) = (false, false, false);
          if let Ok(dtval) = tmp_result {
            if let Value::DateTime(dt) = dtval {
              year = dt.date.year.to_string().into();
              month = dt.date.month.to_string().into();
              day = dt.date.day.to_string().into();
              if let Some(ref time) = dt.time {
                has_time = true;
                hour = time.hour.to_string().into();
                minute = time.minute.to_string().into();
                second = time.second.to_string().into();
                if let Some(ref frac) = time.fraction {
                  has_fraction = true;
                  fraction = frac.to_string().into();
                }
                if let Some(ref offset) = time.offset {
                  if let &TimeOffset::Time(ref amount) = offset {
                    has_offset = true;
                    pos_neg = amount.pos_neg;
                    off_hour = amount.hour.to_string().into();
                    off_minute = amount.minute.to_string().into();
                  }
                }
              }

              let newoffset = if has_offset {
                Some(TimeOffset::Time(TimeOffsetAmount{pos_neg: pos_neg, hour: off_hour, minute: off_minute}))
              } else {
                None
              };
              let newfraction = if has_fraction {
                Some(fraction)
              } else {
                None
              };
              let newtime = if has_time {
                Some(Time{hour: hour, minute: minute, second: second, fraction: newfraction, offset: newoffset})
              } else {
                None
              };
              new_dt = DateTime{
                date: Date{
                  year: year,
                  month: month,
                  day: day,
                },
                time: newtime,
              };
            }
            val_result = Ok(Value::DateTime(new_dt));
          } else {
            return None;
          }
        },
        _ => return None,
      }
      if let Ok(value) = val_result {
        if doc.set_value(key, value) {
          return Some("Success".to_string());
        }
      }
    }
  }
  return None;
}

fn csv_to_vec<'a>(csv: &str) -> Result<Vec<String>, csv::Error> {
  let mut fields = vec![];
  let mut rdr = Reader::from_string(csv).has_headers(false).escape(Some(b'\\')).quote(b'\0');
  while !rdr.done() {
    while let Some(result) = rdr.next_str().into_iter_result() {
      match result {
        Ok(field) => fields.push(field.to_string()),
        Err(err)  => return Err(err),
      }
    }
  }
  Ok(fields)
}

fn print_doc(doc: &TOMLParser) {
  //unimplemented!();
  println!("{}", doc);
}
<|endoftext|>