Collection helper detection: Improve handling of default arguments (#691)

This previously crashed in some cases due to an incomplete refactoring
I had done in a prior commit. Add test coverage to prevent recurrence.

In passing move SQL helper definitions to a dedicated util/helpers
directory and add the helper definition for "get_stat_statements", which
whilst not used much these days, is a good test case here.

Author: lfittl

input/postgres/collection.go

@@ -120,9 +120,11 @@ func (c *Collection) findHelperFunction(name string, inputTypes []string) (state
 			parts := strings.Split(arg, " ")
 
 			// Check if function has more arguments required than we're expecting
-			//
-			// Note we allow extra arguments if they have default values.
-			if idx >= len(inputTypes) && parts[2] != "DEFAULT" {
+			if idx >= len(inputTypes) {
+				// Allow extra arguments if they have default values.
+				if len(parts) >= 3 && parts[2] == "DEFAULT" {
+					break
+				}
 				mismatch = true
 				break
 			}

input/postgres/collection_test.go

@@ -0,0 +1,97 @@
+package postgres_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/pganalyze/collector/input/postgres"
+	"github.com/pganalyze/collector/state"
+	"github.com/pganalyze/collector/util"
+)
+
+type collectionHelperTestpair struct {
+	Name               string
+	InputTypes         []string
+	ExpectedExists     bool
+	ExpectedReturnType string
+}
+
+var collectionHelperTests = []collectionHelperTestpair{
+	{
+		"explain_analyze",
+		[]string{"text", "text[]", "text[]", "text[]"},
+		true,
+		"text",
+	},
+	{
+		"get_stat_statements",
+		[]string{"boolean"},
+		true,
+		"SETOF pg_stat_statements",
+	},
+	// Verify we don't match on a shorter set of arguments
+	{
+		"explain_analyze",
+		[]string{"text", "text[]", "text[]"},
+		false,
+		"",
+	},
+	// Verify we don't match on a longer set of arguments
+	{
+		"explain_analyze",
+		[]string{"text", "text[]", "text[]", "text[]", "text[]"},
+		false,
+		"",
+	},
+	// Verify we don't match on different argument types
+	{
+		"explain_analyze",
+		[]string{"text", "text[]", "text[]", "float"},
+		false,
+		"",
+	},
+	// Verify we allow default arguments to be ommitted
+	{
+		"get_stat_statements",
+		[]string{},
+		true,
+		"SETOF pg_stat_statements",
+	},
+}
+
+func TestCollectionFindHelper(t *testing.T) {
+	db := setupTest(t)
+	defer db.Close()
+
+	_, err := db.Exec("CREATE EXTENSION IF NOT EXISTS pg_stat_statements")
+	if err != nil {
+		t.Fatalf("Could not create pg_stat_statements extension: %s", err)
+	}
+
+	_, err = db.Exec(util.GetStatStatementsHelper)
+	if err != nil {
+		t.Fatalf("Could not load get_stat_statements helper: %s", err)
+	}
+
+	ctx := context.Background()
+	logger := &util.Logger{}
+	server := &state.Server{}
+	opts := state.CollectionOpts{}
+
+	collection, err := postgres.NewCollection(ctx, logger, server, opts, db)
+	if err != nil {
+		t.Fatalf("Could not initialize collection struct: %s", err)
+	}
+
+	for _, pair := range collectionHelperTests {
+		helperExists := collection.HelperExists(pair.Name, pair.InputTypes)
+		if helperExists != pair.ExpectedExists {
+			t.Errorf("Incorrect exists state for helper %s(%+v):\n got: %+v\n expected: %+v", pair.Name, pair.InputTypes, helperExists, pair.ExpectedExists)
+		}
+
+		helperReturnType := collection.HelperReturnType(pair.Name, pair.InputTypes)
+		if helperReturnType != pair.ExpectedReturnType {
+			t.Errorf("Incorrect return type for helper %s(%+v):\n got: %s\n expected: %s", pair.Name, pair.InputTypes, helperReturnType, pair.ExpectedReturnType)
+		}
+	}
+}

util/explain_analyze_helper.sql renamed to util/helpers/explain_analyze.sql

@@ -0,0 +1,4 @@
+CREATE OR REPLACE FUNCTION pganalyze.get_stat_statements(showtext boolean = true) RETURNS SETOF pg_stat_statements AS
+$$
+    /* pganalyze-collector */ SELECT * FROM public.pg_stat_statements(showtext);
+$$ LANGUAGE sql VOLATILE SECURITY DEFINER;

util/helpers/get_stat_statements.sql

@@ -4,5 +4,8 @@ import (
 	_ "embed"
 )
 
-//go:embed explain_analyze_helper.sql
+//go:embed helpers/explain_analyze.sql
 var ExplainAnalyzeHelper string
+
+//go:embed helpers/get_stat_statements.sql
+var GetStatStatementsHelper string