Repository: edgedb/edgedb-cli
Pull request #17 (2020-03-25 16:56:06)

Dump restore implementation
Dump is here. Restore is in To Do.


Patch:
diff --git a/Cargo.lock b/Cargo.lock
index e0c5ec2..f49e845 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -38,15 +38,16 @@ checksum = "cff77d8686867eceff3105329d4698d96c2391c176d5d03adc90c7389162b5b8"
 
 [[package]]
 name = "assert_cmd"
-version = "0.12.1"
+version = "0.12.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "35ad62275a8bda1c2c9a9303aea121eb04204272d3be0735d5dc1f49eb9ff9a9"
+checksum = "936fcf2c692b37c696cd0002c57752b2d9478402450c9ca4a463f6afae16d6f5"
 dependencies = [
  "doc-comment",
  "escargot",
  "predicates",
  "predicates-core",
  "predicates-tree",
+ "wait-timeout",
 ]
 
 [[package]]
@@ -432,6 +433,8 @@ dependencies = [
  "scram",
  "serde",
  "serde_json",
+ "sha1",
+ "shutdown_hooks",
  "snafu",
  "structopt",
  "tempfile",
@@ -446,7 +449,7 @@ dependencies = [
 [[package]]
 name = "edgedb-derive"
 version = "0.1.0"
-source = "git+https://github.com/edgedb/edgedb-rust#34e9d82ba6ebbb32e4a5959aa21221d6ce30f358"
+source = "git+https://github.com/edgedb/edgedb-rust#b7f644ea8fe6333b36af38cd66275102e0d35820"
 dependencies = [
  "quote",
  "syn",
@@ -456,7 +459,7 @@ dependencies = [
 [[package]]
 name = "edgedb-protocol"
 version = "0.1.0"
-source = "git+https://github.com/edgedb/edgedb-rust#34e9d82ba6ebbb32e4a5959aa21221d6ce30f358"
+source = "git+https://github.com/edgedb/edgedb-rust#b7f644ea8fe6333b36af38cd66275102e0d35820"
 dependencies = [
  "bigdecimal",
  "bytes",
@@ -470,7 +473,7 @@ dependencies = [
 [[package]]
 name = "edgeql-parser"
 version = "0.1.0"
-source = "git+https://github.com/edgedb/edgedb#2b8cef056dd82d97556dcfe835eabac6e8ae9d5f"
+source = "git+https://github.com/edgedb/edgedb#722a8e804d3270bc6c3721b1d19920d76f56cdf9"
 dependencies = [
  "combine",
  "snafu",
@@ -1238,6 +1241,12 @@ version = "0.6.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2579985fda508104f7587689507983eadd6a6e84dd35d6d115361f530916fa0d"
 
+[[package]]
+name = "shutdown_hooks"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6057adedbec913419c92996f395ba69931acbd50b7d56955394cd3f7bedbfa45"
+
 [[package]]
 name = "slab"
 version = "0.4.2"
@@ -1572,6 +1581,15 @@ version = "1.0.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6a02e4885ed3bc0f2de90ea6dd45ebcbb66dacffe03547fadbb0eeae2770887d"
 
+[[package]]
+name = "wait-timeout"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9f200f5b12eb75f8c1ed65abd4b2db8a6e1b138a20de009dacee265a2498f3f6"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "wasi"
 version = "0.9.0+wasi-snapshot-preview1"
diff --git a/Cargo.toml b/Cargo.toml
index a867810..2813d6a 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -42,10 +42,12 @@ codespan = "0.8.0"
 codespan-reporting = "0.8.0"
 termcolor = "1.1.0"
 async-listen = "0.2.0"
+sha1 = "0.6.0"
 
 [dev-dependencies]
 assert_cmd = "0.12"
 predicates = "1"
+shutdown_hooks = "0.1.0"
 
 [[bin]]
 name = "edgedb"
diff --git a/src/client.rs b/src/client.rs
index 116edc1..ad90bc1 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -4,11 +4,12 @@ use std::io;
 use std::mem::replace;
 use std::str;
 use std::sync::Arc;
-use std::time::Instant;
+use std::time::{Instant, Duration};
 
 use anyhow::{self, Context};
 use async_std::prelude::StreamExt;
 use async_std::io::{stdin, stdout};
+use async_std::future::timeout;
 use async_std::io::prelude::WriteExt;
 use async_std::net::{TcpStream, ToSocketAddrs};
 use async_listen::ByteStream;
@@ -31,13 +32,15 @@ use crate::commands::backslash;
 use crate::options::Options;
 use crate::print::{self, print_to_stdout, PrintError};
 use crate::prompt;
-use crate::reader::{Reader, ReadError, QueryableDecoder, QueryResponse};
+use crate::reader::{ReadError, QueryableDecoder, QueryResponse};
 use crate::repl;
 use crate::server_params::PostgresAddress;
 use crate::statement::{ReadStatement, EndOfFile};
 use crate::variables::input_variables;
 use crate::error_display::print_query_error;
 
+pub use crate::reader::Reader;
+
 
 const QUERY_OPT_IMPLICIT_LIMIT: u16 = 0xFF01;
 
@@ -46,10 +49,14 @@ pub struct Connection {
     stream: ByteStream,
 }
 
-pub struct Client<'a> {
+pub struct Writer<'a> {
     stream: &'a ByteStream,
     outbuf: BytesMut,
-    reader: Reader<&'a ByteStream>,
+}
+
+pub struct Client<'a> {
+    pub writer: Writer<'a>,
+    pub reader: Reader<&'a ByteStream>,
     pub params: TypeMap<dyn typemap::DebugAny + Send>,
 }
 
@@ -60,8 +67,9 @@ pub struct NoResultExpected {
 
 
 impl Connection {
-    async fn connect_tcp<A: ToSocketAddrs>(addrs: A, options: &Options)
+    async fn connect_tcp<A>(addrs: A, options: &Options)
         -> Result<Connection, anyhow::Error>
+        where A: ToSocketAddrs+fmt::Debug,
     {
         let start = Instant::now();
         let conn = loop {
@@ -77,10 +85,14 @@ impl Connection {
                                                     wait))?
                         }
                     } else {
-                        Err(e)?
+                        Err(e).with_context(
+                            || format!("Can't connect to {:?}", addrs))?;
                     }
                 }
-                Err(e) => Err(e)?,
+                Err(e) => {
+                    Err(e).with_context(
+                        || format!("Can't connect to {:?}", addrs))?;
+                }
                 Ok(conn) => break conn,
             }
         };
@@ -127,10 +139,14 @@ impl Connection {
                                                         wait))?
                             }
                         } else {
-                            Err(e)?
+                            Err(e).with_context(|| format!(
+                                "Can't connect to unix socket {:?}", path))?
                         }
                     }
-                    Err(e) => Err(e)?,
+                    Err(e) => {
+                        Err(e).with_context(|| format!(
+                            "Can't connect to unix socket {:?}", path))?
+                    }
                     Ok(conn) => break conn,
                 }
             };
@@ -164,8 +180,11 @@ impl Connection {
         let (rd, stream) = (&self.stream, &self.stream);
         let reader = Reader::new(rd);
         let mut cli = Client {
-            stream, reader,
-            outbuf: BytesMut::with_capacity(8912),
+            reader,
+            writer: Writer {
+                outbuf: BytesMut::with_capacity(8912),
+                stream,
+            },
             params: TypeMap::custom(),
         };
         let mut params = HashMap::new();
@@ -500,6 +519,22 @@ async fn _interactive_main(
     }
 }
 
+impl<'a> Writer<'a> {
+
+    pub async fn send_messages<'x, I>(&mut self, msgs: I)
+        -> Result<(), anyhow::Error>
+        where I: IntoIterator<Item=&'x ClientMessage>
+    {
+        self.outbuf.truncate(0);
+        for msg in msgs {
+            msg.encode(&mut self.outbuf)?;
+        }
+        self.stream.write_all(&self.outbuf[..]).await?;
+        Ok(())
+    }
+
+}
+
 impl<'a> Client<'a> {
     pub async fn scram(&mut self, options: &Options)
         -> Result<(), anyhow::Error>
@@ -582,15 +617,16 @@ impl<'a> Client<'a> {
         Ok(())
     }
 
-    async fn send_messages<'x, I>(&mut self, msgs: I)
+    pub async fn send_messages<'x, I>(&mut self, msgs: I)
         -> Result<(), anyhow::Error>
         where I: IntoIterator<Item=&'x ClientMessage>
     {
-        self.outbuf.truncate(0);
-        for msg in msgs {
-            msg.encode(&mut self.outbuf)?;
-        }
-        self.stream.write_all(&self.outbuf[..]).await?;
+        self.writer.send_messages(msgs).await
+    }
+
+    pub async fn err_sync(&mut self) -> Result<(), anyhow::Error> {
+        self.writer.send_messages(&[ClientMessage::Sync]).await?;
+        timeout(Duration::from_secs(10), self.reader.wait_ready()).await??;
         Ok(())
     }
 
diff --git a/src/commands/backslash.rs b/src/commands/backslash.rs
index 01eb8b2..a693c39 100644
--- a/src/commands/backslash.rs
+++ b/src/commands/backslash.rs
@@ -37,6 +37,7 @@ Introspection
   \li[IS+] [PATTERN]       list indexes
                            (alias: \list-indexes)
   \list-ports              list ports
+  \dump FILENAME           dump current database into a file
 
 Editing
   \s, \history             show history
@@ -79,6 +80,7 @@ pub const HINTS: &'static [&'static str] = &[
     r"\d+ NAME",
     r"\describe NAME",
     r"\describe+ NAME",
+    r"\dump FILENAME",
     r"\e [N]",
     r"\edit [N]",
     r"\emacs",
@@ -150,6 +152,7 @@ pub const COMMAND_NAMES: &'static [&'static str] = &[
     r"\d+",
     r"\describe",
     r"\describe+",
+    r"\dump",
     r"\e",
     r"\edit",
     r"\emacs",
@@ -277,6 +280,7 @@ pub enum Command {
     ShowLimit,
     SetOutput { mode: OutputMode },
     ShowOutput,
+    Dump { filename: String },
 }
 
 pub struct ParseError {
@@ -428,6 +432,7 @@ pub fn parse(s: &str) -> Result<Command, ParseError> {
         ("no-introspect-types", None) => Ok(Command::NoIntrospectTypes),
         ("verbose-errors", None) => Ok(Command::VerboseErrors),
         ("no-verbose-errors", None) => Ok(Command::NoVerboseErrors),
+        ("dump", Some(param)) => Ok(Command::Dump { filename: param.into() }),
         ("output", None) => Ok(Command::ShowOutput),
         ("output", Some(param)) => {
             Ok(Command::SetOutput {
@@ -626,6 +631,10 @@ pub async fn execute<'x>(cli: &mut Client<'x>, cmd: Command,
             });
             Ok(Skip)
         }
+        Dump { filename } => {
+            commands::dump(cli, &options, filename.as_ref()).await?;
+            Ok(Skip)
+        }
     }
 }
 
diff --git a/src/commands/cli.rs b/src/commands/cli.rs
index 44a37de..e438179 100644
--- a/src/commands/cli.rs
+++ b/src/commands/cli.rs
@@ -195,5 +195,25 @@ pub fn main(options: Options) -> Result<(), anyhow::Error> {
                 Ok(())
             }).into()
         },
+        Command::Dump(dump) => {
+            task::block_on(async {
+                let mut conn = Connection::from_options(&options).await?;
+                let mut cli = conn.authenticate(
+                    &options, &options.database).await?;
+                commands::dump(&mut cli, &cmdopt, &dump.file.as_ref()).await?;
+                Ok(())
+            }).into()
+        },
+        Command::Restore(restore) => {
+            task::block_on(async {
+                let mut conn = Connection::from_options(&options).await?;
+                let mut cli = conn.authenticate(
+                    &options, &options.database).await?;
+                commands::restore(&mut cli, &cmdopt,
+                    &restore.file.as_ref(),
+                    restore.allow_non_empty).await?;
+                Ok(())
+            }).into()
+        },
     }
 }
diff --git a/src/commands/dump.rs b/src/commands/dump.rs
new file mode 100644
index 0000000..1351b10
--- /dev/null
+++ b/src/commands/dump.rs
@@ -0,0 +1,107 @@
+use std::ffi::OsString;
+use std::default::Default;
+
+use anyhow::Context;
+use async_std::path::Path;
+use async_std::fs;
+use async_std::io::{self, Write, prelude::WriteExt};
+
+use edgedb_protocol::client_message::{ClientMessage, Dump};
+use edgedb_protocol::server_message::ServerMessage;
+use crate::commands::Options;
+use crate::client::Client;
+
+type Output = Box<dyn Write + Unpin + Send>;
+
+
+pub async fn dump<'x>(cli: &mut Client<'x>, _options: &Options,
+    filename: &Path)
+    -> Result<(), anyhow::Error>
+{
+    let (mut output, tmp_filename) = if filename.to_str() == Some("-") {
+        (Box::new(io::stdout()) as Output, None)
+    } else if cfg!(windows)
+        || filename.starts_with("/dev/")
+        || filename.file_name().is_none()
+    {
+        let file = fs::File::create(&filename).await
+            .context(filename.display().to_string())?;
+        (Box::new(file) as Output, None)
+    } else {
+        let name = filename.file_name().unwrap();
+        let mut tmp_name = OsString::with_capacity(name.len() + 10);
+        tmp_name.push(".");
+        tmp_name.push(name);
+        tmp_name.push(".edb.tmp");
+        let tmp_filename = filename.with_file_name(tmp_name);
+        if tmp_filename.exists().await {
+            fs::remove_file(&tmp_filename).await.ok();
+        }
+        let tmpfile = fs::File::create(&tmp_filename).await
+            .context(tmp_filename.display().to_string())?;
+        (Box::new(tmpfile) as Output, Some(tmp_filename))
+    };
+    output.write_all(
+        b"\xFF\xD8\x00\x00\xD8EDGEDB\x00DUMP\x00\
+          \x00\x00\x00\x00\x00\x00\x00\x01"
+        ).await?;
+
+    cli.send_messages(&[
+        ClientMessage::Dump(Dump {
+            headers: Default::default(),
+        }),
+        ClientMessage::Sync,
+    ]).await?;
+
+    let mut header_buf = Vec::with_capacity(25);
+    let msg = cli.reader.message().await?;
+    match msg {
+        ServerMessage::DumpHeader(packet) => {
+            // this is ensured because length in the protocol is u32 too
+            assert!(packet.data.len() <= u32::max_value() as usize);
+
+            header_buf.truncate(0);
+            header_buf.push(b'H');
+            header_buf.extend(
+                &sha1::Sha1::from(&packet.data).digest().bytes()[..]);
+            header_buf.extend(
+                &(packet.data.len() as u32).to_be_bytes()[..]);
+            output.write_all(&header_buf).await?;
+            output.write_all(&packet.data).await?;
+        }
+        _ => {
+            return Err(anyhow::anyhow!(
+                "WARNING: unsolicited message {:?}", msg));
+        }
+    }
+    loop {
+        let msg = cli.reader.message().await?;
+        match msg {
+            ServerMessage::CommandComplete(..) => {
+                cli.reader.wait_ready().await?;
+                break;
+            }
+            ServerMessage::DumpBlock(packet) => {
+                // this is ensured because length in the protocol is u32 too
+                assert!(packet.data.len() <= u32::max_value() as usize);
+
+                header_buf.truncate(0);
+                header_buf.push(b'D');
+                header_buf.extend(
+                    &sha1::Sha1::from(&packet.data).digest().bytes()[..]);
+                header_buf.extend(
+                    &(packet.data.len() as u32).to_be_bytes()[..]);
+                output.write_all(&header_buf).await?;
+                output.write_all(&packet.data).await?;
+            }
+            _ => {
+                return Err(anyhow::anyhow!(
+                    "WARNING: unsolicited message {:?}", msg));
+            }
+        }
+    }
+    if let Some(tmp_filename) = tmp_filename {
+        fs::rename(tmp_filename, filename).await?;
+    }
+    Ok(())
+}
diff --git a/src/commands/mod.rs b/src/commands/mod.rs
index 825563a..da48209 100644
--- a/src/commands/mod.rs
+++ b/src/commands/mod.rs
@@ -2,6 +2,7 @@ pub mod cli;
 pub mod backslash;
 pub mod options;
 mod configure;
+mod dump;
 mod filter;
 mod helpers;
 mod list;
@@ -15,12 +16,14 @@ mod list_ports;
 mod list_roles;
 mod list_scalar_types;
 mod psql;
+mod restore;
 mod roles;
 mod table;
 mod type_names;
 mod describe;
 
 pub use self::configure::configure;
+pub use self::dump::dump;
 pub use self::describe::describe;
 pub use self::list_aliases::list_aliases;
 pub use self::list_casts::list_casts;
@@ -32,4 +35,5 @@ pub use self::list_ports::list_ports;
 pub use self::list_roles::list_roles;
 pub use self::list_scalar_types::list_scalar_types;
 pub use self::options::Options;
+pub use self::restore::restore;
 pub use self::psql::psql;
diff --git a/src/commands/restore.rs b/src/commands/restore.rs
new file mode 100644
index 0000000..4bba7ff
--- /dev/null
+++ b/src/commands/restore.rs
@@ -0,0 +1,219 @@
+use std::collections::HashMap;
+use std::convert::TryInto;
+use std::time::{Instant, Duration};
+use std::mem::transmute;
+
+use anyhow::Context;
+use async_std::path::Path;
+use async_std::fs;
+use async_std::io::{self, Read, prelude::ReadExt};
+use async_std::future::{timeout, pending};
+use async_std::prelude::{FutureExt, StreamExt};
+use async_listen::ByteStream;
+use bytes::{Bytes, BytesMut, BufMut};
+
+use edgedb_protocol::client_message::{ClientMessage, Restore, RestoreBlock};
+use edgedb_protocol::server_message::ServerMessage;
+use edgedb_protocol::value::Value;
+use crate::commands::Options;
+use crate::commands::helpers::print_result;
+use crate::client::{Client, Reader, Writer};
+
+type Input = Box<dyn Read + Unpin + Send>;
+
+const MAX_SUPPORTED_DUMP_VER: i64 = 1;
+
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum PacketType {
+    Header,
+    Block,
+}
+
+
+async fn read_packet(input: &mut Input, expected: PacketType)
+    -> Result<Option<Bytes>, anyhow::Error>
+{
+    let mut buf = [0u8; 1+20+4];
+    let mut read = 0;
+    while read < buf.len() {
+        let n = input.read(&mut buf[read..]).await
+            .context("Cannot read packet header")?;
+        if n == 0 {  // EOF
+            if read == 0 {
+                return Ok(None);
+            }
+            Err(io::Error::from(io::ErrorKind::UnexpectedEof))
+            .context("Cannot read packet header")?
+        }
+        read += n;
+    }
+    let typ = match buf[0] {
+        b'H' => PacketType::Header,
+        b'D' => PacketType::Block,
+        _ => return Err(anyhow::anyhow!("Invalid block type {:x}", buf[0])),
+    };
+    if typ != expected {
+        return Err(anyhow::anyhow!("Expected block {:?} got {:?}",
+                    expected, typ));
+    }
+    let len = u32::from_be_bytes(buf[1+20..].try_into().unwrap()) as usize;
+    let mut buf = BytesMut::with_capacity(len);
+    unsafe {
+        // this is safe because we use read_exact which initializes whole slice
+        let dest: &mut [u8] = transmute(buf.bytes_mut());
+        input.read_exact(dest).await
+            .with_context(|| format!("Error reading block of {} bytes", len))?;
+        buf.advance_mut(dest.len());
+    }
+    return Ok(Some(buf.freeze()));
+}
+
+
+async fn is_empty_db(cli: &mut Client<'_>) -> Result<bool, anyhow::Error> {
+    let mut query = cli.query::<i64>(r###"SELECT
+            count(
+                schema::Module
+                FILTER NOT .builtin AND NOT .name = "default"
+            ) + count(
+                schema::Object
+                FILTER .name LIKE "default::%"
+            )
+        "###, &Value::empty_tuple()).await?;
+    let mut non_empty = false;
+    while let Some(num) = query.next().await.transpose()? {
+        if num > 0 {
+            non_empty = true;
+        }
+    }
+    return Ok(non_empty);
+}
+
+pub async fn restore<'x>(cli: &mut Client<'x>, options: &Options,
+    filename: &Path, allow_non_empty: bool)
+    -> Result<(), anyhow::Error>
+{
+    use PacketType::*;
+    if !allow_non_empty {
+        if is_empty_db(cli).await.context("Error checking DB emptyness")? {
+            if options.command_line {
+                return Err(anyhow::anyhow!("\
+                    cannot restore: the database is not empty; \
+                    consider using the --allow-non-empty option"));
+            } else {
+                return Err(anyhow::anyhow!(
+                    "cannot restore: the database is not empty"));
+            }
+        }
+    }
+
+    // TODO(tailhook) check that DB is empty
+    let file_ctx = &|| format!("Failed to read dump {}", filename.display());
+    let mut input = if filename.to_str() == Some("-") {
+        Box::new(io::stdin()) as Input
+    } else {
+        fs::File::open(filename).await
+        .map(Box::new)
+        .with_context(file_ctx)?
+        as Input
+    };
+    let mut buf = [0u8; 17+8];
+    input.read_exact(&mut buf).await
+        .context("Cannot read header")
+        .with_context(file_ctx)?;
+    if &buf[..17] != b"\xFF\xD8\x00\x00\xD8EDGEDB\x00DUMP\x00" {
+        Err(anyhow::anyhow!("File is not an edgedb dump"))
+        .with_context(file_ctx)?
+    }
+    let version = i64::from_be_bytes(buf[17..].try_into().unwrap());
+    if version == 0 || version > MAX_SUPPORTED_DUMP_VER {
+        Err(anyhow::anyhow!("Unsupported dump version {}", version))
+        .with_context(file_ctx)?
+    }
+    let header = read_packet(&mut input, Header).await.with_context(file_ctx)?
+        .ok_or_else(|| anyhow::anyhow!("Dump is empty"))
+                       .with_context(file_ctx)?;
+    let start_headers = Instant::now();
+    cli.send_messages(&[
+        ClientMessage::Restore(Restore {
+            headers: HashMap::new(),
+            jobs: 1,
+            data: header,
+        })
+    ]).await?;
+    loop {
+        let msg = cli.reader.message().await?;
+        match msg {
+            ServerMessage::RestoreReady(_) => {
+                eprintln!("Schema applied in {:?}", start_headers.elapsed());
+                break;
+            }
+            ServerMessage::ErrorResponse(err) => {
+                cli.err_sync().await.ok();
+                return Err(anyhow::anyhow!(err)
+                    .context("Error initiating restore protocol"));
+            }
+            _ => {
+                return Err(anyhow::anyhow!(
+                    "WARNING: unsolicited message {:?}", msg));
+            }
+        }
+    }
+    let result = send_blocks(&mut cli.writer, &mut input, filename)
+        .race(wait_response(&mut cli.reader, start_headers))
+        .await;
+    if let Err(..) = result {
+        cli.err_sync().await.ok();
+    }
+    result
+}
+
+async fn send_blocks(writer: &mut Writer<'_>, input: &mut Input,
+    filename: &Path)
+    -> Result<(), anyhow::Error>
+{
+    use PacketType::*;
+
+    let start_blocks = Instant::now();
+    while
+        let Some(data) = read_packet(input, Block).await
+            .with_context(|| format!("Failed to read dump {}",
+                                     filename.display()))?
+    {
+        writer.send_messages(&[
+            ClientMessage::RestoreBlock(RestoreBlock { data })
+        ]).await?;
+    }
+    writer.send_messages(&[ClientMessage::RestoreEof]).await?;
+    eprintln!("Blocks sent in {:?}", start_blocks.elapsed());
+
+    // This future should be canceled by wait_response() receiving
+    // CommandComplete
+    let start_waiting = Instant::now();
+    loop {
+        timeout(Duration::from_secs(60), pending()).await?;
+        eprintln!("Waiting for complete {:?}", start_waiting.elapsed());
+    }
+}
+
+async fn wait_response(reader: &mut Reader<&'_ ByteStream>, start: Instant)
+    -> Result<(), anyhow::Error>
+{
+    loop {
+        let msg = reader.message().await?;
+        match msg {
+            ServerMessage::CommandComplete(c) => {
+                eprintln!("Complete in {:?}", start.elapsed());
+                print_result(c.status_data);
+                break;
+            }
+            ServerMessage::ErrorResponse(err) => {
+                return Err(anyhow::anyhow!(err));
+            }
+            _ => {
+                return Err(anyhow::anyhow!(
+                    "WARNING: unsolicited message {:?}", msg));
+            }
+        }
+    }
+    Ok(())
+}
diff --git a/src/options.rs b/src/options.rs
index bfda94d..18879b8 100644
--- a/src/options.rs
+++ b/src/options.rs
@@ -1,6 +1,7 @@
 use std::env;
 use std::time::Duration;
 use std::process::exit;
+use std::path::PathBuf;
 
 use atty;
 use structopt::StructOpt;
@@ -100,6 +101,10 @@ pub enum Command {
     Configure(Configure),
     Describe(Describe),
     Query(Query),
+    /// Create a database backup
+    Dump(Dump),
+    /// Restore a database backup from file
+    Restore(Restore),
 }
 
 #[derive(StructOpt, Clone, Debug)]
@@ -180,6 +185,22 @@ pub struct Query {
     pub queries: Vec<String>,
 }
 
+#[derive(StructOpt, Clone, Debug)]
+#[structopt(setting=AppSettings::DisableVersion)]
+pub struct Dump {
+    pub file: PathBuf,
+}
+
+#[derive(StructOpt, Clone, Debug)]
+#[structopt(setting=AppSettings::DisableVersion)]
+pub struct Restore {
+    pub file: PathBuf,
+
+    /// Allow restoring the database dump into a non-empty database
+    #[structopt(long)]
+    pub allow_non_empty: bool,
+}
+
 #[derive(StructOpt, Clone, Debug)]
 #[structopt(setting=AppSettings::DisableVersion)]
 pub struct Configure {
diff --git a/src/reader.rs b/src/reader.rs
index 3aa15c3..d9b7d5e 100644
--- a/src/reader.rs
+++ b/src/reader.rs
@@ -137,7 +137,7 @@ impl<T: AsyncRead + Unpin> Reader<T> {
 
             buf.reserve(next_read);
             unsafe {
-                // this is save because the underlying ByteStream always
+                // this is safe because the underlying ByteStream always
                 // initializes read bytes
                 let dest: &mut [u8] = transmute(buf.bytes_mut());
                 match Pin::new(&mut *stream).poll_read(cx, dest) {


Test patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index accbbc9..7b86833 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -5,6 +5,7 @@ name: Tests
 jobs:
   test:
     runs-on: ${{ matrix.os }}
+    timeout-minutes: 25
     strategy:
       max-parallel: 4
       matrix:
diff --git a/tests/func/dump_restore.rs b/tests/func/dump_restore.rs
new file mode 100644
index 0000000..c592da4
--- /dev/null
+++ b/tests/func/dump_restore.rs
@@ -0,0 +1,21 @@
+use crate::SERVER;
+
+#[test]
+fn dump_restore_cycle() {
+    SERVER.admin_cmd().arg("create-database").arg("dump_01")
+        .assert().success();
+    SERVER.database_cmd("dump_01").arg("query")
+        .arg("CREATE TYPE Hello { CREATE REQUIRED PROPERTY name -> str; }")
+        .arg("INSERT Hello { name := 'world' }")
+        .assert().success();
+    SERVER.database_cmd("dump_01").arg("dump").arg("dump_01.dump")
+        .assert().success();
+    SERVER.admin_cmd().arg("create-database").arg("restore_01")
+        .assert().success();
+    SERVER.database_cmd("restore_01").arg("restore").arg("dump_01.dump")
+        .assert().success();
+    SERVER.database_cmd("restore_01").arg("query")
+        .arg("SELECT Hello.name")
+        .assert().success()
+        .stdout("\"world\"\n");
+}
diff --git a/tests/func/main.rs b/tests/func/main.rs
index aac4b44..3ea7dab 100644
--- a/tests/func/main.rs
+++ b/tests/func/main.rs
@@ -1,3 +1,4 @@
+use std::sync::Mutex;
 use std::convert::TryInto;
 use std::io::{BufReader, BufRead};
 use std::process::Child;
@@ -7,6 +8,8 @@ use std::thread::{self, JoinHandle};
 use assert_cmd::Command;
 use serde_json::from_str;
 
+mod dump_restore;
+
 
 lazy_static::lazy_static! {
     pub static ref SERVER: ServerGuard = ServerGuard::start();
@@ -14,19 +17,17 @@ lazy_static::lazy_static! {
 
 #[test]
 fn simple_query() {
-    let guard = &*SERVER;
-    let cmd = Command::cargo_bin("edgedb").expect("binary found")
-        .arg("--admin")
-        .arg("--port").arg(guard.port.to_string())
-        .arg("query").arg("SELECT 1+7")
-        .env("EDGEDB_HOST", &guard.runstate_dir)
-        .assert();
+    let cmd = SERVER.admin_cmd().arg("query").arg("SELECT 1+7").assert();
     cmd.success().stdout("8\n");
 }
 
-pub struct ServerGuard {
+pub struct ShutdownInfo {
     process: Child,
     thread: Option<JoinHandle<()>>,
+}
+
+pub struct ServerGuard {
+    shutdown_info: Mutex<ShutdownInfo>,
     port: u16,
     runstate_dir: String,
 }
@@ -40,8 +41,8 @@ impl ServerGuard {
         cmd.arg("--testmode");
         cmd.arg("--echo-runtime-info");
         cmd.arg("--port=auto");
-        cmd.arg("--auto-shutdown");
         cmd.stdout(Stdio::piped());
+
         let mut process = cmd.spawn().expect("Can run edgedb-server");
         let process_in = process.stdout.take().expect("stdout is pipe");
         let (tx, rx) = sync_channel(1);
@@ -75,18 +76,41 @@ impl ServerGuard {
             }
         });
         let (port, runstate_dir) = rx.recv().expect("valid port received");
+
+        shutdown_hooks::add_shutdown_hook(stop_process);
+
         ServerGuard {
-            process,
-            thread: Some(thread),
+            shutdown_info: Mutex::new(ShutdownInfo {
+                process,
+                thread: Some(thread),
+            }),
             port,
             runstate_dir,
         }
     }
-}
 
-impl Drop for ServerGuard {
-    fn drop(&mut self) {
-        self.process.kill().ok();
-        self.thread.take().expect("not yet joined").join().ok();
+    pub fn admin_cmd(&self) -> Command {
+        let mut cmd = Command::cargo_bin("edgedb").expect("binary found");
+        cmd.arg("--admin");
+        cmd.arg("--port").arg(self.port.to_string());
+        cmd.env("EDGEDB_HOST", &self.runstate_dir);
+        return cmd
+    }
+
+    pub fn database_cmd(&self, database_name: &str) -> Command {
+        let mut cmd = Command::cargo_bin("edgedb").expect("binary found");
+        cmd.arg("--admin");
+        cmd.arg("--port").arg(self.port.to_string());
+        cmd.arg("--database").arg(database_name);
+        cmd.env("EDGEDB_HOST", &self.runstate_dir);
+        return cmd
     }
 }
+
+
+extern fn stop_process() {
+    let mut sinfo = SERVER.shutdown_info.lock().expect("shutdown mutex works");
+    sinfo.process.kill().ok();
+    sinfo.process.wait().ok();
+    sinfo.thread.take().expect("not yet joined").join().ok();
+}