// Copyright 2023 Dolthub, Inc. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package _go import ( "bufio" "context" "encoding/json" goerrors "errors" "fmt" "math" "os" "path/filepath" "strings" "testing" "time" "github.com/cockroachdb/apd/v3" "github.com/cockroachdb/errors" "github.com/dolthub/dolt/go/libraries/doltcore/doltdb" "github.com/dolthub/dolt/go/libraries/doltcore/env" "github.com/dolthub/dolt/go/libraries/utils/filesys" "github.com/dolthub/dolt/go/libraries/utils/svcs" "github.com/dolthub/go-mysql-server/sql" gmstypes "github.com/dolthub/go-mysql-server/sql/types" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgconn" "github.com/jackc/pgx/v5/pgtype" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/dolthub/doltgresql/core/id" "github.com/dolthub/doltgresql/postgres/parser/duration" pgtree "github.com/dolthub/doltgresql/postgres/parser/sem/tree" "github.com/dolthub/doltgresql/postgres/parser/timeofday" "github.com/dolthub/doltgresql/postgres/parser/uuid" dserver "github.com/dolthub/doltgresql/server" "github.com/dolthub/doltgresql/server/auth" "github.com/dolthub/doltgresql/server/functions" "github.com/dolthub/doltgresql/server/types" "github.com/dolthub/doltgresql/servercfg" "github.com/dolthub/doltgresql/servercfg/cfgdetails" ) // runOnPostgres is a debug setting to redirect the test framework to a local running postgres server, // rather than starting a doltgres server. const runOnPostgres = false // serverHost is the host of the local Doltgres server used for testing; set to IPv4 loopback address. var serverHost = "127.0.0.1" // ScriptTest defines a consistent structure for testing queries. type ScriptTest struct { // Name of the script. Name string // The database to create and use. If not provided, then it defaults to "postgres". Database string // The SQL statements to execute as setup, in order. Results are not checked, but statements must not error. SetUpScript []string // The set of assertions to make after setup, in order Assertions []ScriptTestAssertion // When using RunScripts, setting this on one (or more) tests causes RunScripts to ignore all tests that have this // set to false (which is the default value). This allows a developer to easily "focus" on a specific test without // having to comment out other tests, pull it into a different function, etc. In addition, CI ensures that this is // false before passing, meaning this prevents the commented-out situation where the developer forgets to uncomment // their code. Focus bool // Skip is used to completely skip a test including setup Skip bool // UseLocalFileSystem determines if the test should use the local filesystem UseLocalFileSystem bool } // ExpectedNotice specifies what notices are expected during a script test assertion. type ExpectedNotice struct { Severity string Message string } // ScriptTestAssertion are the assertions upon which the script executes its main "testing" logic. type ScriptTestAssertion struct { Query string Expected []sql.Row // Expected or ExpectedRaw should be used, but not both at the same time ExpectedRaw [][][]byte // ExpectedRaw or Expected should be used, but not both at the same time ExpectedErr string ExpectedNotices []ExpectedNotice Focus bool BindVars []any // SkipResultsCheck is used to skip assertions on the expected rows returned from a query. For now, this is // included as some messages do not have a full logical implementation. Skipping the results check allows us to // force the test client to not send of those messages. SkipResultsCheck bool // Skip is used to completely skip a test, not execute its query at all, and record it as a skipped test // in the test suite results. Skip bool // Username specifies the user's name to use for the command. This creates a new connection, using the given name. // By default (when the string is empty), the `postgres` superuser account is used. Any consecutive queries that // have the same username and password will reuse the same connection. The `postgres` superuser account will always // reuse the same connection. Do note that specifying the `postgres` account manually will create a connection // that is different from the primary one. Username string // Password specifies the password that will be used alongside the given username. This field is essentially ignored // when no username is given. If a username is given and the password is empty, then it is assumed that the password // is the empty string. Password string // ExpectedTag is used to check the command tag returned from the server. // This is checked only if no Expected is defined ExpectedTag string // ExpectedColNames are used to check the column names returned from the server. ExpectedColNames []string // ExpectedColTypes are used to check the column types returned from the server. ExpectedColTypes []id.Type // CopyFromSTDIN is used to test the COPY FROM STDIN command. CopyFromStdInFile string } // EmptyCommandTag is special command tag placeholder to check for the empty string const EmptyCommandTag = "EMPTY_COMMAND_TAG" // Connection contains the default and current connections. type Connection struct { Default *pgx.Conn Current *pgx.Conn Username string Password string } // receivedNotices tracks the NOTICE messages received over the connection to the Doltgres server, so that tests // can assert what notices are expected to be sent to the client. var receivedNotices []*pgconn.Notice // RunScript runs the given script. func RunScript(t *testing.T, script ScriptTest, normalizeRows bool) { scriptDatabase := script.Database if len(scriptDatabase) == 0 { scriptDatabase = "postgres" } var ctx context.Context var conn *Connection if runOnPostgres { ctx = context.Background() pgxConn, err := pgx.Connect(ctx, fmt.Sprintf("postgres://postgres:password@127.0.0.1:%d/%s?sslmode=disable", 5432, scriptDatabase)) require.NoError(t, err) conn = &Connection{ Default: pgxConn, Current: pgxConn, } require.NoError(t, pgxConn.Ping(ctx)) defer func() { conn.Close(ctx) }() } else { var controller *svcs.Controller if script.UseLocalFileSystem { port, err := sql.GetEmptyPort() require.NoError(t, err) ctx, conn, controller = CreateServerLocalWithPort(t, scriptDatabase, port) } else { ctx, conn, controller = CreateServer(t, scriptDatabase) } defer func() { conn.Close(ctx) controller.Stop() err := controller.WaitForStop() require.NoError(t, err) }() } t.Run(script.Name, func(t *testing.T) { runScript(t, ctx, script, conn, normalizeRows) }) } // runScript runs the script given on the postgres connection provided func runScript(t *testing.T, ctx context.Context, script ScriptTest, conn *Connection, normalizeRows bool) { if script.Skip { t.Skip("Skip has been set in the script") } // Run the setup for _, query := range script.SetUpScript { _, err := conn.Exec(ctx, query) require.NoError(t, err, "error running setup query: %s", query) } assertions := script.Assertions // First, we'll run through the scripts to check for the Focus variable. If it's true, then append it to the new slice. focusAssertions := make([]ScriptTestAssertion, 0, len(script.Assertions)) for _, assertion := range script.Assertions { if assertion.Focus { // If this is running in GitHub Actions, then we'll panic, because someone forgot to disable it before committing if _, ok := os.LookupEnv("GITHUB_ACTION"); ok { panic(fmt.Sprintf("The assertion `%s` has Focus set to `true`. GitHub Actions requires that "+ "all tests are run, which Focus circumvents, leading to this error. Please disable Focus on "+ "all tests.", assertion.Query)) } focusAssertions = append(focusAssertions, assertion) } } // If we have assertions with Focus set, then we replace the normal script slice with the new slice. if len(focusAssertions) > 0 { assertions = focusAssertions } // Run the assertions for _, assertion := range assertions { t.Run(assertion.Query, func(t *testing.T) { if assertion.Skip { t.Skip("Skip has been set in the assertion") } // Clear out any previously received notices receivedNotices = nil // Use the provided username and password to create a new connection (if a username has been specified). // This will automatically handle connection reuse, using the default connection is no user is specified, etc. if err := conn.Connect(ctx, assertion.Username, assertion.Password); err != nil { if assertion.ExpectedErr != "" { require.Error(t, err) assert.Contains(t, err.Error(), assertion.ExpectedErr) } else { require.NoError(t, err) } return } // If we're skipping the results check, then we call Execute, as it uses a simplified message model. if assertion.CopyFromStdInFile != "" { copyFromStdin(t, conn.Current, assertion.Query, assertion.CopyFromStdInFile) } else if assertion.SkipResultsCheck || assertion.ExpectedErr != "" { _, err := conn.Exec(ctx, assertion.Query, assertion.BindVars...) if assertion.ExpectedErr != "" { require.Error(t, err) assert.Contains(t, err.Error(), assertion.ExpectedErr) } else { require.NoError(t, err) } } else if assertion.ExpectedTag != "" { commandTag, err := conn.Exec(ctx, assertion.Query) require.NoError(t, err) tag := assertion.ExpectedTag if tag == EmptyCommandTag { tag = "" } assert.Equal(t, tag, commandTag.String()) } else { rows, err := conn.Query(ctx, assertion.Query, assertion.BindVars...) require.NoError(t, err) readRows, readRawRows, err := ReadRows(rows, normalizeRows) require.NoError(t, err) if assertion.ExpectedColNames != nil { fields := rows.FieldDescriptions() if assert.Len(t, fields, len(assertion.ExpectedColNames), "expected length of columns") { for i, col := range assertion.ExpectedColNames { assert.Equal(t, col, fields[i].Name) } } } if assertion.ExpectedColTypes != nil { fields := rows.FieldDescriptions() if assert.Len(t, fields, len(assertion.ExpectedColTypes), "columns returned and types expected are not the same length") { for i, colId := range assertion.ExpectedColTypes { assert.Equal(t, id.Cache().ToOID(colId.AsId()), fields[i].DataTypeOID, `"%s" expected type "%s" but received "%s"`, fields[i].Name, colId.TypeName(), id.Type(id.Cache().ToInternal(fields[i].DataTypeOID)).TypeName()) } } } // not an exact match but works well enough for our tests orderBy := strings.Contains(strings.ToLower(assertion.Query), "order by") if assertion.ExpectedRaw != nil { if orderBy { assert.Equal(t, assertion.ExpectedRaw, readRawRows, "wrong result for query %s", assertion.Query) } else { assert.ElementsMatch(t, assertion.ExpectedRaw, readRawRows, "wrong result for query %s", assertion.Query) } } else { if normalizeRows { if orderBy { assert.Equal(t, NormalizeExpectedRow(rows.FieldDescriptions(), assertion.Expected), readRows, "wrong result for query %s", assertion.Query) } else { assert.ElementsMatch(t, NormalizeExpectedRow(rows.FieldDescriptions(), assertion.Expected), readRows, "wrong result for query %s", assertion.Query) } } else { if orderBy { assert.Equal(t, assertion.Expected, readRows, "wrong result for query %s", assertion.Query) } else { assert.ElementsMatch(t, assertion.Expected, readRows, "wrong result for query %s", assertion.Query) } } } if len(assertion.ExpectedNotices) > 0 { if len(assertion.ExpectedNotices) == len(receivedNotices) { for i, notice := range receivedNotices { assert.Equal(t, assertion.ExpectedNotices[i].Severity, notice.Severity) assert.Equal(t, assertion.ExpectedNotices[i].Message, notice.Message) } } else { if len(receivedNotices) == 0 { receivedNotices = []*pgconn.Notice{} } assert.Fail(t, "Received notices do not match expected notices", "Expected %d notices, but received %d. Expected: %v, Received: %v", len(assertion.ExpectedNotices), len(receivedNotices), assertion.ExpectedNotices, receivedNotices) } } } }) } } func copyFromStdin(t *testing.T, conn *pgx.Conn, query string, filename string) { filePath := filepath.Join("testdata", filename) file, err := os.Open(filePath) if err != nil { t.Fatalf("Failed to open file: %v", err) } defer file.Close() reader := bufio.NewReader(file) _, err = conn.PgConn().CopyFrom(context.Background(), reader, query) require.NoError(t, err) } // RunScripts runs the given collection of scripts. This normalizes all rows before comparing them. func RunScripts(t *testing.T, scripts []ScriptTest) { runScripts(t, scripts, true) } // RunScriptsWithoutNormalization runs the given collection of scripts, without normalizing any rows. func RunScriptsWithoutNormalization(t *testing.T, scripts []ScriptTest) { runScripts(t, scripts, false) } // runScripts is the implementation of both RunScripts and RunScriptsWithoutNormalization. func runScripts(t *testing.T, scripts []ScriptTest, normalizeRows bool) { // First, we'll run through the scripts to check for the Focus variable. If it's true, then append it to the new slice. focusScripts := make([]ScriptTest, 0, len(scripts)) for _, script := range scripts { if script.Focus { // If this is running in GitHub Actions, then we'll panic, because someone forgot to disable it before committing if _, ok := os.LookupEnv("GITHUB_ACTION"); ok { panic(fmt.Sprintf("The script `%s` has Focus set to `true`. GitHub Actions requires that "+ "all tests are run, which Focus circumvents, leading to this error. Please disable Focus on "+ "all tests.", script.Name)) } focusScripts = append(focusScripts, script) } } // If we have scripts with Focus set, then we replace the normal script slice with the new slice. if len(focusScripts) > 0 { scripts = focusScripts } for _, script := range scripts { RunScript(t, script, normalizeRows) } } func ptr[T any](val T) *T { return &val } var testServerLogLevel = "warn" func init() { if logLevel, ok := os.LookupEnv("TEST_SERVER_LOG_LEVEL"); ok { testServerLogLevel = logLevel } } // CreateServer creates a server with the given database, returning a connection to the server. The server will close // when the connection is closed (or loses its connection to the server). The accompanying [svcs.Controller] may be used // to wait until the server has closed. func CreateServer(t *testing.T, database string) (context.Context, *Connection, *svcs.Controller) { port, err := sql.GetEmptyPort() require.NoError(t, err) return CreateServerWithPort(t, database, port) } // CreateServerWithPort creates a server with the given database and port, returning a connection to the server. The server will close // when the connection is closed (or loses its connection to the server). The accompanying [svcs.Controller] may be used // to wait until the server has closed. func CreateServerWithPort(t *testing.T, database string, port int) (context.Context, *Connection, *svcs.Controller) { require.NotEmpty(t, database) controller, err := dserver.RunInMemory(&servercfg.DoltgresConfig{ DoltgresConfig: cfgdetails.DoltgresConfig{ ListenerConfig: &cfgdetails.DoltgresListenerConfig{ PortNumber: &port, HostStr: &serverHost, }, LogLevelStr: &testServerLogLevel, }, }, dserver.NewListener) require.NoError(t, err) auth.ClearDatabase() fmt.Printf("port is %d\n", port) ctx := context.Background() connection := newTestDatabaseConnection(t, ctx, database, serverHost, port) return ctx, connection, controller } // CreateServerLocalWithPort creates a server using the local file system at [os.TempDir]. A Connection is returned to // |database| at 127.0.0.1:|port|. The server will close when the connection is closed or lost. The returned // [svcs.Controller] may be used to wait for the server to stop. func CreateServerLocalWithPort(t *testing.T, database string, port int) (context.Context, *Connection, *svcs.Controller) { // We avoid using [T.TempDir] because it results in a file lock conflict on Windows. [T.TempDir] registers a // [T.Cleanup] function that runs without checking the [svcs.Controller] and it cannot be overwritten. // TODO(elianddb): Setup an optional [T.Cleanup] function for the temporary directory. Our default setup for now is // preferable for debugging the database after a failure. // t.Name() contains "/" for subtests, which os.MkdirTemp rejects as a path separator in the pattern. safeName := strings.ReplaceAll(t.Name(), "/", "_") dbDir, err := os.MkdirTemp(os.TempDir(), safeName) require.NoError(t, err) fileSys, err := filesys.LocalFilesysWithWorkingDir(dbDir) require.NoError(t, err) ctx := context.Background() doltEnv := env.Load(ctx, env.GetCurrentUserHomeDir, fileSys, doltdb.LocalDirDoltDB, dserver.Version) controller, err := dserver.RunOnDisk(ctx, &servercfg.DoltgresConfig{ DoltgresConfig: cfgdetails.DoltgresConfig{ ListenerConfig: &cfgdetails.DoltgresListenerConfig{ PortNumber: &port, HostStr: &serverHost, }, LogLevelStr: &testServerLogLevel, }, }, doltEnv) require.NoError(t, err) auth.ClearDatabase() fmt.Printf("port is %d\n", port) connection := newTestDatabaseConnection(t, ctx, database, serverHost, port) return ctx, connection, controller } // newTestDatabaseConnection returns a Connection to the test |database| at |host|:|port|. If the |database| provided // does not exist, it will be automatically created. func newTestDatabaseConnection(t *testing.T, ctx context.Context, database, host string, port int) *Connection { const connectionUrlFmt = "postgres://postgres:password@%s:%d/%s?DateStyle=ISO%%2C%%20MDY" func() { var conn *pgx.Conn var err error // Connections can happen before the server has a chance to grab the port so we retry. for range 3 { conn, err = pgx.Connect(ctx, fmt.Sprintf(connectionUrlFmt, host, port, "")) if err == nil { break } time.Sleep(time.Second) } require.NoError(t, err) _, err = conn.Exec(ctx, fmt.Sprintf("CREATE DATABASE IF NOT EXISTS %s", database)) require.NoError(t, err) defer require.NoError(t, conn.Close(ctx)) }() config, err := pgx.ParseConfig(fmt.Sprintf(connectionUrlFmt, host, port, database)) require.NoError(t, err) config.OnNotice = func(conn *pgconn.PgConn, notice *pgconn.Notice) { receivedNotices = append(receivedNotices, notice) } // pgx v5.9.1+ skips Describe(portal) on statement-cache hits, so stale field descriptions // from before a DDL change (e.g. DROP COLUMN) cause a field-count/value-count mismatch and // panic in rows.Values(). DescribeExec re-describes every query, keeping schema info fresh. // It also preserves binary wire format (needed for record type decoding). config.DefaultQueryExecMode = pgx.QueryExecModeDescribeExec conn, err := pgx.ConnectConfig(ctx, config) require.NoError(t, err) // Ping tests that Doltgres can handle empty queries, and makes sure connection is alive. require.NoError(t, conn.Ping(ctx)) return &Connection{ Default: conn, Current: conn, } } // ReadRows reads all of the given rows into a slice, then closes the rows. If `normalizeRows` is true, then the rows // will be normalized such that all integers are int64, etc. Normalization does not affect the raw returned bytes. func ReadRows(rows pgx.Rows, normalizeRows bool) (readRows []sql.Row, readRawRows [][][]byte, err error) { defer func() { err = goerrors.Join(err, rows.Err()) }() var slices []sql.Row var rawSlices [][][]byte for rows.Next() { var rawSlice [][]byte for _, rawValue := range rows.RawValues() { rawSlice = append(rawSlice, append([]byte{}, rawValue...)) } rawSlices = append(rawSlices, rawSlice) row, err := rows.Values() if err != nil { return nil, nil, err } slices = append(slices, row) } return NormalizeRows(rows.FieldDescriptions(), slices, normalizeRows), rawSlices, nil } // NormalizeRows normalizes each value's type within each row, as the tests only want to compare values. Returns a new // set of rows in the same order. func NormalizeRows(fds []pgconn.FieldDescription, rows []sql.Row, normalize bool) []sql.Row { newRows := make([]sql.Row, len(rows)) for i := range rows { newRows[i] = NormalizeRow(fds, rows[i], normalize) } return newRows } // NormalizeRow normalizes each value's type, as the tests only want to compare values. // Returns a new row. func NormalizeRow(fds []pgconn.FieldDescription, row sql.Row, normalize bool) sql.Row { if len(row) == 0 { return nil } newRow := make(sql.Row, len(row)) for i := range row { dt, ok := types.IDToBuiltInDoltgresType[id.Type(id.Cache().ToInternal(fds[i].DataTypeOID))] if !ok { // try using text type dt = types.Text } newRow[i] = NormalizeValToString(dt, row[i]) if normalize { newRow[i] = NormalizeIntsAndFloats(newRow[i]) } } return newRow } // NormalizeExpectedRow normalizes each value's type, as the tests only want to compare values. Returns a new row. func NormalizeExpectedRow(fds []pgconn.FieldDescription, rows []sql.Row) []sql.Row { newRows := make([]sql.Row, len(rows)) for ri, row := range rows { if len(row) == 0 { newRows[ri] = nil } else if len(row) != len(fds) { // Return if the expected row count does not match the field description count, we'll error elsewhere return rows } else { newRow := make(sql.Row, len(row)) for i := range row { dt, ok := types.IDToBuiltInDoltgresType[id.Type(id.Cache().ToInternal(fds[i].DataTypeOID))] if !ok { // try using text type dt = types.Text } if dt.ID == types.Json.ID && row[i] != nil { newRow[i] = UnmarshalAndMarshalJsonString(row[i].(string)) } else if dt.IsArrayType() && dt.ArrayBaseType().ID == types.Json.ID { v, err := dt.IoInput(nil, row[i].(string)) if err != nil { panic(err) } arr := v.([]any) newArr := make([]any, len(arr)) for j, el := range arr { switch e := el.(type) { case string: newArr[j] = UnmarshalAndMarshalJsonString(e) case sql.JSONWrapper: iface, err := e.ToInterface(context.Background()) if err != nil { panic(err) } b, err := json.Marshal(iface) if err != nil { panic(err) } newArr[j] = string(b) default: newArr[j] = el } } ret, err := dt.FormatValue(newArr) if err != nil { panic(err) } newRow[i] = ret } else if dt.ID == types.Date.ID { newRow[i] = row[i] if row[i] != nil { if t, _, err := pgtree.ParseDTimestampTZ(nil, row[i].(string), pgtree.TimeFamilyPrecisionToRoundDuration(6), time.UTC); err == nil { newRow[i] = functions.FormatDateTimeWithBC(t.Time.UTC(), "2006-01-02", dt.ID == types.TimestampTZ.ID) } } } else if dt.ID == types.Timestamp.ID || dt.ID == types.TimestampTZ.ID { newRow[i] = row[i] if row[i] != nil { if t, _, err := pgtree.ParseDTimestampTZ(nil, row[i].(string), pgtree.TimeFamilyPrecisionToRoundDuration(6), time.UTC); err == nil { newRow[i] = functions.FormatDateTimeWithBC(t.Time.UTC(), "2006-01-02 15:04:05.999999", dt.ID == types.TimestampTZ.ID) } } } else { newRow[i] = NormalizeIntsAndFloats(row[i]) } } newRows[ri] = newRow } } return newRows } // UnmarshalAndMarshalJsonString is used to normalize expected json type value to compare the actual value. // JSON type value is in string format, and since Postrges JSON type preserves the input string if valid, // it cannot be compared to the returned map as json.Marshal method space padded key value pair. // To allow result matching, we unmarshal and marshal the expected string. This causes missing check // for the identical format as the input of the json string. func UnmarshalAndMarshalJsonString(val string) string { var decoded any err := json.Unmarshal([]byte(val), &decoded) if err != nil { panic(err) } ret, err := json.Marshal(decoded) if err != nil { panic(err) } return string(ret) } // NormalizeValToString normalizes values into types that can be compared. // JSON types, any pg types and time and decimal type values are converted into string value. // |normalizeNumeric| defines whether to normalize Numeric values into either Numeric type or string type. // There are an infinite number of ways to represent the same value in-memory, // so we must at least normalize Numeric values. func NormalizeValToString(dt *types.DoltgresType, v any) any { if v == nil { return nil } switch dt.ID { case types.Json.ID: jsonBytes, err := json.Marshal(v) if err != nil { panic(err) } return string(jsonBytes) case types.JsonB.ID: // JSONB normalizes JSON with spaces after ':' and ',' (PostgreSQL JSONB format) s, err := gmstypes.JSONDocument{Val: v}.JSONString() if err != nil { panic(err) } return s case types.InternalChar.ID: if v == nil { return nil } var b []byte if v.(int32) == 0 { b = []byte{} } else { b = []byte{uint8(v.(int32))} } val, err := dt.FormatValue(string(b)) if err != nil { panic(err) } return val case types.Interval.ID, types.Uuid.ID: // These values need to be normalized into the appropriate types // before being converted to string type using the Doltgres // IoOutput method. if v == nil { return nil } tVal, err := dt.FormatValue(NormalizeVal(dt, v)) if err != nil { panic(err) } return tVal case types.Date.ID: if v == nil { return nil } return functions.FormatDateTimeWithBC(v.(time.Time), "2006-01-02", false) case types.Timestamp.ID, types.TimestampTZ.ID: if v == nil { return nil } return functions.FormatDateTimeWithBC(v.(time.Time).UTC(), "2006-01-02 15:04:05.999999", dt.ID == types.TimestampTZ.ID) } switch val := v.(type) { case bool: if val { return "t" } else { return "f" } case pgtype.Time: return timeofday.TimeOfDay(val.Microseconds).String() case []any: if dt.IsArrayType() { return NormalizeArrayType(dt, val) } } return v } // NormalizeArrayType normalizes array types by normalizing its elements first, // then to a string using the type IoOutput method. func NormalizeArrayType(dt *types.DoltgresType, arr []any) any { newVal := make([]any, len(arr)) for i, el := range arr { newVal[i] = NormalizeVal(dt.ArrayBaseType(), el) } baseType := dt.ArrayBaseType() if baseType.ID == types.Bool.ID { sqlVal, err := dt.SQL(sql.NewEmptyContext(), nil, newVal) if err != nil { panic(err) } return sqlVal.ToString() } else { ret, err := dt.FormatValue(newVal) if err != nil { panic(err) } return ret } } // NormalizeVal normalizes values to the Doltgres type expects, so it can be used to // convert the values using the given Doltgres type. This is used to normalize array // types as the type conversion expects certain type values. func NormalizeVal(dt *types.DoltgresType, v any) any { switch dt.ID { case types.Json.ID: str, err := json.Marshal(v) if err != nil { panic(err) } return string(str) case types.JsonB.ID: return gmstypes.JSONDocument{Val: v} case types.Oid.ID, types.Regclass.ID, types.Regproc.ID, types.Regtype.ID: if uval, ok := v.(uint32); ok { if internalID := id.Cache().ToInternal(uval); internalID.IsValid() { return internalID } return id.NewOID(uval).AsId() } } switch val := v.(type) { case pgtype.Numeric: if val.NaN { return math.NaN() } else if val.InfinityModifier != pgtype.Finite { return math.Inf(int(val.InfinityModifier)) } else if !val.Valid { return nil } else { return apd.New(val.Int.Int64(), val.Exp) } case pgtype.Time: // This value type is used for TIME type. return timeofday.FromInt(val.Microseconds) case pgtype.Interval: // This value type is used for INTERVAL type. return duration.MakeDuration(val.Microseconds*functions.NanosPerMicro, int64(val.Days), int64(val.Months)) case [16]byte: // This value type is used for UUID type. u, err := uuid.FromBytes(val[:]) if err != nil { panic(err) } return u case []any: baseType := dt if baseType.IsArrayType() { baseType = baseType.ArrayBaseType() } newVal := make([]any, len(val)) for i, el := range val { newVal[i] = NormalizeVal(baseType, el) } return newVal } return v } // NormalizeIntsAndFloats normalizes all int and float types // to int64 and float64, respectively. func NormalizeIntsAndFloats(v any) any { switch val := v.(type) { case int: return int64(val) case int8: return int64(val) case int16: return int64(val) case int32: return int64(val) case uint: return int64(val) case uint8: return int64(val) case uint16: return int64(val) case uint32: return int64(val) case uint64: // PostgreSQL does not support an uint64 type, so we can always convert this to an int64 safely. return int64(val) case float32: return float64(val) case []any: for i := range val { val[i] = NormalizeIntsAndFloats(val[i]) } return val default: return val } } // Numeric creates a numeric value from a string. func Numeric(str string) pgtype.Numeric { numeric := pgtype.Numeric{} if err := numeric.Scan(str); err != nil { panic(err) } return numeric } // Timestamp is a helper function to convert timestamp strings to pgtype.Timestamp instances. If // the string cannot be converted, this function will panic. func Timestamp(timestampStr string) (ret pgtype.Timestamp) { t, err := time.Parse("2006-01-02 15:04:05", timestampStr) if err != nil { panic("invalid timestamp format: " + err.Error()) } if err := ret.Scan(t); err != nil { panic("failed to set pgtype.Timestamp: " + err.Error()) } return ret } // Date is a helper function to convert date strings to pgtype.Date instances. If the string // cannot be converted, this function will panic. func Date(dateStr string) (d pgtype.Date) { t, err := time.Parse("2006-01-02", dateStr) if err != nil { panic("invalid date format: " + err.Error()) } if err := d.Scan(t); err != nil { panic("failed to set pgtype.Date: " + err.Error()) } return d } // UUID is a helper function to convert UUID strings to pgtype.UUID instances. If the string // cannot be converted, this function will panic. func UUID(s string) (u pgtype.UUID) { if err := u.Scan(s); err != nil { panic(err) } return u } // Connect replaces the Current connection with a new one, using the given username and password. If the username is // empty, then the default connection is used. If the username and password match the existing connection, then no new // connection is made. func (conn *Connection) Connect(ctx context.Context, username string, password string) error { // Reuse the existing connection if it's the same username and password if username == conn.Username && password == conn.Password && conn.Current != nil { return nil } // Username or password has changed, so we'll close the Current connection only if it's not the default if conn.Username != "" && conn.Current != nil { _ = conn.Current.Close(ctx) } conn.Username = username conn.Password = password if username == "" { conn.Current = conn.Default return nil } else { var err error config := conn.Default.Config() conn.Current, err = pgx.Connect(ctx, fmt.Sprintf("postgres://%s:%s@127.0.0.1:%d/%s", username, password, config.Port, config.Database)) return err } } // Exec calls Exec on the current connection. func (conn *Connection) Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error) { if conn.Current == nil { return pgconn.CommandTag{}, errors.New("EXEC: current connection is nil") } return conn.Current.Exec(ctx, sql, args...) } // Query calls Query on the current connection. func (conn *Connection) Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) { if conn.Current == nil { return nil, errors.New("QUERY: current connection is nil") } return conn.Current.Query(ctx, sql, args...) } // Close closes the connections. func (conn *Connection) Close(ctx context.Context) { if conn.Default != nil { _ = conn.Default.Close(ctx) } if conn.Current != nil && conn.Current != conn.Default { _ = conn.Current.Close(ctx) } conn.Default = nil conn.Current = nil }