<commit_before>// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package postgres

import (
	"context"
	"database/sql"
	"fmt"
	"strings"
	"sync/atomic"
	"time"
	"unicode"

	"golang.org/x/discovery/internal/config"
	"golang.org/x/discovery/internal/derrors"
	"golang.org/x/discovery/internal/log"
)

// DB wraps a sql.DB to provide an API for interacting with discovery data
// stored in Postgres.
type DB struct {
	db *sql.DB
}

func (db *DB) exec(ctx context.Context, query string, args ...interface{}) (res sql.Result, err error) {
	defer logQuery(query, args)(&err)

	return db.db.ExecContext(ctx, query, args...)
}

func execTx(ctx context.Context, tx *sql.Tx, query string, args ...interface{}) (res sql.Result, err error) {
	defer logQuery(query, args)(&err)

	return tx.ExecContext(ctx, query, args...)
}

func (db *DB) query(ctx context.Context, query string, args ...interface{}) (_ *sql.Rows, err error) {
	defer logQuery(query, args)(&err)
	return db.db.QueryContext(ctx, query, args...)
}

func (db *DB) queryRow(ctx context.Context, query string, args ...interface{}) *sql.Row {
	defer logQuery(query, args)(nil)
	return db.db.QueryRowContext(ctx, query, args...)
}

var (
	queryCounter         int64 // atomic: per-process counter for unique query IDs
	queryLoggingDisabled bool  // For use in tests only: not concurrency-safe.
)

type queryEndLogEntry struct {
	ID              string
	Query           string
	Args            []interface{}
	ArgsTruncated   bool
	DurationSeconds float64
	Error           string `json:",omitempty"`
}

func logQuery(query string, args []interface{}) func(*error) {
	if queryLoggingDisabled {
		return func(*error) {}
	}
	const maxlen = 300 // maximum length of displayed query

	// To make the query more compact and readable, replace newlines with spaces
	// and collapse adjacent whitespace.
	var r []rune
	for _, c := range query {
		if c == '\n' {
			c = ' '
		}
		if len(r) == 0 || !unicode.IsSpace(r[len(r)-1]) || !unicode.IsSpace(c) {
			r = append(r, c)
		}
	}
	query = string(r)
	if len(query) > maxlen {
		query = query[:maxlen] + "..."
	}

	instanceID := config.InstanceID()
	if instanceID == "" {
		instanceID = "local"
	} else {
		// Instance IDs are long strings. The low-order part seems quite random, so
		// shortening the ID will still likely result in something unique.
		instanceID = instanceID[len(instanceID)-4:]
	}
	n := atomic.AddInt64(&queryCounter, 1)
	uid := fmt.Sprintf("%s-%d", instanceID, n)

	const maxargs = 20 // maximum displayed args
	var moreargs string
	if len(args) > maxargs {
		args = args[:maxargs]
		moreargs = "..."
	}

	log.Debugf("%s %s %v%s", uid, query, args, moreargs)
	start := time.Now()
	return func(errp *error) {
		dur := time.Since(start)
		if errp == nil { // happens with queryRow
			log.Debugf("%s done", uid)
		} else {
			derrors.Wrap(errp, "DB running query %s", uid)
			entry := queryEndLogEntry{
				ID:              uid,
				Query:           query,
				Args:            args,
				ArgsTruncated:   moreargs != "",
				DurationSeconds: dur.Seconds(),
			}
			if *errp == nil {
				log.Debug(entry)
			} else {
				entry.Error = (*errp).Error()
				log.Error(entry)
			}
		}
	}
}

// Open creates a new DB for the given Postgres connection string.
func Open(driverName, dbinfo string) (*DB, error) {
	db, err := sql.Open(driverName, dbinfo)
	if err != nil {
		return nil, err
	}
<commit_msg>internal/postgres: shorten query args in logs<commit_after>// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package postgres

import (
	"context"
	"database/sql"
	"fmt"
	"strings"
	"sync/atomic"
	"time"
	"unicode"

	"golang.org/x/discovery/internal/config"
	"golang.org/x/discovery/internal/derrors"
	"golang.org/x/discovery/internal/log"
)

// DB wraps a sql.DB to provide an API for interacting with discovery data
// stored in Postgres.
type DB struct {
	db *sql.DB
}

func (db *DB) exec(ctx context.Context, query string, args ...interface{}) (res sql.Result, err error) {
	defer logQuery(query, args)(&err)

	return db.db.ExecContext(ctx, query, args...)
}

func execTx(ctx context.Context, tx *sql.Tx, query string, args ...interface{}) (res sql.Result, err error) {
	defer logQuery(query, args)(&err)

	return tx.ExecContext(ctx, query, args...)
}

func (db *DB) query(ctx context.Context, query string, args ...interface{}) (_ *sql.Rows, err error) {
	defer logQuery(query, args)(&err)
	return db.db.QueryContext(ctx, query, args...)
}

func (db *DB) queryRow(ctx context.Context, query string, args ...interface{}) *sql.Row {
	defer logQuery(query, args)(nil)
	return db.db.QueryRowContext(ctx, query, args...)
}

var (
	queryCounter         int64 // atomic: per-process counter for unique query IDs
	queryLoggingDisabled bool  // For use in tests only: not concurrency-safe.
)

type queryEndLogEntry struct {
	ID              string
	Query           string
	Args            string
	DurationSeconds float64
	Error           string `json:",omitempty"`
}

func logQuery(query string, args []interface{}) func(*error) {
	if queryLoggingDisabled {
		return func(*error) {}
	}
	const maxlen = 300 // maximum length of displayed query

	// To make the query more compact and readable, replace newlines with spaces
	// and collapse adjacent whitespace.
	var r []rune
	for _, c := range query {
		if c == '\n' {
			c = ' '
		}
		if len(r) == 0 || !unicode.IsSpace(r[len(r)-1]) || !unicode.IsSpace(c) {
			r = append(r, c)
		}
	}
	query = string(r)
	if len(query) > maxlen {
		query = query[:maxlen] + "..."
	}

	instanceID := config.InstanceID()
	if instanceID == "" {
		instanceID = "local"
	} else {
		// Instance IDs are long strings. The low-order part seems quite random, so
		// shortening the ID will still likely result in something unique.
		instanceID = instanceID[len(instanceID)-4:]
	}
	n := atomic.AddInt64(&queryCounter, 1)
	uid := fmt.Sprintf("%s-%d", instanceID, n)

	// Construct a short string of the args.
	const (
		maxArgs   = 20
		maxArgLen = 50
	)
	var argStrings []string
	for i := 0; i < len(args) && i < maxArgs; i++ {
		s := fmt.Sprint(args[i])
		if len(s) > maxArgLen {
			s = s[:maxArgLen] + "..."
		}
		argStrings = append(argStrings, s)
	}
	if len(args) > maxArgs {
		argStrings = append(argStrings, "...")
	}
	argString := strings.Join(argStrings, ", ")

	log.Debugf("%s %s args=%s", uid, query, argString)
	start := time.Now()
	return func(errp *error) {
		dur := time.Since(start)
		if errp == nil { // happens with queryRow
			log.Debugf("%s done", uid)
		} else {
			derrors.Wrap(errp, "DB running query %s", uid)
			entry := queryEndLogEntry{
				ID:              uid,
				Query:           query,
				Args:            argString,
				DurationSeconds: dur.Seconds(),
			}
			if *errp == nil {
				log.Debug(entry)
			} else {
				entry.Error = (*errp).Error()
				log.Error(entry)
			}
		}
	}
}

// Open creates a new DB for the given Postgres connection string.
func Open(driverName, dbinfo string) (*DB, error) {
	db, err := sql.Open(driverName, dbinfo)
	if err != nil {
		return nil, err
	}

{'fraction_non_alphanumeric': 0.14195061387777624, 'fraction_numerical': 0.005655952545178645, 'mean_word_length': 3.1169789892106756, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '493500', 'n_tokens_mistral': 2615, 'n_tokens_neox': 2410, 'n_words': 751}