From 3a32f531a0c404aec4de26869875a7c4c6565c29 Mon Sep 17 00:00:00 2001
From: Dan Sosedoff <dan.sosedoff@gmail.com>
Date: Sat, 25 May 2024 10:59:23 -0700
Subject: [PATCH] Allow database stats downloads (#738)

* Add button to download database stats
* Return xml as well
---
 pkg/api/api.go     | 39 +++++++++++++++++++++++++++++++++++++--
 pkg/api/helpers.go |  7 ++++---
 static/index.html  |  4 +++-
 static/js/app.js   | 13 ++++++++++---
 4 files changed, 54 insertions(+), 9 deletions(-)

diff --git a/pkg/api/api.go b/pkg/api/api.go
index 6d89581..c1d8b4f 100644
--- a/pkg/api/api.go
+++ b/pkg/api/api.go
@@ -506,8 +506,43 @@ func GetTableConstraints(c *gin.Context) {
 
 // GetTablesStats renders data sizes and estimated rows for all tables in the database
 func GetTablesStats(c *gin.Context) {
-	res, err := DB(c).TablesStats()
-	serveResult(c, res, err)
+	db := DB(c)
+
+	connCtx, err := db.GetConnContext()
+	if err != nil {
+		badRequest(c, err)
+		return
+	}
+
+	res, err := db.TablesStats()
+	if err != nil {
+		badRequest(c, err)
+		return
+	}
+
+	format := getQueryParam(c, "format")
+	if format == "" {
+		format = "json"
+	}
+
+	// Save as attachment if exporting parameter is set
+	if getQueryParam(c, "export") == "true" {
+		ts := time.Now().Format(time.DateOnly)
+
+		filename := fmt.Sprintf("pgweb-dbstats-%s-%s.%s", connCtx.Database, ts, format)
+		c.Writer.Header().Set("Content-disposition", "attachment;filename="+filename)
+	}
+
+	switch format {
+	case "json":
+		c.JSON(http.StatusOK, res)
+	case "csv":
+		c.Data(http.StatusOK, "text/csv", res.CSV())
+	case "xml":
+		c.XML(200, res)
+	default:
+		badRequest(c, "invalid format")
+	}
 }
 
 // HandleQuery runs the database query
diff --git a/pkg/api/helpers.go b/pkg/api/helpers.go
index ed7be10..8e0fcb9 100644
--- a/pkg/api/helpers.go
+++ b/pkg/api/helpers.go
@@ -155,11 +155,12 @@ func assetContentType(name string) string {
 
 // Send a query result to client
 func serveResult(c *gin.Context, result interface{}, err interface{}) {
-	if err == nil {
-		successResponse(c, result)
-	} else {
+	if err != nil {
 		badRequest(c, err)
+		return
 	}
+
+	successResponse(c, result)
 }
 
 // Send successful response back to client
diff --git a/static/index.html b/static/index.html
index 0697b3d..d92fa7d 100644
--- a/static/index.html
+++ b/static/index.html
@@ -324,7 +324,9 @@
   </div>
   <div id="current_database_context_menu">
     <ul class="dropdown-menu" role="menu">
-      <li><a href="#" data-action="show_tables_stats">Show Tables Stats</a></li>
+      <li><a href="#" data-action="show_db_stats">Show Database Stats</a></li>
+      <li><a href="#" data-action="download_db_stats">Download Database Stats</a></li>
+      <li class="divider"></li>
       <li><a href="#" data-action="export">Export SQL dump</a></li>
     </ul>
   </div>
diff --git a/static/js/app.js b/static/js/app.js
index b8ad384..50e3fb5 100644
--- a/static/js/app.js
+++ b/static/js/app.js
@@ -660,7 +660,7 @@ function showPaginatedTableContent() {
   showTableContent(sortColumn, sortOrder);
 }
 
-function showTablesStats() {
+function showDatabaseStats() {
   getTablesStats(function(data) {
     buildTable(data);
 
@@ -671,6 +671,10 @@ function showTablesStats() {
   });
 }
 
+function downloadDatabaseStats() {
+  openInNewWindow("api/tables_stats", { format: "csv", export: "true" });
+}
+
 function showTableStructure() {
   var name = getCurrentObject().name;
 
@@ -1267,8 +1271,11 @@ function bindCurrentDatabaseMenu() {
       var menuItem = $(e.target);
 
       switch(menuItem.data("action")) {
-        case "show_tables_stats":
-          showTablesStats();
+        case "show_db_stats":
+          showDatabaseStats();
+          break;
+        case "download_db_stats":
+          downloadDatabaseStats();
           break;
         case "export":
           openInNewWindow("api/export");