<commit_before>package pgx

import (
	"database/sql"
	"errors"
	"fmt"
	"time"
)

// Row is a convenience wrapper over Rows that is returned by QueryRow.
type Row Rows

// Scan works the same as (*Rows Scan) with the following exceptions. If no
// rows were found it returns ErrNoRows. If multiple rows are returned it
// ignores all but the first.
func (r *Row) Scan(dest ...interface{}) (err error) {
	rows := (*Rows)(r)

	if rows.Err() != nil {
		return rows.Err()
	}

	if !rows.Next() {
		if rows.Err() == nil {
			return ErrNoRows
		}
		return rows.Err()
	}

	rows.Scan(dest...)
	rows.Close()
	return rows.Err()
}

// Rows is the result set returned from *Conn.Query. Rows must be closed before
// the *Conn can be used again. Rows are closed by explicitly calling Close(),
// calling Next() until it returns false, or when a fatal error occurs.
type Rows struct {
	conn       *Conn
	mr         *msgReader
	fields     []FieldDescription
	vr         ValueReader
	rowCount   int
	columnIdx  int
	err        error
	startTime  time.Time
	sql        string
	args       []interface{}
	log        func(lvl int, msg string, ctx ...interface{})
	shouldLog  func(lvl int) bool
	afterClose func(*Rows)
	unlockConn bool
	closed     bool
}

func (rows *Rows) FieldDescriptions() []FieldDescription {
	return rows.fields
}

func (rows *Rows) close() {
	if rows.closed {
		return
	}

	if rows.unlockConn {
		rows.conn.unlock()
		rows.unlockConn = false
	}

	rows.closed = true

	if rows.err == nil {
		if rows.shouldLog(LogLevelInfo) {
			endTime := time.Now()
			rows.log(LogLevelInfo, "Query", "sql", rows.sql, "args", logQueryArgs(rows.args), "time", endTime.Sub(rows.startTime), "rowCount", rows.rowCount)
		}
	} else if rows.shouldLog(LogLevelError) {
		rows.log(LogLevelError, "Query", "sql", rows.sql, "args", logQueryArgs(rows.args))
	}

	if rows.afterClose != nil {
		rows.afterClose(rows)
	}
}

func (rows *Rows) readUntilReadyForQuery() {
	for {
		t, r, err := rows.conn.rxMsg()
		if err != nil {
			rows.close()
			return
		}

		switch t {
		case readyForQuery:
			rows.conn.rxReadyForQuery(r)
			rows.close()
			return
		case rowDescription:
		case dataRow:
		case commandComplete:
		case bindComplete:
		case errorResponse:
			err = rows.conn.rxErrorResponse(r)
			if rows.err == nil {
				rows.err = err
			}
		default:
			err = rows.conn.processContextFreeMsg(t, r)
			if err != nil {
				rows.close()
				return
			}
		}
	}
}

// Close closes the rows, making the connection ready for use again. It is safe
// to call Close after rows is already closed.
func (rows *Rows) Close() {
	if rows.closed {
		return
	}
	rows.readUntilReadyForQuery()
	rows.close()
}

func (rows *Rows) Err() error {
	return rows.err
}

// abort signals that the query was not successfully sent to the server.
// This differs from Fatal in that it is not necessary to readUntilReadyForQuery
func (rows *Rows) abort(err error) {
	if rows.err != nil {
		return
	}

	rows.err = err
	rows.close()
}

// Fatal signals an error occurred after the query was sent to the server. It
// closes the rows automatically.
func (rows *Rows) Fatal(err error) {
	if rows.err != nil {
		return
	}

	rows.err = err
	rows.Close()
}

// Next prepares the next row for reading. It returns true if there is another
// row and false if no more rows are available. It automatically closes rows
// when all rows are read.
func (rows *Rows) Next() bool {
	if rows.closed {
		return false
	}

	rows.rowCount++
	rows.columnIdx = 0
	rows.vr = ValueReader{}

	for {
		t, r, err := rows.conn.rxMsg()
		if err != nil {
			rows.Fatal(err)
			return false
		}

		switch t {
		case readyForQuery:
			rows.conn.rxReadyForQuery(r)
			rows.close()
			return false
		case dataRow:
			fieldCount := r.readInt16()
			if int(fieldCount) != len(rows.fields) {
				rows.Fatal(ProtocolError(fmt.Sprintf("Row description field count (%v) and data row field count (%v) do not match", len(rows.fields), fieldCount)))
				return false
			}

			rows.mr = r
			return true
		case commandComplete:
		case bindComplete:
		default:
			err = rows.conn.processContextFreeMsg(t, r)
			if err != nil {
				rows.Fatal(err)
				return false
			}
		}
	}
}

// Conn returns the *Conn this *Rows is using.
func (rows *Rows) Conn() *Conn {
	return rows.conn
}

func (rows *Rows) nextColumn() (*ValueReader, bool) {
	if rows.closed {
		return nil, false
	}
	if len(rows.fields) <= rows.columnIdx {
		rows.Fatal(ProtocolError("No next column available"))
		return nil, false
	}

	if rows.vr.Len() > 0 {
		rows.mr.readBytes(rows.vr.Len())
	}

	fd := &rows.fields[rows.columnIdx]
	rows.columnIdx++
	size := rows.mr.readInt32()
	rows.vr = ValueReader{mr: rows.mr, fd: fd, valueBytesRemaining: size}
	return &rows.vr, true
}

type scanArgError struct {
	col int
	err error
}

func (e scanArgError) Error() string {
	return fmt.Sprintf("can't scan into dest[%d]: %v", e.col, e.err)
}

// Scan reads the values from the current row into dest values positionally.
// dest can include pointers to core types, values implementing the Scanner
// interface, []byte, and nil. []byte will skip the decoding process and directly
// copy the raw bytes received from PostgreSQL. nil will skip the value entirely.
func (rows *Rows) Scan(dest ...interface{}) (err error) {
	if len(rows.fields) != len(dest) {
		err = fmt.Errorf("Scan received wrong number of arguments, got %d but expected %d", len(dest), len(rows.fields))
		rows.Fatal(err)
		return err
	}

	for i, d := range dest {
		vr, _ := rows.nextColumn()

		if d == nil {
			continue
		}

		// Check for []byte first as we allow sidestepping the decoding process and retrieving the raw bytes
		if b, ok := d.(*[]byte); ok {
			// If it actually is a bytea then pass it through decodeBytea (so it can be decoded if it is in text format)
			// Otherwise read the bytes directly regardless of what the actual type is.
			if vr.Type().DataType == ByteaOid {
				*b = decodeBytea(vr)
			} else {
				if vr.Len() != -1 {
					*b = vr.ReadBytes(vr.Len())
				} else {
					*b = nil
				}
			}
		} else if s, ok := d.(Scanner); ok {
			err = s.Scan(vr)
			if err != nil {
				rows.Fatal(scanArgError{col: i, err: err})
			}
		} else if s, ok := d.(sql.Scanner); ok {
			var val interface{}
			if 0 <= vr.Len() {
				switch vr.Type().DataType {
				case BoolOid:
					val = decodeBool(vr)
				case Int8Oid:
					val = int64(decodeInt8(vr))
				case Int2Oid:
					val = int64(decodeInt2(vr))
				case Int4Oid:
					val = int64(decodeInt4(vr))
				case TextOid, VarcharOid:
					val = decodeText(vr)
				case OidOid:
					val = int64(decodeOid(vr))
				case Float4Oid:
					val = float64(decodeFloat4(vr))
				case Float8Oid:
					val = decodeFloat8(vr)
				case DateOid:
					val = decodeDate(vr)
				case TimestampOid:
					val = decodeTimestamp(vr)
				case TimestampTzOid:
					val = decodeTimestampTz(vr)
				default:
					val = vr.ReadBytes(vr.Len())
				}
			}
			err = s.Scan(val)
			if err != nil {
				rows.Fatal(scanArgError{col: i, err: err})
			}
		} else if vr.Type().DataType == JsonOid || vr.Type().DataType == JsonbOid {
			// Because the argument passed to decodeJSON will escape the heap.
			// This allows d to be stack allocated and only copied to the heap when
			// we actually are decoding JSON. This saves one memory allocation per
			// row.
			d2 := d
			decodeJSON(vr, &d2)
		} else {
			if err := Decode(vr, d); err != nil {
				rows.Fatal(scanArgError{col: i, err: err})
			}
		}
		if vr.Err() != nil {
			rows.Fatal(scanArgError{col: i, err: vr.Err()})
		}

		if rows.Err() != nil {
			return rows.Err()
		}
	}

	return nil
}

// Values returns an array of the row values
func (rows *Rows) Values() ([]interface{}, error) {
	if rows.closed {
		return nil, errors.New("rows is closed")
	}

	values := make([]interface{}, 0, len(rows.fields))

	for range rows.fields {
		vr, _ := rows.nextColumn()

		if vr.Len() == -1 {
			values = append(values, nil)
			continue
		}

		switch vr.Type().FormatCode {
		// All intrinsic types (except string) are encoded with binary
		// encoding so anything else should be treated as a string
		case TextFormatCode:
			values = append(values, vr.ReadString(vr.Len()))
		case BinaryFormatCode:
			switch vr.Type().DataType {
			case BoolOid:
				values = append(values, decodeBool(vr))
			case ByteaOid:
				values = append(values, decodeBytea(vr))
			case Int8Oid:
				values = append(values, decodeInt8(vr))
			case Int2Oid:
				values = append(values, decodeInt2(vr))
			case Int4Oid:
				values = append(values, decodeInt4(vr))
			case OidOid:
				values = append(values, decodeOid(vr))
			case Float4Oid:
				values = append(values, decodeFloat4(vr))
			case Float8Oid:
				values = append(values, decodeFloat8(vr))
			case BoolArrayOid:
				values = append(values, decodeBoolArray(vr))
			case Int2ArrayOid:
				values = append(values, decodeInt2Array(vr))
			case Int4ArrayOid:
				values = append(values, decodeInt4Array(vr))
			case Int8ArrayOid:
				values = append(values, decodeInt8Array(vr))
			case Float4ArrayOid:
				values = append(values, decodeFloat4Array(vr))
			case Float8ArrayOid:
				values = append(values, decodeFloat8Array(vr))
			case TextArrayOid, VarcharArrayOid:
				values = append(values, decodeTextArray(vr))
			case TimestampArrayOid, TimestampTzArrayOid:
				values = append(values, decodeTimestampArray(vr))
			case DateOid:
				values = append(values, decodeDate(vr))
			case TimestampTzOid:
				values = append(values, decodeTimestampTz(vr))
			case TimestampOid:
				values = append(values, decodeTimestamp(vr))
			case InetOid, CidrOid:
				values = append(values, decodeInet(vr))
			case JsonOid:
				var d interface{}
				decodeJSON(vr, &d)
				values = append(values, d)
			case JsonbOid:
				var d interface{}
				decodeJSON(vr, &d)
				values = append(values, d)
			default:
				rows.Fatal(errors.New("Values cannot handle binary format non-intrinsic types"))
			}
		default:
			rows.Fatal(errors.New("Unknown format code"))
		}

		if vr.Err() != nil {
			rows.Fatal(vr.Err())
		}

		if rows.Err() != nil {
			return nil, rows.Err()
		}
	}

	return values, rows.Err()
}

// AfterClose adds f to a LILO queue of functions that will be called when
// rows is closed.
func (rows *Rows) AfterClose(f func(*Rows)) {
	if rows.afterClose == nil {
		rows.afterClose = f
	} else {
		prevFn := rows.afterClose
		rows.afterClose = func(rows *Rows) {
			f(rows)
			prevFn(rows)
		}
	}
}

// Query executes sql with args. If there is an error the returned *Rows will
// be returned in an error state. So it is allowed to ignore the error returned
// from Query and handle it in *Rows.
func (c *Conn) Query(sql string, args ...interface{}) (*Rows, error) {
	c.lastActivityTime = time.Now()

	rows := c.getRows(sql, args)

	if err := c.lock(); err != nil {
		rows.abort(err)
		return rows, err
	}
	rows.unlockConn = true

	ps, ok := c.preparedStatements[sql]
	if !ok {
		var err error
		ps, err = c.Prepare("", sql)
		if err != nil {
			rows.abort(err)
			return rows, rows.err
		}
	}

	rows.fields = ps.FieldDescriptions
	err := c.sendPreparedQuery(ps, args...)
	if err != nil {
		rows.abort(err)
	}
	return rows, rows.err
}

func (c *Conn) getRows(sql string, args []interface{}) *Rows {
	if len(c.preallocatedRows) == 0 {
		c.preallocatedRows = make([]Rows, 64)
	}

	r := &c.preallocatedRows[len(c.preallocatedRows)-1]
	c.preallocatedRows = c.preallocatedRows[0 : len(c.preallocatedRows)-1]

	r.conn = c
	r.startTime = c.lastActivityTime
	r.sql = sql
	r.args = args
	r.log = c.log
	r.shouldLog = c.shouldLog

	return r
}

// QueryRow is a convenience wrapper over Query. Any error that occurs while
// querying is deferred until calling Scan on the returned *Row. That *Row will
// error with ErrNoRows if no rows are returned.
func (c *Conn) QueryRow(sql string, args ...interface{}) *Row {
	rows, _ := c.Query(sql, args...)
	return (*Row)(rows)
}
<commit_msg>Avoid two memory allocations per query<commit_after>package pgx

import (
	"database/sql"
	"errors"
	"fmt"
	"time"
)

// Row is a convenience wrapper over Rows that is returned by QueryRow.
type Row Rows

// Scan works the same as (*Rows Scan) with the following exceptions. If no
// rows were found it returns ErrNoRows. If multiple rows are returned it
// ignores all but the first.
func (r *Row) Scan(dest ...interface{}) (err error) {
	rows := (*Rows)(r)

	if rows.Err() != nil {
		return rows.Err()
	}

	if !rows.Next() {
		if rows.Err() == nil {
			return ErrNoRows
		}
		return rows.Err()
	}

	rows.Scan(dest...)
	rows.Close()
	return rows.Err()
}

// Rows is the result set returned from *Conn.Query. Rows must be closed before
// the *Conn can be used again. Rows are closed by explicitly calling Close(),
// calling Next() until it returns false, or when a fatal error occurs.
type Rows struct {
	conn       *Conn
	mr         *msgReader
	fields     []FieldDescription
	vr         ValueReader
	rowCount   int
	columnIdx  int
	err        error
	startTime  time.Time
	sql        string
	args       []interface{}
	afterClose func(*Rows)
	unlockConn bool
	closed     bool
}

func (rows *Rows) FieldDescriptions() []FieldDescription {
	return rows.fields
}

func (rows *Rows) close() {
	if rows.closed {
		return
	}

	if rows.unlockConn {
		rows.conn.unlock()
		rows.unlockConn = false
	}

	rows.closed = true

	if rows.err == nil {
		if rows.conn.shouldLog(LogLevelInfo) {
			endTime := time.Now()
			rows.conn.log(LogLevelInfo, "Query", "sql", rows.sql, "args", logQueryArgs(rows.args), "time", endTime.Sub(rows.startTime), "rowCount", rows.rowCount)
		}
	} else if rows.conn.shouldLog(LogLevelError) {
		rows.conn.log(LogLevelError, "Query", "sql", rows.sql, "args", logQueryArgs(rows.args))
	}

	if rows.afterClose != nil {
		rows.afterClose(rows)
	}
}

func (rows *Rows) readUntilReadyForQuery() {
	for {
		t, r, err := rows.conn.rxMsg()
		if err != nil {
			rows.close()
			return
		}

		switch t {
		case readyForQuery:
			rows.conn.rxReadyForQuery(r)
			rows.close()
			return
		case rowDescription:
		case dataRow:
		case commandComplete:
		case bindComplete:
		case errorResponse:
			err = rows.conn.rxErrorResponse(r)
			if rows.err == nil {
				rows.err = err
			}
		default:
			err = rows.conn.processContextFreeMsg(t, r)
			if err != nil {
				rows.close()
				return
			}
		}
	}
}

// Close closes the rows, making the connection ready for use again. It is safe
// to call Close after rows is already closed.
func (rows *Rows) Close() {
	if rows.closed {
		return
	}
	rows.readUntilReadyForQuery()
	rows.close()
}

func (rows *Rows) Err() error {
	return rows.err
}

// abort signals that the query was not successfully sent to the server.
// This differs from Fatal in that it is not necessary to readUntilReadyForQuery
func (rows *Rows) abort(err error) {
	if rows.err != nil {
		return
	}

	rows.err = err
	rows.close()
}

// Fatal signals an error occurred after the query was sent to the server. It
// closes the rows automatically.
func (rows *Rows) Fatal(err error) {
	if rows.err != nil {
		return
	}

	rows.err = err
	rows.Close()
}

// Next prepares the next row for reading. It returns true if there is another
// row and false if no more rows are available. It automatically closes rows
// when all rows are read.
func (rows *Rows) Next() bool {
	if rows.closed {
		return false
	}

	rows.rowCount++
	rows.columnIdx = 0
	rows.vr = ValueReader{}

	for {
		t, r, err := rows.conn.rxMsg()
		if err != nil {
			rows.Fatal(err)
			return false
		}

		switch t {
		case readyForQuery:
			rows.conn.rxReadyForQuery(r)
			rows.close()
			return false
		case dataRow:
			fieldCount := r.readInt16()
			if int(fieldCount) != len(rows.fields) {
				rows.Fatal(ProtocolError(fmt.Sprintf("Row description field count (%v) and data row field count (%v) do not match", len(rows.fields), fieldCount)))
				return false
			}

			rows.mr = r
			return true
		case commandComplete:
		case bindComplete:
		default:
			err = rows.conn.processContextFreeMsg(t, r)
			if err != nil {
				rows.Fatal(err)
				return false
			}
		}
	}
}

// Conn returns the *Conn this *Rows is using.
func (rows *Rows) Conn() *Conn {
	return rows.conn
}

func (rows *Rows) nextColumn() (*ValueReader, bool) {
	if rows.closed {
		return nil, false
	}
	if len(rows.fields) <= rows.columnIdx {
		rows.Fatal(ProtocolError("No next column available"))
		return nil, false
	}

	if rows.vr.Len() > 0 {
		rows.mr.readBytes(rows.vr.Len())
	}

	fd := &rows.fields[rows.columnIdx]
	rows.columnIdx++
	size := rows.mr.readInt32()
	rows.vr = ValueReader{mr: rows.mr, fd: fd, valueBytesRemaining: size}
	return &rows.vr, true
}

type scanArgError struct {
	col int
	err error
}

func (e scanArgError) Error() string {
	return fmt.Sprintf("can't scan into dest[%d]: %v", e.col, e.err)
}

// Scan reads the values from the current row into dest values positionally.
// dest can include pointers to core types, values implementing the Scanner
// interface, []byte, and nil. []byte will skip the decoding process and directly
// copy the raw bytes received from PostgreSQL. nil will skip the value entirely.
func (rows *Rows) Scan(dest ...interface{}) (err error) {
	if len(rows.fields) != len(dest) {
		err = fmt.Errorf("Scan received wrong number of arguments, got %d but expected %d", len(dest), len(rows.fields))
		rows.Fatal(err)
		return err
	}

	for i, d := range dest {
		vr, _ := rows.nextColumn()

		if d == nil {
			continue
		}

		// Check for []byte first as we allow sidestepping the decoding process and retrieving the raw bytes
		if b, ok := d.(*[]byte); ok {
			// If it actually is a bytea then pass it through decodeBytea (so it can be decoded if it is in text format)
			// Otherwise read the bytes directly regardless of what the actual type is.
			if vr.Type().DataType == ByteaOid {
				*b = decodeBytea(vr)
			} else {
				if vr.Len() != -1 {
					*b = vr.ReadBytes(vr.Len())
				} else {
					*b = nil
				}
			}
		} else if s, ok := d.(Scanner); ok {
			err = s.Scan(vr)
			if err != nil {
				rows.Fatal(scanArgError{col: i, err: err})
			}
		} else if s, ok := d.(sql.Scanner); ok {
			var val interface{}
			if 0 <= vr.Len() {
				switch vr.Type().DataType {
				case BoolOid:
					val = decodeBool(vr)
				case Int8Oid:
					val = int64(decodeInt8(vr))
				case Int2Oid:
					val = int64(decodeInt2(vr))
				case Int4Oid:
					val = int64(decodeInt4(vr))
				case TextOid, VarcharOid:
					val = decodeText(vr)
				case OidOid:
					val = int64(decodeOid(vr))
				case Float4Oid:
					val = float64(decodeFloat4(vr))
				case Float8Oid:
					val = decodeFloat8(vr)
				case DateOid:
					val = decodeDate(vr)
				case TimestampOid:
					val = decodeTimestamp(vr)
				case TimestampTzOid:
					val = decodeTimestampTz(vr)
				default:
					val = vr.ReadBytes(vr.Len())
				}
			}
			err = s.Scan(val)
			if err != nil {
				rows.Fatal(scanArgError{col: i, err: err})
			}
		} else if vr.Type().DataType == JsonOid || vr.Type().DataType == JsonbOid {
			// Because the argument passed to decodeJSON will escape the heap.
			// This allows d to be stack allocated and only copied to the heap when
			// we actually are decoding JSON. This saves one memory allocation per
			// row.
			d2 := d
			decodeJSON(vr, &d2)
		} else {
			if err := Decode(vr, d); err != nil {
				rows.Fatal(scanArgError{col: i, err: err})
			}
		}
		if vr.Err() != nil {
			rows.Fatal(scanArgError{col: i, err: vr.Err()})
		}

		if rows.Err() != nil {
			return rows.Err()
		}
	}

	return nil
}

// Values returns an array of the row values
func (rows *Rows) Values() ([]interface{}, error) {
	if rows.closed {
		return nil, errors.New("rows is closed")
	}

	values := make([]interface{}, 0, len(rows.fields))

	for range rows.fields {
		vr, _ := rows.nextColumn()

		if vr.Len() == -1 {
			values = append(values, nil)
			continue
		}

		switch vr.Type().FormatCode {
		// All intrinsic types (except string) are encoded with binary
		// encoding so anything else should be treated as a string
		case TextFormatCode:
			values = append(values, vr.ReadString(vr.Len()))
		case BinaryFormatCode:
			switch vr.Type().DataType {
			case BoolOid:
				values = append(values, decodeBool(vr))
			case ByteaOid:
				values = append(values, decodeBytea(vr))
			case Int8Oid:
				values = append(values, decodeInt8(vr))
			case Int2Oid:
				values = append(values, decodeInt2(vr))
			case Int4Oid:
				values = append(values, decodeInt4(vr))
			case OidOid:
				values = append(values, decodeOid(vr))
			case Float4Oid:
				values = append(values, decodeFloat4(vr))
			case Float8Oid:
				values = append(values, decodeFloat8(vr))
			case BoolArrayOid:
				values = append(values, decodeBoolArray(vr))
			case Int2ArrayOid:
				values = append(values, decodeInt2Array(vr))
			case Int4ArrayOid:
				values = append(values, decodeInt4Array(vr))
			case Int8ArrayOid:
				values = append(values, decodeInt8Array(vr))
			case Float4ArrayOid:
				values = append(values, decodeFloat4Array(vr))
			case Float8ArrayOid:
				values = append(values, decodeFloat8Array(vr))
			case TextArrayOid, VarcharArrayOid:
				values = append(values, decodeTextArray(vr))
			case TimestampArrayOid, TimestampTzArrayOid:
				values = append(values, decodeTimestampArray(vr))
			case DateOid:
				values = append(values, decodeDate(vr))
			case TimestampTzOid:
				values = append(values, decodeTimestampTz(vr))
			case TimestampOid:
				values = append(values, decodeTimestamp(vr))
			case InetOid, CidrOid:
				values = append(values, decodeInet(vr))
			case JsonOid:
				var d interface{}
				decodeJSON(vr, &d)
				values = append(values, d)
			case JsonbOid:
				var d interface{}
				decodeJSON(vr, &d)
				values = append(values, d)
			default:
				rows.Fatal(errors.New("Values cannot handle binary format non-intrinsic types"))
			}
		default:
			rows.Fatal(errors.New("Unknown format code"))
		}

		if vr.Err() != nil {
			rows.Fatal(vr.Err())
		}

		if rows.Err() != nil {
			return nil, rows.Err()
		}
	}

	return values, rows.Err()
}

// AfterClose adds f to a LILO queue of functions that will be called when
// rows is closed.
func (rows *Rows) AfterClose(f func(*Rows)) {
	if rows.afterClose == nil {
		rows.afterClose = f
	} else {
		prevFn := rows.afterClose
		rows.afterClose = func(rows *Rows) {
			f(rows)
			prevFn(rows)
		}
	}
}

// Query executes sql with args. If there is an error the returned *Rows will
// be returned in an error state. So it is allowed to ignore the error returned
// from Query and handle it in *Rows.
func (c *Conn) Query(sql string, args ...interface{}) (*Rows, error) {
	c.lastActivityTime = time.Now()

	rows := c.getRows(sql, args)

	if err := c.lock(); err != nil {
		rows.abort(err)
		return rows, err
	}
	rows.unlockConn = true

	ps, ok := c.preparedStatements[sql]
	if !ok {
		var err error
		ps, err = c.Prepare("", sql)
		if err != nil {
			rows.abort(err)
			return rows, rows.err
		}
	}

	rows.fields = ps.FieldDescriptions
	err := c.sendPreparedQuery(ps, args...)
	if err != nil {
		rows.abort(err)
	}
	return rows, rows.err
}

func (c *Conn) getRows(sql string, args []interface{}) *Rows {
	if len(c.preallocatedRows) == 0 {
		c.preallocatedRows = make([]Rows, 64)
	}

	r := &c.preallocatedRows[len(c.preallocatedRows)-1]
	c.preallocatedRows = c.preallocatedRows[0 : len(c.preallocatedRows)-1]

	r.conn = c
	r.startTime = c.lastActivityTime
	r.sql = sql
	r.args = args

	return r
}

// QueryRow is a convenience wrapper over Query. Any error that occurs while
// querying is deferred until calling Scan on the returned *Row. That *Row will
// error with ErrNoRows if no rows are returned.
func (c *Conn) QueryRow(sql string, args ...interface{}) *Row {
	rows, _ := c.Query(sql, args...)
	return (*Row)(rows)
}
<|endoftext|>


<commit_before><commit_msg>update<commit_after><|endoftext|>


<commit_before>package gojenkins

type Queue struct {
	Items []Item `json:"items"`
}

type Item struct {
	Actions                    []Action `json:"actions"`
	Blocked                    bool     `json:"blocked"`
	Buildable                  bool     `json:"buildable"`
	Id                         int      `json:"id"`
	InQueueSince               int64    `json:"inQueueSince"`
	Params                     string   `json:"params"`
	Stuck                      bool     `json:"stuck"`
	Task                       Task     `json:"task"`
	URL                        string   `json:"url"`
	Why                        string   `json:"why"`
	BuildableStartMilliseconds int64    `json:"buildableStartMilliseconds"`
	Pending                    bool     `json:"pending"`
}

type Action struct {
	Causes               []Cause               `json:"causes"`
	ParameterDefinitions []ParameterDefinition `json:"parameterDefinitions"`
}

type Cause struct {
	ShortDescription string `json:"shortDescription"`
	UserId           string `json:"userId"`
	UserName         string `json:"userName"`
	UpstreamCause
}

type ParameterDefinition struct {
	Name string `json:"name"`
}

type Task struct {
	Name  string `json:"name"`
	Url   string `json:"url"`
	Color string `json:"color"`
}
<commit_msg>GetJob deserializes `parameters` from Build response<commit_after>package gojenkins

type Queue struct {
	Items []Item `json:"items"`
}

type Item struct {
	Actions                    []Action `json:"actions"`
	Blocked                    bool     `json:"blocked"`
	Buildable                  bool     `json:"buildable"`
	Id                         int      `json:"id"`
	InQueueSince               int64    `json:"inQueueSince"`
	Params                     string   `json:"params"`
	Stuck                      bool     `json:"stuck"`
	Task                       Task     `json:"task"`
	URL                        string   `json:"url"`
	Why                        string   `json:"why"`
	BuildableStartMilliseconds int64    `json:"buildableStartMilliseconds"`
	Pending                    bool     `json:"pending"`
}

type Action struct {
	Causes               []Cause               `json:"causes"`
	Parameter            []Parameter           `json:"parameters"`
	ParameterDefinitions []ParameterDefinition `json:"parameterDefinitions"`
}

type Cause struct {
	ShortDescription string `json:"shortDescription"`
	UserId           string `json:"userId"`
	UserName         string `json:"userName"`
	UpstreamCause
}

type ParameterDefinition struct {
	Name string `json:"name"`
}

// Parameter for a build
type Parameter struct {
	Name  string      `json:"name"`
	Value interface{} `json:"value"`
}

type Task struct {
	Name  string `json:"name"`
	Url   string `json:"url"`
	Color string `json:"color"`
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"html/template"
	"net/url"
	"strings"
	"time"
)

// A Quote stores all the info needed to render a quote on the site.
type Quote struct {
	ID         uint
	Text       string
	Notes      string
	CreatedAt  int64
	Rating     uint
	Up         uint
	Down       uint
	TotalVotes uint
	Score      uint // Unseen. This is generated for sorting.
}

func NewQuote(text string) Quote {
	return Quote{
		Text:      text,
		CreatedAt: time.Now().Unix(),
	}
}

func nl2br(text string) template.HTML {
	return template.HTML(strings.Replace(template.HTMLEscapeString(text), "\n", "
", -1))
}

func AddQuote(v url.Values) (int64, error) {
	stmt, err := db.Prepare("INSERT INTO quotes (body, notes) VALUES (?, ?)")
	checkErr(err)

	result, err := stmt.Exec(v.Get("quote"), v.Get("comment"))
	checkErr(err)

	return result.LastInsertId()
}

func GetQuote(sid string) []Quote {

	q := make([]Quote, 0)

	rows, err := db.Query("SELECT id,body,notes,rating FROM quotes WHERE id = ?", sid)
	checkErr(err)

	for rows.Next() {
		var body, notes string
		var id, rating uint
		err = rows.Scan(&id, &body, &notes, &rating)
		checkErr(err)
		newQuote := Quote{ID: id, Text: body, Notes: notes, Rating: rating}
		q = append(q, newQuote)
	}

	return q
}

func Browse(page int) []Quote {
	q := make([]Quote, 0)

	rows, err := db.Query("SELECT id,body,notes,rating FROM quotes ORDER BY id ASC LIMIT 20 OFFSET ?", page*20)
	checkErr(err)

	fmt.Println("iterating")

	for rows.Next() {
		var body, notes string
		var id, rating uint
		err = rows.Scan(&id, &body, &notes, &rating)
		checkErr(err)

		newQuote := Quote{ID: id, Text: body, Notes: notes, Rating: rating}
		q = append(q, newQuote)
	}
	fmt.Println("returning browse")

	return q
}

func Latest() []Quote {
	q := make([]Quote, 0)

	rows, err := db.Query("SELECT id,body,notes,rating FROM quotes ORDER BY id DESC LIMIT 20")
	checkErr(err)

	for rows.Next() {
		var body, notes string
		var id, rating uint
		err = rows.Scan(&id, &body, &notes, &rating)
		checkErr(err)

		newQuote := Quote{ID: id, Text: body, Notes: notes, Rating: rating}
		q = append(q, newQuote)
	}
	return q
}

func Top() []Quote {
	q := make([]Quote, 0)

	rows, err := db.Query("SELECT id,body,notes,rating FROM quotes ORDER BY rating DESC LIMIT 20")
	checkErr(err)

	for rows.Next() {
		var body, notes string
		var id, rating uint
		err = rows.Scan(&id, &body, &notes, &rating)
		checkErr(err)

		newQuote := Quote{ID: id, Text: body, Notes: notes, Rating: rating}
		q = append(q, newQuote)
	}
	return q
}

func Search(searchText string) []Quote {
	q := make([]Quote, 0)

	query := "SELECT id,body,notes,rating FROM quotes WHERE 1=1"

	terms := strings.Split(searchText, " ")

	for i := 0; i < len(terms); i++ {
		// This took WAY too long for what it was.
		// Note to future self: Go doesn't like '%?%'. It takes it literally and
		// ignores the question mark as a binding parameter.
		query += " AND body LIKE '%' || ? || '%'"
	}
	query += " ORDER BY id DESC"

	// We have to cast `terms` to []interface{} because Go sucks
	args := make([]interface{}, len(terms))

	for i := range terms {
		args[i] = terms[i]
	}

	rows, err := db.Query(query, args...)

	checkErr(err)

	for rows.Next() {
		var body, notes string
		var id, rating uint
		err = rows.Scan(&id, &body, &notes, &rating)
		checkErr(err)

		newQuote := Quote{ID: id, Text: body, Notes: notes, Rating: rating}
		q = append(q, newQuote)
	}
	return q

}

func (q *Quote) String() string {
	return fmt.Sprint("%d: %s", q.ID, q.Text)
}
<commit_msg>Refactor database calls in quote.go<commit_after>package main

import (
	"fmt"
	"html/template"
	"net/url"
	"strings"
	"time"
)

// A Quote stores all the info needed to render a quote on the site.
type Quote struct {
	ID         uint
	Text       string
	Notes      string
	CreatedAt  int64
	Rating     uint
	Up         uint
	Down       uint
	TotalVotes uint
	Score      uint // Unseen. This is generated for sorting.
}

func NewQuote(text string) Quote {
	return Quote{
		Text:      text,
		CreatedAt: time.Now().Unix(),
	}
}

func nl2br(text string) template.HTML {
	return template.HTML(strings.Replace(template.HTMLEscapeString(text), "\n", "
", -1))
}

func AddQuote(v url.Values) (int64, error) {
	stmt, err := db.Prepare("INSERT INTO quotes (body, notes) VALUES (?, ?)")
	checkErr(err)

	result, err := stmt.Exec(v.Get("quote"), v.Get("comment"))
	checkErr(err)

	return result.LastInsertId()
}

func GetQuote(sid string) []Quote {
	return getQuotesFromDatabase("SELECT id,body,notes,rating FROM quotes WHERE id = ?", sid)
}

func Browse(page int) []Quote {
	return getQuotesFromDatabase("SELECT id,body,notes,rating FROM quotes ORDER BY id ASC LIMIT 20 OFFSET ?", page*20)
}

func Latest() []Quote {
	return getQuotesFromDatabase("SELECT id,body,notes,rating FROM quotes ORDER BY id DESC LIMIT 20")
}

func Top() []Quote {
	return getQuotesFromDatabase("SELECT id,body,notes,rating FROM quotes ORDER BY rating DESC LIMIT 20")
}

func Search(searchText string) []Quote {
	query := "SELECT id,body,notes,rating FROM quotes WHERE 1=1"
	terms := strings.Split(searchText, " ")
	for i := 0; i < len(terms); i++ {
		// This took WAY too long for what it was.
		// Note to future self: Go doesn't like '%?%'. It takes it literally and
		// ignores the question mark as a binding parameter.
		query += " AND body LIKE '%' || ? || '%'"
	}
	query += " ORDER BY id DESC"

	// We have to cast `terms` to []interface{} because Go sucks
	args := make([]interface{}, len(terms))

	for i := range terms {
		args[i] = terms[i]
	}

	return getQuotesFromDatabase(query, args...)
}

func (q *Quote) String() string {
	return fmt.Sprint("%d: %s", q.ID, q.Text)
}

func getQuotesFromDatabase(statement string, args ...interface{}) []Quote {
	var q []Quote

	rows, err := db.Query(statement, args...)
	checkErr(err)

	for rows.Next() {
		var body, notes string
		var id, rating uint
		err = rows.Scan(&id, &body, &notes, &rating)
		checkErr(err)
		newQuote := Quote{
			ID:     id,
			Text:   body,
			Notes:  notes,
			Rating: rating,
		}
		q = append(q, newQuote)
	}

	return q
}
<|endoftext|>


<commit_before>// xgettext-soy is a tool to extract messages from Soy templates in the PO
// (gettext) file format.
package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/robfig/gettext/po"
	"github.com/robfig/soy/ast"
	"github.com/robfig/soy/parse"
	"github.com/robfig/soy/parsepasses"
	"github.com/robfig/soy/soymsg/pomsg"
	"github.com/robfig/soy/template"
)

func usage() {
	fmt.Println(`xgettext-soy is a tool to extract messages from Soy templates.

Usage:

	./xgettext-soy [INPUTPATH]...

INPUTPATH elements may be files or directories. Input directories will be
recursively searched for *.soy files.

The resulting POT (PO template) file is written to STDOUT`)
}

var registry = template.Registry{}

func main() {
	if len(os.Args) < 2 || strings.HasSuffix(os.Args[1], "help") {
		usage()
		os.Exit(1)
	}

	// Add all the sources to the registry.
	for _, src := range os.Args[1:] {
		err := filepath.Walk(src, walkSource)
		if err != nil {
			exit(err)
		}
	}
	parsepasses.ProcessMessages(registry)

	var e = extractor{&po.File{}}
	for _, t := range registry.Templates {
		e.extract(t.Node)
	}
	e.file.WriteTo(os.Stdout)
}

func walkSource(path string, info os.FileInfo, err error) error {
	if err != nil {
		return err
	}
	if !strings.HasSuffix(path, ".soy") {
		return nil
	}

	content, err := ioutil.ReadFile(path)
	if err != nil {
		return err
	}

	tree, err := parse.SoyFile(path, string(content))
	if err != nil {
		return err
	}
	if err = registry.Add(tree); err != nil {
		return err
	}
	return nil
}

type extractor struct {
	file *po.File
}

func (e extractor) extract(node ast.Node) {
	switch node := node.(type) {
	case *ast.MsgNode:
		if err := pomsg.Validate(node); err != nil {
			exit(err)
		}
		var pluralVar = ""
		if plural, ok := node.Body.Children()[0].(*ast.MsgPluralNode); ok {
			pluralVar = " var=" + plural.VarName
		}
		e.file.Messages = append(e.file.Messages, po.Message{
			Comment: po.Comment{
				ExtractedComments: []string{node.Desc},
				References:        []string{fmt.Sprintf("id=%d%v", node.ID, pluralVar)},
			},
			Ctxt:     node.Meaning,
			Id:       pomsg.Msgid(node),
			IdPlural: pomsg.MsgidPlural(node),
		})
	default:
		if parent, ok := node.(ast.ParentNode); ok {
			for _, child := range parent.Children() {
				e.extract(child)
			}
		}
	}
}

func exit(err error) {
	fmt.Fprintln(os.Stderr, err)
	os.Exit(1)
}
<commit_msg>xgettext-soy: write usage to stderr<commit_after>// xgettext-soy is a tool to extract messages from Soy templates in the PO
// (gettext) file format.
package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/robfig/gettext/po"
	"github.com/robfig/soy/ast"
	"github.com/robfig/soy/parse"
	"github.com/robfig/soy/parsepasses"
	"github.com/robfig/soy/soymsg/pomsg"
	"github.com/robfig/soy/template"
)

func usage() {
	fmt.Fprintln(os.Stderr, `xgettext-soy is a tool to extract messages from Soy templates.

Usage:

	./xgettext-soy [INPUTPATH]...

INPUTPATH elements may be files or directories. Input directories will be
recursively searched for *.soy files.

The resulting POT (PO template) file is written to STDOUT`)
}

var registry = template.Registry{}

func main() {
	if len(os.Args) < 2 || strings.HasSuffix(os.Args[1], "help") {
		usage()
		os.Exit(1)
	}

	// Add all the sources to the registry.
	for _, src := range os.Args[1:] {
		err := filepath.Walk(src, walkSource)
		if err != nil {
			exit(err)
		}
	}
	parsepasses.ProcessMessages(registry)

	var e = extractor{&po.File{}}
	for _, t := range registry.Templates {
		e.extract(t.Node)
	}
	e.file.WriteTo(os.Stdout)
}

func walkSource(path string, info os.FileInfo, err error) error {
	if err != nil {
		return err
	}
	if !strings.HasSuffix(path, ".soy") {
		return nil
	}

	content, err := ioutil.ReadFile(path)
	if err != nil {
		return err
	}

	tree, err := parse.SoyFile(path, string(content))
	if err != nil {
		return err
	}
	if err = registry.Add(tree); err != nil {
		return err
	}
	return nil
}

type extractor struct {
	file *po.File
}

func (e extractor) extract(node ast.Node) {
	switch node := node.(type) {
	case *ast.MsgNode:
		if err := pomsg.Validate(node); err != nil {
			exit(err)
		}
		var pluralVar = ""
		if plural, ok := node.Body.Children()[0].(*ast.MsgPluralNode); ok {
			pluralVar = " var=" + plural.VarName
		}
		e.file.Messages = append(e.file.Messages, po.Message{
			Comment: po.Comment{
				ExtractedComments: []string{node.Desc},
				References:        []string{fmt.Sprintf("id=%d%v", node.ID, pluralVar)},
			},
			Ctxt:     node.Meaning,
			Id:       pomsg.Msgid(node),
			IdPlural: pomsg.MsgidPlural(node),
		})
	default:
		if parent, ok := node.(ast.ParentNode); ok {
			for _, child := range parent.Children() {
				e.extract(child)
			}
		}
	}
}

func exit(err error) {
	fmt.Fprintln(os.Stderr, err)
	os.Exit(1)
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"errors"
	"fmt"
	"io"
	"path/filepath"

	"github.com/eiblog/eiblog/setting"
	"github.com/qiniu/api.v7/auth/qbox"
	"github.com/qiniu/api.v7/storage"
)

// 进度条
func onProgress(fsize, uploaded int64) {
	d := int(float64(uploaded) / float64(fsize) * 100)
	if fsize == uploaded {
		fmt.Printf("\rUpload completed!          \n")
	} else {
		fmt.Printf("\r%02d%% uploaded              ", int(d))
	}
}

// 上传文件
func FileUpload(name string, size int64, data io.Reader) (string, error) {
	if setting.Conf.Qiniu.AccessKey == "" || setting.Conf.Qiniu.SecretKey == "" {
		return "", errors.New("qiniu config error")
	}

	key := getKey(name)
	mac := qbox.NewMac(setting.Conf.Qiniu.AccessKey, setting.Conf.Qiniu.SecretKey)
	// 设置上传的策略
	putPolicy := &storage.PutPolicy{
		Scope:      setting.Conf.Qiniu.Bucket,
		Expires:    3600,
		InsertOnly: 1,
	}
	// 上传token
	upToken := putPolicy.UploadToken(mac)

	// 上传配置
	cfg := &storage.Config{
		Zone:     &storage.ZoneHuadong,
		UseHTTPS: true,
	}
	// uploader
	uploader := storage.NewFormUploader(cfg)
	ret := new(storage.PutRet)
	putExtra := &storage.PutExtra{}

	err := uploader.Put(context.Background(), ret, upToken, key, data, size, putExtra)
	if err != nil {
		return "", err
	}

	url := "https://" + setting.Conf.Qiniu.Domain + "/" + key
	return url, nil
}

// 删除文件
func FileDelete(name string) error {
	key := getKey(name)

	mac := qbox.NewMac(setting.Conf.Qiniu.AccessKey, setting.Conf.Qiniu.SecretKey)
	// 上传配置
	cfg := &storage.Config{
		Zone:     &storage.ZoneHuadong,
		UseHTTPS: true,
	}
	// manager
	bucketManager := storage.NewBucketManager(mac, cfg)
	// Delete
	err := bucketManager.Delete(setting.Conf.Qiniu.Bucket, key)
	if err != nil {
		return err
	}
	return nil
}

// 修复路径
func getKey(name string) string {
	ext := filepath.Ext(name)
	var key string
	switch ext {
	case ".bmp", ".png", ".jpg", ".gif", ".ico":
		key = "blog/img/" + name
	case ".mov", ".mp4":
		key = "blog/video/" + name
	case ".go", ".js", ".css", ".cpp", ".php", ".rb",
		".java", ".py", ".sql", ".lua", ".html",
		".sh", ".xml", ".cs":
		key = "blog/code/" + name
	case ".txt", ".md", ".ini", ".yaml", ".yml",
		".doc", ".ppt", ".pdf":
		key = "blog/document/" + name
	case ".zip", ".rar", ".tar", ".gz":
		key = "blog/archive/" + name
	}
	return key
}
<commit_msg>fix(qiniu): empty file name<commit_after>package main

import (
	"context"
	"errors"
	"fmt"
	"io"
	"path/filepath"

	"github.com/eiblog/eiblog/setting"
	"github.com/qiniu/api.v7/auth/qbox"
	"github.com/qiniu/api.v7/storage"
)

// 进度条
func onProgress(fsize, uploaded int64) {
	d := int(float64(uploaded) / float64(fsize) * 100)
	if fsize == uploaded {
		fmt.Printf("\rUpload completed!          \n")
	} else {
		fmt.Printf("\r%02d%% uploaded              ", int(d))
	}
}

// 上传文件
func FileUpload(name string, size int64, data io.Reader) (string, error) {
	if setting.Conf.Qiniu.AccessKey == "" || setting.Conf.Qiniu.SecretKey == "" {
		return "", errors.New("qiniu config error")
	}

	key := getKey(name)
	mac := qbox.NewMac(setting.Conf.Qiniu.AccessKey, setting.Conf.Qiniu.SecretKey)
	// 设置上传的策略
	putPolicy := &storage.PutPolicy{
		Scope:      setting.Conf.Qiniu.Bucket,
		Expires:    3600,
		InsertOnly: 1,
	}
	// 上传token
	upToken := putPolicy.UploadToken(mac)

	// 上传配置
	cfg := &storage.Config{
		Zone:     &storage.ZoneHuadong,
		UseHTTPS: true,
	}
	// uploader
	uploader := storage.NewFormUploader(cfg)
	ret := new(storage.PutRet)
	putExtra := &storage.PutExtra{}

	err := uploader.Put(context.Background(), ret, upToken, key, data, size, putExtra)
	if err != nil {
		return "", err
	}

	url := "https://" + setting.Conf.Qiniu.Domain + "/" + key
	return url, nil
}

// 删除文件
func FileDelete(name string) error {
	key := getKey(name)

	mac := qbox.NewMac(setting.Conf.Qiniu.AccessKey, setting.Conf.Qiniu.SecretKey)
	// 上传配置
	cfg := &storage.Config{
		Zone:     &storage.ZoneHuadong,
		UseHTTPS: true,
	}
	// manager
	bucketManager := storage.NewBucketManager(mac, cfg)
	// Delete
	err := bucketManager.Delete(setting.Conf.Qiniu.Bucket, key)
	if err != nil {
		return err
	}
	return nil
}

// 修复路径
func getKey(name string) string {
	ext := filepath.Ext(name)
	var key string
	switch ext {
	case ".bmp", ".png", ".jpg", ".gif", ".ico", ".jpeg":
		key = "blog/img/" + name
	case ".mov", ".mp4":
		key = "blog/video/" + name
	case ".go", ".js", ".css", ".cpp", ".php", ".rb",
		".java", ".py", ".sql", ".lua", ".html",
		".sh", ".xml", ".cs":
		key = "blog/code/" + name
	case ".txt", ".md", ".ini", ".yaml", ".yml",
		".doc", ".ppt", ".pdf":
		key = "blog/document/" + name
	case ".zip", ".rar", ".tar", ".gz":
		key = "blog/archive/" + name
	default:
		key = "blog/other/" + name
	}
	return key
}
<|endoftext|>


<commit_before>package jobs

import (
	"fmt"
	"strings"
	"sync"
	"time"

	log "github.com/meifamily/logrus"

	"github.com/meifamily/ptt-alertor/models/ptt/article"
	board "github.com/meifamily/ptt-alertor/models/ptt/board/redis"
	user "github.com/meifamily/ptt-alertor/models/user/redis"
	"github.com/meifamily/ptt-alertor/myutil"
)

const checkBoardDuration = 200 * time.Millisecond
const checkHighBoardDuration = 1 * time.Second

var boardCh = make(chan *board.Board)
var ckerCh = make(chan Checker)

type Checker struct {
	email        string
	line         string
	lineNotify   string
	messenger    string
	telegram     string
	telegramChat int64
	board        string
	keyword      string
	author       string
	articles     article.Articles
	subType      string
	word         string
}

func (cker Checker) String() string {
	subType := "關鍵字"
	if cker.author != "" {
		subType = "作者"
	}
	return fmt.Sprintf("%s@%s\r\n看板：%s；%s：%s%s", cker.word, cker.board, cker.board, subType, cker.word, cker.articles.String())
}

// Self return Checker itself
func (cker Checker) Self() Checker {
	return cker
}

// Run is main in Job
func (cker Checker) Run() {
	highBoards := highBoards()
	var wgHigh sync.WaitGroup
	var wg sync.WaitGroup
	go func(highBoards []*board.Board) {
		for {
			checkBoards(&wgHigh, highBoards, checkHighBoardDuration)
			wgHigh.Wait()
		}
	}(highBoards)
	go func() {
		for {
			checkBoards(&wg, new(board.Board).All(), checkBoardDuration)
			wg.Wait()
		}
	}()

	for {
		select {
		case bd := <-boardCh:
			checkSubscriber(bd, cker)
		case cker := <-ckerCh:
			ckCh <- cker
		}
	}
}

func highBoards() (highBoards []*board.Board) {
	boardcfg := myutil.Config("board")
	highBoardNames := strings.Split(boardcfg["high"], ",")
	for _, name := range highBoardNames {
		bd := new(board.Board)
		bd.Name = name
		highBoards = append(highBoards, bd)
	}
	return highBoards
}

func checkBoards(wg *sync.WaitGroup, bds []*board.Board, duration time.Duration) {
	wg.Add(len(bds))
	for _, bd := range bds {
		time.Sleep(duration)
		go checkNewArticle(wg, bd, boardCh)
	}
}

func checkNewArticle(wg *sync.WaitGroup, bd *board.Board, boardCh chan *board.Board) {
	defer wg.Done()
	bd.WithNewArticles()
	if bd.NewArticles == nil {
		bd.Articles = bd.OnlineArticles
		log.WithField("board", bd.Name).Info("Created Articles")
		bd.Save()
	}
	if len(bd.NewArticles) != 0 {
		bd.Articles = bd.OnlineArticles
		log.WithField("board", bd.Name).Info("Updated Articles")
		err := bd.Save()
		if err == nil {
			boardCh <- bd
		}
	}
}

func checkSubscriber(bd *board.Board, cker Checker) {
	users := new(user.User).All()
	for _, user := range users {
		if user.Enable {
			cker.email = user.Profile.Email
			cker.line = user.Profile.Line
			cker.lineNotify = user.Profile.LineAccessToken
			cker.messenger = user.Profile.Messenger
			cker.telegram = user.Profile.Telegram
			cker.telegramChat = user.Profile.TelegramChat
			go subscribeChecker(user, bd, cker)
		}
	}
}

func subscribeChecker(user *user.User, bd *board.Board, cker Checker) {
	for _, sub := range user.Subscribes {
		if bd.Name == sub.Board {
			cker.board = sub.Board
			for _, keyword := range sub.Keywords {
				go keywordChecker(keyword, bd, cker)
			}
			for _, author := range sub.Authors {
				go authorChecker(author, bd, cker)
			}
		}
	}
}

func keywordChecker(keyword string, bd *board.Board, cker Checker) {
	keywordArticles := make(article.Articles, 0)
	for _, newAtcl := range bd.NewArticles {
		if newAtcl.MatchKeyword(keyword) {
			newAtcl.Author = ""
			keywordArticles = append(keywordArticles, newAtcl)
		}
	}
	if len(keywordArticles) != 0 {
		cker.keyword = keyword
		cker.articles = keywordArticles
		cker.subType = "keyword"
		cker.word = keyword
		ckerCh <- cker
	}
}

func authorChecker(author string, bd *board.Board, cker Checker) {
	authorArticles := make(article.Articles, 0)
	for _, newAtcl := range bd.NewArticles {
		if strings.EqualFold(newAtcl.Author, author) {
			authorArticles = append(authorArticles, newAtcl)
		}
	}
	if len(authorArticles) != 0 {
		cker.author = author
		cker.articles = authorArticles
		cker.subType = "author"
		cker.word = author
		ckerCh <- cker
	}
}
<commit_msg>increase duration to testing cpu usage<commit_after>package jobs

import (
	"fmt"
	"strings"
	"sync"
	"time"

	log "github.com/meifamily/logrus"

	"github.com/meifamily/ptt-alertor/models/ptt/article"
	board "github.com/meifamily/ptt-alertor/models/ptt/board/redis"
	user "github.com/meifamily/ptt-alertor/models/user/redis"
	"github.com/meifamily/ptt-alertor/myutil"
)

const checkBoardDuration = 500 * time.Millisecond
const checkHighBoardDuration = 1 * time.Second

var boardCh = make(chan *board.Board)
var ckerCh = make(chan Checker)

type Checker struct {
	email        string
	line         string
	lineNotify   string
	messenger    string
	telegram     string
	telegramChat int64
	board        string
	keyword      string
	author       string
	articles     article.Articles
	subType      string
	word         string
}

func (cker Checker) String() string {
	subType := "關鍵字"
	if cker.author != "" {
		subType = "作者"
	}
	return fmt.Sprintf("%s@%s\r\n看板：%s；%s：%s%s", cker.word, cker.board, cker.board, subType, cker.word, cker.articles.String())
}

// Self return Checker itself
func (cker Checker) Self() Checker {
	return cker
}

// Run is main in Job
func (cker Checker) Run() {
	highBoards := highBoards()
	var wgHigh sync.WaitGroup
	var wg sync.WaitGroup
	go func(highBoards []*board.Board) {
		for {
			checkBoards(&wgHigh, highBoards, checkHighBoardDuration)
			wgHigh.Wait()
		}
	}(highBoards)
	go func() {
		for {
			checkBoards(&wg, new(board.Board).All(), checkBoardDuration)
			wg.Wait()
		}
	}()

	for {
		select {
		case bd := <-boardCh:
			checkSubscriber(bd, cker)
		case cker := <-ckerCh:
			ckCh <- cker
		}
	}
}

func highBoards() (highBoards []*board.Board) {
	boardcfg := myutil.Config("board")
	highBoardNames := strings.Split(boardcfg["high"], ",")
	for _, name := range highBoardNames {
		bd := new(board.Board)
		bd.Name = name
		highBoards = append(highBoards, bd)
	}
	return highBoards
}

func checkBoards(wg *sync.WaitGroup, bds []*board.Board, duration time.Duration) {
	wg.Add(len(bds))
	for _, bd := range bds {
		time.Sleep(duration)
		go checkNewArticle(wg, bd, boardCh)
	}
}

func checkNewArticle(wg *sync.WaitGroup, bd *board.Board, boardCh chan *board.Board) {
	defer wg.Done()
	bd.WithNewArticles()
	if bd.NewArticles == nil {
		bd.Articles = bd.OnlineArticles
		log.WithField("board", bd.Name).Info("Created Articles")
		bd.Save()
	}
	if len(bd.NewArticles) != 0 {
		bd.Articles = bd.OnlineArticles
		log.WithField("board", bd.Name).Info("Updated Articles")
		err := bd.Save()
		if err == nil {
			boardCh <- bd
		}
	}
}

func checkSubscriber(bd *board.Board, cker Checker) {
	users := new(user.User).All()
	for _, user := range users {
		if user.Enable {
			cker.email = user.Profile.Email
			cker.line = user.Profile.Line
			cker.lineNotify = user.Profile.LineAccessToken
			cker.messenger = user.Profile.Messenger
			cker.telegram = user.Profile.Telegram
			cker.telegramChat = user.Profile.TelegramChat
			go subscribeChecker(user, bd, cker)
		}
	}
}

func subscribeChecker(user *user.User, bd *board.Board, cker Checker) {
	for _, sub := range user.Subscribes {
		if bd.Name == sub.Board {
			cker.board = sub.Board
			for _, keyword := range sub.Keywords {
				go keywordChecker(keyword, bd, cker)
			}
			for _, author := range sub.Authors {
				go authorChecker(author, bd, cker)
			}
		}
	}
}

func keywordChecker(keyword string, bd *board.Board, cker Checker) {
	keywordArticles := make(article.Articles, 0)
	for _, newAtcl := range bd.NewArticles {
		if newAtcl.MatchKeyword(keyword) {
			newAtcl.Author = ""
			keywordArticles = append(keywordArticles, newAtcl)
		}
	}
	if len(keywordArticles) != 0 {
		cker.keyword = keyword
		cker.articles = keywordArticles
		cker.subType = "keyword"
		cker.word = keyword
		ckerCh <- cker
	}
}

func authorChecker(author string, bd *board.Board, cker Checker) {
	authorArticles := make(article.Articles, 0)
	for _, newAtcl := range bd.NewArticles {
		if strings.EqualFold(newAtcl.Author, author) {
			authorArticles = append(authorArticles, newAtcl)
		}
	}
	if len(authorArticles) != 0 {
		cker.author = author
		cker.articles = authorArticles
		cker.subType = "author"
		cker.word = author
		ckerCh <- cker
	}
}
<|endoftext|>


<commit_before>// Package journal provides write bindings to the systemd journal
package journal

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"strconv"
	"strings"
	"syscall"
	"encoding/binary"
)

// Priority of a journal message
type Priority int

const (
	PriEmerg Priority = iota
	PriAlert
	PriCrit
	PriErr
	PriWarning
	PriNotice
	PriInfo
	PriDebug
)

var conn net.Conn

func init() {
	var err error
	conn, err = net.Dial("unixgram", "/run/systemd/journal/socket")
}

// Enabled returns true iff the systemd journal is available for logging
func Enabled() bool {
	return conn != nil
}

// Send a message to the systemd journal. vars is a map of journald fields to
// values.  Fields must be composed of uppercase letters, numbers, and
// underscores, but must not start with an underscore. Within these
// restrictions, any arbitrary field name may be used.  Some names have special
// significance: see the journalctl documentation
// (http://www.freedesktop.org/software/systemd/man/systemd.journal-fields.html)
// for more details.  vars may be nil.
func Send(message string, priority Priority, vars map[string]string) error {
	if conn == nil {
		return journalError("could not connect to journald socket")
	}

	data := new(bytes.Buffer)
	appendVariable(data, "PRIORITY", strconv.Itoa(int(priority)))
	appendVariable(data, "MESSAGE", message)
	for k, v := range vars {
		appendVariable(data, k, v)
	}

	_, err := io.Copy(conn, data)
	if err != nil && isSocketSpaceError(err) {
		file, err := tempFd()
		if err != nil {
			return journalError(err.Error())
		}
		_, err = io.Copy(file, data)
		if err != nil {
			return journalError(err.Error())
		}

		rights := syscall.UnixRights(int(file.Fd()))

		/* this connection should always be a UnixConn, but better safe than sorry */
		unixConn, ok := conn.(*net.UnixConn)
		if !ok {
			return journalError("can't send file through non-Unix connection")
		}
		unixConn.WriteMsgUnix([]byte{}, rights, nil)
	} else if err != nil {
		return journalError(err.Error())
	}
	return nil
}

func appendVariable(w io.Writer, name, value string) {
	if !validVarName(name) {
		journalError("variable name contains invalid character, ignoring")
	}
	if strings.ContainsRune(value, '\n') {
		/* When the value contains a newline, we write:
		 * - the variable name, followed by a newline
		 * - the size (in 64bit little endian format)
		 * - the data, followed by a newline
		 */
		fmt.Fprintln(w, name)
		binary.Write(w, binary.LittleEndian, uint64(len(value)))
		fmt.Fprintln(w, value)
	} else {
		/* just write the variable and value all on one line */
		fmt.Fprintf(w, "%s=%s\n", name, value)
	}
}

func validVarName(name string) bool {
	/* The variable name must be in uppercase and consist only of characters,
	 * numbers and underscores, and may not begin with an underscore. (from the docs)
	 */
	valid := true
	valid = valid && name[0] != '_'
	for _, c := range name {
		valid = valid && ('A' <= c && c <= 'Z') || ('0' <= c && c <= '9') || c == '_'
	}
	return valid
}

func isSocketSpaceError(err error) bool {
	opErr, ok := err.(*net.OpError)
	if !ok {
		return false
	}

	sysErr, ok := opErr.Err.(syscall.Errno)
	if !ok {
		return false
	}

	return sysErr == syscall.EMSGSIZE || sysErr == syscall.ENOBUFS
}

func tempFd() (*os.File, error) {
	file, err := ioutil.TempFile("/dev/shm/", "journal.XXXXX")
	if err != nil {
		return nil, err
	}
	syscall.Unlink(file.Name())
	if err != nil {
		return nil, err
	}
	return file, nil
}

func journalError(s string) error {
	s = "journal error: " + s
	fmt.Fprintln(os.Stderr, s)
	return errors.New(s)
}
<commit_msg>Fix compile error<commit_after>// Package journal provides write bindings to the systemd journal
package journal

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"strconv"
	"strings"
	"syscall"
	"encoding/binary"
)

// Priority of a journal message
type Priority int

const (
	PriEmerg Priority = iota
	PriAlert
	PriCrit
	PriErr
	PriWarning
	PriNotice
	PriInfo
	PriDebug
)

var conn net.Conn

func init() {
	conn, _ = net.Dial("unixgram", "/run/systemd/journal/socket")
}

// Enabled returns true iff the systemd journal is available for logging
func Enabled() bool {
	return conn != nil
}

// Send a message to the systemd journal. vars is a map of journald fields to
// values.  Fields must be composed of uppercase letters, numbers, and
// underscores, but must not start with an underscore. Within these
// restrictions, any arbitrary field name may be used.  Some names have special
// significance: see the journalctl documentation
// (http://www.freedesktop.org/software/systemd/man/systemd.journal-fields.html)
// for more details.  vars may be nil.
func Send(message string, priority Priority, vars map[string]string) error {
	if conn == nil {
		return journalError("could not connect to journald socket")
	}

	data := new(bytes.Buffer)
	appendVariable(data, "PRIORITY", strconv.Itoa(int(priority)))
	appendVariable(data, "MESSAGE", message)
	for k, v := range vars {
		appendVariable(data, k, v)
	}

	_, err := io.Copy(conn, data)
	if err != nil && isSocketSpaceError(err) {
		file, err := tempFd()
		if err != nil {
			return journalError(err.Error())
		}
		_, err = io.Copy(file, data)
		if err != nil {
			return journalError(err.Error())
		}

		rights := syscall.UnixRights(int(file.Fd()))

		/* this connection should always be a UnixConn, but better safe than sorry */
		unixConn, ok := conn.(*net.UnixConn)
		if !ok {
			return journalError("can't send file through non-Unix connection")
		}
		unixConn.WriteMsgUnix([]byte{}, rights, nil)
	} else if err != nil {
		return journalError(err.Error())
	}
	return nil
}

func appendVariable(w io.Writer, name, value string) {
	if !validVarName(name) {
		journalError("variable name contains invalid character, ignoring")
	}
	if strings.ContainsRune(value, '\n') {
		/* When the value contains a newline, we write:
		 * - the variable name, followed by a newline
		 * - the size (in 64bit little endian format)
		 * - the data, followed by a newline
		 */
		fmt.Fprintln(w, name)
		binary.Write(w, binary.LittleEndian, uint64(len(value)))
		fmt.Fprintln(w, value)
	} else {
		/* just write the variable and value all on one line */
		fmt.Fprintf(w, "%s=%s\n", name, value)
	}
}

func validVarName(name string) bool {
	/* The variable name must be in uppercase and consist only of characters,
	 * numbers and underscores, and may not begin with an underscore. (from the docs)
	 */
	valid := true
	valid = valid && name[0] != '_'
	for _, c := range name {
		valid = valid && ('A' <= c && c <= 'Z') || ('0' <= c && c <= '9') || c == '_'
	}
	return valid
}

func isSocketSpaceError(err error) bool {
	opErr, ok := err.(*net.OpError)
	if !ok {
		return false
	}

	sysErr, ok := opErr.Err.(syscall.Errno)
	if !ok {
		return false
	}

	return sysErr == syscall.EMSGSIZE || sysErr == syscall.ENOBUFS
}

func tempFd() (*os.File, error) {
	file, err := ioutil.TempFile("/dev/shm/", "journal.XXXXX")
	if err != nil {
		return nil, err
	}
	syscall.Unlink(file.Name())
	if err != nil {
		return nil, err
	}
	return file, nil
}

func journalError(s string) error {
	s = "journal error: " + s
	fmt.Fprintln(os.Stderr, s)
	return errors.New(s)
}
<|endoftext|>


<commit_before>package wc

import "testing"

func assertEqual(t *testing.T, actual Histogram, expected Histogram) {
	if !actual.Equal(expected) {
		t.Errorf("Expected %v to equal %v", actual, expected)
	}
}

func assertNotEqual(t *testing.T, h Histogram, other Histogram) {
	if h.Equal(other) || other.Equal(h) {
		t.Errorf("Expected %v to not equal %v", h, other)
	}
}

func TestEqual(t *testing.T) {
	h1 := Histogram{"hello": 1, "world": 2}
	h2 := Histogram{"hello": 1, "world": 2}
	assertEqual(t, h1, h2)
}

func TestNotEqual(t *testing.T) {
	t.SkipNow()
	h1 := Histogram{"word": 1}
	h2 := Histogram{"word": 1, "games": 2}
	assertNotEqual(t, h1, h2)
}

func TestCountOneWord(t *testing.T) {
	t.SkipNow()
	h := Histogram{"word": 1}
	assertEqual(t, WordCount("word"), h)
}

func TestCountOneOfEach(t *testing.T) {
	t.SkipNow()
	h := Histogram{"one": 1, "of": 1, "each": 1}
	assertEqual(t, WordCount("one of each"), h)
}

func TestCountMultipleOccurrences(t *testing.T) {
	t.SkipNow()
	actual := WordCount("one fish two fish red fish blue fish")
	expected := Histogram{"one": 1, "fish": 4, "two": 1, "red": 1, "blue": 1}
	assertEqual(t, actual, expected)
}

func TestIgnorePunctuation(t *testing.T) {
	t.SkipNow()
	actual := WordCount("car : carpet as java : javascript!!&@$%^&")
	expected := Histogram{"car": 1, "carpet": 1, "as": 1, "java": 1, "javascript": 1}
	assertEqual(t, actual, expected)
}

func TestIncludeNumbers(t *testing.T) {
	t.SkipNow()
	actual := WordCount("testing, 1, 2 testing")
	expected := Histogram{"testing": 2, "1": 1, "2": 1}
	assertEqual(t, actual, expected)
}

func TestNormalizeCase(t *testing.T) {
	t.SkipNow()
	assertEqual(t, WordCount("go Go GO"), Histogram{"go": 3})
}
<commit_msg>Remove SkipNow() in go/word-count. See #1304<commit_after>package wc

import (
	"fmt"
	"testing"
)

var testCases = []struct {
	description string
	input       string
	output      Histogram
}{
	{
		description: "a single word",
		input:       "word",
		output:      Histogram{"word": 1},
	},
	{
		description: "one of each",
		input:       "one of each",
		output:      Histogram{"one": 1, "of": 1, "each": 1},
	},
	{
		description: "multiple occurrences",
		input:       "one fish two fish red fish blue fish",
		output:      Histogram{"one": 1, "fish": 4, "two": 1, "red": 1, "blue": 1},
	},
	{
		description: "ignore punctuation",
		input:       "car : carpet as java : javascript!!&@$%^&",
		output:      Histogram{"car": 1, "carpet": 1, "as": 1, "java": 1, "javascript": 1},
	},
	{
		description: "including numbers",
		input:       "testing, 1, 2 testing",
		output:      Histogram{"testing": 2, "1": 1, "2": 1},
	},
	{
		description: "normalises case",
		input:       "go Go GO",
		output:      Histogram{"go": 3},
	},
}

func TestWordCount(t *testing.T) {
	for _, tt := range testCases {
		expected := fmt.Sprintf("%v", tt.output)
		actual := fmt.Sprintf("%v", WordCount(tt.input))

		if expected != actual {
			t.Fatalf("%s\n\tExpected: %v\n\tGot: %v", tt.description, expected, actual)
		} else {
			t.Logf("PASS: %s - WordCount(%s)", tt.description, tt.input)
		}
	}
}
<|endoftext|>


<commit_before>package server

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"time"

	protolion "go.pedge.io/lion"
	protorpclog "go.pedge.io/proto/rpclog"
	"golang.org/x/net/context"
	"golang.org/x/sync/errgroup"

	"github.com/cenkalti/backoff"
	"github.com/gogo/protobuf/types"
	"github.com/golang/groupcache"
	"github.com/pachyderm/pachyderm/src/client"
	pfsclient "github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/server/pkg/obj"
)

type objBlockAPIServer struct {
	protorpclog.Logger
	dir         string
	localServer *localBlockAPIServer
	objClient   obj.Client
	cache       *groupcache.Group
}

func newObjBlockAPIServer(dir string, cacheBytes int64, objClient obj.Client) (*objBlockAPIServer, error) {
	localServer, err := newLocalBlockAPIServer(dir)
	if err != nil {
		return nil, err
	}
	return &objBlockAPIServer{
		Logger:      protorpclog.NewLogger("pfs.BlockAPI.Obj"),
		dir:         dir,
		localServer: localServer,
		objClient:   objClient,
		cache: groupcache.NewGroup("block", 1024*1024*1024*10,
			groupcache.GetterFunc(func(ctx groupcache.Context, key string, dest groupcache.Sink) (retErr error) {
				var reader io.ReadCloser
				var err error
				backoff.RetryNotify(func() error {
					reader, err = objClient.Reader(localServer.blockPath(client.NewBlock(key)), 0, 0)
					if err != nil && objClient.IsRetryable(err) {
						return err
					}
					return nil
				}, obj.NewExponentialBackOffConfig(), func(err error, d time.Duration) {
					protolion.Infof("Error creating reader; retrying in %s: %#v", d, obj.RetryError{
						Err:               err.Error(),
						TimeTillNextRetry: d.String(),
					})
				})
				if err != nil {
					return err
				}
				defer func() {
					if err := reader.Close(); err != nil && retErr == nil {
						retErr = err
					}
				}()
				block, err := ioutil.ReadAll(reader)
				if err != nil {
					return err
				}
				return dest.SetBytes(block)
			})),
	}, nil
}

func newMinioBlockAPIServer(dir string, cacheBytes int64) (*objBlockAPIServer, error) {
	objClient, err := obj.NewMinioClientFromSecret("")
	if err != nil {
		return nil, err
	}
	return newObjBlockAPIServer(dir, cacheBytes, objClient)
}

func newAmazonBlockAPIServer(dir string, cacheBytes int64) (*objBlockAPIServer, error) {
	objClient, err := obj.NewAmazonClientFromSecret("")
	if err != nil {
		return nil, err
	}
	return newObjBlockAPIServer(dir, cacheBytes, objClient)
}

func newGoogleBlockAPIServer(dir string, cacheBytes int64) (*objBlockAPIServer, error) {
	objClient, err := obj.NewGoogleClientFromSecret(context.Background(), "")
	if err != nil {
		return nil, err
	}
	return newObjBlockAPIServer(dir, cacheBytes, objClient)
}

func newMicrosoftBlockAPIServer(dir string, cacheBytes int64) (*objBlockAPIServer, error) {
	objClient, err := obj.NewMicrosoftClientFromSecret("")
	if err != nil {
		return nil, err
	}
	return newObjBlockAPIServer(dir, cacheBytes, objClient)
}

func (s *objBlockAPIServer) PutBlock(putBlockServer pfsclient.BlockAPI_PutBlockServer) (retErr error) {
	func() { s.Log(nil, nil, nil, 0) }()
	result := &pfsclient.BlockRefs{}
	defer func(start time.Time) { s.Log(nil, result, retErr, time.Since(start)) }(time.Now())
	defer drainBlockServer(putBlockServer)
	putBlockRequest, err := putBlockServer.Recv()
	if err != nil {
		if err != io.EOF {
			return err
		}
		return putBlockServer.SendAndClose(result)
	}
	reader := bufio.NewReader(&putBlockReader{
		server: putBlockServer,
		buffer: bytes.NewBuffer(putBlockRequest.Value),
	})
	var eg errgroup.Group
	decoder := json.NewDecoder(reader)
	for {
		blockRef, data, err := readBlock(putBlockRequest.Delimiter, reader, decoder)
		if err != nil {
			return err
		}
		result.BlockRef = append(result.BlockRef, blockRef)
		eg.Go(func() (retErr error) {
			backoff.RetryNotify(func() error {
				path := s.localServer.blockPath(blockRef.Block)
				// We don't want to overwrite blocks that already exist, since:
				// 1) blocks are content-addressable, so it will be the same block
				// 2) we risk exceeding the object store's rate limit
				if s.objClient.Exists(path) {
					return nil
				}
				writer, err := s.objClient.Writer(path)
				if err != nil {
					retErr = err
					return nil
				}
				if _, err := writer.Write(data); err != nil {
					retErr = err
					return nil
				}
				if err := writer.Close(); err != nil {
					if s.objClient.IsRetryable(err) {
						return err
					}
					retErr = err
					return nil
				}
				return nil
			}, obj.NewExponentialBackOffConfig(), func(err error, d time.Duration) {
				protolion.Infof("Error writing; retrying in %s: %#v", d, obj.RetryError{
					Err:               err.Error(),
					TimeTillNextRetry: d.String(),
				})
			})
			return
		})
		if (blockRef.Range.Upper - blockRef.Range.Lower) < uint64(blockSize) {
			break
		}
	}
	if err := eg.Wait(); err != nil {
		return err
	}
	return putBlockServer.SendAndClose(result)
}

func (s *objBlockAPIServer) GetBlock(request *pfsclient.GetBlockRequest, getBlockServer pfsclient.BlockAPI_GetBlockServer) (retErr error) {
	func() { s.Log(nil, nil, nil, 0) }()
	defer func(start time.Time) { s.Log(request, nil, retErr, time.Since(start)) }(time.Now())
	var data []byte
	sink := groupcache.AllocatingByteSliceSink(&data)
	if err := s.cache.Get(getBlockServer.Context(), request.Block.Hash, sink); err != nil {
		return err
	}
	if request.SizeBytes != 0 && request.SizeBytes+request.OffsetBytes < uint64(len(data)) {
		data = data[request.OffsetBytes : request.OffsetBytes+request.SizeBytes]
	} else if request.OffsetBytes < uint64(len(data)) {
		data = data[request.OffsetBytes:]
	} else {
		data = nil
	}
	return getBlockServer.Send(&types.BytesValue{Value: data})
}

func (s *objBlockAPIServer) DeleteBlock(ctx context.Context, request *pfsclient.DeleteBlockRequest) (response *types.Empty, retErr error) {
	func() { s.Log(nil, nil, nil, 0) }()
	defer func(start time.Time) { s.Log(request, response, retErr, time.Since(start)) }(time.Now())
	backoff.RetryNotify(func() error {
		if err := s.objClient.Delete(s.localServer.blockPath(request.Block)); err != nil && !s.objClient.IsNotExist(err) {
			return err
		}
		return nil
	}, obj.NewExponentialBackOffConfig(), func(err error, d time.Duration) {
		protolion.Infof("Error deleting block; retrying in %s: %#v", d, obj.RetryError{
			Err:               err.Error(),
			TimeTillNextRetry: d.String(),
		})
	})
	return &types.Empty{}, nil
}

func (s *objBlockAPIServer) InspectBlock(ctx context.Context, request *pfsclient.InspectBlockRequest) (response *pfsclient.BlockInfo, retErr error) {
	return nil, fmt.Errorf("not implemented")
}

func (s *objBlockAPIServer) ListBlock(ctx context.Context, request *pfsclient.ListBlockRequest) (response *pfsclient.BlockInfos, retErr error) {
	return nil, fmt.Errorf("not implemented")
}
<commit_msg>Insulate us against more obj store errors.<commit_after>package server

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"time"

	protolion "go.pedge.io/lion"
	protorpclog "go.pedge.io/proto/rpclog"
	"golang.org/x/net/context"
	"golang.org/x/sync/errgroup"

	"github.com/cenkalti/backoff"
	"github.com/gogo/protobuf/types"
	"github.com/golang/groupcache"
	"github.com/pachyderm/pachyderm/src/client"
	pfsclient "github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/server/pkg/obj"
)

type objBlockAPIServer struct {
	protorpclog.Logger
	dir         string
	localServer *localBlockAPIServer
	objClient   obj.Client
	cache       *groupcache.Group
}

func newObjBlockAPIServer(dir string, cacheBytes int64, objClient obj.Client) (*objBlockAPIServer, error) {
	localServer, err := newLocalBlockAPIServer(dir)
	if err != nil {
		return nil, err
	}
	return &objBlockAPIServer{
		Logger:      protorpclog.NewLogger("pfs.BlockAPI.Obj"),
		dir:         dir,
		localServer: localServer,
		objClient:   objClient,
		cache: groupcache.NewGroup("block", 1024*1024*1024*10,
			groupcache.GetterFunc(func(ctx groupcache.Context, key string, dest groupcache.Sink) (retErr error) {
				var reader io.ReadCloser
				var err error
				backoff.RetryNotify(func() error {
					reader, err = objClient.Reader(localServer.blockPath(client.NewBlock(key)), 0, 0)
					if err != nil && objClient.IsRetryable(err) {
						return err
					}
					return nil
				}, obj.NewExponentialBackOffConfig(), func(err error, d time.Duration) {
					protolion.Infof("Error creating reader; retrying in %s: %#v", d, obj.RetryError{
						Err:               err.Error(),
						TimeTillNextRetry: d.String(),
					})
				})
				if err != nil {
					return err
				}
				defer func() {
					if err := reader.Close(); err != nil && retErr == nil {
						retErr = err
					}
				}()
				block, err := ioutil.ReadAll(reader)
				if err != nil {
					return err
				}
				return dest.SetBytes(block)
			})),
	}, nil
}

func newMinioBlockAPIServer(dir string, cacheBytes int64) (*objBlockAPIServer, error) {
	objClient, err := obj.NewMinioClientFromSecret("")
	if err != nil {
		return nil, err
	}
	return newObjBlockAPIServer(dir, cacheBytes, objClient)
}

func newAmazonBlockAPIServer(dir string, cacheBytes int64) (*objBlockAPIServer, error) {
	objClient, err := obj.NewAmazonClientFromSecret("")
	if err != nil {
		return nil, err
	}
	return newObjBlockAPIServer(dir, cacheBytes, objClient)
}

func newGoogleBlockAPIServer(dir string, cacheBytes int64) (*objBlockAPIServer, error) {
	objClient, err := obj.NewGoogleClientFromSecret(context.Background(), "")
	if err != nil {
		return nil, err
	}
	return newObjBlockAPIServer(dir, cacheBytes, objClient)
}

func newMicrosoftBlockAPIServer(dir string, cacheBytes int64) (*objBlockAPIServer, error) {
	objClient, err := obj.NewMicrosoftClientFromSecret("")
	if err != nil {
		return nil, err
	}
	return newObjBlockAPIServer(dir, cacheBytes, objClient)
}

func (s *objBlockAPIServer) PutBlock(putBlockServer pfsclient.BlockAPI_PutBlockServer) (retErr error) {
	func() { s.Log(nil, nil, nil, 0) }()
	result := &pfsclient.BlockRefs{}
	defer func(start time.Time) { s.Log(nil, result, retErr, time.Since(start)) }(time.Now())
	defer drainBlockServer(putBlockServer)
	putBlockRequest, err := putBlockServer.Recv()
	if err != nil {
		if err != io.EOF {
			return err
		}
		return putBlockServer.SendAndClose(result)
	}
	reader := bufio.NewReader(&putBlockReader{
		server: putBlockServer,
		buffer: bytes.NewBuffer(putBlockRequest.Value),
	})
	var eg errgroup.Group
	decoder := json.NewDecoder(reader)
	for {
		blockRef, data, err := readBlock(putBlockRequest.Delimiter, reader, decoder)
		if err != nil {
			return err
		}
		result.BlockRef = append(result.BlockRef, blockRef)
		eg.Go(func() error {
			var outerErr error
			path := s.localServer.blockPath(blockRef.Block)
			backoff.RetryNotify(func() error {
				// We don't want to overwrite blocks that already exist, since:
				// 1) blocks are content-addressable, so it will be the same block
				// 2) we risk exceeding the object store's rate limit
				if s.objClient.Exists(path) {
					return nil
				}
				writer, err := s.objClient.Writer(path)
				if err != nil {
					outerErr = err
					return nil
				}
				if _, err := writer.Write(data); err != nil {
					outerErr = err
					return nil
				}
				if err := writer.Close(); err != nil {
					if s.objClient.IsRetryable(err) {
						return err
					}
					outerErr = err
					return nil
				}
				return nil
			}, obj.NewExponentialBackOffConfig(), func(err error, d time.Duration) {
				protolion.Infof("Error writing; retrying in %s: %#v", d, obj.RetryError{
					Err:               err.Error(),
					TimeTillNextRetry: d.String(),
				})
			})
			// Weird effects can happen with clients racing. Ultimately if the
			// path exists then it doesn't make sense to consider this
			// operation as having errored because we know that it contains the
			// data we want thanks to content addressing.
			if outerErr != nil && !s.objClient.Exists(path) {
				return outerErr
			}
			return nil
		})
		if (blockRef.Range.Upper - blockRef.Range.Lower) < uint64(blockSize) {
			break
		}
	}
	if err := eg.Wait(); err != nil {
		return err
	}
	return putBlockServer.SendAndClose(result)
}

func (s *objBlockAPIServer) GetBlock(request *pfsclient.GetBlockRequest, getBlockServer pfsclient.BlockAPI_GetBlockServer) (retErr error) {
	func() { s.Log(nil, nil, nil, 0) }()
	defer func(start time.Time) { s.Log(request, nil, retErr, time.Since(start)) }(time.Now())
	var data []byte
	sink := groupcache.AllocatingByteSliceSink(&data)
	if err := s.cache.Get(getBlockServer.Context(), request.Block.Hash, sink); err != nil {
		return err
	}
	if request.SizeBytes != 0 && request.SizeBytes+request.OffsetBytes < uint64(len(data)) {
		data = data[request.OffsetBytes : request.OffsetBytes+request.SizeBytes]
	} else if request.OffsetBytes < uint64(len(data)) {
		data = data[request.OffsetBytes:]
	} else {
		data = nil
	}
	return getBlockServer.Send(&types.BytesValue{Value: data})
}

func (s *objBlockAPIServer) DeleteBlock(ctx context.Context, request *pfsclient.DeleteBlockRequest) (response *types.Empty, retErr error) {
	func() { s.Log(nil, nil, nil, 0) }()
	defer func(start time.Time) { s.Log(request, response, retErr, time.Since(start)) }(time.Now())
	backoff.RetryNotify(func() error {
		if err := s.objClient.Delete(s.localServer.blockPath(request.Block)); err != nil && !s.objClient.IsNotExist(err) {
			return err
		}
		return nil
	}, obj.NewExponentialBackOffConfig(), func(err error, d time.Duration) {
		protolion.Infof("Error deleting block; retrying in %s: %#v", d, obj.RetryError{
			Err:               err.Error(),
			TimeTillNextRetry: d.String(),
		})
	})
	return &types.Empty{}, nil
}

func (s *objBlockAPIServer) InspectBlock(ctx context.Context, request *pfsclient.InspectBlockRequest) (response *pfsclient.BlockInfo, retErr error) {
	return nil, fmt.Errorf("not implemented")
}

func (s *objBlockAPIServer) ListBlock(ctx context.Context, request *pfsclient.ListBlockRequest) (response *pfsclient.BlockInfos, retErr error) {
	return nil, fmt.Errorf("not implemented")
}
<|endoftext|>


<commit_before>/*
 * Copyright 2016 Xuyuan Pang
 * Author: Xuyuan Pang
 *
 * 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 hodor

import "net/http"

// BuildRoute creates a route builder from router.
func BuildRoute(router Router) Route {
	return RouteFunc(router.AddRoute).Route()
}

// RouteFunc is a function type implemented Router interface.
type RouteFunc func(method Method, pattern string, handler http.Handler, filters ...Filter)

// Route returns a setter-chain to add a new route step-by-step.
func (f RouteFunc) Route() Route {
	return func(method Method) PatternSetter {
		return func(pattern string) HandlerSetter {
			return func(handler http.Handler, filters ...Filter) {
				f(method, pattern, handler, filters...)
			}
		}
	}
}

// Route easy way to set Method for a route.
type Route func(method Method) PatternSetter

// Method calls MethodSetter function.
func (ms Route) Method(method Method) PatternSetter {
	return ms(method)
}

// Options short for Method(OPTIONS)
func (ms Route) Options() PatternSetter {
	return ms.Method(OPTIONS)
}

// Get short for Method(GET)
func (ms Route) Get() PatternSetter {
	return ms.Method(GET)
}

// Head short for Method(HEAD)
func (ms Route) Head() PatternSetter {
	return ms.Method(HEAD)
}

// Post short for Method(POST)
func (ms Route) Post() PatternSetter {
	return ms.Method(POST)
}

// Put short for Method(PUT)
func (ms Route) Put() PatternSetter {
	return ms.Method(PUT)
}

// Delete short for Method(DELETE)
func (ms Route) Delete() PatternSetter {
	return ms.Method(DELETE)
}

// Trace short for Method(TRACE)
func (ms Route) Trace() PatternSetter {
	return ms.Method(TRACE)
}

// Connect short for Method(CONNECT)
func (ms Route) Connect() PatternSetter {
	return ms.Method(CONNECT)
}

// Patch short for Method(PATCH)
func (ms Route) Patch() PatternSetter {
	return ms.Method(PATCH)
}

// Group creates a Group with root
func (ms Route) Group(root string) Grouper {
	return func(fn func(Route), fs ...Filter) {
		fn(RouteFunc(
			func(method Method, subpattern string, handler http.Handler, subfilters ...Filter) {
				ms.Method(method).
					Pattern(root + subpattern).
					Filters(fs...).
					Filters(subfilters...).
					Handler(handler)
			}).Route())
	}
}

// HandlerSetter easy way to set Handler for a route.
type HandlerSetter func(handler http.Handler, filters ...Filter)

// Handler calls HandlerSetter function.
func (hs HandlerSetter) Handler(handler http.Handler) {
	hs(handler)
}

// HandlerFunc wraps HandlerFunc to Handler
func (hs HandlerSetter) HandlerFunc(hf func(http.ResponseWriter, *http.Request)) {
	hs.Handler(http.HandlerFunc(hf))
}

// Filters returns a new HandlerSetter
func (hs HandlerSetter) Filters(filters ...Filter) HandlerSetter {
	return func(handler http.Handler, fs ...Filter) {
		hs(handler, append(filters, fs...)...)
	}
}

// PatternSetter easy way to set Path for a route.
type PatternSetter func(pattern string) HandlerSetter

// Pattern calls Pattern function.
func (ps PatternSetter) Pattern(pattern string) HandlerSetter {
	return ps(pattern)
}

// Grouper is to add routes grouply.
type Grouper func(func(Route), ...Filter)

// For applies func
func (g Grouper) For(fn func(Route)) {
	g(fn)
}

// Filters creates new Grouper wrapped filters
func (g Grouper) Filters(filters ...Filter) Grouper {
	return func(fn func(Route), fs ...Filter) {
		g(fn, append(filters, fs...)...)
	}
}
<commit_msg>Added FilterFunc method<commit_after>/*
 * Copyright 2016 Xuyuan Pang
 * Author: Xuyuan Pang
 *
 * 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 hodor

import "net/http"

// BuildRoute creates a route builder from router.
func BuildRoute(router Router) Route {
	return RouteFunc(router.AddRoute).Route()
}

// RouteFunc is a function type implemented Router interface.
type RouteFunc func(method Method, pattern string, handler http.Handler, filters ...Filter)

// Route returns a setter-chain to add a new route step-by-step.
func (f RouteFunc) Route() Route {
	return func(method Method) PatternSetter {
		return func(pattern string) HandlerSetter {
			return func(handler http.Handler, filters ...Filter) {
				f(method, pattern, handler, filters...)
			}
		}
	}
}

// Route easy way to set Method for a route.
type Route func(method Method) PatternSetter

// Method calls MethodSetter function.
func (ms Route) Method(method Method) PatternSetter {
	return ms(method)
}

// Options short for Method(OPTIONS)
func (ms Route) Options() PatternSetter {
	return ms.Method(OPTIONS)
}

// Get short for Method(GET)
func (ms Route) Get() PatternSetter {
	return ms.Method(GET)
}

// Head short for Method(HEAD)
func (ms Route) Head() PatternSetter {
	return ms.Method(HEAD)
}

// Post short for Method(POST)
func (ms Route) Post() PatternSetter {
	return ms.Method(POST)
}

// Put short for Method(PUT)
func (ms Route) Put() PatternSetter {
	return ms.Method(PUT)
}

// Delete short for Method(DELETE)
func (ms Route) Delete() PatternSetter {
	return ms.Method(DELETE)
}

// Trace short for Method(TRACE)
func (ms Route) Trace() PatternSetter {
	return ms.Method(TRACE)
}

// Connect short for Method(CONNECT)
func (ms Route) Connect() PatternSetter {
	return ms.Method(CONNECT)
}

// Patch short for Method(PATCH)
func (ms Route) Patch() PatternSetter {
	return ms.Method(PATCH)
}

// Group creates a Group with root
func (ms Route) Group(root string) Grouper {
	return func(fn func(Route), fs ...Filter) {
		fn(RouteFunc(
			func(method Method, subpattern string, handler http.Handler, subfilters ...Filter) {
				ms.Method(method).
					Pattern(root + subpattern).
					Filters(fs...).
					Filters(subfilters...).
					Handler(handler)
			}).Route())
	}
}

// HandlerSetter easy way to set Handler for a route.
type HandlerSetter func(handler http.Handler, filters ...Filter)

// Handler calls HandlerSetter function.
func (hs HandlerSetter) Handler(handler http.Handler) {
	hs(handler)
}

// HandlerFunc wraps HandlerFunc to Handler
func (hs HandlerSetter) HandlerFunc(hf func(http.ResponseWriter, *http.Request)) {
	hs.Handler(http.HandlerFunc(hf))
}

// Filters returns a new HandlerSetter
func (hs HandlerSetter) Filters(filters ...Filter) HandlerSetter {
	return func(handler http.Handler, fs ...Filter) {
		hs(handler, append(filters, fs...)...)
	}
}

// FilterFunc is a sugar method for Filters
func (hs HandlerSetter) FilterFunc(f func(http.Handler) http.Handler) HandlerSetter {
	return hs.Filters(FilterFunc(f))
}

// PatternSetter easy way to set Path for a route.
type PatternSetter func(pattern string) HandlerSetter

// Pattern calls Pattern function.
func (ps PatternSetter) Pattern(pattern string) HandlerSetter {
	return ps(pattern)
}

// Grouper is to add routes grouply.
type Grouper func(func(Route), ...Filter)

// For applies func
func (g Grouper) For(fn func(Route)) {
	g(fn)
}

// Filters creates new Grouper wrapped filters
func (g Grouper) Filters(filters ...Filter) Grouper {
	return func(fn func(Route), fs ...Filter) {
		g(fn, append(filters, fs...)...)
	}
}

// FilterFunc is a sugar method for Filters
func (g Grouper) FilterFunc(f func(http.Handler) http.Handler) Grouper {
	return g.Filters(FilterFunc(f))
}
<|endoftext|>


<commit_before>package atlas

import (
	"fmt"
	"regexp"

	"github.com/hashicorp/atlas-go/v1"
	"github.com/hashicorp/terraform/helper/hashcode"
	"github.com/hashicorp/terraform/helper/schema"
)

var (
	// saneMetaKey is used to sanitize the metadata keys so that
	// they can be accessed as a variable interpolation from TF
	saneMetaKey = regexp.MustCompile("[^a-zA-Z0-9-_]")
)

func resourceArtifact() *schema.Resource {
	return &schema.Resource{
		Create: resourceArtifactRead,
		Read:   resourceArtifactRead,
		Update: resourceArtifactRead,
		Delete: resourceArtifactDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"type": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"build": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},

			"version": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},

			"metadata_keys": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				ForceNew: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Set: func(v interface{}) int {
					return hashcode.String(v.(string))
				},
			},

			"metadata": &schema.Schema{
				Type:     schema.TypeMap,
				Optional: true,
				ForceNew: true,
			},

			"file_url": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},

			"metadata_full": &schema.Schema{
				Type:     schema.TypeMap,
				Computed: true,
			},

			"slug": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},

			"version_real": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceArtifactRead(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*atlas.Client)

	// Parse the slug from the name given of the artifact since the API
	// expects these to be split.
	user, name, err := atlas.ParseSlug(d.Get("name").(string))
	if err != nil {
		return err
	}

	// Filter by version or build if given
	var build, version string
	if v, ok := d.GetOk("version"); ok {
		version = v.(string)
	} else if b, ok := d.GetOk("build"); ok {
		build = b.(string)
	}

	// If we have neither, default to latest version
	if build == "" && version == "" {
		version = "latest"
	}

	// Compile the metadata search params
	md := make(map[string]string)
	for _, v := range d.Get("metadata_keys").(*schema.Set).List() {
		md[v.(string)] = atlas.MetadataAnyValue
	}
	for k, v := range d.Get("metadata").(map[string]interface{}) {
		md[k] = v.(string)
	}

	// Do the search!
	vs, err := client.ArtifactSearch(&atlas.ArtifactSearchOpts{
		User:     user,
		Name:     name,
		Type:     d.Get("type").(string),
		Build:    build,
		Version:  version,
		Metadata: md,
	})
	if err != nil {
		return fmt.Errorf(
			"Error searching for artifact: %s", err)
	}

	if len(vs) == 0 {
		return fmt.Errorf("No matching artifact")
	} else if len(vs) > 1 {
		return fmt.Errorf("Got %d results, only one is allowed", len(vs))
	}
	v := vs[0]

	d.SetId(v.ID)
	if v.ID == "" {
		d.SetId(fmt.Sprintf("%s %d", v.Tag, v.Version))
	}
	d.Set("version_real", v.Version)
	d.Set("metadata_full", cleanMetadata(v.Metadata))
	d.Set("slug", v.Slug)

	d.Set("file_url", "")
	if u, err := client.ArtifactFileURL(v); err != nil {
		return fmt.Errorf(
			"Error reading file URL: %s", err)
	} else if u != nil {
		d.Set("file_url", u.String())
	}

	return nil
}

func resourceArtifactDelete(d *schema.ResourceData, meta interface{}) error {
	// This just always succeeds since this is a readonly element.
	d.SetId("")
	return nil
}

// cleanMetadata is used to ensure the metadata is accessible as
// a variable by doing a simple re-write.
func cleanMetadata(in map[string]string) map[string]string {
	out := make(map[string]string, len(in))
	for k, v := range in {
		sane := saneMetaKey.ReplaceAllString(k, "-")
		out[sane] = v
	}
	return out
}
<commit_msg>providers/atlas: make errors more readable<commit_after>package atlas

import (
	"fmt"
	"regexp"

	"github.com/hashicorp/atlas-go/v1"
	"github.com/hashicorp/terraform/helper/hashcode"
	"github.com/hashicorp/terraform/helper/schema"
)

var (
	// saneMetaKey is used to sanitize the metadata keys so that
	// they can be accessed as a variable interpolation from TF
	saneMetaKey = regexp.MustCompile("[^a-zA-Z0-9-_]")
)

func resourceArtifact() *schema.Resource {
	return &schema.Resource{
		Create: resourceArtifactRead,
		Read:   resourceArtifactRead,
		Update: resourceArtifactRead,
		Delete: resourceArtifactDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"type": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"build": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},

			"version": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},

			"metadata_keys": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				ForceNew: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Set: func(v interface{}) int {
					return hashcode.String(v.(string))
				},
			},

			"metadata": &schema.Schema{
				Type:     schema.TypeMap,
				Optional: true,
				ForceNew: true,
			},

			"file_url": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},

			"metadata_full": &schema.Schema{
				Type:     schema.TypeMap,
				Computed: true,
			},

			"slug": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},

			"version_real": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceArtifactRead(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*atlas.Client)

	// Parse the slug from the name given of the artifact since the API
	// expects these to be split.
	user, name, err := atlas.ParseSlug(d.Get("name").(string))
	if err != nil {
		return err
	}

	// Filter by version or build if given
	var build, version string
	if v, ok := d.GetOk("version"); ok {
		version = v.(string)
	} else if b, ok := d.GetOk("build"); ok {
		build = b.(string)
	}

	// If we have neither, default to latest version
	if build == "" && version == "" {
		version = "latest"
	}

	// Compile the metadata search params
	md := make(map[string]string)
	for _, v := range d.Get("metadata_keys").(*schema.Set).List() {
		md[v.(string)] = atlas.MetadataAnyValue
	}
	for k, v := range d.Get("metadata").(map[string]interface{}) {
		md[k] = v.(string)
	}

	// Do the search!
	vs, err := client.ArtifactSearch(&atlas.ArtifactSearchOpts{
		User:     user,
		Name:     name,
		Type:     d.Get("type").(string),
		Build:    build,
		Version:  version,
		Metadata: md,
	})
	if err != nil {
		return fmt.Errorf(
			"Error searching for artifact '%s/%s': %s",
			user, name, err)
	}

	if len(vs) == 0 {
		return fmt.Errorf("No matching artifact for '%s/%s'", user, name)
	} else if len(vs) > 1 {
		return fmt.Errorf(
			"Got %d results for '%s/%s', only one is allowed",
			user, name, len(vs))
	}
	v := vs[0]

	d.SetId(v.ID)
	if v.ID == "" {
		d.SetId(fmt.Sprintf("%s %d", v.Tag, v.Version))
	}
	d.Set("version_real", v.Version)
	d.Set("metadata_full", cleanMetadata(v.Metadata))
	d.Set("slug", v.Slug)

	d.Set("file_url", "")
	if u, err := client.ArtifactFileURL(v); err != nil {
		return fmt.Errorf(
			"Error reading file URL: %s", err)
	} else if u != nil {
		d.Set("file_url", u.String())
	}

	return nil
}

func resourceArtifactDelete(d *schema.ResourceData, meta interface{}) error {
	// This just always succeeds since this is a readonly element.
	d.SetId("")
	return nil
}

// cleanMetadata is used to ensure the metadata is accessible as
// a variable by doing a simple re-write.
func cleanMetadata(in map[string]string) map[string]string {
	out := make(map[string]string, len(in))
	for k, v := range in {
		sane := saneMetaKey.ReplaceAllString(k, "-")
		out[sane] = v
	}
	return out
}
<|endoftext|>