Add support for user functions (#608)

* Add initial support for functions
* Show functions definitions
* Fix client tests
* Fix schema objects search
* Perform partial matching for functions
* Add function test
* Make sure to close client connections so that database could be dropped in tests
* Fix lint
* Allow to copy the view/functions definitions
* Nits

Makefile

@@ -28,7 +28,7 @@ usage:
 	@echo ""
 
 test:
-	go test -race -cover ./pkg/...
+	go test -v -race -cover ./pkg/...
 
 test-all:
 	@./script/test_all.sh

pkg/api/api.go

@@ -342,13 +342,21 @@ func GetSchemas(c *gin.Context) {
 
 // GetTable renders table information
 func GetTable(c *gin.Context) {
-	var res *client.Result
+	var (
-	var err error
+		res *client.Result
+		err error
+	)
 
-	if c.Request.FormValue("type") == client.ObjTypeMaterializedView {
+	db := DB(c)
-		res, err = DB(c).MaterializedView(c.Params.ByName("table"))
+	tableName := c.Params.ByName("table")
-	} else {
+
-		res, err = DB(c).Table(c.Params.ByName("table"))
+	switch c.Request.FormValue("type") {
+	case client.ObjTypeMaterializedView:
+		res, err = db.MaterializedView(tableName)
+	case client.ObjTypeFunction:
+		res, err = db.Function(tableName)
+	default:
+		res, err = db.Table(tableName)
 	}
 
 	serveResult(c, res, err)
@@ -541,3 +549,9 @@ func DataExport(c *gin.Context) {
 		badRequest(c, err)
 	}
 }
+
+// GetFunction renders function information
+func GetFunction(c *gin.Context) {
+	res, err := DB(c).Function(c.Param("id"))
+	serveResult(c, res, err)
+}

pkg/api/routes.go

@@ -42,6 +42,7 @@ func SetupRoutes(router *gin.Engine) {
 	api.GET("/tables/:table/info", GetTableInfo)
 	api.GET("/tables/:table/indexes", GetTableIndexes)
 	api.GET("/tables/:table/constraints", GetTableConstraints)
+	api.GET("/functions/:id", GetFunction)
 	api.GET("/query", RunQuery)
 	api.POST("/query", RunQuery)
 	api.GET("/explain", ExplainQuery)

pkg/client/client.go

@@ -197,6 +197,10 @@ func (client *Client) MaterializedView(name string) (*Result, error) {
 	return client.query(statements.MaterializedView, name)
 }
 
+func (client *Client) Function(id string) (*Result, error) {
+	return client.query(statements.Function, id)
+}
+
 func (client *Client) TableRows(table string, opts RowsOptions) (*Result, error) {
 	schema, table := getSchemaAndTable(table)
 	sql := fmt.Sprintf(`SELECT * FROM "%s"."%s"`, schema, table)

pkg/client/client_test.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"os/exec"
 	"runtime"
+	"sort"
 	"testing"
 	"time"
 
@@ -32,6 +33,26 @@ func mapKeys(data map[string]*Objects) []string {
 	return result
 }
 
+func objectNames(data []Object) []string {
+	names := make([]string, len(data))
+	for i, obj := range data {
+		names[i] = obj.Name
+	}
+
+	sort.Strings(names)
+	return names
+}
+
+// assertMatches is a helper method to check if src slice contains any elements of expected slice
+func assertMatches(t *testing.T, expected, src []string) {
+	assert.NotEqual(t, 0, len(expected))
+	assert.NotEqual(t, 0, len(src))
+
+	for _, val := range expected {
+		assert.Contains(t, src, val)
+	}
+}
+
 func pgVersion() (int, int) {
 	var major, minor int
 	fmt.Sscanf(os.Getenv("PGVERSION"), "%d.%d", &major, &minor)
@@ -118,12 +139,12 @@ func setupClient() {
 
 func teardownClient() {
 	if testClient != nil {
-		testClient.db.Close()
+		testClient.Close()
 	}
 }
 
 func teardown() {
-	_, err := exec.Command(
+	output, err := exec.Command(
 		testCommands["dropdb"],
 		"-U", serverUser,
 		"-h", serverHost,
@@ -133,31 +154,28 @@ func teardown() {
 
 	if err != nil {
 		fmt.Println("Teardown error:", err)
+		fmt.Printf("%s\n", output)
 	}
 }
 
-func testNewClientFromUrl(t *testing.T) {
+func testNewClientFromURL(t *testing.T) {
-	url := fmt.Sprintf("postgres://%s@%s:%s/%s?sslmode=disable", serverUser, serverHost, serverPort, serverDatabase)
+	t.Run("postgres prefix", func(t *testing.T) {
-	client, err := NewFromUrl(url, nil)
+		url := fmt.Sprintf("postgres://%s@%s:%s/%s?sslmode=disable", serverUser, serverHost, serverPort, serverDatabase)
+		client, err := NewFromUrl(url, nil)
 
-	if err != nil {
+		assert.Equal(t, nil, err)
-		defer client.Close()
+		assert.Equal(t, url, client.ConnectionString)
-	}
+		assert.NoError(t, client.Close())
+	})
 
-	assert.Equal(t, nil, err)
+	t.Run("postgresql prefix", func(t *testing.T) {
-	assert.Equal(t, url, client.ConnectionString)
+		url := fmt.Sprintf("postgresql://%s@%s:%s/%s?sslmode=disable", serverUser, serverHost, serverPort, serverDatabase)
-}
+		client, err := NewFromUrl(url, nil)
 
-func testNewClientFromUrl2(t *testing.T) {
+		assert.Equal(t, nil, err)
-	url := fmt.Sprintf("postgresql://%s@%s:%s/%s?sslmode=disable", serverUser, serverHost, serverPort, serverDatabase)
+		assert.Equal(t, url, client.ConnectionString)
-	client, err := NewFromUrl(url, nil)
+		assert.NoError(t, client.Close())
-
+	})
-	if err != nil {
-		defer client.Close()
-	}
-
-	assert.Equal(t, nil, err)
-	assert.Equal(t, url, client.ConnectionString)
 }
 
 func testClientIdleTime(t *testing.T) {
@@ -202,16 +220,13 @@ func testActivity(t *testing.T) {
 
 	res, err := testClient.Activity()
 	assert.NoError(t, err)
-	for _, val := range expected {
+	assertMatches(t, expected, res.Columns)
-		assert.Contains(t, res.Columns, val)
-	}
 }
 
 func testDatabases(t *testing.T) {
 	res, err := testClient.Databases()
 	assert.NoError(t, err)
-	assert.Contains(t, res, "booktown")
+	assertMatches(t, []string{"booktown", "postgres"}, res)
-	assert.Contains(t, res, "postgres")
 }
 
 func testObjects(t *testing.T) {
@@ -245,16 +260,44 @@ func testObjects(t *testing.T) {
 		"text_sorting",
 	}
 
+	functions := []string{
+		"add_shipment",
+		"add_two_loop",
+		"books_by_subject",
+		"compound_word",
+		"count_by_two",
+		"double_price",
+		"extract_all_titles",
+		"extract_all_titles2",
+		"extract_title",
+		"first",
+		"get_author",
+		"get_author",
+		"get_customer_id",
+		"get_customer_name",
+		"html_linebreaks",
+		"in_stock",
+		"isbn_to_title",
+		"mixed",
+		"raise_test",
+		"ship_item",
+		"stock_amount",
+		"test",
+		"title",
+		"triple_price",
+	}
+
 	assert.NoError(t, err)
-	assert.Equal(t, []string{"schema", "name", "type", "owner", "comment"}, res.Columns)
+	assert.Equal(t, []string{"oid", "schema", "name", "type", "owner", "comment"}, res.Columns)
 	assert.Equal(t, []string{"public"}, mapKeys(objects))
-	assert.Equal(t, tables, objects["public"].Tables)
+	assert.Equal(t, tables, objectNames(objects["public"].Tables))
-	assert.Equal(t, []string{"recent_shipments", "stock_view"}, objects["public"].Views)
+	assertMatches(t, functions, objectNames(objects["public"].Functions))
-	assert.Equal(t, []string{"author_ids", "book_ids", "shipments_ship_id_seq", "subject_ids"}, objects["public"].Sequences)
+	assert.Equal(t, []string{"recent_shipments", "stock_view"}, objectNames(objects["public"].Views))
+	assert.Equal(t, []string{"author_ids", "book_ids", "shipments_ship_id_seq", "subject_ids"}, objectNames(objects["public"].Sequences))
 
 	major, minor := pgVersion()
 	if minor == 0 || minor >= 3 {
-		assert.Equal(t, []string{"m_stock_view"}, objects["public"].MaterializedViews)
+		assert.Equal(t, []string{"m_stock_view"}, objectNames(objects["public"].MaterializedViews))
 	} else {
 		t.Logf("Skipping materialized view on %d.%d\n", major, minor)
 	}
@@ -428,6 +471,33 @@ func testTableRowsOrderEscape(t *testing.T) {
 	assert.Nil(t, rows)
 }
 
+func testFunctions(t *testing.T) {
+	funcName := "get_customer_name"
+	funcID := ""
+
+	res, err := testClient.Objects()
+	assert.NoError(t, err)
+
+	for _, row := range res.Rows {
+		if row[2] == funcName {
+			funcID = row[0].(string)
+			break
+		}
+	}
+
+	res, err = testClient.Function("12345")
+	assert.NoError(t, err)
+	assertMatches(t, []string{"oid", "proname", "functiondef"}, res.Columns)
+	assert.Equal(t, 0, len(res.Rows))
+
+	res, err = testClient.Function(funcID)
+	assert.NoError(t, err)
+	assertMatches(t, []string{"oid", "proname", "functiondef"}, res.Columns)
+	assert.Equal(t, 1, len(res.Rows))
+	assert.Equal(t, funcName, res.Rows[0][1])
+	assert.Contains(t, res.Rows[0][len(res.Columns)-1], "SELECT INTO customer_fname, customer_lname")
+}
+
 func testResult(t *testing.T) {
 	t.Run("json", func(t *testing.T) {
 		result, err := testClient.Query("SELECT * FROM books LIMIT 1")
@@ -466,8 +536,8 @@ func testHistory(t *testing.T) {
 	t.Run("unique queries", func(t *testing.T) {
 		url := fmt.Sprintf("postgres://%s@%s:%s/%s?sslmode=disable", serverUser, serverHost, serverPort, serverDatabase)
 
-		client, err := NewFromUrl(url, nil)
+		client, _ := NewFromUrl(url, nil)
-		assert.NoError(t, err)
+		defer client.Close()
 
 		for i := 0; i < 3; i++ {
 			_, err := client.Query("SELECT * FROM books WHERE id = 1")
@@ -487,6 +557,7 @@ func testReadOnlyMode(t *testing.T) {
 
 	url := fmt.Sprintf("postgres://%s@%s:%s/%s?sslmode=disable", serverUser, serverHost, serverPort, serverDatabase)
 	client, _ := NewFromUrl(url, nil)
+	defer client.Close()
 
 	err := client.SetReadOnlyMode()
 	assert.NoError(t, err)
@@ -522,8 +593,7 @@ func TestAll(t *testing.T) {
 	setup()
 	setupClient()
 
-	testNewClientFromUrl(t)
+	testNewClientFromURL(t)
-	testNewClientFromUrl2(t)
 	testClientIdleTime(t)
 	testTest(t)
 	testInfo(t)
@@ -544,6 +614,7 @@ func TestAll(t *testing.T) {
 	testQueryError(t)
 	testQueryInvalidTable(t)
 	testTableRowsOrderEscape(t)
+	testFunctions(t)
 	testResult(t)
 	testHistory(t)
 	testReadOnlyMode(t)

pkg/client/result.go

@@ -18,6 +18,7 @@ const (
 	ObjTypeView             = "view"
 	ObjTypeMaterializedView = "materialized_view"
 	ObjTypeSequence         = "sequence"
+	ObjTypeFunction         = "function"
 )
 
 type (
@@ -36,11 +37,17 @@ type (
 		Rows       []Row       `json:"rows"`
 	}
 
+	Object struct {
+		OID  string `json:"oid"`
+		Name string `json:"name"`
+	}
+
 	Objects struct {
-		Tables            []string `json:"table"`
+		Tables            []Object `json:"table"`
-		Views             []string `json:"view"`
+		Views             []Object `json:"view"`
-		MaterializedViews []string `json:"materialized_view"`
+		MaterializedViews []Object `json:"materialized_view"`
-		Sequences         []string `json:"sequence"`
+		Functions         []Object `json:"function"`
+		Sequences         []Object `json:"sequence"`
 	}
 )
 
@@ -154,28 +161,34 @@ func ObjectsFromResult(res *Result) map[string]*Objects {
 	objects := map[string]*Objects{}
 
 	for _, row := range res.Rows {
-		schema := row[0].(string)
+		oid := row[0].(string)
-		name := row[1].(string)
+		schema := row[1].(string)
-		objectType := row[2].(string)
+		name := row[2].(string)
+		objectType := row[3].(string)
 
 		if objects[schema] == nil {
 			objects[schema] = &Objects{
-				Tables:            []string{},
+				Tables:            []Object{},
-				Views:             []string{},
+				Views:             []Object{},
-				MaterializedViews: []string{},
+				MaterializedViews: []Object{},
-				Sequences:         []string{},
+				Functions:         []Object{},
+				Sequences:         []Object{},
 			}
 		}
 
+		obj := Object{OID: oid, Name: name}
+
 		switch objectType {
 		case ObjTypeTable:
-			objects[schema].Tables = append(objects[schema].Tables, name)
+			objects[schema].Tables = append(objects[schema].Tables, obj)
 		case ObjTypeView:
-			objects[schema].Views = append(objects[schema].Views, name)
+			objects[schema].Views = append(objects[schema].Views, obj)
 		case ObjTypeMaterializedView:
-			objects[schema].MaterializedViews = append(objects[schema].MaterializedViews, name)
+			objects[schema].MaterializedViews = append(objects[schema].MaterializedViews, obj)
+		case ObjTypeFunction:
+			objects[schema].Functions = append(objects[schema].Functions, obj)
 		case ObjTypeSequence:
-			objects[schema].Sequences = append(objects[schema].Sequences, name)
+			objects[schema].Sequences = append(objects[schema].Sequences, obj)
 		}
 	}
 

pkg/statements/sql.go

@@ -38,6 +38,9 @@ var (
 	//go:embed sql/objects.sql
 	Objects string
 
+	//go:embed sql/function.sql
+	Function string
+
 	// Activity queries for specific PG versions
 	Activity = map[string]string{
 		"default": "SELECT * FROM pg_stat_activity WHERE datname = current_database()",

pkg/statements/sql/function.sql

@@ -0,0 +1,7 @@
+SELECT
+  p.*,
+  pg_get_functiondef(oid) AS functiondef
+FROM
+  pg_catalog.pg_proc p
+WHERE
+  oid = $1::oid

pkg/statements/sql/objects.sql

@@ -1,25 +1,46 @@
-SELECT
+WITH all_objects AS (
-  n.nspname AS schema,
+  SELECT
-  c.relname AS name,
+    c.oid,
-  CASE c.relkind
+    n.nspname AS schema,
-    WHEN 'r' THEN 'table'
+    c.relname AS name,
-    WHEN 'v' THEN 'view'
+    CASE c.relkind
-    WHEN 'm' THEN 'materialized_view'
+      WHEN 'r' THEN 'table'
-    WHEN 'i' THEN 'index'
+      WHEN 'v' THEN 'view'
-    WHEN 'S' THEN 'sequence'
+      WHEN 'm' THEN 'materialized_view'
-    WHEN 's' THEN 'special'
+      WHEN 'i' THEN 'index'
-    WHEN 'f' THEN 'foreign_table'
+      WHEN 'S' THEN 'sequence'
-  END AS type,
+      WHEN 's' THEN 'special'
-  pg_catalog.pg_get_userbyid(c.relowner) AS owner,
+      WHEN 'f' THEN 'foreign_table'
-  pg_catalog.obj_description(c.oid) AS comment
+    END AS type,
-FROM
+    pg_catalog.pg_get_userbyid(c.relowner) AS owner,
-  pg_catalog.pg_class c
+    pg_catalog.obj_description(c.oid) AS comment
-LEFT JOIN
+  FROM
-  pg_catalog.pg_namespace n ON n.oid = c.relnamespace
+    pg_catalog.pg_class c
-WHERE
+  LEFT JOIN
-  c.relkind IN ('r','v','m','S','s','')
+    pg_catalog.pg_namespace n ON n.oid = c.relnamespace
-  AND n.nspname !~ '^pg_toast'
+  WHERE
-  AND n.nspname NOT IN ('information_schema', 'pg_catalog')
+    c.relkind IN ('r','v','m','S','s','')
-  AND has_schema_privilege(n.nspname, 'USAGE')
+    AND n.nspname !~ '^pg_toast'
-ORDER BY
+    AND n.nspname NOT IN ('information_schema', 'pg_catalog')
-  1, 2
+    AND has_schema_privilege(n.nspname, 'USAGE')
+
+  UNION
+
+  SELECT
+    p.oid,
+    n.nspname AS schema,
+    p.proname AS name,
+    'function' AS function,
+    pg_catalog.pg_get_userbyid(p.proowner) AS owner,
+    NULL AS comment
+  FROM
+    pg_catalog.pg_namespace n
+  JOIN
+    pg_catalog.pg_proc p ON p.pronamespace = n.oid
+  WHERE
+    n.nspname !~ '^pg_toast'
+    AND n.nspname NOT IN ('information_schema', 'pg_catalog')
+    AND p.prokind = 'f'
+)
+SELECT * FROM all_objects
+ORDER BY 1, 2

static/css/app.css

@@ -583,6 +583,30 @@
 
 #results_view pre {
   border: 0px none;
+  position: relative;
+}
+
+#results_view .copy {
+  position: absolute;
+  display: none;
+  text-align: center;
+  line-height: 30px;
+  right: 4px;
+  top: 4px;
+  width: 30px;
+  height: 30px;
+  background: #fff;
+  border: 1px solid #ddd;
+  border-radius: 3px;
+  cursor: pointer;
+}
+
+#results_view .copy:hover {
+  border-color: #999;
+}
+
+#results_view pre:hover .copy {
+  display: block;
 }
 
 .full #output {

static/js/app.js

@@ -88,6 +88,7 @@ function getTableRows(table, opts, cb)      { apiCall("get", "/tables/" + table
 function getTableStructure(table, opts, cb) { apiCall("get", "/tables/" + table, opts, cb); }
 function getTableIndexes(table, cb)         { apiCall("get", "/tables/" + table + "/indexes", {}, cb); }
 function getTableConstraints(table, cb)     { apiCall("get", "/tables/" + table + "/constraints", {}, cb); }
+function getFunction(id, cb)                { apiCall("get", "/functions/" + id, {}, cb); }
 function getHistory(cb)                     { apiCall("get", "/history", {}, cb); }
 function getBookmarks(cb)                   { apiCall("get", "/bookmarks", {}, cb); }
 function executeQuery(query, cb)            { apiCall("post", "/query", { query: query }, cb); }
@@ -106,6 +107,7 @@ function buildSchemaSection(name, objects) {
     "table":             "Tables",
     "view":              "Views",
     "materialized_view": "Materialized Views",
+    "function":          "Functions",
     "sequence":          "Sequences"
   };
 
@@ -113,6 +115,7 @@ function buildSchemaSection(name, objects) {
     "table":             '<i class="fa fa-table"></i>',
     "view":              '<i class="fa fa-table"></i>',
     "materialized_view": '<i class="fa fa-table"></i>',
+    "function":          '<i class="fa fa-bolt"></i>',
     "sequence":          '<i class="fa fa-circle-o"></i>'
   };
 
@@ -123,7 +126,7 @@ function buildSchemaSection(name, objects) {
   section += "<div class='schema-name'><i class='fa fa-folder-o'></i><i class='fa fa-folder-open-o'></i> " + name + "</div>";
   section += "<div class='schema-container'>";
 
-  ["table", "view", "materialized_view", "sequence"].forEach(function(group) {
+  ["table", "view", "materialized_view", "function", "sequence"].forEach(function(group) {
     group_klass = "";
     if (name == "public" && group == "table") group_klass = "expanded";
 
@@ -133,8 +136,14 @@ function buildSchemaSection(name, objects) {
 
     if (objects[group]) {
       objects[group].forEach(function(item) {
-        var id = name + "." + item;
+        var id = name + "." + item.name;
-        section += "<li class='schema-item schema-" + group + "' data-type='" + group + "' data-id='" + id + "' data-name='" + item + "'>" + icons[group] + "&nbsp;" + item + "</li>";
+
+        // Use function OID since multiple functions with the same name might exist
+        if (group == "function") {
+          id = item.oid;
+        }
+
+        section += "<li class='schema-item schema-" + group + "' data-type='" + group + "' data-id='" + id + "' data-name='" + item.name + "'>" + icons[group] + "&nbsp;" + item.name + "</li>";
       });
       section += "</ul></div>";
     }
@@ -154,6 +163,7 @@ function loadSchemas() {
         table: [],
         view: [],
         materialized_view: [],
+        function: [],
         sequence: []
       };
     }
@@ -170,13 +180,14 @@ function loadSchemas() {
     autocompleteObjects = [];
     for (schema in data) {
       for (kind in data[schema]) {
-        if (!(kind == "table" || kind == "view" || kind == "materialized_view")) {
+        if (!(kind == "table" || kind == "view" || kind == "materialized_view" || kind == "function")) {
           continue
         }
+
         for (item in data[schema][kind]) {
           autocompleteObjects.push({
-            caption: data[schema][kind][item],
+            caption: data[schema][kind][item].name,
-            value: data[schema][kind][item],
+            value: data[schema][kind][item].name,
             meta: kind
           });
         }
@@ -507,6 +518,11 @@ function showTableContent(sortColumn, sortOrder) {
     return;
   }
 
+  if (getCurrentObject().type == "function") {
+    alert("Cant view rows for a function");
+    return;
+  }
+
   var opts = {
     limit:       getRowsLimit(),
     offset:      getPaginationOffset(),
@@ -569,6 +585,13 @@ function showTableStructure() {
   $("#body").prop("class", "full");
 
   getTableStructure(name, { type: getCurrentObject().type }, function(data) {
+    if (getCurrentObject().type == "function") {
+      var name = data.rows[0][data.columns.indexOf("proname")];
+      var definition = data.rows[0][data.columns.indexOf("functiondef")];
+      showFunctionDefinition(name, definition);
+      return
+    }
+
     buildTable(data);
     $("#results").addClass("no-crop");
   });
@@ -576,14 +599,25 @@ function showTableStructure() {
 
 function showViewDefinition(viewName, viewDefintion) {
   setCurrentTab("table_structure");
+  renderResultsView("View definition for: <strong>" + viewName + "</strong>", viewDefintion);
+}
 
+function showFunctionDefinition(functionName, definition) {
+  setCurrentTab("table_structure");
+  renderResultsView("Function definition for: <strong>" + functionName + "</strong>", definition)
+}
+
+function renderResultsView(title, content) {
   $("#results").addClass("no-crop");
   $("#input").hide();
   $("#body").prop("class", "full");
   $("#results").hide();
 
-  var title = $("<div/>").prop("class", "title").html("View definition for: <strong>" + viewName + "</strong>");
+  var title = $("<div/>").prop("class", "title").html(title);
-  var content = $("<pre/>").text(viewDefintion);
+  var content = $("<pre/>").text(content);
+  console.log(content);
+
+  $("<div/>").html("<i class='fa fa-copy'></i>").addClass("copy").appendTo(content);
 
   $("#results_view").html("");
   title.appendTo("#results_view");
@@ -1138,7 +1172,7 @@ function bindDatabaseObjectsFilter() {
     $(".clear-objects-filter").show();
     $(".schema-group").addClass("expanded");
 
-    filterTimeout = setTimeout(function () {
+    filterTimeout = setTimeout(function() {
       filterObjectsByName(val)
     }, 200);
   });
@@ -1354,6 +1388,10 @@ $(document).ready(function() {
     exportTo("xml");
   });
 
+  $("#results_view").on("click", ".copy", function() {
+    copyToClipboard($(this).parent().text());
+  });
+
   $("#results").on("click", "tr", function(e) {
     $("#results tr.selected").removeClass();
     $(this).addClass("selected");
@@ -1378,7 +1416,11 @@ $(document).ready(function() {
     $(".current-page").data("page", 1);
     $(".filters select, .filters input").val("");
 
-    showTableInfo();
+    if (currentObject.type == "function") {
+      sessionStorage.setItem("tab", "table_structure");
+    } else {
+      showTableInfo();
+    }
 
     switch(sessionStorage.getItem("tab")) {
       case "table_content":