github.com/willyham/dosa@v2.3.1-0.20171024181418-1e446d37ee71+incompatible/connectors/cassandra/config.go (about)

     1  // Copyright (c) 2017 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package cassandra
    22  
    23  import (
    24  	"fmt"
    25  	"time"
    26  
    27  	"github.com/gocql/gocql"
    28  )
    29  
    30  // Config is a wrapper for the gocql.ClusterConfig, adding
    31  // support for yaml
    32  type Config struct {
    33  	gocql.ClusterConfig `yaml:",inline"`
    34  }
    35  
    36  // UnmarshalYAML unmarshals the config into gocql cluster config
    37  func (c *Config) UnmarshalYAML(unmarshal func(interface{}) error) error {
    38  	internal := &internalCluster{}
    39  	if err := unmarshal(internal); err != nil {
    40  		return err
    41  	}
    42  
    43  	c.ClusterConfig = *gocql.NewCluster(internal.Hosts...)
    44  	if internal.CQLVersion != nil {
    45  		c.CQLVersion = *internal.CQLVersion
    46  	}
    47  
    48  	if internal.ProtoVersion != nil {
    49  		c.ProtoVersion = *internal.ProtoVersion
    50  	}
    51  
    52  	if internal.Timeout != nil {
    53  		c.Timeout = *internal.Timeout
    54  	}
    55  
    56  	if internal.ConnectTimeout != nil {
    57  		c.ConnectTimeout = *internal.ConnectTimeout
    58  	}
    59  
    60  	if internal.Port != nil {
    61  		c.Port = *internal.Port
    62  	}
    63  
    64  	if internal.NumConns != nil {
    65  		c.NumConns = *internal.NumConns
    66  	}
    67  
    68  	if internal.Consistency != nil {
    69  		c.Consistency = gocql.ParseConsistency(*internal.Consistency)
    70  	}
    71  
    72  	if internal.RetryPolicy != nil {
    73  		c.RetryPolicy = &gocql.SimpleRetryPolicy{NumRetries: *internal.RetryPolicy}
    74  	}
    75  
    76  	if internal.SocketKeepalive != nil {
    77  		c.ClusterConfig.SocketKeepalive = *internal.SocketKeepalive
    78  	}
    79  
    80  	if internal.MaxPreparedStmts != nil {
    81  		c.ClusterConfig.MaxPreparedStmts = *internal.MaxPreparedStmts
    82  	}
    83  
    84  	if internal.MaxRoutingKeyInfo != nil {
    85  		c.ClusterConfig.MaxRoutingKeyInfo = *internal.MaxRoutingKeyInfo
    86  	}
    87  
    88  	if internal.PageSize != nil {
    89  		c.ClusterConfig.PageSize = *internal.PageSize
    90  	}
    91  
    92  	if internal.SerialConsistency != nil {
    93  		switch *internal.SerialConsistency {
    94  		case gocql.Serial.String():
    95  			c.ClusterConfig.SerialConsistency = gocql.Serial
    96  		case gocql.LocalSerial.String():
    97  			c.ClusterConfig.SerialConsistency = gocql.LocalSerial
    98  		default:
    99  			return fmt.Errorf("invalid serial consistency %q", *internal.SerialConsistency)
   100  		}
   101  	}
   102  
   103  	if internal.HostSelectionPolicy != nil {
   104  		switch *internal.HostSelectionPolicy {
   105  		case "RoundRobinHostPolicy":
   106  			c.ClusterConfig.PoolConfig = gocql.PoolConfig{HostSelectionPolicy: gocql.RoundRobinHostPolicy()}
   107  		case "TokenAwareHostPolicy":
   108  			c.ClusterConfig.PoolConfig = gocql.PoolConfig{HostSelectionPolicy: gocql.TokenAwareHostPolicy(gocql.RoundRobinHostPolicy())}
   109  		default:
   110  			return fmt.Errorf("unrecognized host selection policy: %s", *internal.HostSelectionPolicy)
   111  		}
   112  	}
   113  
   114  	if internal.DataCenter != nil {
   115  		c.ClusterConfig.HostFilter = gocql.DataCentreHostFilter(*internal.DataCenter)
   116  	}
   117  	return nil
   118  }
   119  
   120  type internalCluster struct {
   121  	Hosts               []string       `yaml:"hosts"`
   122  	CQLVersion          *string        `yaml:"cqlVersion"`
   123  	ProtoVersion        *int           `yaml:"protoVersion"`
   124  	Timeout             *time.Duration `yaml:"timeout"`
   125  	ConnectTimeout      *time.Duration `yaml:"connectTimeout"`
   126  	Port                *int           `yaml:"port"`
   127  	NumConns            *int           `yaml:"numConns"`
   128  	Consistency         *string        `yaml:"consistency"`
   129  	RetryPolicy         *int           `yaml:"retryPolicy"`
   130  	SocketKeepalive     *time.Duration `yaml:"socketKeepalive"`
   131  	MaxPreparedStmts    *int           `yaml:"maxPreparedStmts"`
   132  	MaxRoutingKeyInfo   *int           `yaml:"maxRoutingKeyInfo"`
   133  	PageSize            *int           `yaml:"pageSize"`
   134  	SerialConsistency   *string        `yaml:"serialConsistency"`
   135  	HostSelectionPolicy *string        `yaml:"hostSelectionPolicy"`
   136  	DataCenter          *string        `yaml:"dataCenter"`
   137  }