turso-maddy/table/turso.go

package table

import (
	"context"
	"database/sql"
	"fmt"
	"log"
	"os"

	"github.com/foxcpp/maddy/framework/config"
	"github.com/foxcpp/maddy/framework/module"

	"github.com/tursodatabase/go-libsql"
	_ "github.com/tursodatabase/libsql-client-go/libsql"
)

type table struct {
	modName  string
	instName string

	lookup string
	add    string
	list   string
	set    string
	del    string

	dir       string
	sql       *sql.DB
	connector *libsql.Connector
}

func NewTable(modName, instName string, _, _ []string) (module.Module, error) {
	return &table{
		modName:  modName,
		instName: instName,
	}, nil
}

func (s *table) Name() string {
	return s.modName
}

func (s *table) InstanceName() string {
	return s.instName
}

func (s *table) Init(cfg *config.Map) error {
	var (
		initQueries []string
		primaryUrl  string
	)
	cfg.StringList("init", false, false, nil, &initQueries)

	cfg.String("url", false, true, "", &primaryUrl)

	cfg.String("lookup", false, true, "", &s.lookup)
	cfg.String("add", false, false, "", &s.add)
	cfg.String("list", false, false, "", &s.list)
	cfg.String("del", false, false, "", &s.del)
	cfg.String("set", false, false, "", &s.set)

	if _, err := cfg.Process(); err != nil {
		return err
	}

	var err error
	s.sql, err = sql.Open("libsql", primaryUrl)

	if len(initQueries) > 0 {
		tx, err := s.sql.Begin()
		if err != nil {
			return err
		}

		for _, qry := range initQueries {
			res, err := tx.Exec(qry)
			if err != nil {
				return err
			}
			rows, err := res.RowsAffected()
			log.Println("db rows effected: ", rows, err)
		}
		if err = tx.Commit(); err != nil {
			return err
		}
	}

	return err
}

func (s *table) Close() error {
	defer os.RemoveAll(s.dir)
	defer s.connector.Close()

	return s.sql.Close()
}

func (s *table) Lookup(ctx context.Context, val string) (value string, ok bool, err error) {
	defer func() {
		if err != nil {
			err = fmt.Errorf("%s: lookup: %w", s.modName, err)
		}
	}()

	row := s.sql.QueryRowContext(ctx, s.lookup, val)
	err = row.Scan(&value)

	return value, err == nil, err
}

func (s *table) LookupMulti(ctx context.Context, val string) (lis []string, err error) {
	defer func() {
		if err != nil {
			err = fmt.Errorf("%s: lookupMulti: %w", s.modName, err)
		}
	}()

	var rows *sql.Rows
	rows, err = s.sql.QueryContext(ctx, s.lookup, val)
	if err != nil {
		return
	}
	defer rows.Close()

	for rows.Next() {
		var value string
		err = rows.Scan(&value)
		if err != nil {
			return
		}
		lis = append(lis, value)
	}
	err = rows.Err()

	return
}

func (s *table) Keys() (lis []string, err error) {
	if s.list == "" {
		return nil, fmt.Errorf("%s: table is not mutable (no 'list' query)", s.modName)
	}
	defer func() {
		if err != nil {
			err = fmt.Errorf("%s: list: %w", s.modName, err)
		}
	}()

	var rows *sql.Rows
	rows, err = s.sql.Query(s.list)
	if err != nil {
		return
	}
	defer rows.Close()

	for rows.Next() {
		var value string
		err = rows.Scan(&value)
		if err != nil {
			return
		}
		lis = append(lis, value)
	}
	err = rows.Err()

	return
}

func (s *table) RemoveKey(k string) (err error) {
	if s.del == "" {
		return fmt.Errorf("%s: table is not mutable (no 'del' query)", s.modName)
	}

	_, err = s.sql.ExecContext(context.TODO(), s.del, k)
	if err != nil {
		return fmt.Errorf("%s: del %s: %w", s.modName, k, err)
	}
	return nil
}

func (s *table) SetKey(k, v string) (err error) {
	if s.set == "" {
		return fmt.Errorf("%s: table is not mutable (no 'set' query)", s.modName)
	}
	if s.add == "" {
		return fmt.Errorf("%s: table is not mutable (no 'add' query)", s.modName)
	}

	res, err := s.sql.ExecContext(context.TODO(), s.set, k, v)
	if err != nil {
		return fmt.Errorf("%s: set %s: %w", s.modName, k, err)
	}
	var n int64
	if n, err = res.RowsAffected(); err != nil && n == 0 {
		_, err = s.sql.ExecContext(context.TODO(), s.add, k, v)
		if err != nil {
			return fmt.Errorf("%s: add %s: %w", s.modName, k, err)
		}
	}

	return
}

func init() {
	module.Register("table.turso_query", NewTable)
}