Skip to content

Instantly share code, notes, and snippets.

@shibukawa
Created May 7, 2019 17:28
Show Gist options
  • Star 1 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save shibukawa/5ce09c12a8238393a33f443f628796dd to your computer and use it in GitHub Desktop.
Save shibukawa/5ce09c12a8238393a33f443f628796dd to your computer and use it in GitHub Desktop.
package crlite
import (
"database/sql"
"fmt"
"io"
"os/exec"
"syscall"
"time"
"github.com/shibukawa/size"
_ "github.com/lib/pq"
)
type StorageType string
const (
HDD StorageType = "hdd"
SSD = "ssd"
defaultDatabaseFolder = "default-db"
)
type Option struct {
Size size.Size
Percent int
Type StorageType
}
type closer struct {
cmd *exec.Cmd
closed bool
}
func (c *closer) Close() error {
if c.closed {
return fmt.Errorf("Process '%s' is already closed", c.cmd.Path)
}
c.cmd.Process.Signal(syscall.SIGTERM)
return c.cmd.Wait()
}
func OpenInMemoryDB(dbName string, option ...Option) (*sql.DB, io.Closer, error) {
if len(option) > 0 {
return openDB(dbName, true, "", option[0])
} else {
return openDB(dbName, true, "", Option{})
}
}
func OpenLocalDB(path, dbName string, option ...Option) (*sql.DB, io.Closer, error) {
if len(option) > 0 {
return openDB(dbName, false, path, option[0])
} else {
return openDB(dbName, false, path, Option{})
}
}
func openDB(dbName string, memory bool, path string, option Option) (*sql.DB, io.Closer, error) {
var storeOpt string
if memory {
storeOpt = "--store=type=mem"
if option.Size == 0 {
option.Size = size.GB
}
} else {
if path == "" {
path = defaultDatabaseFolder
}
if option.Type == "" {
option.Type = SSD
}
storeOpt = fmt.Sprintf("--store=path=%s,attrs=%s", path, option.Type)
}
if option.Size > 0 {
storeOpt += fmt.Sprintf(",size=%s", option.Size)
} else if option.Percent > 0 {
storeOpt += fmt.Sprintf(",size=%d%%", option.Percent)
}
cmd := exec.Command("cockroach", "start", "--advertise-addr=localhost", "--insecure", storeOpt)
err := cmd.Start()
if err != nil {
return nil, nil, err
}
c := &closer{cmd: cmd, closed: false}
for i := 0; i < 10; i++ {
db, err := sql.Open("postgres", "postgresql://localhost:26257/"+dbName)
if err == nil {
return db, c, nil
}
fmt.Println(err)
if i != 9 {
time.Sleep(500 * time.Millisecond)
}
}
c.Close()
return nil, nil, fmt.Errorf("Timeout to connect local cockroach db")
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment