Refactor api package

pkg/api/api.go

@@ -19,10 +19,14 @@ import (
 )
 
 var (
-	DbClient   *client.Client
+	// DbClient represents the active database connection in a single-session mode
+	DbClient *client.Client
+
+	// DbSessions represents the mapping for client connections
 	DbSessions = map[string]*client.Client{}
 )
 
+// DB returns a database connection from the client context
 func DB(c *gin.Context) *client.Client {
 	if command.Opts.Sessions {
 		return DbSessions[getSessionId(c.Request)]
@@ -63,11 +67,10 @@ func GetSessions(c *gin.Context) {
 	// In debug mode endpoint will return a lot of sensitive information
 	// like full database connection string and all query history.
 	if command.Opts.Debug {
-		c.JSON(200, DbSessions)
+		successResponse(c, DbSessions)
 		return
 	}
-
+	successResponse(c, gin.H{"sessions": len(DbSessions)})
-	c.JSON(200, map[string]int{"sessions": len(DbSessions)})
 }
 
 func ConnectWithBackend(c *gin.Context) {
@@ -81,14 +84,14 @@ func ConnectWithBackend(c *gin.Context) {
 	// Fetch connection credentials
 	cred, err := backend.FetchCredential(c.Param("resource"), c)
 	if err != nil {
-		c.JSON(400, Error{err.Error()})
+		badRequest(c, err)
 		return
 	}
 
 	// Make the new session
 	sessionId, err := securerandom.Uuid()
 	if err != nil {
-		c.JSON(400, Error{err.Error()})
+		badRequest(c, err)
 		return
 	}
 	c.Request.Header.Add("x-session-id", sessionId)
@@ -96,7 +99,7 @@ func ConnectWithBackend(c *gin.Context) {
 	// Connect to the database
 	cl, err := client.NewFromUrl(cred.DatabaseUrl, nil)
 	if err != nil {
-		c.JSON(400, Error{err.Error()})
+		badRequest(c, err)
 		return
 	}
 	cl.External = true
@@ -108,7 +111,7 @@ func ConnectWithBackend(c *gin.Context) {
 	}
 	if err != nil {
 		cl.Close()
-		c.JSON(400, Error{err.Error()})
+		badRequest(c, err)
 		return
 	}
 
@@ -117,7 +120,7 @@ func ConnectWithBackend(c *gin.Context) {
 
 func Connect(c *gin.Context) {
 	if command.Opts.LockSession {
-		c.JSON(400, Error{"Session is locked"})
+		badRequest(c, "Session is locked")
 		return
 	}
 
@@ -125,7 +128,7 @@ func Connect(c *gin.Context) {
 	url := c.Request.FormValue("url")
 
 	if url == "" {
-		c.JSON(400, Error{"Url parameter is required"})
+		badRequest(c, "Url parameter is required")
 		return
 	}
 
@@ -133,7 +136,7 @@ func Connect(c *gin.Context) {
 	url, err := connection.FormatUrl(opts)
 
 	if err != nil {
-		c.JSON(400, Error{err.Error()})
+		badRequest(c, err)
 		return
 	}
 
@@ -143,13 +146,13 @@ func Connect(c *gin.Context) {
 
 	cl, err := client.NewFromUrl(url, sshInfo)
 	if err != nil {
-		c.JSON(400, Error{err.Error()})
+		badRequest(c, err)
 		return
 	}
 
 	err = cl.Test()
 	if err != nil {
-		c.JSON(400, Error{err.Error()})
+		badRequest(c, err)
 		return
 	}
 
@@ -159,16 +162,16 @@ func Connect(c *gin.Context) {
 	}
 	if err != nil {
 		cl.Close()
-		c.JSON(400, Error{err.Error()})
+		badRequest(c, err)
 		return
 	}
 
-	c.JSON(200, info.Format()[0])
+	successResponse(c, info.Format()[0])
 }
 
 func SwitchDb(c *gin.Context) {
 	if command.Opts.LockSession {
-		c.JSON(400, Error{"Session is locked"})
+		badRequest(c, "Session is locked")
 		return
 	}
 
@@ -177,25 +180,25 @@ func SwitchDb(c *gin.Context) {
 		name = c.Request.FormValue("db")
 	}
 	if name == "" {
-		c.JSON(400, Error{"Database name is not provided"})
+		badRequest(c, "Database name is not provided")
 		return
 	}
 
 	conn := DB(c)
 	if conn == nil {
-		c.JSON(400, Error{"Not connected"})
+		badRequest(c, "Not connected")
 		return
 	}
 
 	// Do not allow switching databases for connections from third-party backends
 	if conn.External {
-		c.JSON(400, Error{"Session is locked"})
+		badRequest(c, "Session is locked")
 		return
 	}
 
 	currentUrl, err := neturl.Parse(conn.ConnectionString)
 	if err != nil {
-		c.JSON(400, Error{"Unable to parse current connection string"})
+		badRequest(c, "Unable to parse current connection string")
 		return
 	}
 
@@ -203,13 +206,13 @@ func SwitchDb(c *gin.Context) {
 
 	cl, err := client.NewFromUrl(currentUrl.String(), nil)
 	if err != nil {
-		c.JSON(400, Error{err.Error()})
+		badRequest(c, err)
 		return
 	}
 
 	err = cl.Test()
 	if err != nil {
-		c.JSON(400, Error{err.Error()})
+		badRequest(c, err)
 		return
 	}
 
@@ -219,64 +222,62 @@ func SwitchDb(c *gin.Context) {
 	}
 	if err != nil {
 		cl.Close()
-		c.JSON(400, Error{err.Error()})
+		badRequest(c, err)
 		return
 	}
 
 	conn.Close()
 
-	c.JSON(200, info.Format()[0])
+	successResponse(c, info.Format()[0])
 }
 
 func Disconnect(c *gin.Context) {
 	if command.Opts.LockSession {
-		c.JSON(400, Error{"Session is locked"})
+		badRequest(c, "Session is locked")
 		return
 	}
 
 	conn := DB(c)
 
 	if conn == nil {
-		c.JSON(400, Error{"Not connected"})
+		badRequest(c, "Not connected")
 		return
 	}
 
 	err := conn.Close()
 	if err != nil {
-		c.JSON(400, Error{err.Error()})
+		badRequest(c, err)
 		return
 	}
 
-	c.JSON(200, map[string]bool{"success": true})
+	successResponse(c, gin.H{"success": true})
 }
 
 func GetDatabases(c *gin.Context) {
 	conn := DB(c)
 	if conn.External {
-		c.JSON(403, Error{"Not permitted"})
+		errorResponse(c, 403, "Not permitted")
 		return
 	}
 
 	names, err := DB(c).Databases()
-	serveResult(names, err, c)
+	serveResult(c, names, err)
 }
 
 func GetObjects(c *gin.Context) {
 	result, err := DB(c).Objects()
 	if err != nil {
-		c.JSON(400, NewError(err))
+		badRequest(c, err)
 		return
 	}
-
+	successResponse(c, client.ObjectsFromResult(result))
-	objects := client.ObjectsFromResult(result)
-	c.JSON(200, objects)
 }
 
 func RunQuery(c *gin.Context) {
 	query := cleanQuery(c.Request.FormValue("query"))
 
 	if query == "" {
-		c.JSON(400, NewError(errors.New("Query parameter is missing")))
+		badRequest(c, "Query parameter is missing")
 		return
 	}
 
@@ -287,7 +288,7 @@ func ExplainQuery(c *gin.Context) {
 	query := cleanQuery(c.Request.FormValue("query"))
 
 	if query == "" {
-		c.JSON(400, NewError(errors.New("Query parameter is missing")))
+		badRequest(c, "Query parameter is missing")
 		return
 	}
 
@@ -296,7 +297,7 @@ func ExplainQuery(c *gin.Context) {
 
 func GetSchemas(c *gin.Context) {
 	res, err := DB(c).Schemas()
-	serveResult(res, err, c)
+	serveResult(c, res, err)
 }
 
 func GetTable(c *gin.Context) {
@@ -309,19 +310,19 @@ func GetTable(c *gin.Context) {
 		res, err = DB(c).Table(c.Params.ByName("table"))
 	}
 
-	serveResult(res, err, c)
+	serveResult(c, res, err)
 }
 
 func GetTableRows(c *gin.Context) {
 	offset, err := parseIntFormValue(c, "offset", 0)
 	if err != nil {
-		c.JSON(400, NewError(err))
+		badRequest(c, err)
 		return
 	}
 
 	limit, err := parseIntFormValue(c, "limit", 100)
 	if err != nil {
-		c.JSON(400, NewError(err))
+		badRequest(c, err)
 		return
 	}
 
@@ -335,13 +336,13 @@ func GetTableRows(c *gin.Context) {
 
 	res, err := DB(c).TableRows(c.Params.ByName("table"), opts)
 	if err != nil {
-		c.JSON(400, NewError(err))
+		badRequest(c, err)
 		return
 	}
 
 	countRes, err := DB(c).TableRowsCount(c.Params.ByName("table"), opts)
 	if err != nil {
-		c.JSON(400, NewError(err))
+		badRequest(c, err)
 		return
 	}
 
@@ -361,51 +362,49 @@ func GetTableRows(c *gin.Context) {
 		PerPage: numFetch,
 	}
 
-	serveResult(res, err, c)
+	serveResult(c, res, err)
 }
 
 func GetTableInfo(c *gin.Context) {
 	res, err := DB(c).TableInfo(c.Params.ByName("table"))
-
+	if err == nil {
-	if err != nil {
+		successResponse(c, res.Format()[0])
-		c.JSON(400, NewError(err))
+	} else {
-		return
+		badRequest(c, err)
 	}
-
-	c.JSON(200, res.Format()[0])
 }
 
 func GetHistory(c *gin.Context) {
-	c.JSON(200, DB(c).History)
+	successResponse(c, DB(c).History)
 }
 
 func GetConnectionInfo(c *gin.Context) {
 	res, err := DB(c).Info()
 
 	if err != nil {
-		c.JSON(400, NewError(err))
+		badRequest(c, err)
 		return
 	}
 
 	info := res.Format()[0]
 	info["session_lock"] = command.Opts.LockSession
 
-	c.JSON(200, info)
+	successResponse(c, info)
 }
 
 func GetActivity(c *gin.Context) {
 	res, err := DB(c).Activity()
-	serveResult(res, err, c)
+	serveResult(c, res, err)
 }
 
 func GetTableIndexes(c *gin.Context) {
 	res, err := DB(c).TableIndexes(c.Params.ByName("table"))
-	serveResult(res, err, c)
+	serveResult(c, res, err)
 }
 
 func GetTableConstraints(c *gin.Context) {
 	res, err := DB(c).TableConstraints(c.Params.ByName("table"))
-	serveResult(res, err, c)
+	serveResult(c, res, err)
 }
 
 func HandleQuery(query string, c *gin.Context) {
@@ -416,7 +415,7 @@ func HandleQuery(query string, c *gin.Context) {
 
 	result, err := DB(c).Query(query)
 	if err != nil {
-		c.JSON(400, NewError(err))
+		badRequest(c, err)
 		return
 	}
 
@@ -445,17 +444,15 @@ func HandleQuery(query string, c *gin.Context) {
 
 func GetBookmarks(c *gin.Context) {
 	bookmarks, err := bookmarks.ReadAll(bookmarks.Path(command.Opts.BookmarksDir))
-	serveResult(bookmarks, err, c)
+	serveResult(c, bookmarks, err)
 }
 
 func GetInfo(c *gin.Context) {
-	info := map[string]string{
+	successResponse(c, gin.H{
 		"version":    command.VERSION,
 		"git_sha":    command.GitCommit,
 		"build_time": command.BuildTime,
-	}
+	})
-
-	c.JSON(200, info)
 }
 
 // Export database or table data
@@ -464,7 +461,7 @@ func DataExport(c *gin.Context) {
 
 	info, err := db.Info()
 	if err != nil {
-		c.JSON(400, Error{err.Error()})
+		badRequest(c, err)
 		return
 	}
 
@@ -475,7 +472,7 @@ func DataExport(c *gin.Context) {
 	// If pg_dump is not available the following code will not show an error in browser.
 	// This is due to the headers being written first.
 	if !dump.CanExport() {
-		c.JSON(400, Error{"pg_dump is not found"})
+		badRequest(c, "pg_dump is not found")
 		return
 	}
 
@@ -485,11 +482,13 @@ func DataExport(c *gin.Context) {
 		filename = filename + "_" + dump.Table
 	}
 
-	attachment := fmt.Sprintf(`attachment; filename="%s.sql.gz"`, filename)
+	c.Header(
-	c.Header("Content-Disposition", attachment)
+		"Content-Disposition",
+		fmt.Sprintf(`attachment; filename="%s.sql.gz"`, filename),
+	)
 
 	err = dump.Export(db.ConnectionString, c.Writer)
 	if err != nil {
-		c.JSON(400, Error{err.Error()})
+		badRequest(c, err)
 	}
 }

pkg/api/helpers.go

@@ -14,30 +14,33 @@ import (
 	"github.com/sosedoff/pgweb/pkg/shared"
 )
 
-var extraMimeTypes = map[string]string{
+var (
-	".icon": "image-x-icon",
+	// Mime types definitions
-	".ttf":  "application/x-font-ttf",
+	extraMimeTypes = map[string]string{
-	".woff": "application/x-font-woff",
+		".icon": "image-x-icon",
-	".eot":  "application/vnd.ms-fontobject",
+		".ttf":  "application/x-font-ttf",
-	".svg":  "image/svg+xml",
+		".woff": "application/x-font-woff",
-	".html": "text/html; charset-utf-8",
+		".eot":  "application/vnd.ms-fontobject",
-}
+		".svg":  "image/svg+xml",
+		".html": "text/html; charset-utf-8",
+	}
 
-// Paths that dont require database connection
+	// Paths that dont require database connection
-var allowedPaths = map[string]bool{
+	allowedPaths = map[string]bool{
-	"/api/sessions":  true,
+		"/api/sessions":  true,
-	"/api/info":      true,
+		"/api/info":      true,
-	"/api/connect":   true,
+		"/api/connect":   true,
-	"/api/bookmarks": true,
+		"/api/bookmarks": true,
-	"/api/history":   true,
+		"/api/history":   true,
-}
+	}
 
-// List of characters replaced by javascript code to make queries url-safe.
+	// List of characters replaced by javascript code to make queries url-safe.
-var base64subs = map[string]string{
+	base64subs = map[string]string{
-	"-": "+",
+		"-": "+",
-	"_": "/",
+		"_": "/",
-	".": "=",
+		".": "=",
-}
+	}
+)
 
 type Error struct {
 	Message string `json:"error"`
@@ -151,11 +154,37 @@ func serveStaticAsset(path string, c *gin.Context) {
 	c.Data(200, assetContentType(path), data)
 }
 
-func serveResult(result interface{}, err error, c *gin.Context) {
+// Send a query result to client
-	if err != nil {
+func serveResult(c *gin.Context, result interface{}, err interface{}) {
-		c.JSON(400, NewError(err))
+	if err == nil {
-		return
+		successResponse(c, result)
+	} else {
+		badRequest(c, err)
+	}
+}
+
+// Send successful response back to client
+func successResponse(c *gin.Context, data interface{}) {
+	c.JSON(200, data)
+}
+
+// Send an error response back to client
+func errorResponse(c *gin.Context, status int, err interface{}) {
+	var message interface{}
+
+	switch v := err.(type) {
+	case error:
+		message = v.Error()
+	case string:
+		message = v
+	default:
+		message = v
 	}
 
-	c.JSON(200, result)
+	c.AbortWithStatusJSON(status, gin.H{"status": status, "error": message})
+}
+
+// Send a bad request (http 400) back to client
+func badRequest(c *gin.Context, err interface{}) {
+	errorResponse(c, 400, err)
 }

pkg/api/helpers_test.go

@@ -43,13 +43,13 @@ func Test_getSessionId(t *testing.T) {
 func Test_serveResult(t *testing.T) {
 	server := gin.Default()
 	server.GET("/good", func(c *gin.Context) {
-		serveResult(gin.H{"foo": "bar"}, nil, c)
+		serveResult(c, gin.H{"foo": "bar"}, nil)
 	})
 	server.GET("/bad", func(c *gin.Context) {
-		serveResult(nil, errors.New("message"), c)
+		serveResult(c, nil, errors.New("message"))
 	})
 	server.GET("/nodata", func(c *gin.Context) {
-		serveResult(nil, nil, c)
+		serveResult(c, nil, nil)
 	})
 
 	w := httptest.NewRecorder()
@@ -62,7 +62,7 @@ func Test_serveResult(t *testing.T) {
 	req, _ = http.NewRequest("GET", "/bad", nil)
 	server.ServeHTTP(w, req)
 	assert.Equal(t, 400, w.Code)
-	assert.Equal(t, `{"error":"message"}`, w.Body.String())
+	assert.Equal(t, `{"error":"message","status":400}`, w.Body.String())
 
 	w = httptest.NewRecorder()
 	req, _ = http.NewRequest("GET", "/nodata", nil)

pkg/api/middleware.go

@@ -9,21 +9,21 @@ import (
 	"github.com/sosedoff/pgweb/pkg/command"
 )
 
-// Middleware function to check database connection status before running queries
+// Middleware to check database connection status before running queries
 func dbCheckMiddleware() gin.HandlerFunc {
 	return func(c *gin.Context) {
 		path := strings.Replace(c.Request.URL.Path, command.Opts.Prefix, "", -1)
 
-		if allowedPaths[path] == true {
+		// Allow whitelisted paths
+		if allowedPaths[path] {
 			c.Next()
 			return
 		}
 
-		// We dont care about sessions unless they're enabled
+		// Check if session exists in single-session mode
 		if !command.Opts.Sessions {
 			if DbClient == nil {
-				c.JSON(400, Error{"Not connected"})
+				badRequest(c, "Not connected")
-				c.Abort()
 				return
 			}
 
@@ -31,17 +31,17 @@ func dbCheckMiddleware() gin.HandlerFunc {
 			return
 		}
 
+		// Determine session ID from the client request
 		sessionId := getSessionId(c.Request)
 		if sessionId == "" {
-			c.JSON(400, Error{"Session ID is required"})
+			badRequest(c, "Session ID is required")
-			c.Abort()
 			return
 		}
 
+		// Determine the database connection handle for the session
 		conn := DbSessions[sessionId]
 		if conn == nil {
-			c.JSON(400, Error{"Not connected"})
+			badRequest(c, "Not connected")
-			c.Abort()
 			return
 		}
 
@@ -49,7 +49,7 @@ func dbCheckMiddleware() gin.HandlerFunc {
 	}
 }
 
-// Middleware function to print out request parameters and body for debugging
+// Middleware to print out request parameters and body for debugging
 func requestInspectMiddleware() gin.HandlerFunc {
 	return func(c *gin.Context) {
 		err := c.Request.ParseForm()
@@ -57,6 +57,7 @@ func requestInspectMiddleware() gin.HandlerFunc {
 	}
 }
 
+// Middleware to inject CORS headers
 func corsMiddleware() gin.HandlerFunc {
 	return func(c *gin.Context) {
 		c.Header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")

pkg/api/session_cleanup.go

@@ -7,6 +7,16 @@ import (
 	"github.com/sosedoff/pgweb/pkg/command"
 )
 
+// StartSessionCleanup starts a goroutine to cleanup idle database sessions
+func StartSessionCleanup() {
+	for range time.Tick(time.Minute) {
+		if command.Opts.Debug {
+			log.Println("Triggering idle session deletion")
+		}
+		cleanupIdleSessions()
+	}
+}
+
 func cleanupIdleSessions() {
 	ids := []string{}
 
@@ -16,11 +26,11 @@ func cleanupIdleSessions() {
 			ids = append(ids, id)
 		}
 	}
-
-	// Close and delete idle sessions
 	if len(ids) == 0 {
 		return
 	}
+
+	// Close and delete idle sessions
 	log.Println("Closing", len(ids), "idle sessions")
 	for _, id := range ids {
 		// TODO: concurrent map edit will trigger panic
@@ -29,17 +39,3 @@ func cleanupIdleSessions() {
 		}
 	}
 }
-
-func StartSessionCleanup() {
-	ticker := time.NewTicker(time.Minute)
-
-	for {
-		<-ticker.C
-
-		if command.Opts.Debug {
-			log.Println("Triggering idle session deletion")
-		}
-
-		cleanupIdleSessions()
-	}
-}