// +build go1.8

package xray

import (
	"context"
	"database/sql"
)

func (db *DB) Begin(ctx context.Context, opts *sql.TxOptions) (*Tx, error) {
	tx, err := db.db.BeginTx(ctx, opts)
	return &Tx{db, tx}, err
}

func (db *DB) Prepare(ctx context.Context, query string) (*Stmt, error) {
	stmt, err := db.db.PrepareContext(ctx, enhanceQuery(ctx, query))
	return &Stmt{db, stmt, query}, err
}

func (db *DB) Ping(ctx context.Context) error {
	return Capture(ctx, db.dbname, func(ctx context.Context) error {
		db.populate(ctx, "PING")
		return db.db.PingContext(ctx)
	})
}

func (db *DB) Exec(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
	var res sql.Result

	err := Capture(ctx, db.dbname, func(ctx context.Context) error {
		db.populate(ctx, query)

		var err error
		res, err = db.db.ExecContext(ctx, enhanceQuery(ctx, query), args...)
		return err
	})

	return res, err
}

func (db *DB) Query(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) {
	var res *sql.Rows

	err := Capture(ctx, db.dbname, func(ctx context.Context) error {
		db.populate(ctx, query)

		var err error
		res, err = db.db.QueryContext(ctx, enhanceQuery(ctx, query), args...)
		return err
	})

	return res, err
}

func (db *DB) QueryRow(ctx context.Context, query string, args ...interface{}) *sql.Row {
	var res *sql.Row

	Capture(ctx, db.dbname, func(ctx context.Context) error {
		db.populate(ctx, query)

		res = db.db.QueryRowContext(ctx, enhanceQuery(ctx, query), args...)
		return nil
	})

	return res
}

func (tx *Tx) Exec(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
	var res sql.Result

	err := Capture(ctx, tx.db.dbname, func(ctx context.Context) error {
		tx.db.populate(ctx, query)

		var err error
		res, err = tx.tx.ExecContext(ctx, enhanceQuery(ctx, query), args...)
		return err
	})

	return res, err
}

func (tx *Tx) Query(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) {
	var res *sql.Rows

	err := Capture(ctx, tx.db.dbname, func(ctx context.Context) error {
		tx.db.populate(ctx, query)

		var err error
		res, err = tx.tx.QueryContext(ctx, enhanceQuery(ctx, query), args...)
		return err
	})

	return res, err
}

func (tx *Tx) QueryRow(ctx context.Context, query string, args ...interface{}) *sql.Row {
	var res *sql.Row

	Capture(ctx, tx.db.dbname, func(ctx context.Context) error {
		tx.db.populate(ctx, query)

		res = tx.tx.QueryRowContext(ctx, enhanceQuery(ctx, query), args...)
		return nil
	})

	return res
}

func (tx *Tx) Stmt(ctx context.Context, stmt *Stmt) *Stmt {
	return &Stmt{stmt.db, tx.tx.StmtContext(ctx, stmt.stmt), stmt.query}
}

func (stmt *Stmt) Exec(ctx context.Context, args ...interface{}) (sql.Result, error) {
	var res sql.Result

	err := Capture(ctx, stmt.db.dbname, func(ctx context.Context) error {
		stmt.populate(ctx, stmt.query)

		var err error
		res, err = stmt.stmt.ExecContext(ctx, args...)
		return err
	})

	return res, err
}

func (stmt *Stmt) Query(ctx context.Context, args ...interface{}) (*sql.Rows, error) {
	var res *sql.Rows

	err := Capture(ctx, stmt.db.dbname, func(ctx context.Context) error {
		stmt.populate(ctx, stmt.query)

		var err error
		res, err = stmt.stmt.QueryContext(ctx, args...)
		return err
	})

	return res, err
}

func (stmt *Stmt) QueryRow(ctx context.Context, args ...interface{}) *sql.Row {
	var res *sql.Row

	Capture(ctx, stmt.db.dbname, func(ctx context.Context) error {
		stmt.populate(ctx, stmt.query)

		res = stmt.stmt.QueryRowContext(ctx, args...)
		return nil
	})

	return res
}
