Handle support/permissions errors in info call (sosedoff#635)

* Handle support/permissions errors in info call * Fix linting
jonaz · Apr 10, 2024 · e2d8991 · e2d8991
1 parent feee40f
commit e2d8991
Show file tree

Hide file tree

Showing 4 changed files with 68 additions and 15 deletions.
diff --git a/pkg/client/client.go b/pkg/client/client.go
@@ -206,7 +206,15 @@ func (client *Client) TestWithTimeout(timeout time.Duration) (result error) {
 }
 
 func (client *Client) Info() (*Result, error) {
-	return client.query(statements.Info)
+	result, err := client.query(statements.Info)
+	if err != nil {
+		msg := err.Error()
+		if strings.Contains(msg, "inet_") && (strings.Contains(msg, "not supported") || strings.Contains(msg, "permission denied")) {
+			// Fetch client information without inet_ function calls
+			result, err = client.query(statements.InfoSimple)
+		}
+	}
+	return result, err
 }
 
 func (client *Client) Databases() ([]string, error) {

diff --git a/pkg/client/client_test.go b/pkg/client/client_test.go
@@ -203,21 +203,57 @@ func testTest(t *testing.T) {
 }
 
 func testInfo(t *testing.T) {
-	expected := []string{
-		"session_user",
-		"current_user",
-		"current_database",
-		"current_schemas",
-		"inet_client_addr",
-		"inet_client_port",
-		"inet_server_addr",
-		"inet_server_port",
-		"version",
-	}
+	t.Run("normal", func(t *testing.T) {
+		expected := []string{
+			"session_user",
+			"current_user",
+			"current_database",
+			"current_schemas",
+			"inet_client_addr",
+			"inet_client_port",
+			"inet_server_addr",
+			"inet_server_port",
+			"version",
+		}
 
-	res, err := testClient.Info()
-	assert.NoError(t, err)
-	assert.Equal(t, expected, res.Columns)
+		res, err := testClient.Info()
+		assert.NoError(t, err)
+		assert.Equal(t, expected, res.Columns)
+	})
+
+	t.Run("with restrictions", func(t *testing.T) {
+		expected := []string{
+			"session_user",
+			"current_user",
+			"current_database",
+			"current_schemas",
+			"version",
+		}
+
+		// Prepare a new user and database
+		testClient.db.MustExec("DROP DATABASE IF EXISTS testdb")
+		testClient.db.Exec("DROP OWNED BY IF EXISTS testuser") //nolint:all
+		testClient.db.MustExec("DROP ROLE IF EXISTS testuser")
+		testClient.db.MustExec("CREATE ROLE testuser WITH PASSWORD 'secret' LOGIN NOSUPERUSER NOINHERIT")
+		testClient.db.MustExec("CREATE DATABASE testdb OWNER testuser")
+
+		// Disable access to inet_ calls for new user
+		url := fmt.Sprintf("postgres://%s:@%s:%s/testdb?sslmode=disable", serverUser, serverHost, serverPort)
+		client, err := NewFromUrl(url, nil)
+		assert.NoError(t, err)
+		client.db.MustExec("REVOKE EXECUTE ON FUNCTION inet_client_addr() FROM PUBLIC")
+		assert.NoError(t, client.Close())
+
+		// Connect using new user
+		url = fmt.Sprintf("postgres://testuser:secret@%s:%s/testdb?sslmode=disable", serverHost, serverPort)
+		client, err = NewFromUrl(url, nil)
+		assert.NoError(t, err)
+		defer client.Close()
+
+		res, err := client.Info()
+		assert.NoError(t, err)
+		assert.Equal(t, expected, res.Columns)
+	})
 }
 
 func testActivity(t *testing.T) {

diff --git a/pkg/statements/sql.go b/pkg/statements/sql.go
@@ -14,6 +14,9 @@ var (
 	//go:embed sql/info.sql
 	Info string
 
+	//go:embed sql/info_simple.sql
+	InfoSimple string
+
 	//go:embed sql/estimated_row_count.sql
 	EstimatedTableRowCount string
 

diff --git a/pkg/statements/sql/info_simple.sql b/pkg/statements/sql/info_simple.sql
@@ -0,0 +1,6 @@
+SELECT
+  session_user,
+  current_user,
+  current_database(),
+  current_schemas(false),
+  version()