pgweb/pkg/client/client_test.go
2024-01-12 21:17:14 -06:00

763 lines
20 KiB
Go

package client
import (
"fmt"
"log"
"os"
"os/exec"
"runtime"
"sort"
"strings"
"testing"
"time"
"github.com/sosedoff/pgweb/pkg/command"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
var (
testClient *Client
testCommands map[string]string
serverHost string
serverPort string
serverUser string
serverPassword string
serverDatabase string
)
func mapKeys(data map[string]*Objects) []string {
result := []string{}
for k := range data {
result = append(result, k)
}
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)
return major, minor
}
func getVar(name, def string) string {
val := os.Getenv(name)
if val == "" {
return def
}
return val
}
func initVars() {
// We need to load default options to make sure all stuff works
if err := command.SetDefaultOptions(); err != nil {
log.Fatal(err)
}
serverHost = getVar("PGHOST", "localhost")
serverPort = getVar("PGPORT", "5432")
serverUser = getVar("PGUSER", "postgres")
serverPassword = getVar("PGPASSWORD", "postgres")
serverDatabase = getVar("PGDATABASE", "booktown")
}
func setupCommands() {
testCommands = map[string]string{
"createdb": "createdb",
"psql": "psql",
"dropdb": "dropdb",
}
if onWindows() {
for k, v := range testCommands {
testCommands[k] = v + ".exe"
}
}
}
func onWindows() bool {
return runtime.GOOS == "windows"
}
func setup() {
// No pretty JSON for tests
command.Opts.DisablePrettyJSON = true
out, err := exec.Command(
testCommands["createdb"],
"-U", serverUser,
"-h", serverHost,
"-p", serverPort,
serverDatabase,
).CombinedOutput()
if err != nil {
fmt.Println("Database creation failed:", string(out))
fmt.Println("Error:", err)
os.Exit(1)
}
out, err = exec.Command(
testCommands["psql"],
"-U", serverUser,
"-h", serverHost,
"-p", serverPort,
"-f", "../../data/booktown.sql",
serverDatabase,
).CombinedOutput()
if err != nil {
fmt.Println("Database import failed:", string(out))
fmt.Println("Error:", err)
os.Exit(1)
}
}
func setupClient() {
url := fmt.Sprintf("postgres://%s@%s:%s/%s?sslmode=disable", serverUser, serverHost, serverPort, serverDatabase)
testClient, _ = NewFromUrl(url, nil)
}
func teardownClient() {
if testClient != nil {
testClient.Close()
}
}
func teardown(t *testing.T, allowFail bool) {
output, err := exec.Command(
testCommands["dropdb"],
"-U", serverUser,
"-h", serverHost,
"-p", serverPort,
serverDatabase,
).CombinedOutput()
if err != nil && strings.Contains(err.Error(), "does not exist") {
t.Log("Teardown error:", err)
t.Logf("%s\n", output)
if !allowFail {
assert.NoError(t, err)
}
}
}
func testNewClientFromURL(t *testing.T) {
t.Run("postgres prefix", func(t *testing.T) {
url := fmt.Sprintf("postgres://%s@%s:%s/%s?sslmode=disable", serverUser, serverHost, serverPort, serverDatabase)
client, err := NewFromUrl(url, nil)
assert.Equal(t, nil, err)
assert.Equal(t, url, client.ConnectionString)
assert.NoError(t, client.Close())
})
t.Run("postgresql prefix", func(t *testing.T) {
url := fmt.Sprintf("postgresql://%s@%s:%s/%s?sslmode=disable", serverUser, serverHost, serverPort, serverDatabase)
client, err := NewFromUrl(url, nil)
assert.Equal(t, nil, err)
assert.Equal(t, url, client.ConnectionString)
assert.NoError(t, client.Close())
})
}
func testClientIdleTime(t *testing.T) {
examples := map[time.Time]bool{
time.Now(): false, // Current time
time.Now().Add(time.Minute * -30): false, // 30 minutes ago
time.Now().Add(time.Minute * -240): true, // 240 minutes ago
time.Now().Add(time.Minute * 30): false, // 30 minutes in future
time.Now().Add(time.Minute * 128): false, // 128 minutes in future
}
for ts, expected := range examples {
testClient.lastQueryTime = ts
assert.Equal(t, expected, testClient.IsIdle())
}
}
func testTest(t *testing.T) {
examples := []struct {
name string
input string
err error
}{
{
name: "success",
input: fmt.Sprintf("postgres://%s@%s:%s/%s?sslmode=disable", serverUser, serverHost, serverPort, serverDatabase),
err: nil,
},
{
name: "connection refused",
input: "postgresql://localhost:5433/dbname",
err: ErrConnectionRefused,
},
{
name: "invalid user",
input: fmt.Sprintf("postgres://%s:%s@%s:%s/%s?sslmode=disable", "foo", serverPassword, serverHost, serverPort, serverDatabase),
err: ErrAuthFailed,
},
{
name: "invalid password",
input: fmt.Sprintf("postgres://%s:%s@%s:%s/%s?sslmode=disable", serverUser, "foo", serverHost, serverPort, serverDatabase),
err: ErrAuthFailed,
},
{
name: "invalid database",
input: fmt.Sprintf("postgres://%s@%s:%s/%s?sslmode=disable", serverUser, serverHost, serverPort, "foo"),
err: ErrDatabaseNotExist,
},
}
for _, ex := range examples {
t.Run(ex.name, func(t *testing.T) {
conn, err := NewFromUrl(ex.input, nil)
require.NoError(t, err)
require.Equal(t, ex.err, conn.Test())
})
}
}
func testInfo(t *testing.T) {
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)
})
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) {
expected := []string{"datid", "pid", "query", "query_start", "state", "client_addr"}
res, err := testClient.Activity()
assert.NoError(t, err)
assertMatches(t, expected, res.Columns)
}
func testDatabases(t *testing.T) {
res, err := testClient.Databases()
assert.NoError(t, err)
assertMatches(t, []string{"booktown", "postgres"}, res)
}
func testSchemas(t *testing.T) {
res, err := testClient.Schemas()
assert.NoError(t, err)
assert.Equal(t, []string{"public"}, res)
}
func testObjects(t *testing.T) {
res, err := testClient.Objects()
objects := ObjectsFromResult(res)
tables := []string{
"alternate_stock",
"authors",
"book_backup",
"book_queue",
"books",
"customers",
"daily_inventory",
"distinguished_authors",
"dummies",
"editions",
"employees",
"favorite_authors",
"favorite_books",
"money_example",
"my_list",
"numeric_values",
"publishers",
"schedules",
"shipments",
"states",
"stock",
"stock_backup",
"subjects",
"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{"oid", "schema", "name", "type", "owner", "comment"}, res.Columns)
assert.Equal(t, []string{"public"}, mapKeys(objects))
assert.Equal(t, tables, objectNames(objects["public"].Tables))
assertMatches(t, functions, objectNames(objects["public"].Functions))
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"}, objectNames(objects["public"].MaterializedViews))
} else {
t.Logf("Skipping materialized view on %d.%d\n", major, minor)
}
}
func testTable(t *testing.T) {
columns := []string{
"column_name",
"data_type",
"is_nullable",
"character_maximum_length",
"character_set_catalog",
"column_default",
"comment",
}
res, err := testClient.Table("books")
assert.NoError(t, err)
assert.Equal(t, columns, res.Columns)
assert.Equal(t, 4, len(res.Rows))
}
func testTableRows(t *testing.T) {
res, err := testClient.TableRows("books", RowsOptions{})
assert.NoError(t, err)
assert.Equal(t, 4, len(res.Columns))
assert.Equal(t, 15, len(res.Rows))
}
func testTableInfo(t *testing.T) {
res, err := testClient.TableInfo("books")
assert.NoError(t, err)
assert.Equal(t, 4, len(res.Columns))
assert.Equal(t, 1, len(res.Rows))
}
func testEstimatedTableRowsCount(t *testing.T) {
res, err := testClient.EstimatedTableRowsCount("books", RowsOptions{})
assert.NoError(t, err)
assert.Equal(t, []string{"reltuples"}, res.Columns)
assert.Equal(t, []Row{{int64(15)}}, res.Rows)
}
func testTableRowsCount(t *testing.T) {
res, err := testClient.TableRowsCount("books", RowsOptions{})
assert.NoError(t, err)
assert.Equal(t, []string{"count"}, res.Columns)
assert.Equal(t, []Row{{int64(15)}}, res.Rows)
}
func testTableRowsCountWithLargeTable(t *testing.T) {
testClient.db.MustExec(`CREATE TABLE large_table AS SELECT s FROM generate_series(1,1000000) s;`)
testClient.db.MustExec(`VACUUM large_table;`)
res, err := testClient.TableRowsCount("large_table", RowsOptions{})
assert.Equal(t, nil, err)
assert.Equal(t, []string{"reltuples"}, res.Columns)
assert.Equal(t, []Row{{int64(1000000)}}, res.Rows)
}
func testTableIndexes(t *testing.T) {
res, err := testClient.TableIndexes("books")
assert.NoError(t, err)
assert.Equal(t, []string{"index_name", "index_size", "index_definition"}, res.Columns)
assert.Equal(t, 2, len(res.Rows))
}
func testTableConstraints(t *testing.T) {
res, err := testClient.TableConstraints("editions")
assert.NoError(t, err)
assert.Equal(t, []string{"name", "definition"}, res.Columns)
assert.Equal(t, Row{"pkey", "PRIMARY KEY (isbn)"}, res.Rows[0])
assert.Equal(t, Row{"integrity", "CHECK (book_id IS NOT NULL AND edition IS NOT NULL)"}, res.Rows[1])
}
func testTableNameWithCamelCase(t *testing.T) {
testClient.db.MustExec(`CREATE TABLE "exampleTable" (id int, name varchar);`)
testClient.db.MustExec(`INSERT INTO "exampleTable" (id, name) VALUES (1, 'foo'), (2, 'bar');`)
_, err := testClient.Table("exampleTable")
assert.NoError(t, err)
_, err = testClient.TableInfo("exampleTable")
assert.NoError(t, err)
_, err = testClient.TableConstraints("exampleTable")
assert.NoError(t, err)
_, err = testClient.TableIndexes("exampleTable")
assert.NoError(t, err)
_, err = testClient.TableRowsCount("exampleTable", RowsOptions{})
assert.NoError(t, err)
_, err = testClient.EstimatedTableRowsCount("exampleTable", RowsOptions{})
assert.NoError(t, err)
}
func testQuery(t *testing.T) {
t.Run("basic query", func(t *testing.T) {
res, err := testClient.Query("SELECT * FROM books")
assert.NoError(t, err)
assert.Equal(t, 4, len(res.Columns))
assert.Equal(t, 15, len(res.Rows))
})
t.Run("error", func(t *testing.T) {
res, err := testClient.Query("SELCT * FROM books")
assert.NotNil(t, err)
assert.Equal(t, "pq: syntax error at or near \"SELCT\"", err.Error())
assert.Nil(t, res)
})
t.Run("invalid table", func(t *testing.T) {
res, err := testClient.Query("SELECT * FROM books2")
assert.NotNil(t, err)
assert.Equal(t, "pq: relation \"books2\" does not exist", err.Error())
assert.Nil(t, res)
})
t.Run("timeout", func(t *testing.T) {
testClient.queryTimeout = time.Millisecond * 100
defer func() {
testClient.queryTimeout = 0
}()
res, err := testClient.query("SELECT pg_sleep(1);")
assert.Equal(t, "pq: canceling statement due to user request", err.Error())
assert.Nil(t, res)
})
}
func testUpdateQuery(t *testing.T) {
t.Run("updating data", func(t *testing.T) {
// Add new row
testClient.db.MustExec("INSERT INTO books (id, title) VALUES (8888, 'Test Book'), (8889, 'Test Book 2')")
// Update without return values
res, err := testClient.Query("UPDATE books SET title = 'Foo' WHERE id >= 8888 AND id <= 8889")
assert.NoError(t, err)
assert.Equal(t, "Rows Affected", res.Columns[0])
assert.Equal(t, int64(2), res.Rows[0][0])
// Update with return values
res, err = testClient.Query("UPDATE books SET title = 'Foo2' WHERE id >= 8888 AND id <= 8889 RETURNING id, title")
assert.NoError(t, err)
assert.Equal(t, []string{"id", "title"}, res.Columns)
assert.Equal(t, Row{int64(8888), "Foo2"}, res.Rows[0])
assert.Equal(t, Row{int64(8889), "Foo2"}, res.Rows[1])
})
t.Run("deleting data", func(t *testing.T) {
// Add new row
testClient.db.MustExec("INSERT INTO books (id, title) VALUES (9999, 'Test Book')")
// Delete the existing row
res, err := testClient.Query("DELETE FROM books WHERE id = 9999")
assert.NoError(t, err)
assert.Equal(t, "Rows Affected", res.Columns[0])
assert.Equal(t, int64(1), res.Rows[0][0])
// Deleting already deleted row
res, err = testClient.Query("DELETE FROM books WHERE id = 9999")
assert.NoError(t, err)
assert.Equal(t, int64(0), res.Rows[0][0])
// Delete with returning value
testClient.db.MustExec("INSERT INTO books (id, title) VALUES (9999, 'Test Book')")
res, err = testClient.Query("DELETE FROM books WHERE id = 9999 RETURNING id")
assert.NoError(t, err)
assert.Equal(t, int64(9999), res.Rows[0][0])
})
}
func testTableRowsOrderEscape(t *testing.T) {
rows, err := testClient.TableRows("dummies", RowsOptions{SortColumn: "isDummy"})
assert.NoError(t, err)
assert.Equal(t, 2, len(rows.Rows))
rows, err = testClient.TableRows("dummies", RowsOptions{SortColumn: "isdummy"})
assert.NotNil(t, err)
assert.Equal(t, `pq: column "isdummy" does not exist`, err.Error())
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")
assert.NoError(t, err)
assert.Equal(t, `[{"author_id":4156,"id":7808,"subject_id":9,"title":"The Shining"}]`, string(result.JSON()))
result, err = testClient.Query("SELECT 'NaN'::float AS value;")
assert.NoError(t, err)
assert.Equal(t, `[{"value":null}]`, string(result.JSON()))
})
t.Run("csv", func(t *testing.T) {
expected := "id,title,author_id,subject_id\n156,The Tell-Tale Heart,115,9\n"
res, err := testClient.Query("SELECT * FROM books ORDER BY id ASC LIMIT 1")
assert.NoError(t, err)
assert.Equal(t, expected, string(res.CSV()))
})
}
func testHistory(t *testing.T) {
t.Run("success", func(t *testing.T) {
_, err := testClient.Query("SELECT * FROM books WHERE id = 12345")
query := testClient.History[len(testClient.History)-1].Query
assert.NoError(t, err)
assert.Equal(t, "SELECT * FROM books WHERE id = 12345", query)
})
t.Run("failed query", func(t *testing.T) {
_, err := testClient.Query("SELECT * FROM books123")
query := testClient.History[len(testClient.History)-1].Query
assert.NotNil(t, err)
assert.NotEqual(t, "SELECT * FROM books123", query)
})
t.Run("unique queries", func(t *testing.T) {
url := fmt.Sprintf("postgres://%s@%s:%s/%s?sslmode=disable", serverUser, serverHost, serverPort, serverDatabase)
client, _ := NewFromUrl(url, nil)
defer client.Close()
for i := 0; i < 3; i++ {
_, err := client.Query("SELECT * FROM books WHERE id = 1")
assert.NoError(t, err)
}
assert.Equal(t, 1, len(client.History))
assert.Equal(t, "SELECT * FROM books WHERE id = 1", client.History[0].Query)
})
}
func testReadOnlyMode(t *testing.T) {
command.Opts.ReadOnly = true
defer func() {
command.Opts.ReadOnly = false
}()
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)
_, err = client.Query("\nCREATE TABLE foobar(id integer);\n")
assert.NotNil(t, err)
assert.Error(t, err, "query contains keywords not allowed in read-only mode")
// Turn off guard
_, err = client.db.Exec("SET default_transaction_read_only=off;")
assert.NoError(t, err)
_, err = client.Query("\nCREATE TABLE foobar(id integer);\n")
assert.NotNil(t, err)
assert.Contains(t, err.Error(), "query contains keywords not allowed in read-only mode")
_, err = client.Query("-- CREATE TABLE foobar(id integer);\nSELECT 'foo';")
assert.NoError(t, err)
_, err = client.Query("/* CREATE TABLE foobar(id integer); */ SELECT 'foo';")
assert.NoError(t, err)
t.Run("with local readonly flag", func(t *testing.T) {
command.Opts.ReadOnly = false
client.readonly = true
_, err := client.Query("INSERT INTO foobar(id) VALUES(1)")
assert.Error(t, err, "query contains keywords not allowed in read-only mode")
})
}
func testTablesStats(t *testing.T) {
columns := []string{
"schema_name",
"table_name",
"total_size",
"data_size",
"index_size",
"estimated_rows_count",
"estimated_rows",
"index_to_data_ratio",
"indexes_count",
"columns_count",
}
result, err := testClient.TablesStats()
assert.NoError(t, err)
assert.Equal(t, columns, result.Columns)
}
func testConnContext(t *testing.T) {
result, err := testClient.GetConnContext()
assert.NoError(t, err)
assert.Equal(t, "localhost", result.Host)
assert.Equal(t, "postgres", result.User)
assert.Equal(t, "booktown", result.Database)
assert.Equal(t, "default", result.Mode)
}
func TestAll(t *testing.T) {
if onWindows() {
t.Log("Unit testing on Windows platform is not supported.")
return
}
initVars()
setupCommands()
teardown(t, false)
setup()
setupClient()
testNewClientFromURL(t)
testClientIdleTime(t)
testTest(t)
testInfo(t)
testActivity(t)
testDatabases(t)
testSchemas(t)
testObjects(t)
testTable(t)
testTableRows(t)
testTableInfo(t)
testEstimatedTableRowsCount(t)
testTableRowsCount(t)
testTableRowsCountWithLargeTable(t)
testTableIndexes(t)
testTableConstraints(t)
testTableNameWithCamelCase(t)
testQuery(t)
testUpdateQuery(t)
testTableRowsOrderEscape(t)
testFunctions(t)
testResult(t)
testHistory(t)
testReadOnlyMode(t)
testDumpExport(t)
testTablesStats(t)
testConnContext(t)
teardownClient()
teardown(t, true)
}