package rethinkdb import ( "fmt" "strings" "github.com/dancannon/gorethink" ) func makeDB(session *gorethink.Session, name string) error { _, err := gorethink.DBCreate(name).RunWrite(session) if err != nil { if strings.Contains(err.Error(), "already exists") { return nil } return err } // Note: Wait() not supported by RethinkDB 2.3+ resp, err := gorethink.DB(name).Wait().Run(session) if resp != nil { resp.Close() } return err } // Table holds the configuration for setting up a RethinkDB table type Table struct { Name string PrimaryKey interface{} // Keys are the index names. If len(value) is 0, it is a simple index // on the field matching the key. Otherwise, it is a compound index // on the list of fields in the corrensponding slice value. SecondaryIndexes map[string][]string Config map[string]string } func (t Table) term(dbName string) gorethink.Term { return gorethink.DB(dbName).Table(t.Name) } func (t Table) wait(session *gorethink.Session, dbName string) error { // Note: Wait() not supported by RethinkDB 2.3+ resp, err := t.term(dbName).Wait().Run(session) if resp != nil { resp.Close() } return err } func (t Table) create(session *gorethink.Session, dbName string, numReplicas uint) error { createOpts := gorethink.TableCreateOpts{ PrimaryKey: t.PrimaryKey, Durability: "hard", } if _, err := gorethink.DB(dbName).TableCreate(t.Name, createOpts).RunWrite(session); err != nil { if !strings.Contains(err.Error(), "already exists") { return fmt.Errorf("unable to run table creation: %s", err) } } reconfigureOpts := gorethink.ReconfigureOpts{ Shards: 1, Replicas: numReplicas, } if _, err := t.term(dbName).Reconfigure(reconfigureOpts).RunWrite(session); err != nil { return fmt.Errorf("unable to reconfigure table replication: %s", err) } if err := t.wait(session, dbName); err != nil { return fmt.Errorf("unable to wait for table to be ready after reconfiguring replication: %s", err) } if _, err := t.term(dbName).Config().Update(t.Config).RunWrite(session); err != nil { return fmt.Errorf("unable to configure table linearizability: %s", err) } if err := t.wait(session, dbName); err != nil { return fmt.Errorf("unable to wait for table to be ready after configuring linearizability: %s", err) } for indexName, fieldNames := range t.SecondaryIndexes { if len(fieldNames) == 0 { // The field name is the index name. fieldNames = []string{indexName} } if _, err := t.term(dbName).IndexCreateFunc(indexName, func(row gorethink.Term) interface{} { fields := make([]interface{}, len(fieldNames)) for i, fieldName := range fieldNames { term := row for _, subfield := range strings.Split(fieldName, ".") { term = term.Field(subfield) } fields[i] = term } if len(fields) == 1 { return fields[0] } return fields }).RunWrite(session); err != nil { if !strings.Contains(err.Error(), "already exists") { return fmt.Errorf("unable to create secondary index %q: %s", indexName, err) } } } if err := t.wait(session, dbName); err != nil { return fmt.Errorf("unable to wait for table to be ready after creating secondary indexes: %s", err) } return nil } // SetupDB handles creating the database and creating all tables and indexes. func SetupDB(session *gorethink.Session, dbName string, tables []Table) error { if err := makeDB(session, dbName); err != nil { return fmt.Errorf("unable to create database: %s", err) } cursor, err := gorethink.DB("rethinkdb").Table("server_config").Count().Run(session) if err != nil { return fmt.Errorf("unable to query db server config: %s", err) } var replicaCount uint if err := cursor.One(&replicaCount); err != nil { return fmt.Errorf("unable to scan db server config count: %s", err) } for _, table := range tables { if err = table.create(session, dbName, replicaCount); err != nil { return fmt.Errorf("unable to create table %q: %s", table.Name, err) } } return nil }