Skip to content

Commit

Permalink
feat: return column counts on overview
Browse files Browse the repository at this point in the history
  • Loading branch information
@frectonz
frectonz committed Jun 25, 2024
1 parent 1724fc2 commit 0a3c445
Showing 1 changed file with 147 additions and 15 deletions.
162 changes: 147 additions & 15 deletions src/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -316,7 +316,7 @@ mod sqlite {
let modified = Some(metadata.modified()?.into());
let created = metadata.created().ok().map(Into::into);

let (tables, indexes, triggers, views, row_counts) = self
let (tables, indexes, triggers, views, row_counts, column_counts) = self
.conn
.call(move |conn| {
let tables = conn.query_row(
Expand Down Expand Up @@ -358,16 +358,16 @@ mod sqlite {
let mut stmt =
conn.prepare(r#"SELECT name FROM sqlite_master WHERE type="table""#)?;
let table_names = stmt.query_map([], |row| row.get::<_, String>(0))?;
let table_names = table_names.collect::<Result<Vec<_>, _>>()?;

let mut row_counts = HashMap::with_capacity(tables as usize);
for name in table_names {
let name = name?;
for name in table_names.iter() {
let count =
conn.query_row(&format!("SELECT count(*) FROM '{name}'"), (), |r| {
r.get::<_, i32>(0)
})?;

row_counts.insert(name, count);
row_counts.insert(name.to_owned(), count);
}

let mut row_counts = row_counts
Expand All @@ -377,7 +377,23 @@ mod sqlite {

row_counts.sort_by(|a, b| b.count.cmp(&a.count));

Ok((tables, indexes, triggers, views, row_counts))
let mut column_counts = HashMap::with_capacity(tables as usize);
for name in table_names.iter() {
let mut columns = conn.prepare(&format!("PRAGMA table_info('{name}')"))?;
let count =
columns.query_map((), |r| r.get::<_, String>(1))?.count() as i32;

column_counts.insert(name.to_owned(), count);
}

let mut column_counts = column_counts
.into_iter()
.map(|(name, count)| responses::Count { name, count })
.collect::<Vec<_>>();

column_counts.sort_by(|a, b| b.count.cmp(&a.count));

Ok((tables, indexes, triggers, views, row_counts, column_counts))
})
.await?;

Expand All @@ -392,6 +408,7 @@ mod sqlite {
triggers,
views,
row_counts,
column_counts,
})
}

Expand Down Expand Up @@ -741,7 +758,7 @@ mod libsql {
.collect::<Vec<_>>();

let mut row_counts = HashMap::with_capacity(table_names.len());
for name in table_names {
for name in table_names.iter() {
let count = conn
.query(&format!("SELECT count(*) FROM '{name}'"), ())
.await?
Expand All @@ -750,7 +767,7 @@ mod libsql {
.ok_or_eyre("no row returned from db")?
.get::<i32>(0)?;

row_counts.insert(name, count);
row_counts.insert(name.to_owned(), count);
}

let mut row_counts = row_counts
Expand All @@ -760,6 +777,30 @@ mod libsql {

row_counts.sort_by(|a, b| b.count.cmp(&a.count));

let mut column_counts = HashMap::with_capacity(table_names.len());
for name in table_names.iter() {
let count = conn
.query(&format!("PRAGMA table_info('{name}')"), ())
.await?
.into_stream()
.map_ok(|r| r.get::<String>(1))
.collect::<Vec<_>>()
.await
.into_iter()
.filter_map(|r| r.ok())
.filter_map(|r| r.ok())
.count() as i32;

column_counts.insert(name.to_owned(), count);
}

let mut column_counts = column_counts
.into_iter()
.map(|(name, count)| responses::Count { name, count })
.collect::<Vec<_>>();

column_counts.sort_by(|a, b| b.count.cmp(&a.count));

Ok(responses::Overview {
file_name,
sqlite_version: Some(sqlite_version),
Expand All @@ -771,6 +812,7 @@ mod libsql {
triggers,
views,
row_counts,
column_counts,
})
}

Expand Down Expand Up @@ -1137,7 +1179,7 @@ mod postgres {
.await?
.get(0);

let mut table_counts = self
let mut row_counts = self
.client
.query(
r#"
Expand All @@ -1155,7 +1197,7 @@ mod postgres {
})
.collect::<Vec<_>>();

for table in table_counts.iter_mut() {
for table in row_counts.iter_mut() {
let count: i64 = self
.client
.query_one(&format!(r#"SELECT count(*) FROM "{}""#, table.name), &[])
Expand All @@ -1164,7 +1206,48 @@ mod postgres {
table.count = count as i32;
}

table_counts.sort_by(|a, b| b.count.cmp(&a.count));
row_counts.sort_by(|a, b| b.count.cmp(&a.count));

let mut column_counts = self
.client
.query(
r#"
SELECT relname
FROM pg_stat_user_tables
WHERE schemaname = 'public'
"#,
&[],
)
.await?
.into_iter()
.map(|r| Count {
name: r.get(0),
count: 0,
})
.collect::<Vec<_>>();

for table in column_counts.iter_mut() {
let count: i64 = self
.client
.query_one(
&format!(
r#"
SELECT count(*)
FROM information_schema.columns
WHERE table_schema = 'public'
AND table_name = '{}'
"#,
table.name
),
&[],
)
.await?
.get(0);

table.count = count as i32;
}

column_counts.sort_by(|a, b| b.count.cmp(&a.count));

Ok(responses::Overview {
file_name,
Expand All @@ -1176,7 +1259,8 @@ mod postgres {
indexes: indexes as i32,
triggers: triggers as i32,
views: views as i32,
row_counts: table_counts,
row_counts,
column_counts,
})
}

Expand Down Expand Up @@ -1504,6 +1588,32 @@ mod mysql {

row_counts.sort_by(|a, b| b.count.cmp(&a.count));

let mut column_counts = r#"
SELECT TABLE_NAME AS name
FROM information_schema.tables
WHERE table_schema = database()
"#
.with(())
.map(&mut conn, |name| Count { name, count: 0 })
.await?;

for count in column_counts.iter_mut() {
count.count = r#"
SELECT count(*) AS count
FROM information_schema.columns
WHERE table_schema = database() AND table_name = :table_name
"#
.with(params! {
"table_name" => &count.name
})
.first(&mut conn)
.await?
.map(|count: i32| count)
.ok_or_eyre("couldn't count columns")?;
}

column_counts.sort_by(|a, b| b.count.cmp(&a.count));

Ok(responses::Overview {
file_name,
sqlite_version: None,
Expand All @@ -1515,6 +1625,7 @@ mod mysql {
triggers,
views,
row_counts,
column_counts,
})
}

Expand Down Expand Up @@ -1793,7 +1904,7 @@ mod duckdb {
let created = metadata.created().ok().map(Into::into);

let c = self.conn.clone();
let (tables, indexes, triggers, views, row_counts) =
let (tables, indexes, triggers, views, row_counts, column_counts) =
tokio::task::spawn_blocking(move || {
let c = c.lock().expect("could not get lock on connection");

Expand Down Expand Up @@ -1843,18 +1954,37 @@ mod duckdb {
.collect::<Vec<String>>();

let mut row_counts = Vec::with_capacity(table_names.len());
for name in table_names {
for name in table_names.iter() {
let count: i32 =
c.query_row(&format!(r#"SELECT count(*) FROM "{name}""#), [], |row| {
row.get(0)
})?;

row_counts.push(Count { name, count });
row_counts.push(Count {
name: name.to_owned(),
count,
});
}

row_counts.sort_by(|a, b| b.count.cmp(&a.count));

eyre::Ok((tables, indexes, triggers, views, row_counts))
let mut column_counts = Vec::with_capacity(table_names.len());
for name in table_names.iter() {
let count: i32 = c.query_row(
"SELECT column_count FROM duckdb_tables WHERE table_name = ?",
[&name],
|row| row.get(0),
)?;

column_counts.push(Count {
name: name.to_owned(),
count,
});
}

column_counts.sort_by(|a, b| b.count.cmp(&a.count));

eyre::Ok((tables, indexes, triggers, views, row_counts, column_counts))
})
.await??;

Expand All @@ -1869,6 +1999,7 @@ mod duckdb {
triggers,
views,
row_counts,
column_counts,
})
}

Expand Down Expand Up @@ -2120,6 +2251,7 @@ mod responses {
pub triggers: i32,
pub views: i32,
pub row_counts: Vec<Count>,
pub column_counts: Vec<Count>,
}

#[derive(Serialize)]
Expand Down

0 comments on commit 0a3c445

Please sign in to comment.