Created
May 7, 2019 17:28
-
-
Save shibukawa/5ce09c12a8238393a33f443f628796dd to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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