<commit_before>package db

import (
	"bytes"
	"database/sql"
	"fmt"
	"io"
	"io/ioutil"
	"net/url"
	"strings"
	"sync"

	"github.com/koron/nvgd/config"
	"github.com/koron/nvgd/ltsv"
	"github.com/koron/nvgd/protocol"
)

var NullReplacement = "(null)"

// Param is connection parameter for the database.
type Param struct {
	// Driver represents driver name for database.
	Driver string `yaml:"driver"`

	// Name represents driver-specific data source name.
	Name string `yaml:"name"`
}

// Config represents configuration for Handler.
type Config map[string]Param

// Handler is database protocol handler.
type Handler struct {
	Config *Config

	l         sync.Mutex
	databases map[string]*sql.DB
}

var dbconfig Config

func init() {
	protocol.MustRegister("db", &Handler{
		Config:    &dbconfig,
		databases: make(map[string]*sql.DB),
	})
	config.RegisterProtocol("db", &dbconfig)
}

// Open creates a database handler.
func (h *Handler) Open(u *url.URL) (io.ReadCloser, error) {
	var (
		name  = u.Host
		query = u.Path
	)
	db, err := h.openDB(name)
	if err != nil {
		return nil, err
	}
	// TODO: sanitize query!
	if strings.HasPrefix(query, "/") {
		query = query[1:]
	}
	fmt.Printf("query=%s\n", query)
	rows, err := db.Query(query)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return h.rows2ltsv(rows)
}

func (h *Handler) openDB(name string) (*sql.DB, error) {
	h.l.Lock()
	defer h.l.Unlock()
	if db, ok := h.databases[name]; ok {
		return db, nil
	}
	p, ok := (*h.Config)[name]
	if !ok {
		return nil, fmt.Errorf("unknown database: %q", name)
	}<commit_msg>wrap queries by transaction for safety<commit_after>package db

import (
	"bytes"
	"database/sql"
	"fmt"
	"io"
	"io/ioutil"
	"net/url"
	"strings"
	"sync"

	"github.com/koron/nvgd/config"
	"github.com/koron/nvgd/ltsv"
	"github.com/koron/nvgd/protocol"
)

var NullReplacement = "(null)"

// Param is connection parameter for the database.
type Param struct {
	// Driver represents driver name for database.
	Driver string `yaml:"driver"`

	// Name represents driver-specific data source name.
	Name string `yaml:"name"`
}

// Config represents configuration for Handler.
type Config map[string]Param

// Handler is database protocol handler.
type Handler struct {
	Config *Config

	l         sync.Mutex
	databases map[string]*sql.DB
}

var dbconfig Config

func init() {
	protocol.MustRegister("db", &Handler{
		Config:    &dbconfig,
		databases: make(map[string]*sql.DB),
	})
	config.RegisterProtocol("db", &dbconfig)
}

// Open creates a database handler.
func (h *Handler) Open(u *url.URL) (io.ReadCloser, error) {
	var (
		name  = u.Host
		query = u.Path
	)
	db, err := h.openDB(name)
	if err != nil {
		return nil, err
	}
	if strings.HasPrefix(query, "/") {
		query = query[1:]
	}
	// TODO: sanitize query!
	//fmt.Printf("query=%s\n", query)
	return h.execQuery(db, query)
}

// execQuery executes a query in a transaction which will be rollbacked.
func (h *Handler) execQuery(db *sql.DB, q string) (io.ReadCloser, error) {
	tx, err := db.Begin()
	if err != nil {
		return nil, err
	}
	defer tx.Rollback()
	rows, err := tx.Query(q)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return h.rows2ltsv(rows)
}

func (h *Handler) openDB(name string) (*sql.DB, error) {
	h.l.Lock()
	defer h.l.Unlock()
	if db, ok := h.databases[name]; ok {
		return db, nil
	}
	p, ok := (*h.Config)[name]
	if !ok {
		return nil, fmt.Errorf("unknown database: %q", name)
	}
{'fraction_non_alphanumeric': 0.15140947399011914, 'fraction_numerical': 0.0011624527753560012, 'mean_word_length': 3.5832223701731025, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '2094322', 'n_tokens_mistral': 1308, 'n_tokens_neox': 1285, 'n_words': 353}