// CreateKeyspace creates a keyspace if necessary
// ks -> keyspace name
// rs -> replication strategy class
// rf -> replication factor
func CreateKeyspace(c *gocql.ClusterConfig, ks string, rs string, rf int) error {
kis, err := GetKeyspaces(c)
if err != nil {
return err
}
kss := set.NewStringSet(kis)
if !kss.Contains(ks) {
log.Printf("Creating keyspace: %v\n", ks)
c.Keyspace = ""
s, err := c.CreateSession()
if err != nil {
return err
}
defer s.Close()
if rs == "" {
rs = "SimpleStrategy"
}
err = s.Query(fmt.Sprintf("CREATE KEYSPACE IF NOT EXISTS %v WITH REPLICATION = {'class': '%v', 'replication_factor': %v};", ks, rs, rf)).Exec()
if err != nil {
return err
}
}
c.Keyspace = ks
return nil
}
// CreateRequiredTypes ensures all the types passed in are created if necessary
func CreateRequiredTypes(c *gocql.ClusterConfig, rt []UDT) error {
rtn := []string{}
etn := []string{}
rtm := map[string]UDT{}
for _, u := range rt {
rtn = append(rtn, u.Name)
rtm[u.Name] = u
}
rts := set.NewStringSet(rtn)
s, err := c.CreateSession()
if err != nil {
return err
}
q := `SELECT type_name FROM system.schema_usertypes WHERE keyspace_name = '%v';`
q = fmt.Sprintf(q, c.Keyspace)
iter := s.Query(q).Iter()
for n := ""; iter.Scan(&n); {
etn = append(etn, n)
}
if err := iter.Close(); err != nil {
return err
}
ets := set.NewStringSet(etn)
missing := rts.Difference(ets).Items()
if len(missing) > 0 {
for _, mt := range missing {
err := CreateUDT(c, rtm[mt])
if err != nil {
return err
}
}
}
return nil
}
// CreateKeyspaceWithNetworkTopologyStrategy creates a keyspace if necesary with NetworkTopologyStrategy
// ks -> keyspace name
// rfmap -> map of datacenter name to replication factor for that DC
func CreateKeyspaceWithNetworkTopologyStrategy(c *gocql.ClusterConfig, ks string, rfmap map[string]uint) error {
kis, err := GetKeyspaces(c)
if err != nil {
return err
}
kss := set.NewStringSet(kis)
if !kss.Contains(ks) {
log.Printf("Creating keyspace: %v\n", ks)
c.Keyspace = ""
s, err := c.CreateSession()
if err != nil {
return err
}
defer s.Close()
q := fmt.Sprintf("CREATE KEYSPACE IF NOT EXISTS %v WITH REPLICATION = {'class': '%v', ", ks, "NetworkTopologyStrategy")
rfsl := []string{}
for dc, rf := range rfmap {
rfsl = append(rfsl, fmt.Sprintf("'%v' : %v", dc, rf))
}
q = fmt.Sprintf("%v%v};", q, strings.Join(rfsl, ", "))
err = s.Query(q).Exec()
if err != nil {
return err
}
}
c.Keyspace = ks
return nil
}
// Output each metric in the given registry to cassandra periodically by creating new session
// the given cassandrasyslogger.
func Cassandra(r Registry, d time.Duration, cassandra_cluster *gocql.ClusterConfig, query string, server string) {
for {
session, _ := cassandra_cluster.CreateSession()
r.Each(func(name string, i interface{}) {
switch metric := i.(type) {
case Counter:
if err := session.Query(query,
server, name, time.Now(), float32(metric.Count())).Exec(); err != nil {
panic(err)
}
if metric.Get_reset_on_submit() {
metric.Clear()
}
case Gauge:
if err := session.Query(`INSERT INTO metrics2(server,metric, time, v) VALUES
(?, ?, ?, ?) USING TTL 1209600`,
server, name, time.Now(), float32(metric.Value())).Exec(); err != nil {
panic(err)
}
}
})
session.Close()
time.Sleep(d)
}
}
func prepareCassandraCluster(cluster *gocql.ClusterConfig) {
cluster.Keyspace = "system"
session, err := cluster.CreateSession()
if err != nil {
panic(err)
}
defer session.Close()
err = session.Query("DROP KEYSPACE IF EXISTS ?", cassandraKeyspace).RetryPolicy(nil).Exec()
if err != nil {
panic(err)
}
err = session.Query("DROP KEYSPACE IF EXISTS ?", cassandraKeyspace).RetryPolicy(nil).Exec()
if err != nil {
panic(err)
}
err = session.Query("CREATE KEYSPACE ? WITH replication = { 'class' : 'SimpleStrategy', 'replication_factor' : 1 }", cassandraKeyspace).RetryPolicy(nil).Exec()
if err != nil {
panic(err)
}
err = session.Query("CREATE TABLE ?.kvbench (id UUID primary key, value varchar)", cassandraKeyspace).RetryPolicy(nil).Exec()
if err != nil {
panic(err)
}
}
//NewStorageInstance returns a new StorageDatastore containing a connection to the specified cluster
func NewStorageInstance(cluster *gocql.ClusterConfig) (*StorageDatastore, error) {
session, err := cluster.CreateSession()
if err != nil {
return nil, err
}
return &StorageDatastore{session}, nil
}
// NewCassandraDatabase creates an instance of database, backed by Cassandra.
func NewCassandraDatabase(clusterConfig *gocql.ClusterConfig) (cassandraDatabase, error) {
session, err := clusterConfig.CreateSession()
if err != nil {
return cassandraDatabase{}, err
}
return cassandraDatabase{
session: session,
}, nil
}
// CreateUDT creates a user-defined type
func CreateUDT(c *gocql.ClusterConfig, u UDT) error {
qs := fmt.Sprintf("CREATE TYPE IF NOT EXISTS %v ( %v );", u.Name, strings.Join(u.Columns, ", "))
s, err := c.CreateSession()
if err != nil {
return err
}
defer s.Close()
log.Printf("Creating UDT %v\n", u.Name)
return s.Query(qs).Exec()
}
// CreateTable creates a table
func CreateTable(c *gocql.ClusterConfig, t CTable) error {
qs := fmt.Sprintf("CREATE TABLE IF NOT EXISTS %v ( %v ) %v ;", t.Name, strings.Join(t.Columns, ", "), t.Options)
s, err := c.CreateSession()
if err != nil {
return err
}
defer s.Close()
log.Printf("Creating table %v\n", t.Name)
return s.Query(qs).Exec()
}
// NewCassandraDatabase creates an instance of database, backed by Cassandra.
func NewCassandraDatabase(clusterConfig *gocql.ClusterConfig) (Database, error) {
session, err := clusterConfig.CreateSession()
if err != nil {
return nil, err
}
return &defaultDatabase{
session: session,
allMetricsCache: make(map[api.MetricKey]bool),
allMetricsMutex: &sync.Mutex{},
tagIndexCache: make(map[tagIndexCacheKey]bool),
tagIndexMutex: &sync.Mutex{},
}, nil
}
// DropKeyspace deletes a keyspace and all data associated with it
func DropKeyspace(c *gocql.ClusterConfig, ks string) error {
c.Keyspace = ""
s, err := c.CreateSession()
if err != nil {
return err
}
defer s.Close()
log.Printf("Dropping keyspace: %v\n", ks)
err = s.Query(fmt.Sprintf("DROP KEYSPACE IF EXISTS %v\n", ks)).Exec()
if err != nil {
return err
}
return nil
}
// GetTables returns all tables in configured keyspace
func GetTables(c *gocql.ClusterConfig) ([]string, error) {
tables := []string{}
s, err := c.CreateSession()
if err != nil {
return tables, err
}
defer s.Close()
q := s.Query("SELECT columnfamily_name FROM system.schema_columnfamilies WHERE keyspace_name = ?;", c.Keyspace).Iter()
var tn string
for q.Scan(&tn) {
tables = append(tables, tn)
}
return tables, q.Close()
}
// GetKeyspaces returns all extant keyspaces
func GetKeyspaces(c *gocql.ClusterConfig) ([]string, error) {
kss := []string{}
s, err := c.CreateSession()
if err != nil {
return kss, err
}
defer s.Close()
q := s.Query("SELECT keyspace_name FROM system.schema_keyspaces;").Iter()
var kn string
for q.Scan(&kn) {
kss = append(kss, kn)
}
return kss, q.Close()
}
/*
Run cassandra task
*/
func (c *CassandraTask) Run(r *http.Request, data map[string]interface{}) (response *Response) {
response = NewResponse(http.StatusOK)
queries := []*Response{}
for _, query := range c.config.Queries {
args := []interface{}{}
var (
Result []map[string]interface{}
cluster *gocql.ClusterConfig
session *gocql.Session
err error
)
qr := NewResponse(http.StatusOK).StripStatusData()
chosts := []string{}
for _, i := range query.Cluster {
var chost string
if chost, err = c.Interpolate(i, data); err != nil {
qr.Error(err)
goto Append
}
chosts = append(chosts, chost)
}
// instantiate cluster
cluster = gocql.NewCluster(chosts...)
if cluster.Keyspace, err = c.Interpolate(query.Keyspace, data); err != nil {
qr.Error(err)
goto Append
}
if session, err = cluster.CreateSession(); err != nil {
qr.Error(err)
goto Append
}
if c.config.ReturnQueries {
qr.AddValue("query", query.Query)
}
for _, arg := range query.Args {
final, err := c.Interpolate(arg, data)
if err != nil {
qr.Error(err)
goto Append
} else {
args = append(args, final)
}
}
if c.config.ReturnQueries {
qr.AddValue("args", args)
}
// slicemap to result
if Result, err = session.Query(query.Query, args...).Iter().SliceMap(); err != nil {
qr.Error(err)
goto Append
} else {
qr.Result(Result)
goto Append
}
Append:
queries = append(queries, qr)
}
// single result
if c.config.singleResultIndex != -1 {
response.Result(queries[c.config.singleResultIndex])
} else {
response.Result(queries)
}
return
}
// WithSession adds a cassandra Session to the mix ;)
func WithSession(cassandra *gocql.ClusterConfig) (*gocql.Session, *cqlc.Context, error) {
ctx := cqlc.NewContext()
session, err := cassandra.CreateSession()
return session, ctx, err
}
