From 3701c866377ca50a43f0fbbacd314a18552cb7b6 Mon Sep 17 00:00:00 2001
From: Billy Chan <30400950+billy1624@users.noreply.github.com>
Date: Sat, 7 Aug 2021 18:23:52 +0800
Subject: [PATCH] CLI Generate Entity for Postgres (#61)

* cli generate entity for Postgres

* Fixup
---
 sea-orm-cli/Cargo.toml     |  3 ++-
 sea-orm-cli/src/main.rs    | 37 +++++++++++++++++++++++++++----------
 sea-orm-codegen/Cargo.toml |  2 +-
 3 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/sea-orm-cli/Cargo.toml b/sea-orm-cli/Cargo.toml
index c5ae5823..85fc693f 100644
--- a/sea-orm-cli/Cargo.toml
+++ b/sea-orm-cli/Cargo.toml
@@ -20,8 +20,9 @@ clap = { version = "^2.33.3" }
 dotenv = { version = "^0.15" }
 async-std = { version = "^1.9", features = [ "attributes" ] }
 sea-orm-codegen = { path = "../sea-orm-codegen" }
-sea-schema = { version = "^0.2", default-features = false, features = [
+sea-schema = { version = "^0.2", git = "https://github.com/SeaQL/sea-schema.git", branch = "pg-discovery-writer", default-features = false, features = [
     "sqlx-mysql",
+    "sqlx-postgres",
     "discovery",
     "writer",
 ] }
diff --git a/sea-orm-cli/src/main.rs b/sea-orm-cli/src/main.rs
index ecfa7892..9a512050 100644
--- a/sea-orm-cli/src/main.rs
+++ b/sea-orm-cli/src/main.rs
@@ -1,8 +1,6 @@
 use clap::ArgMatches;
 use dotenv::dotenv;
 use sea_orm_codegen::{EntityTransformer, OutputFile};
-use sea_schema::mysql::discovery::SchemaDiscovery;
-use sqlx::MySqlPool;
 use std::{error::Error, fmt::Display, fs, io::Write, path::Path, process::Command};
 
 mod cli;
@@ -28,14 +26,33 @@ async fn run_generate_command(matches: &ArgMatches<'_>) -> Result<(), Box<dyn Er
             let schema = args.value_of("DATABASE_SCHEMA").unwrap();
             let output_dir = args.value_of("OUTPUT_DIR").unwrap();
 
-            let connection = MySqlPool::connect(url).await?;
-            let schema_discovery = SchemaDiscovery::new(connection, schema);
-            let schema = schema_discovery.discover().await;
-            let table_stmts = schema
-                .tables
-                .into_iter()
-                .map(|schema| schema.write())
-                .collect();
+            let table_stmts = if url.starts_with("mysql://") {
+                use sea_schema::mysql::discovery::SchemaDiscovery;
+                use sqlx::MySqlPool;
+
+                let connection = MySqlPool::connect(url).await?;
+                let schema_discovery = SchemaDiscovery::new(connection, schema);
+                let schema = schema_discovery.discover().await;
+                schema
+                    .tables
+                    .into_iter()
+                    .map(|schema| schema.write())
+                    .collect()
+            } else if url.starts_with("postgres://") {
+                use sea_schema::postgres::discovery::SchemaDiscovery;
+                use sqlx::PgPool;
+
+                let connection = PgPool::connect(url).await?;
+                let schema_discovery = SchemaDiscovery::new(connection, schema);
+                let schema = schema_discovery.discover().await;
+                schema
+                    .tables
+                    .into_iter()
+                    .map(|schema| schema.write())
+                    .collect()
+            } else {
+                panic!("This database is not supported ({})", url)
+            };
 
             let output = EntityTransformer::transform(table_stmts)?.generate();
 
diff --git a/sea-orm-codegen/Cargo.toml b/sea-orm-codegen/Cargo.toml
index 6c67bbf1..86ccf937 100644
--- a/sea-orm-codegen/Cargo.toml
+++ b/sea-orm-codegen/Cargo.toml
@@ -16,7 +16,7 @@ name = "sea_orm_codegen"
 path = "src/lib.rs"
 
 [dependencies]
-sea-query = { version = "~0.12.8" }
+sea-query = { version = "~0.12.8", git = "https://github.com/SeaQL/sea-query.git" }
 syn = { version = "^1", default-features = false, features = [
     "derive",
     "parsing",