From 51042aa536af8bda0cc2e789bbc36d9ba8c58caa Mon Sep 17 00:00:00 2001 From: Chris Tsang Date: Mon, 9 May 2022 17:44:54 +0800 Subject: [PATCH] More probes --- src/mysql/probe.rs | 8 ++++---- src/postgres/probe.rs | 8 ++++---- src/probe.rs | 34 +++++++++++++++++++++++++++++++++- src/sqlite/def/table.rs | 7 ++++--- src/sqlite/probe.rs | 22 ++++++++++++++++++---- 5 files changed, 63 insertions(+), 16 deletions(-) diff --git a/src/mysql/probe.rs b/src/mysql/probe.rs index 826d3c55..099694cb 100644 --- a/src/mysql/probe.rs +++ b/src/mysql/probe.rs @@ -9,7 +9,6 @@ impl SchemaProbe for MySql { } fn query_tables() -> SelectStatement { - let mut stmt = Query::select(); let (expr, tbl_ref, condition) = ( Expr::col(Alias::new("table_name")), (Alias::new("information_schema"), Alias::new("tables")).into_table_ref(), @@ -18,9 +17,10 @@ impl SchemaProbe for MySql { .equals(Alias::new("tables"), Alias::new("table_schema")), ), ); - stmt.expr_as(expr, Alias::new("table_name")) + Query::select() + .expr_as(expr, Alias::new("table_name")) .from(tbl_ref) - .cond_where(condition); - stmt + .cond_where(condition) + .take() } } diff --git a/src/postgres/probe.rs b/src/postgres/probe.rs index 18a8699a..cbf9916f 100644 --- a/src/postgres/probe.rs +++ b/src/postgres/probe.rs @@ -9,7 +9,6 @@ impl SchemaProbe for Postgres { } fn query_tables() -> SelectStatement { - let mut stmt = Query::select(); let (expr, tbl_ref, condition) = ( Expr::col(Alias::new("table_name")), (Alias::new("information_schema"), Alias::new("tables")).into_table_ref(), @@ -20,9 +19,10 @@ impl SchemaProbe for Postgres { ) .add(Expr::col(Alias::new("table_type")).eq("BASE TABLE")), ); - stmt.expr_as(expr, Alias::new("table_name")) + Query::select() + .expr_as(expr, Alias::new("table_name")) .from(tbl_ref) - .cond_where(condition); - stmt + .cond_where(condition) + .take() } } diff --git a/src/probe.rs b/src/probe.rs index 6e8de16e..e488050a 100644 --- a/src/probe.rs +++ b/src/probe.rs @@ -1,7 +1,39 @@ -use sea_query::{SelectStatement, SimpleExpr}; +use sea_query::{Alias, Condition, Expr, Query, SelectStatement, SimpleExpr}; pub trait SchemaProbe { fn get_current_schema() -> SimpleExpr; fn query_tables() -> SelectStatement; + + fn has_table(&self, table: T) -> SelectStatement + where + T: AsRef, + { + let mut subquery = Self::query_tables(); + subquery.cond_where(Expr::col(Alias::new("table_name")).eq(table.as_ref())); + Query::select() + .expr_as(Expr::cust("COUNT(*) > 0"), Alias::new("has_table")) + .from_subquery(subquery, Alias::new("subquery")) + .take() + } + + fn has_column(table: T, column: C) -> SelectStatement + where + T: AsRef, + C: AsRef, + { + Query::select() + .expr_as(Expr::cust("COUNT(*) > 0"), Alias::new("has_column")) + .from((Alias::new("information_schema"), Alias::new("columns"))) + .cond_where( + Condition::all() + .add( + Expr::expr(Self::get_current_schema()) + .equals(Alias::new("columns"), Alias::new("table_schema")), + ) + .add(Expr::col(Alias::new("table_name")).eq(table.as_ref())) + .add(Expr::col(Alias::new("column_name")).eq(column.as_ref())), + ) + .take() + } } diff --git a/src/sqlite/def/table.rs b/src/sqlite/def/table.rs index a81a4a80..4bd6bc1a 100644 --- a/src/sqlite/def/table.rs +++ b/src/sqlite/def/table.rs @@ -61,6 +61,7 @@ impl TableDef { Ok(self) } + /// Get a list of all the indexes in the table. /// Note that this does not get the column name mapped by the index. /// To get the column name mapped by the index, the `self.get_single_indexinfo` method is invoked @@ -70,7 +71,7 @@ impl TableDef { indexes: &mut Vec, ) -> DiscoveryResult<()> { let mut index_query = String::default(); - index_query.push_str("pragma index_list('"); + index_query.push_str("PRAGMA index_list('"); index_query.push_str(&self.name); index_query.push_str("')"); @@ -107,7 +108,7 @@ impl TableDef { /// Get a list of all the foreign keys in the table pub async fn get_foreign_keys(&mut self, executor: &Executor) -> DiscoveryResult<&mut Self> { let mut index_query = String::default(); - index_query.push_str("pragma foreign_key_list('"); + index_query.push_str("PRAGMA foreign_key_list('"); index_query.push_str(&self.name); index_query.push_str("')"); @@ -125,7 +126,7 @@ impl TableDef { /// Get a list of all the columns in the table mapped as [ColumnInfo] pub async fn get_column_info(&mut self, executor: &Executor) -> DiscoveryResult<&TableDef> { let mut index_query = String::default(); - index_query.push_str("pragma table_info('"); + index_query.push_str("PRAGMA table_info('"); index_query.push_str(&self.name); index_query.push_str("')"); diff --git a/src/sqlite/probe.rs b/src/sqlite/probe.rs index e43826a2..c202f993 100644 --- a/src/sqlite/probe.rs +++ b/src/sqlite/probe.rs @@ -9,7 +9,6 @@ impl SchemaProbe for Sqlite { } fn query_tables() -> SelectStatement { - let mut stmt = Query::select(); let (expr, tbl_ref, condition) = ( Expr::col(Alias::new("name")), Alias::new("sqlite_master").into_table_ref(), @@ -17,9 +16,24 @@ impl SchemaProbe for Sqlite { .add(Expr::col(Alias::new("type")).eq("table")) .add(Expr::col(Alias::new("name")).ne("sqlite_sequence")), ); - stmt.expr_as(expr, Alias::new("table_name")) + Query::select() + .expr_as(expr, Alias::new("table_name")) .from(tbl_ref) - .cond_where(condition); - stmt + .cond_where(condition) + .take() + } + + fn has_column(table: T, column: C) -> SelectStatement + where + T: AsRef, + C: AsRef, + { + Query::select() + .expr(Expr::cust_with_values( + "COUNT(*) > 0 AS \"has_column\" FROM pragma_table_info(?)", + [table.as_ref()], + )) + .and_where(Expr::col(Alias::new("name")).eq(column.as_ref())) + .take() } }