Generator

Cargo.toml

@@ -5,24 +5,22 @@ edition = "2021"
 rust-version = "1.70"
 license = "Apache-2.0 OR MIT"
 
-[[bin]]
-name = "cli"
-required-features = ["postgres", "clap"]
+[features]
+default = ["postgres"]
+postgres = ["tokio-postgres"]
 
 [dependencies]
 anyhow = "1"
-clap = { version = "4", features = ["derive"], optional = true }
+async-recursion = "1"
 heck = "0.4.0"
 indexmap = "2"
-postgres = { version = "0.19.3", optional = true }
-postgres-types = { version = "0.2.3", features = ["derive"] }
 thiserror = "1.0.31"
 time = { version = "0.3.9", features = ["parsing"] }
+tokio = { version = "1", features = ["rt-multi-thread", "macros"] }
+tokio-postgres = { version = "0.7", features = ["with-chrono-0_4"], optional = true }
 
 [dev-dependencies]
 chrono = "0.4"
-postgres = { version = "0.19.3", features = ["with-chrono-0_4"] }
+clap = { version = "4", features = ["derive"] }
+similar-asserts = "1.4.2"
 tempfile = "3.3"
-
-[features]
-default = ["postgres"]

examples/generate.rs

@@ -0,0 +1,22 @@
+use clap::Parser;
+use instant_models::Schema;
+use tokio_postgres::NoTls;
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    let opts = Options::parse();
+    let (client, conn) = tokio_postgres::connect(&opts.db, NoTls).await?;
+    tokio::spawn(conn);
+
+    let schema = Schema::from_postgres(&client).await?;
+    println!("{schema}");
+
+    Ok(())
+}
+
+/// Generate Rust code from Postgres database
+#[derive(Parser, Debug)]
+struct Options {
+    /// Database to inspect (as URL)
+    db: String,
+}

src/column.rs

@@ -1,27 +1,25 @@
-use crate::{Type, TypeAsRef};
+use std::fmt;
+use std::sync::Arc;
+
 use heck::AsSnakeCase;
-use std::borrow::Cow;
+use tokio_postgres::Client;
 
-#[derive(Debug, PartialEq)]
-pub struct ForeignKey {
-    to_table: Cow<'static, str>,
-    columns: Vec<Cow<'static, str>>,
-}
+use crate::types::{Type, TypeAsRef};
 
 #[derive(Debug, PartialEq)]
 pub struct Column {
-    pub name: Cow<'static, str>,
+    pub name: Arc<str>,
     pub r#type: Type,
     pub null: bool,
     pub primary_key: bool,
     pub foreign_key: Option<ForeignKey>,
     pub unique: bool,
-    pub default: Option<Cow<'static, str>>,
-    pub type_def: Option<Cow<'static, str>>,
+    pub default: Option<String>,
+    pub type_def: Option<String>,
 }
 
 impl Column {
-    pub fn new(name: Cow<'static, str>, r#type: Type) -> Self {
+    pub fn new(name: Arc<str>, r#type: Type) -> Self {
         Self {
             name,
             r#type,
@@ -34,39 +32,16 @@ impl Column {
         }
     }
 
-    pub fn set_null(mut self, value: bool) -> Self {
-        self.null = value;
-        self
-    }
-
-    pub fn set_primary_key(mut self, value: bool) -> Self {
-        self.primary_key = value;
-        self
-    }
-
-    pub fn set_foreign_key(mut self, value: Option<ForeignKey>) -> Self {
-        self.foreign_key = value;
-        self
-    }
-
-    pub fn set_unique(mut self, value: bool) -> Self {
-        self.unique = value;
-        self
-    }
-
-    pub fn set_default(mut self, value: Option<Cow<'static, str>>) -> Self {
-        self.default = value;
-        self
-    }
-
-    pub fn set_type_def(mut self, value: Option<Cow<'static, str>>) -> Self {
-        self.type_def = value;
-        self
+    pub fn new_field(&self) -> NewField {
+        NewField {
+            lifetime: Some("a"),
+            val: self,
+        }
     }
 }
 
-impl std::fmt::Display for Column {
-    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
+impl fmt::Display for Column {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
         if self.null {
             write!(fmt, "{}: Option<{}>", AsSnakeCase(&self.name), self.r#type)
         } else {
@@ -75,16 +50,16 @@ impl std::fmt::Display for Column {
     }
 }
 
-pub struct NewValue<'a> {
+pub struct NewField<'a> {
     pub lifetime: Option<&'a str>,
     pub val: &'a Column,
 }
 
-impl std::fmt::Display for NewValue<'_> {
-    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
+impl fmt::Display for NewField<'_> {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
         if self.val.null && self.val.default.is_some() {
             panic!(
-                "Column `{}` is both NULL and takes a default value `{}`",
+                "column `{}` is both NULL and takes a default value `{}`",
                 self.val.name,
                 self.val.default.as_ref().unwrap()
             );
@@ -116,13 +91,46 @@ impl std::fmt::Display for NewValue<'_> {
 #[derive(Debug, PartialEq)]
 pub enum Constraint {
     ForeignKey {
-        name: Cow<'static, str>,
-        columns: Cow<'static, [Cow<'static, str>]>,
-        ref_table: Cow<'static, str>,
-        ref_columns: Cow<'static, [Cow<'static, str>]>,
+        name: String,
+        columns: Vec<String>,
+        ref_table: String,
+        ref_columns: Vec<String>,
     },
     PrimaryKey {
-        name: Cow<'static, str>,
-        columns: Vec<Cow<'static, str>>,
+        name: Vec<String>,
+        columns: Vec<String>,
     },
 }
+
+#[derive(Debug, PartialEq)]
+pub struct ForeignKey {
+    to_table: String,
+    columns: Vec<String>,
+}
+
+impl ForeignKey {
+    pub(crate) async fn from_postgres(
+        name: &str,
+        client: &Client,
+    ) -> Result<Self, tokio_postgres::Error> {
+        let sql = r#"
+            SELECT tc.table_name, kcu.column_name
+            FROM information_schema.table_constraints AS tc
+                JOIN information_schema.key_column_usage AS kcu
+                    ON tc.constraint_name = kcu.constraint_name
+            WHERE tc.constraint_type = 'FOREIGN KEY' AND tc.constraint_name = $1
+        "#;
+
+        let rows = client.query(sql, &[&name]).await?;
+        let row = match &rows[..] {
+            [row] => row,
+            [] => panic!("foreign key constraint {name} not found"),
+            _ => todo!("unsupported composite foreign key {name}"),
+        };
+
+        Ok(Self {
+            to_table: row.get::<_, &str>(0).to_owned(),
+            columns: vec![row.get::<_, &str>(1).to_owned()],
+        })
+    }
+}

src/lib.rs

@@ -1,8 +1,64 @@
-mod struct_builder;
-pub use struct_builder::*;
+use std::fmt;
 
-mod column;
-pub use column::*;
+use tokio_postgres::Client;
 
+mod table;
+pub use table::Table;
+mod column;
+pub use column::Column;
 mod types;
-pub use types::*;
+pub use types::{Enum, Type, TypeDefinition, Composite};
+
+pub struct Schema {
+    pub tables: Vec<Table>,
+    pub types: Vec<TypeDefinition>,
+}
+
+impl Schema {
+    #[cfg(feature = "postgres")]
+    pub async fn from_postgres(client: &Client) -> anyhow::Result<Self> {
+        let sql = r#"
+            SELECT table_name
+            FROM information_schema.tables
+            WHERE table_schema = 'public'
+        "#;
+
+        let (mut tables, mut types) = (Vec::new(), Vec::new());
+        for row in client.query(sql, &[]).await? {
+            let name = row.get::<_, &str>(0);
+            let table = Table::from_postgres(name, client).await?;
+
+            for column in table.columns.values() {
+                match &column.r#type {
+                    Type::Composite(c) => types.push(TypeDefinition::Composite(
+                        Composite::from_postgres(&c.name, client).await?,
+                    )),
+                    Type::Enum(e) => types.push(TypeDefinition::Enum(
+                        Enum::from_postgres(&e.name, client).await?,
+                    )),
+                    _ => {}
+                }
+            }
+
+            tables.push(table);
+        }
+
+        Ok(Self { tables, types })
+    }
+}
+
+impl fmt::Display for Schema {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        for table in &self.tables {
+            f.write_fmt(format_args!("{table}"))?;
+            f.write_str("\n")?;
+        }
+
+        for type_def in &self.types {
+            f.write_fmt(format_args!("{type_def}"))?;
+            f.write_str("\n")?;
+        }
+
+        Ok(())
+    }
+}