package dbconnect import ( "context" "database/sql" "encoding/json" "fmt" "net/url" "reflect" "strings" "github.com/jmoiron/sqlx" "github.com/pkg/errors" "github.com/xo/dburl" // SQL drivers self-register with the database/sql package. // https://github.com/golang/go/wiki/SQLDrivers _ "github.com/denisenkom/go-mssqldb" _ "github.com/go-sql-driver/mysql" _ "github.com/mattn/go-sqlite3" "github.com/kshvakov/clickhouse" "github.com/lib/pq" ) // SQLClient is a Client that talks to a SQL database. type SQLClient struct { Dialect string driver *sqlx.DB } // NewSQLClient creates a SQL client based on its URL scheme. func NewSQLClient(ctx context.Context, originURL *url.URL) (Client, error) { res, err := dburl.Parse(originURL.String()) if err != nil { helpText := fmt.Sprintf("supported drivers: %+q, see documentation for more details: %s", sql.Drivers(), "https://godoc.org/github.com/xo/dburl") return nil, fmt.Errorf("could not parse sql database url '%s': %s\n%s", originURL, err.Error(), helpText) } // Establishes the driver, but does not test the connection. driver, err := sqlx.Open(res.Driver, res.DSN) if err != nil { return nil, fmt.Errorf("could not open sql driver %s: %s\n%s", res.Driver, err.Error(), res.DSN) } // Closes the driver, will occur when the context finishes. go func() { <-ctx.Done() _ = driver.Close() }() return &SQLClient{driver.DriverName(), driver}, nil } // Ping verifies a connection to the database is still alive. func (client *SQLClient) Ping(ctx context.Context) error { return client.driver.PingContext(ctx) } // Submit queries or executes a command to the SQL database. func (client *SQLClient) Submit(ctx context.Context, cmd *Command) (interface{}, error) { txx, err := cmd.ValidateSQL(client.Dialect) if err != nil { return nil, err } ctx, cancel := context.WithTimeout(ctx, cmd.Timeout) defer cancel() var res interface{} // Get the next available sql.Conn and submit the Command. err = sqlConn(ctx, client.driver, txx, func(conn *sql.Conn) error { stmt := cmd.Statement args := cmd.Arguments.Positional if cmd.Mode == "query" { res, err = sqlQuery(ctx, conn, stmt, args) } else { res, err = sqlExec(ctx, conn, stmt, args) } return err }) return res, err } // ValidateSQL extends the contract of Command for SQL dialects: // mode is conformed, arguments are []sql.NamedArg, and isolation is a sql.IsolationLevel. // // When the command should not be wrapped in a transaction, *sql.TxOptions and error will both be nil. func (cmd *Command) ValidateSQL(dialect string) (*sql.TxOptions, error) { err := cmd.Validate() if err != nil { return nil, err } mode, err := sqlMode(cmd.Mode) if err != nil { return nil, err } // Mutates Arguments to only use positional arguments with the type sql.NamedArg. // This is a required by the sql.Driver before submitting arguments. cmd.Arguments.sql(dialect) iso, err := sqlIsolation(cmd.Isolation) if err != nil { return nil, err } // When isolation is out-of-range, this is indicative that no // transaction should be executed and sql.TxOptions should be nil. if iso < sql.LevelDefault { return nil, nil } // In query mode, execute the transaction in read-only, unless it's Microsoft SQL // which does not support that type of transaction. readOnly := mode == "query" && dialect != "mssql" return &sql.TxOptions{Isolation: iso, ReadOnly: readOnly}, nil } // sqlConn gets the next available sql.Conn in the connection pool and runs a function to use it. // // If the transaction options are nil, run the useIt function outside a transaction. // This is potentially an unsafe operation if the command does not clean up its state. func sqlConn(ctx context.Context, driver *sqlx.DB, txx *sql.TxOptions, useIt func(*sql.Conn) error) error { conn, err := driver.Conn(ctx) if err != nil { return err } defer conn.Close() // If transaction options are specified, begin and defer a rollback to catch errors. var tx *sql.Tx if txx != nil { tx, err = conn.BeginTx(ctx, txx) if err != nil { return err } defer tx.Rollback() } err = useIt(conn) // Check if useIt was successful and a transaction exists before committing. if err == nil && tx != nil { err = tx.Commit() } return err } // sqlQuery queries rows on a sql.Conn and returns an array of result objects. func sqlQuery(ctx context.Context, conn *sql.Conn, stmt string, args []interface{}) ([]map[string]interface{}, error) { rows, err := conn.QueryContext(ctx, stmt, args...) if err == nil { return sqlRows(rows) } return nil, err } // sqlExec executes a command on a sql.Conn and returns the result of the operation. func sqlExec(ctx context.Context, conn *sql.Conn, stmt string, args []interface{}) (sqlResult, error) { exec, err := conn.ExecContext(ctx, stmt, args...) if err == nil { return sqlResultFrom(exec), nil } return sqlResult{}, err } // sql mutates Arguments to contain a positional []sql.NamedArg. // // The actual return type is []interface{} due to the native Golang // function signatures for sql.Exec and sql.Query being generic. func (args *Arguments) sql(dialect string) { result := args.Positional for i, val := range result { result[i] = sqlArg("", val, dialect) } for key, val := range args.Named { result = append(result, sqlArg(key, val, dialect)) } args.Positional = result args.Named = map[string]interface{}{} } // sqlArg creates a sql.NamedArg from a key-value pair and an optional dialect. // // Certain dialects will need to wrap objects, such as arrays, to conform its driver requirements. func sqlArg(key, val interface{}, dialect string) sql.NamedArg { switch reflect.ValueOf(val).Kind() { // PostgreSQL and Clickhouse require arrays to be wrapped before // being inserted into the driver interface. case reflect.Slice, reflect.Array: switch dialect { case "postgres": val = pq.Array(val) case "clickhouse": val = clickhouse.Array(val) } } return sql.Named(fmt.Sprint(key), val) } // sqlIsolation tries to match a string to a sql.IsolationLevel. func sqlIsolation(str string) (sql.IsolationLevel, error) { if str == "none" { return sql.IsolationLevel(-1), nil } for iso := sql.LevelDefault; ; iso++ { if iso > sql.LevelLinearizable { return -1, fmt.Errorf("cannot provide an invalid sql isolation level: '%s'", str) } if str == "" || strings.EqualFold(iso.String(), strings.ReplaceAll(str, "_", " ")) { return iso, nil } } } // sqlMode tries to match a string to a command mode: 'query' or 'exec' for now. func sqlMode(str string) (string, error) { switch str { case "query", "exec": return str, nil default: return "", fmt.Errorf("cannot provide invalid sql mode: '%s'", str) } } // sqlRows scans through a SQL result set and returns an array of objects. func sqlRows(rows *sql.Rows) ([]map[string]interface{}, error) { columns, err := rows.Columns() if err != nil { return nil, errors.Wrap(err, "could not extract columns from result") } defer rows.Close() types, err := rows.ColumnTypes() if err != nil { // Some drivers do not support type extraction, so fail silently and continue. types = make([]*sql.ColumnType, len(columns)) } values := make([]interface{}, len(columns)) pointers := make([]interface{}, len(columns)) var results []map[string]interface{} for rows.Next() { for i := range columns { pointers[i] = &values[i] } _ = rows.Scan(pointers...) // Convert a row, an array of values, into an object where // each key is the name of its respective column. entry := make(map[string]interface{}) for i, col := range columns { entry[col] = sqlValue(values[i], types[i]) } results = append(results, entry) } return results, nil } // sqlValue handles special cases where sql.Rows does not return a "human-readable" object. func sqlValue(val interface{}, col *sql.ColumnType) interface{} { bytes, ok := val.([]byte) if ok { // Opportunistically check for embeded JSON and convert it to a first-class object. var embeded interface{} if json.Unmarshal(bytes, &embeded) == nil { return embeded } // STOR-604: investigate a way to coerce PostgreSQL arrays '{a, b, ...}' into JSON. // Although easy with strings, it becomes more difficult with special types like INET[]. return string(bytes) } return val } // sqlResult is a thin wrapper around sql.Result. type sqlResult struct { LastInsertId int64 `json:"last_insert_id"` RowsAffected int64 `json:"rows_affected"` } // sqlResultFrom converts sql.Result into a JSON-marshable sqlResult. func sqlResultFrom(res sql.Result) sqlResult { insertID, errID := res.LastInsertId() rowsAffected, errRows := res.RowsAffected() // If an error occurs when extracting the result, it is because the // driver does not support that specific field. Instead of passing this // to the user, omit the field in the response. if errID != nil { insertID = -1 } if errRows != nil { rowsAffected = -1 } return sqlResult{insertID, rowsAffected} }