database stream patches

Author: krukah

src/clustering/lookup.rs

@@ -164,9 +164,9 @@ impl Upload for Lookup {
         while reader.read_exact(buffer).is_ok() {
             match u16::from_be_bytes(buffer.clone()) {
                 2 => {
-                    reader.read_u32::<BE>().expect("observation length");
+                    assert!(8 == reader.read_u32::<BE>().expect("observation length"));
                     let iso = reader.read_i64::<BE>().expect("read observation");
-                    reader.read_u32::<BE>().expect("abstraction length");
+                    assert!(8 == reader.read_u32::<BE>().expect("abstraction length"));
                     let abs = reader.read_i64::<BE>().expect("read abstraction");
                     let observation = Isomorphism::from(iso);
                     let abstraction = Abstraction::from(abs);

src/lib.rs

@@ -100,7 +100,8 @@ pub fn init() {
 pub async fn db() -> std::sync::Arc<tokio_postgres::Client> {
     log::info!("connecting to database");
     let tls = tokio_postgres::tls::NoTls;
-    let ref url = std::env::var("DB_URL").expect("DB_URL not set");
+    let ref url =
+        std::env::var("DB_URL").unwrap_or("postgresql://localhost:5432/postgres".to_string());
     let (client, connection) = tokio_postgres::connect(url, tls)
         .await
         .expect("database connection failed");

src/save/upload.rs

@@ -1,9 +1,4 @@
 use crate::cards::street::Street;
-use byteorder::ReadBytesExt;
-use byteorder::BE;
-use std::fs::File;
-use std::io::BufReader;
-use tokio_postgres::types::ToSql;
 use tokio_postgres::types::Type;
 
 // blueprint    ~ 154M, (grows with number of CFR iterations)
@@ -35,18 +30,6 @@ pub trait Upload {
     fn load(street: Street) -> Self;
     /// write to disk
     fn save(&self);
-    /// given a BufReader<File>, read a row according to Self::columns()
-    fn read(reader: &mut BufReader<File>) -> Vec<Box<dyn ToSql + Sync>> {
-        Self::columns()
-            .iter()
-            .cloned()
-            .map(|ty| match ty {
-                Type::FLOAT4 => Box::new(reader.read_f32::<BE>().unwrap()) as Box<dyn ToSql + Sync>,
-                Type::INT8 => Box::new(reader.read_i64::<BE>().unwrap()) as Box<dyn ToSql + Sync>,
-                _ => panic!("unsupported type: {}", ty),
-            })
-            .collect::<Vec<Box<dyn ToSql + Sync>>>()
-    }
 
     /// query to nuke table in Postgres
     fn nuke() -> String {

src/save/writer.rs

@@ -6,12 +6,16 @@ use crate::clustering::metric::Metric;
 use crate::clustering::transitions::Decomp;
 use crate::mccfr::encoder::Encoder;
 use crate::mccfr::profile::Profile;
+use byteorder::ReadBytesExt;
+use byteorder::BE;
 use std::fs::File;
 use std::io::BufReader;
 use std::io::Read;
 use std::io::Seek;
 use std::sync::Arc;
 use tokio_postgres::binary_copy::BinaryCopyInWriter;
+use tokio_postgres::types::ToSql;
+use tokio_postgres::types::Type;
 use tokio_postgres::Client;
 use tokio_postgres::Error as E;
 
@@ -28,59 +32,22 @@ impl Writer {
         let postgres = Self(crate::db().await);
         postgres.upload::<Metric>().await?;
         postgres.upload::<Decomp>().await?;
-        postgres.upload::<Encoder>().await?;
+        postgres.upload::<Encoder>().await?; // Lookup ?
         postgres.upload::<Profile>().await?;
         postgres.derive::<Abstraction>().await?;
         postgres.derive::<Street>().await?;
         postgres.vacuum().await?;
         Ok(())
     }
 
-    async fn upload<U>(&self) -> Result<(), E>
-    where
-        U: Upload,
-    {
-        let ref name = U::name();
-        if self.has_rows(name).await? {
-            log::info!("tables data already uploaded ({})", name);
-            Ok(())
-        } else {
-            log::info!("copying {}", name);
-            self.0.batch_execute(&U::prepare()).await?;
-            self.0.batch_execute(&U::nuke()).await?;
-            let sink = self.0.copy_in(&U::copy()).await?;
-            let writer = BinaryCopyInWriter::new(sink, U::columns());
-            futures::pin_mut!(writer);
-            let ref mut count = [0u8; 2];
-            for ref mut reader in U::sources()
-                .iter()
-                .map(|s| File::open(s).expect("file not found"))
-                .map(|f| BufReader::new(f))
-            {
-                reader.seek(std::io::SeekFrom::Start(19)).unwrap();
-                while let Ok(_) = reader.read_exact(count) {
-                    match u16::from_be_bytes(count.clone()) {
-                        0xFFFF => break,
-                        length => {
-                            assert!(length == U::columns().len() as u16);
-                            let row = U::read(reader);
-                            let row = row.iter().map(|b| &**b).collect::<Vec<_>>();
-                            writer.as_mut().write(&row).await?;
-                        }
-                    }
-                }
-            }
-            writer.finish().await?;
-            self.0.batch_execute(&U::indices()).await?;
-            Ok(())
-        }
-    }
-
     async fn derive<D>(&self) -> Result<(), E>
     where
         D: Derive,
     {
         let ref name = D::name();
+        // if !does_exist(name).await? {
+        // create
+        // }
         if self.has_rows(name).await? {
             log::info!("tables data already uploaded ({})", name);
             Ok(())
@@ -103,6 +70,73 @@ impl Writer {
         }
     }
 
+    async fn upload<U>(&self) -> Result<(), E>
+    where
+        U: Upload,
+    {
+        let ref name = U::name();
+        // if !does_exist(name).await? {
+        // create
+        // }
+        if self.has_rows(name).await? {
+            log::info!("tables data already uploaded ({})", name);
+            Ok(())
+        } else {
+            log::info!("copying {}", name);
+            self.prepare::<U>().await?;
+            self.stream::<U>().await?;
+            self.index::<U>().await?;
+            Ok(())
+        }
+    }
+
+    async fn prepare<T>(&self) -> Result<(), E>
+    where
+        T: Upload,
+    {
+        self.0.batch_execute(&T::prepare()).await?;
+        self.0.batch_execute(&T::nuke()).await?;
+        Ok(())
+    }
+
+    async fn index<T>(&self) -> Result<(), E>
+    where
+        T: Upload,
+    {
+        self.0.batch_execute(&T::indices()).await?;
+        Ok(())
+    }
+
+    async fn stream<T>(&self) -> Result<(), E>
+    where
+        T: Upload,
+    {
+        let sink = self.0.copy_in(&T::copy()).await?;
+        let writer = BinaryCopyInWriter::new(sink, T::columns());
+        futures::pin_mut!(writer);
+        let ref mut count = [0u8; 2];
+        for ref mut reader in T::sources()
+            .iter()
+            .map(|s| File::open(s).expect("file not found"))
+            .map(|f| BufReader::new(f))
+        {
+            reader.seek(std::io::SeekFrom::Start(19)).unwrap();
+            while let Ok(_) = reader.read_exact(count) {
+                match u16::from_be_bytes(count.clone()) {
+                    0xFFFF => break,
+                    length => {
+                        assert!(length == T::columns().len() as u16);
+                        let row = Self::read::<T>(reader);
+                        let row = row.iter().map(|b| &**b).collect::<Vec<_>>();
+                        writer.as_mut().write(&row).await?;
+                    }
+                }
+            }
+        }
+        writer.finish().await?;
+        Ok(())
+    }
+
     async fn vacuum(&self) -> Result<(), E> {
         self.0.batch_execute("VACUUM ANALYZE;").await
     }
@@ -116,7 +150,7 @@ impl Writer {
                 ",
                 table
             );
-            Ok(0 != self.0.query_one(sql, &[]).await?.get::<_, i64>(0))
+            Ok(!self.0.query(sql, &[]).await?.is_empty())
         } else {
             Ok(false)
         }
@@ -130,6 +164,27 @@ impl Writer {
             ",
             table
         );
-        Ok(1 == self.0.query_one(sql, &[]).await?.get::<_, i64>(0))
+        Ok(!self.0.query(sql, &[]).await?.is_empty())
+    }
+
+    fn read<T>(reader: &mut BufReader<File>) -> Vec<Box<dyn ToSql + Sync>>
+    where
+        T: Upload,
+    {
+        T::columns()
+            .iter()
+            .cloned()
+            .map(|ty| match ty {
+                Type::FLOAT4 => {
+                    assert!(reader.read_u32::<BE>().expect("length") == 4);
+                    Box::new(reader.read_f32::<BE>().expect("data")) as Box<dyn ToSql + Sync>
+                }
+                Type::INT8 => {
+                    assert!(reader.read_u32::<BE>().expect("length") == 8);
+                    Box::new(reader.read_i64::<BE>().expect("data")) as Box<dyn ToSql + Sync>
+                }
+                _ => panic!("unsupported type: {}", ty),
+            })
+            .collect::<Vec<Box<dyn ToSql + Sync>>>()
     }
 }