Notebook query elapsed time (#644)

SilasMarvin · web-flow · commit 04f7e26c6514 · 2023-05-22T18:20:06.000-07:00
diff --git a/pgml-dashboard/src/lib.rs b/pgml-dashboard/src/lib.rs
@@ -19,6 +19,7 @@ pub mod guards;
 pub mod models;
 mod responses;
 mod templates;
+mod utils;
 
 use guards::Cluster;
 use responses::{BadRequest, ResponseOk};
@@ -558,6 +559,7 @@ pub async fn uploaded_index(cluster: Cluster, table_name: &str) -> ResponseOk {
     let sql = templates::Sql::new(
         cluster.pool(),
         &format!("SELECT * FROM {} LIMIT 10", table_name),
+        true, 
     )
     .await
     .unwrap();
diff --git a/pgml-dashboard/src/models.rs b/pgml-dashboard/src/models.rs
@@ -296,18 +296,19 @@ impl Cell {
     pub async fn render(&mut self, pool: &PgPool) -> anyhow::Result<()> {
         let cell_type: CellType = self.cell_type.into();
 
-        let rendering = match cell_type {
+        let (rendering, execution_time) = match cell_type {
             CellType::Sql => {
-                let queries = self.contents.split(";");
+                let queries: Vec<&str> = self.contents.split(';').filter(|q| !q.trim().is_empty()).collect();
                 let mut rendering = String::new();
+                let mut total_execution_duration = std::time::Duration::default(); 
+                let render_individual_execution_duration = queries.len() > 1;
 
                 for query in queries {
-                    if query.trim().is_empty() {
-                        continue;
-                    }
-
-                    let result = match templates::Sql::new(pool, query).await {
-                        Ok(sql) => sql.render_once()?,
+                    let result = match templates::Sql::new(pool, query, render_individual_execution_duration).await {
+                        Ok(sql) => {
+                            total_execution_duration += sql.execution_duration;
+                            sql.render_once()?
+                        },
                         Err(err) => templates::SqlError {
                             error: format!("{:?}", err),
                         }
@@ -317,7 +318,12 @@ impl Cell {
                     rendering.push_str(&result);
                 }
 
-                rendering
+                let execution_time = PgInterval{
+                    months: 0,
+                    days: 0,
+                    microseconds: total_execution_duration.as_micros().try_into().unwrap_or(0)
+                };
+                (rendering, Some(execution_time))
             }
 
             CellType::Markdown => {
@@ -335,21 +341,23 @@ impl Cell {
                     front_matter_delimiter: None,
                 };
 
-                format!(
+                (format!(
                     "<div class=\"markdown-body\">{}</div>",
                     markdown_to_html(&self.contents, &options)
-                )
+                ), None)
             }
         };
 
         sqlx::query!(
-            "UPDATE pgml.notebook_cells SET rendering = $1 WHERE id = $2",
+            "UPDATE pgml.notebook_cells SET rendering = $1, execution_time = $2 WHERE id = $3",
             rendering,
+            execution_time,
             self.id
         )
         .execute(pool)
         .await?;
 
+        self.execution_time = execution_time;
         self.rendering = Some(rendering);
 
         Ok(())
diff --git a/pgml-dashboard/src/templates.rs b/pgml-dashboard/src/templates.rs
@@ -55,10 +55,12 @@ pub struct Undo {
 pub struct Sql {
     pub columns: Vec<String>,
     pub rows: Vec<Vec<String>>,
+    pub execution_duration: std::time::Duration,
+    pub render_execution_duration: bool,
 }
 
 impl Sql {
-    pub async fn new(pool: &PgPool, query: &str) -> anyhow::Result<Sql> {
+    pub async fn new(pool: &PgPool, query: &str, render_execution_duration: bool) -> anyhow::Result<Sql> {
         let prepared_stmt = pool.prepare(query).await?;
         let cols = prepared_stmt.columns();
 
@@ -67,7 +69,9 @@ impl Sql {
 
         cols.iter().for_each(|c| columns.push(c.name().to_string()));
 
+        let now = std::time::Instant::now();
         let result = prepared_stmt.query().fetch_all(pool).await?;
+        let execution_duration = now.elapsed();
 
         for row in result.iter() {
             let mut values = Vec::new();
@@ -199,7 +203,7 @@ impl Sql {
             rows.push(values);
         }
 
-        Ok(Sql { columns, rows })
+        Ok(Sql { columns, rows, execution_duration, render_execution_duration })
     }
 }
 
diff --git a/pgml-dashboard/src/utils.rs b/pgml-dashboard/src/utils.rs
@@ -0,0 +1,9 @@
+pub fn format_microseconds(microseconds: f64) -> String {
+    if microseconds >= 1000000. {
+        format!("{}s", microseconds / 1000000.)
+    } else if microseconds >= 1000. {
+        format!("{}ms", microseconds / 1000.)
+    } else {
+        format!("{}μs", microseconds)
+    }
+}
diff --git a/pgml-dashboard/templates/cell.html b/pgml-dashboard/templates/cell.html
@@ -75,7 +75,7 @@
       <div class="notebook-duration">
         <div class="markdown-body">
           <div class="flex flex-row-reverse">
-            <code>TODO</code>
+            <code>Time: <%= crate::utils::format_microseconds(cell.execution_time.unwrap().microseconds as f64) %></code>
           </div>
         </div>
       </div>
diff --git a/pgml-dashboard/templates/sql.html b/pgml-dashboard/templates/sql.html
@@ -24,6 +24,9 @@
       <% } %>
     </tbody>
   </table>
+  <% if render_execution_duration { %>
+    <code>Time: <%= crate::utils::format_microseconds(execution_duration.as_micros() as f64) %></code>
+  <% } %>
 </div>
 
 <% } %>

Original file line number	Diff line number	Diff line change
`@@ -55,10 +55,12 @@ pub struct Undo {`
`55`	`55`	`pub struct Sql {`
`56`	`56`	`pub columns: Vec<String>,`
`57`	`57`	`pub rows: Vec<Vec<String>>,`
	`58`	`+ pub execution_duration: std::time::Duration,`
	`59`	`+ pub render_execution_duration: bool,`
`58`	`60`	`}`
`59`	`61`
`60`	`62`	`impl Sql {`
`61`		`- pub async fn new(pool: &PgPool, query: &str) -> anyhow::Result<Sql> {`
	`63`	`+ pub async fn new(pool: &PgPool, query: &str, render_execution_duration: bool) -> anyhow::Result<Sql> {`
`62`	`64`	`let prepared_stmt = pool.prepare(query).await?;`
`63`	`65`	`let cols = prepared_stmt.columns();`
`64`	`66`
`@@ -67,7 +69,9 @@ impl Sql {`
`67`	`69`
`68`	`70`	`cols.iter().for_each(\|c\| columns.push(c.name().to_string()));`
`69`	`71`
	`72`	`+ let now = std::time::Instant::now();`
`70`	`73`	`let result = prepared_stmt.query().fetch_all(pool).await?;`
	`74`	`+ let execution_duration = now.elapsed();`
`71`	`75`
`72`	`76`	`for row in result.iter() {`
`73`	`77`	`let mut values = Vec::new();`
`@@ -199,7 +203,7 @@ impl Sql {`
`199`	`203`	`rows.push(values);`
`200`	`204`	`}`
`201`	`205`
`202`		`- Ok(Sql { columns, rows })`
	`206`	`+ Ok(Sql { columns, rows, execution_duration, render_execution_duration })`
`203`	`207`	`}`
`204`	`208`	`}`
`205`	`209`