feat: limit query

Author: l1xnan

src-tauri/src/dialect/clickhouse.rs

@@ -69,6 +69,18 @@ impl Connection for ClickhouseDialect {
     log::info!("show columns: {}", &sql);
     self.query(&sql, 0, 0).await
   }
+
+  #[allow(clippy::unused_async)]
+  async fn query_count(&self, sql: &str) -> anyhow::Result<usize> {
+    let mut client = self.client().await?;
+    let block = client.query(sql).fetch_all().await?;
+    if let Some(row) = block.rows().next() {
+      let total = row.get::<u32, _>(0)?;
+      Ok(total as usize)
+    } else {
+      Err(anyhow::anyhow!("null"))
+    }
+  }
 }
 
 impl ClickhouseDialect {

src-tauri/src/dialect/mysql.rs

@@ -70,6 +70,16 @@ impl Connection for MySqlDialect {
     log::info!("show columns: {}", &sql);
     self.query(&sql, 0, 0).await
   }
+
+  #[allow(clippy::unused_async)]
+  async fn query_count(&self, sql: &str) -> anyhow::Result<usize> {
+    let mut conn = self.get_conn()?;
+    if let Some(total) = conn.query_first::<usize, _>(sql)? {
+      Ok(total)
+    } else {
+      Err(anyhow::anyhow!("null"))
+    }
+  }
 }
 
 impl MySqlDialect {

src-tauri/src/dialect/postgres.rs

@@ -72,6 +72,13 @@ impl Connection for PostgresDialect {
     log::info!("show columns: {}", &sql);
     self.query(&sql, 0, 0).await
   }
+
+  async fn query_count(&self, sql: &str) -> anyhow::Result<usize> {
+    let conn = self.get_conn(&self.database()).await?;
+    let row = conn.query_one(sql, &[]).await?;
+    let total: u32 = row.get(0);
+    Ok(total as usize)
+  }
 }
 
 impl PostgresDialect {
@@ -109,15 +116,15 @@ impl PostgresDialect {
   pub async fn get_tables(&self, db: &str) -> anyhow::Result<Vec<Table>> {
     let client = self.get_conn(db).await?;
 
-    let sql = r#"
-    select
-      table_catalog as db_name,
-      table_schema as table_schema,
-      table_name as table_name,
-      table_type as table_type,
-      CASE WHEN table_type='BASE TABLE' THEN 'table' ELSE 'view' END as type
-    from information_schema.tables WHERE table_schema='public'
-    "#;
+    let sql = "
+      select
+        table_catalog as db_name,
+        table_schema as table_schema,
+        table_name as table_name,
+        table_type as table_type,
+        CASE WHEN table_type='BASE TABLE' THEN 'table' ELSE 'view' END as type
+      from information_schema.tables WHERE table_schema='public'
+      ";
     let mut tables = vec![];
     for row in client.query(sql, &[]).await? {
       tables.push(Table {