Discover unique indexes

src/postgres/discovery/mod.rs

@@ -2,10 +2,12 @@
 
 use crate::debug_print;
 use crate::postgres::def::*;
-use crate::postgres::parser::parse_table_constraint_query_results;
+use crate::postgres::parser::{
+    parse_table_constraint_query_results, parse_unique_index_query_results,
+};
 use crate::postgres::query::{
     ColumnQueryResult, EnumQueryResult, SchemaQueryBuilder, TableConstraintsQueryResult,
-    TableQueryResult,
+    TableQueryResult, UniqueIndexQueryResult,
 };
 use crate::sqlx_types::SqlxError;
 use futures::future;
@@ -101,32 +103,28 @@ impl SchemaDiscovery {
         let (
             check_constraints,
             not_null_constraints,
-            unique_constraints,
             primary_key_constraints,
             reference_constraints,
             exclusion_constraints,
         ) = constraints.into_iter().fold(
-            (
-                Vec::new(),
-                Vec::new(),
-                Vec::new(),
-                Vec::new(),
-                Vec::new(),
-                Vec::new(),
-            ),
+            (Vec::new(), Vec::new(), Vec::new(), Vec::new(), Vec::new()),
             |mut acc, constraint| {
                 match constraint {
                     Constraint::Check(check) => acc.0.push(check),
                     Constraint::NotNull(not_null) => acc.1.push(not_null),
-                    Constraint::Unique(unique) => acc.2.push(unique),
-                    Constraint::PrimaryKey(primary_key) => acc.3.push(primary_key),
-                    Constraint::References(references) => acc.4.push(references),
-                    Constraint::Exclusion(exclusion) => acc.5.push(exclusion),
+                    Constraint::Unique(_) => (),
+                    Constraint::PrimaryKey(primary_key) => acc.2.push(primary_key),
+                    Constraint::References(references) => acc.3.push(references),
+                    Constraint::Exclusion(exclusion) => acc.4.push(exclusion),
                 }
                 acc
             },
         );
 
+        let unique_constraints = self
+            .discover_unique_indexes(self.schema.clone(), table.clone())
+            .await?;
+
         Ok(TableDef {
             info,
             columns,
@@ -189,6 +187,33 @@ impl SchemaDiscovery {
             .collect())
     }
 
+    pub async fn discover_unique_indexes(
+        &self,
+        schema: SeaRc<dyn Iden>,
+        table: SeaRc<dyn Iden>,
+    ) -> Result<Vec<Unique>, SqlxError> {
+        let rows = self
+            .executor
+            .fetch_all(
+                self.query
+                    .query_table_unique_indexes(schema.clone(), table.clone()),
+            )
+            .await?;
+
+        let results = rows.into_iter().map(|row| {
+            let result: UniqueIndexQueryResult = (&row).into();
+            debug_print!("{:?}", result);
+            result
+        });
+
+        Ok(parse_unique_index_query_results(Box::new(results))
+            .map(|index| {
+                debug_print!("{:?}", index);
+                index
+            })
+            .collect())
+    }
+
     pub async fn discover_enums(&self) -> Result<Vec<EnumDef>, SqlxError> {
         let rows = self.executor.fetch_all(self.query.query_enums()).await?;
 

src/postgres/parser/mod.rs

@@ -1,8 +1,10 @@
 mod column;
+mod pg_indexes;
 mod table;
 mod table_constraints;
 
 pub use column::*;
+pub use pg_indexes::*;
 pub use table::*;
 pub use table_constraints::*;
 

src/postgres/parser/pg_indexes.rs

@@ -0,0 +1,47 @@
+use crate::postgres::{def::*, query::UniqueIndexQueryResult};
+
+pub struct UniqueIndexQueryResultParser {
+    curr: Option<UniqueIndexQueryResult>,
+    results: Box<dyn Iterator<Item = UniqueIndexQueryResult>>,
+}
+
+pub fn parse_unique_index_query_results(
+    results: Box<dyn Iterator<Item = UniqueIndexQueryResult>>,
+) -> impl Iterator<Item = Unique> {
+    UniqueIndexQueryResultParser {
+        curr: None,
+        results,
+    }
+}
+
+impl Iterator for UniqueIndexQueryResultParser {
+    type Item = Unique;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let result = if let Some(result) = self.curr.take() {
+            result
+        } else {
+            self.results.next()?
+        };
+
+        let index_name = result.index_name;
+        let mut columns = vec![result.column_name];
+
+        for result in self.results.by_ref() {
+            if result.index_name != index_name {
+                self.curr = Some(result);
+                return Some(Unique {
+                    name: index_name,
+                    columns,
+                });
+            }
+
+            columns.push(result.column_name);
+        }
+
+        Some(Unique {
+            name: index_name,
+            columns,
+        })
+    }
+}

src/postgres/query/pg_indexes.rs

@@ -1,4 +1,6 @@
-use sea_query::Iden;
+use super::SchemaQueryBuilder;
+use crate::sqlx_types::postgres::PgRow;
+use sea_query::{Alias, Condition, Expr, Iden, JoinType, Order, Query, SeaRc, SelectStatement};
 
 #[derive(Debug, Iden)]
 pub enum PgIndexes {
@@ -10,3 +12,136 @@ pub enum PgIndexes {
     #[iden = "indexname"]
     IndexName,
 }
+
+#[derive(Debug, Iden)]
+pub enum PgIndex {
+    Table,
+    #[iden = "indexrelid"]
+    IndexRelId,
+    #[iden = "indrelid"]
+    IndRelId,
+    #[iden = "indisunique"]
+    IndIsUnique,
+    #[iden = "indisprimary"]
+    IndIsPrimary,
+}
+
+#[derive(Debug, Iden)]
+pub enum PgClass {
+    Table,
+    Oid,
+    #[iden = "relnamespace"]
+    RelNamespace,
+    #[iden = "relname"]
+    RelName,
+}
+
+#[derive(Debug, Iden)]
+pub enum PgNamespace {
+    Table,
+    Oid,
+    #[iden = "nspname"]
+    NspName,
+}
+
+#[derive(Debug, Iden)]
+pub enum PgAttribute {
+    Table,
+    Oid,
+    #[iden = "attrelid"]
+    AttRelId,
+    #[iden = "attname"]
+    AttName,
+}
+
+#[derive(Debug, Default)]
+pub struct UniqueIndexQueryResult {
+    pub index_name: String,
+    pub table_schema: String,
+    pub table_name: String,
+    pub column_name: String,
+}
+
+impl SchemaQueryBuilder {
+    pub fn query_table_unique_indexes(
+        &self,
+        schema: SeaRc<dyn Iden>,
+        table: SeaRc<dyn Iden>,
+    ) -> SelectStatement {
+        let idx = Alias::new("idx");
+        let insp = Alias::new("insp");
+        let tbl = Alias::new("tbl");
+        let tnsp = Alias::new("tnsp");
+        let col = Alias::new("col");
+
+        Query::select()
+            .column((idx.clone(), PgClass::RelName))
+            .column((insp.clone(), PgNamespace::NspName))
+            .column((tbl.clone(), PgClass::RelName))
+            .column((col.clone(), PgAttribute::AttName))
+            .from(PgIndex::Table)
+            .join_as(
+                JoinType::Join,
+                PgClass::Table,
+                idx.clone(),
+                Expr::col((idx.clone(), PgClass::Oid))
+                    .equals((PgIndex::Table, PgIndex::IndexRelId)),
+            )
+            .join_as(
+                JoinType::Join,
+                PgNamespace::Table,
+                insp.clone(),
+                Expr::col((insp.clone(), PgNamespace::Oid))
+                    .equals((idx.clone(), PgClass::RelNamespace)),
+            )
+            .join_as(
+                JoinType::Join,
+                PgClass::Table,
+                tbl.clone(),
+                Expr::col((tbl.clone(), PgClass::Oid)).equals((PgIndex::Table, PgIndex::IndRelId)),
+            )
+            .join_as(
+                JoinType::Join,
+                PgNamespace::Table,
+                tnsp.clone(),
+                Expr::col((tnsp.clone(), PgNamespace::Oid))
+                    .equals((tbl.clone(), PgClass::RelNamespace)),
+            )
+            .join_as(
+                JoinType::Join,
+                PgAttribute::Table,
+                col.clone(),
+                Expr::col((col.clone(), PgAttribute::AttRelId))
+                    .equals((idx.clone(), PgAttribute::Oid)),
+            )
+            .cond_where(
+                Condition::all()
+                    .add(Expr::col((PgIndex::Table, PgIndex::IndIsUnique)).eq(true))
+                    .add(Expr::col((PgIndex::Table, PgIndex::IndIsPrimary)).eq(false))
+                    .add(Expr::col((tbl.clone(), PgClass::RelName)).eq(table.to_string()))
+                    .add(Expr::col((tnsp.clone(), PgNamespace::NspName)).eq(schema.to_string())),
+            )
+            .order_by((PgIndex::Table, PgIndex::IndexRelId), Order::Asc)
+            .take()
+    }
+}
+
+#[cfg(feature = "sqlx-postgres")]
+impl From<&PgRow> for UniqueIndexQueryResult {
+    fn from(row: &PgRow) -> Self {
+        use crate::sqlx_types::Row;
+        Self {
+            index_name: row.get(0),
+            table_schema: row.get(1),
+            table_name: row.get(2),
+            column_name: row.get(3),
+        }
+    }
+}
+
+#[cfg(not(feature = "sqlx-postgres"))]
+impl From<&PgRow> for UniqueIndexQueryResult {
+    fn from(_: &PgRow) -> Self {
+        Self::default()
+    }
+}

src/sqlite/def/schema.rs

@@ -5,3 +5,16 @@ pub struct Schema {
     pub tables: Vec<TableDef>,
     pub indexes: Vec<IndexInfo>,
 }
+
+impl Schema {
+    pub fn merge_indexes_into_table(mut self) -> Self {
+        for table in self.tables.iter_mut() {
+            for index in self.indexes.iter() {
+                if index.unique && index.table_name == table.name {
+                    table.constraints.push(index.clone());
+                }
+            }
+        }
+        self
+    }
+}

tests/live/mysql/src/main.rs

@@ -287,6 +287,19 @@ fn create_lineitem_table() -> TableCreateStatement {
                 .col(Alias::new("order_id")),
         )
         .primary_key(Index::create().col(Alias::new("id")))
+        .index(
+            Index::create()
+                .unique()
+                .name("UNI_lineitem_cake_id")
+                .col(Alias::new("cake_id")),
+        )
+        .index(
+            Index::create()
+                .unique()
+                .name("UNI_lineitem_order_id_cake_id")
+                .col(Alias::new("order_id"))
+                .col(Alias::new("cake_id")),
+        )
         .foreign_key(
             ForeignKey::create()
                 .name("FK_lineitem_cake")

tests/live/postgres/src/main.rs

@@ -64,15 +64,6 @@ async fn main() {
         println!("{};", sql);
         println!();
         sqlx::query(&sql).execute(&mut *executor).await.unwrap();
-        if sql.starts_with(r#"CREATE TABLE "fkey_parent_table""#) {
-            let sql = Index::create()
-                .table(Alias::new("fkey_parent_table"))
-                .name("IDX_fkey_parent_table_unique_u")
-                .col(Alias::new("u"))
-                .unique()
-                .to_string(PostgresQueryBuilder);
-            sqlx::query(&sql).execute(&mut *executor).await.unwrap();
-        }
     }
 
     let schema_discovery = SchemaDiscovery::new(connection, "public");
@@ -302,6 +293,19 @@ fn create_lineitem_table() -> TableCreateStatement {
                 .name("lineitem_pkey")
                 .col(Alias::new("id")),
         )
+        .index(
+            Index::create()
+                .unique()
+                .name("UNI_lineitem_cake_id")
+                .col(Alias::new("cake_id")),
+        )
+        .index(
+            Index::create()
+                .unique()
+                .name("UNI_lineitem_cake_id_order_id")
+                .col(Alias::new("cake_id"))
+                .col(Alias::new("order_id")),
+        )
         .foreign_key(
             ForeignKey::create()
                 .name("FK_lineitem_cake")
@@ -478,6 +482,12 @@ fn create_fkey_parent_table() -> TableCreateStatement {
                 .auto_increment(),
         )
         .col(ColumnDef::new(Alias::new("u")).integer().not_null())
+        .index(
+            Index::create()
+                .unique()
+                .name("IDX_fkey_parent_table_unique_u")
+                .col(Alias::new("u")),
+        )
         .to_owned()
 }