pquery.go

package pquery

import (
	"database/sql"
	"strings"
	"sync"

	log "github.com/sirupsen/logrus"
	"github.com/xyproto/algernon/lua/convert"
	lua "github.com/xyproto/gopher-lua"

	// Using the PostgreSQL database engine
	_ "github.com/lib/pq"
)

const (
	defaultQuery            = "SELECT version()"
	defaultConnectionString = "host=localhost port=5432 user=postgres dbname=test sslmode=disable"
)

var (
	// global map from connection string to database connection, to reuse connections, protected by a mutex
	reuseDB  = make(map[string]*sql.DB)
	reuseMut = &sync.RWMutex{}
)

// Load makes functions related to building a library of Lua code available
func Load(L *lua.LState) {
	// Register the PQ function
	L.SetGlobal("PQ", L.NewFunction(func(L *lua.LState) int {
		// Check if the optional argument is given
		query := defaultQuery
		if L.GetTop() >= 1 {
			query = L.ToString(1)
			if query == "" {
				query = defaultQuery
			}
		}
		connectionString := defaultConnectionString
		if L.GetTop() >= 2 {
			connectionString = L.ToString(2)
		}

		// Check if there is a connection that can be reused
		var db *sql.DB
		reuseMut.RLock()
		conn, ok := reuseDB[connectionString]
		reuseMut.RUnlock()

		if ok {
			// It exists, but is it still alive?
			err := conn.Ping()
			if err != nil {
				// no
				// log.Info("did not reuse the connection")
				reuseMut.Lock()
				delete(reuseDB, connectionString)
				reuseMut.Unlock()
			} else {
				// yes
				// log.Info("reused the connection")
				db = conn
			}
		}
		// Create a new connection, if needed
		var err error
		if db == nil {
			db, err = sql.Open("postgres", connectionString)
			if err != nil {
				log.Error("Could not connect to database using " + connectionString + ": " + err.Error())
				return 0 // No results
			}
			// Save the connection for later
			reuseMut.Lock()
			reuseDB[connectionString] = db
			reuseMut.Unlock()
		}
		// log.Info(fmt.Sprintf("PostgreSQL database: %v (%T)\n", db, db))
		reuseMut.Lock()
		rows, err := db.Query(query)
		reuseMut.Unlock()
		if err != nil {
			errMsg := err.Error()
			if strings.Contains(errMsg, ": connect: connection refused") {
				log.Info("PostgreSQL connection string: " + connectionString)
				log.Info("PostgreSQL query: " + query)
				log.Error("Could not connect to database: " + errMsg)
			} else if strings.Contains(errMsg, "missing") && strings.Contains(errMsg, "in connection info string") {
				log.Info("PostgreSQL connection string: " + connectionString)
				log.Info("PostgreSQL query: " + query)
				log.Error(errMsg)
			} else {
				log.Info("PostgreSQL query: " + query)
				log.Error("Query failed: " + errMsg)
			}
			return 0 // No results
		}
		if rows == nil {
			// Return an empty table
			L.Push(L.NewTable())
			return 1 // number of results
		}
		// Return the rows as a table
		var (
			values []string
			value  string
		)
		for rows.Next() {
			err = rows.Scan(&value)
			if err != nil {
				break
			}
			values = append(values, value)
		}
		// Convert the strings to a Lua table
		table := convert.Strings2table(L, values)
		// Return the table
		L.Push(table)
		return 1 // number of results
	}))
}