<commit_before>// Copyright 2014 The zhgo 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 db

import (
    "database/sql"
    _ "github.com/go-sql-driver/mysql"
    _ "github.com/mxk/go-sqlite/sqlite3"
    _ "github.com/lib/pq"
    "log"
    "reflect"
    "regexp"
    "strconv"
    "strings"
)

// Server struct
type Server struct {
    // Name
    Name string

    // Database type: mysql or sqlite3
    Type string

    // Data Source Name
    DSN string

    // Follow
    Follow string
}

// Server instance
var dbObjects map[string]*sql.DB = make(map[string]*sql.DB)

// Execute query, only return sql.Result
func (e *Server) Exec(sql string, args ...interface{}) (sql.Result, error) {
    sql, args = e.parseSQL(sql, args)
    stmt, err := e.prepare(sql)
    if err != nil {
        return nil, err
    }

    defer stmt.Close()

    result, err := stmt.Exec(args...)
    if err != nil {
        return nil, err
    }

    return result, nil
}

// Get row.
func (e *Server) Row(ptr interface{}, sql string, args ...interface{}) error {
    rows, columns, err := e.rows(sql, args)
    if err != nil {
        log.Printf("%s\n", err)
        return err
    }

    defer rows.Close()

    columnsLen := len(columns)

    kind, ptrRow, scan, err := scanVariables(ptr, columnsLen, false)
    if err != nil {
        log.Printf("%s\n", err)
        return err
    }

    // Return data
    val := reflect.ValueOf(ptr).Elem()

    if rows.Next() {
        if err := rows.Scan(scan...); err != nil {
            log.Printf("%s\n", err)
            return err
        }

        switch kind {
            case reflect.Struct: // struct
            val.Set(reflect.ValueOf(ptrRow).Elem())

            case reflect.Map: //map
            row := make(map[string]interface{}, columnsLen)
            for i := 0; i < columnsLen; i++ {
                row[columns[i]] = typeAssertion(*(scan[i].(*interface{})))
            }
            val.Set(reflect.ValueOf(row))

            case reflect.Slice: //slice
            row := make([]interface{}, columnsLen)
            for i := 0; i < columnsLen; i++ {
                row[i] = typeAssertion(*(scan[i].(*interface{})))
            }
            val.Set(reflect.ValueOf(row))
        }
    }

    if err = rows.Err(); err != nil {
        log.Printf("%s\n", err)
        return err
    }

    return nil
}

// Get all rows
func (e *Server) Rows(ptr interface{}, sql string, args ...interface{}) error {
    rows, columns, err := e.rows(sql, args)
    if err != nil {
        log.Printf("%s\n", err)
        return err
    }

    defer rows.Close()

    columnsLen := len(columns)

    kind, ptrRow, scan, err := scanVariables(ptr, columnsLen, true)
    if err != nil {
        log.Printf("%s\n", err)
        return err
    }

    //return data
    val := reflect.ValueOf(ptr).Elem()

    for rows.Next() {
        if err := rows.Scan(scan...); err != nil {
            log.Printf("%s\n", err)
            return err
        }

        switch kind {
            case reflect.Struct: // struct
            val.Set(reflect.Append(val, reflect.ValueOf(ptrRow).Elem()))

            case reflect.Map: // map
            row := make(map[string]interface{}, columnsLen)
            for i := 0; i < columnsLen; i++ {
                row[columns[i]] = typeAssertion(*(scan[i].(*interface{})))
            }
            val.Set(reflect.Append(val, reflect.ValueOf(row)))

            case reflect.Slice: // slice
            row := make([]interface{}, columnsLen)
            for i := 0; i < columnsLen; i++ {
                row[i] = typeAssertion(*(scan[i].(*interface{})))
            }
            val.Set(reflect.Append(val, reflect.ValueOf(row)))
        }
    }

    if err = rows.Err(); err != nil {
        log.Printf("%s\n", err)
        return err
    }

    return nil
}

// Execute query, return sql.Rows, rows.Columns
func (e *Server) rows(sql string, args []interface{}) (*sql.Rows, []string, error) {
    sql, args = e.parseSQL(sql, args)
    stmt, err := e.prepare(sql)
    if err != nil {
        return nil, nil, err
    }

    defer stmt.Close()

    rows, err := stmt.Query(args...)
    if err != nil {
        return nil, nil, err
    }

    columns, err := rows.Columns()
    if err != nil {
        return nil, nil, err
    }

    return rows, columns, nil
}

// Prepare SQL
func (e *Server) prepare(sql string) (*sql.Stmt, error) {
    if err := e.connect(); err != nil {
        return nil, err
    }

    stmt, err := dbObjects[e.Name].Prepare(sql)
    if err != nil {
        return nil, err
    }

    return stmt, nil
}

// Connect to database
func (e *Server) connect() error {
    db, s := dbObjects[e.Name]
    // err := db.Ping(); err != nil
    if s == false || db == nil {
        var err error
        dbObjects[e.Name], err = sql.Open(e.Type, e.DSN)
        if err != nil {
            return err
        }
    }

    return nil
}

// sql compatibility
func (e *Server) parseSQL(str string, args []interface{}) (string, []interface{}) {
    switch e.Type {
        case "postgres":
        return str, args
        case "mysql":
        str = e.parseQuotes(str)
        return e.parseParameters(str, args)
        case "sqlite3":
        return e.parseParameters(str, args)
    }
    return str, args
}

// " to `
func (e *Server) parseQuotes(str string) string {
    return strings.Replace(str, `"`, "`", -1)
}

// $1, $2, $3 to ?, ?, ?
func (e *Server) parseParameters(str string, args []interface{}) (string, []interface{}) {
    re := regexp.MustCompile(`\$(\d+)`)
    sli := re.FindAllStringSubmatch(str, -1)
    newArgs := make([]interface{}, len(sli))
    for i, v := range sli {
        vi, err := strconv.ParseInt(v[1], 10, 0)
        if err != nil {
            log.Printf("%s\n", err)
        }
        newArgs[i] = args[vi-1]
    }
    return re.ReplaceAllString(str, "?"), newArgs
}

// New Server
func NewServer(name string, typ string, dsn string) *Server {
    return &Server{Name: typ+name, Type: typ, DSN: dsn}
}
<commit_msg>Change reference modules to zhgo.<commit_after>// Copyright 2014 The zhgo 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 db

import (
    "database/sql"
    _ "github.com/zhgo/mysql"
    _ "github.com/zhgo/postgresql"
    _ "github.com/zhgo/sqlite/sqlite3"
    "log"
    "reflect"
    "regexp"
    "strconv"
    "strings"
)

// Server struct
type Server struct {
    // Name
    Name string

    // Database type: mysql or sqlite3
    Type string

    // Data Source Name
    DSN string

    // Follow
    Follow string
}

// Server instance
var dbObjects map[string]*sql.DB = make(map[string]*sql.DB)

// Execute query, only return sql.Result
func (e *Server) Exec(sql string, args ...interface{}) (sql.Result, error) {
    sql, args = e.parseSQL(sql, args)
    stmt, err := e.prepare(sql)
    if err != nil {
        return nil, err
    }

    defer stmt.Close()

    result, err := stmt.Exec(args...)
    if err != nil {
        return nil, err
    }

    return result, nil
}

// Get row.
func (e *Server) Row(ptr interface{}, sql string, args ...interface{}) error {
    rows, columns, err := e.rows(sql, args)
    if err != nil {
        log.Printf("%s\n", err)
        return err
    }

    defer rows.Close()

    columnsLen := len(columns)

    kind, ptrRow, scan, err := scanVariables(ptr, columnsLen, false)
    if err != nil {
        log.Printf("%s\n", err)
        return err
    }

    // Return data
    val := reflect.ValueOf(ptr).Elem()

    if rows.Next() {
        if err := rows.Scan(scan...); err != nil {
            log.Printf("%s\n", err)
            return err
        }

        switch kind {
            case reflect.Struct: // struct
            val.Set(reflect.ValueOf(ptrRow).Elem())

            case reflect.Map: //map
            row := make(map[string]interface{}, columnsLen)
            for i := 0; i < columnsLen; i++ {
                row[columns[i]] = typeAssertion(*(scan[i].(*interface{})))
            }
            val.Set(reflect.ValueOf(row))

            case reflect.Slice: //slice
            row := make([]interface{}, columnsLen)
            for i := 0; i < columnsLen; i++ {
                row[i] = typeAssertion(*(scan[i].(*interface{})))
            }
            val.Set(reflect.ValueOf(row))
        }
    }

    if err = rows.Err(); err != nil {
        log.Printf("%s\n", err)
        return err
    }

    return nil
}

// Get all rows
func (e *Server) Rows(ptr interface{}, sql string, args ...interface{}) error {
    rows, columns, err := e.rows(sql, args)
    if err != nil {
        log.Printf("%s\n", err)
        return err
    }

    defer rows.Close()

    columnsLen := len(columns)

    kind, ptrRow, scan, err := scanVariables(ptr, columnsLen, true)
    if err != nil {
        log.Printf("%s\n", err)
        return err
    }

    //return data
    val := reflect.ValueOf(ptr).Elem()

    for rows.Next() {
        if err := rows.Scan(scan...); err != nil {
            log.Printf("%s\n", err)
            return err
        }

        switch kind {
            case reflect.Struct: // struct
            val.Set(reflect.Append(val, reflect.ValueOf(ptrRow).Elem()))

            case reflect.Map: // map
            row := make(map[string]interface{}, columnsLen)
            for i := 0; i < columnsLen; i++ {
                row[columns[i]] = typeAssertion(*(scan[i].(*interface{})))
            }
            val.Set(reflect.Append(val, reflect.ValueOf(row)))

            case reflect.Slice: // slice
            row := make([]interface{}, columnsLen)
            for i := 0; i < columnsLen; i++ {
                row[i] = typeAssertion(*(scan[i].(*interface{})))
            }
            val.Set(reflect.Append(val, reflect.ValueOf(row)))
        }
    }

    if err = rows.Err(); err != nil {
        log.Printf("%s\n", err)
        return err
    }

    return nil
}

// Execute query, return sql.Rows, rows.Columns
func (e *Server) rows(sql string, args []interface{}) (*sql.Rows, []string, error) {
    sql, args = e.parseSQL(sql, args)
    stmt, err := e.prepare(sql)
    if err != nil {
        return nil, nil, err
    }

    defer stmt.Close()

    rows, err := stmt.Query(args...)
    if err != nil {
        return nil, nil, err
    }

    columns, err := rows.Columns()
    if err != nil {
        return nil, nil, err
    }

    return rows, columns, nil
}

// Prepare SQL
func (e *Server) prepare(sql string) (*sql.Stmt, error) {
    if err := e.connect(); err != nil {
        return nil, err
    }

    stmt, err := dbObjects[e.Name].Prepare(sql)
    if err != nil {
        return nil, err
    }

    return stmt, nil
}

// Connect to database
func (e *Server) connect() error {
    db, s := dbObjects[e.Name]
    // err := db.Ping(); err != nil
    if s == false || db == nil {
        var err error
        dbObjects[e.Name], err = sql.Open(e.Type, e.DSN)
        if err != nil {
            return err
        }
    }

    return nil
}

// sql compatibility
func (e *Server) parseSQL(str string, args []interface{}) (string, []interface{}) {
    switch e.Type {
        case "postgres":
        return str, args
        case "mysql":
        str = e.parseQuotes(str)
        return e.parseParameters(str, args)
        case "sqlite3":
        return e.parseParameters(str, args)
    }
    return str, args
}

// " to `
func (e *Server) parseQuotes(str string) string {
    return strings.Replace(str, `"`, "`", -1)
}

// $1, $2, $3 to ?, ?, ?
func (e *Server) parseParameters(str string, args []interface{}) (string, []interface{}) {
    re := regexp.MustCompile(`\$(\d+)`)
    sli := re.FindAllStringSubmatch(str, -1)
    newArgs := make([]interface{}, len(sli))
    for i, v := range sli {
        vi, err := strconv.ParseInt(v[1], 10, 0)
        if err != nil {
            log.Printf("%s\n", err)
        }
        newArgs[i] = args[vi-1]
    }
    return re.ReplaceAllString(str, "?"), newArgs
}

// New Server
func NewServer(name string, typ string, dsn string) *Server {
    return &Server{Name: typ+name, Type: typ, DSN: dsn}
}

{'fraction_non_alphanumeric': 0.15441176470588236, 'fraction_numerical': 0.003431372549019608, 'mean_word_length': 1.930572180991142, 'pattern_counts': {'":': 6, '<': 11, '<?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': '7071917', 'n_tokens_mistral': 3984, 'n_tokens_neox': 3878, 'n_words': 1086}