diff --git a/pkg/client/client.go b/pkg/client/client.go index 242d593..bb33d17 100644 --- 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 index eaa4209..6f796f2 100644 --- 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 index d6b2f3f..fb2f4fe 100644 --- 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 new file mode 100644 index 0000000..5ebd0b9 --- /dev/null +++ b/pkg/statements/sql/info_simple.sql @@ -0,0 +1,6 @@ +SELECT + session_user, + current_user, + current_database(), + current_schemas(false), + version()