connector.go

// Copyright 2018 Eryx <evorui аt gmail dοt cοm>, All rights reserved.
//
// 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 pgsqlgo

import (
	"database/sql"
	"errors"
	"fmt"
	"strings"

	_ "github.com/lib/pq"
	"github.com/lynkdb/lynkapi/go/lynktable"
)

// NewConnector opens a PostgreSQL connection pool.
//
// Supported option keys:
//   - "host", "port", "user", "pass", "dbname" for a TCP connection
//   - "socket", "user", "pass", "dbname" for a unix domain socket connection
func NewConnector(opts map[string]string) (lynktable.Connector, error) {

	var dsn string

	if v := opts["host"]; v != "" {

		dsn = fmt.Sprintf(
			"host=%s port=%s user=%s password=%s dbname=%s sslmode=disable",
			v, opts["port"], opts["user"], opts["pass"], opts["dbname"],
		)

	} else if v := opts["socket"]; v != "" {

		dsn = fmt.Sprintf(
			"host=%s user=%s password=%s dbname=%s sslmode=disable",
			v, opts["user"], opts["pass"], opts["dbname"],
		)

	} else {
		return nil, errors.New("Incorrect configuration")
	}

	db, err := sql.Open("postgres", dsn)
	if err != nil {
		return nil, err
	}

	base, err := lynktable.NewBase(db)
	if err != nil {
		return nil, err
	}
	base.BindVar = dialectStmtBindVar
	base.QuoteStr = dialectQuoteStr

	for k, v := range dialectStmts {
		base.StmtSet(k, v)
	}

	return &Dialect{
		Base:   base,
		dbName: opts["dbname"],
	}, nil
}

// Upsert inserts a row, or on conflict of the given keys updates the
// remaining columns from the EXCLUDED row (INSERT ... ON CONFLICT DO UPDATE).
func (dc *Dialect) Upsert(tableName string, item map[string]any, keys map[string]any) lynktable.Result {

	if len(item) == 0 || len(keys) == 0 {
		return lynktable.NewResult(nil, errors.New("invalid args"))
	}

	var (
		cols, vars, vals, sets, pkeys = []string{}, []string{}, []any{}, []string{}, []string{}
	)

	for key, val := range item {
		cols = append(cols, dialectQuoteStr(key))
		vars = append(vars, "?")
		vals = append(vals, val)
		sets = append(sets, fmt.Sprintf("%s = EXCLUDED.%s", dialectQuoteStr(key), dialectQuoteStr(key)))
	}

	for key, val := range keys {
		cols = append(cols, dialectQuoteStr(key))
		vars = append(vars, "?")
		vals = append(vals, val)
		pkeys = append(pkeys, dialectQuoteStr(key))
	}

	q := fmt.Sprintf("INSERT INTO %s (%s) VALUES (%s) ON CONFLICT (%s) DO UPDATE SET %s",
		dialectQuoteStr(tableName),
		strings.Join(cols, ","),
		strings.Join(vars, ","),
		strings.Join(pkeys, ","),
		strings.Join(sets, ","))

	if dc.BindVar != nil {
		q, vals = dc.BindVar(q, vals)
	}

	stmt, err := dc.DB().Prepare(q)
	if err != nil {
		return lynktable.NewResult(nil, err)
	}
	defer stmt.Close()

	res, err := stmt.Exec(vals...)
	if err != nil {
		return lynktable.NewResult(nil, err)
	}

	return lynktable.NewResult(res, nil)
}