session.go 32 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212
  1. // Copyright (c) 2012 The gocql Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style
  3. // license that can be found in the LICENSE file.
  4. package gocql
  5. import (
  6. "bytes"
  7. "encoding/binary"
  8. "errors"
  9. "fmt"
  10. "io"
  11. "log"
  12. "net"
  13. "strconv"
  14. "strings"
  15. "sync"
  16. "time"
  17. "unicode"
  18. "github.com/gocql/gocql/internal/lru"
  19. )
  20. // Session is the interface used by users to interact with the database.
  21. //
  22. // It's safe for concurrent use by multiple goroutines and a typical usage
  23. // scenario is to have one global session object to interact with the
  24. // whole Cassandra cluster.
  25. //
  26. // This type extends the Node interface by adding a convinient query builder
  27. // and automatically sets a default consinstency level on all operations
  28. // that do not have a consistency level set.
  29. type Session struct {
  30. pool *policyConnPool
  31. cons Consistency
  32. pageSize int
  33. prefetch float64
  34. routingKeyInfoCache routingKeyInfoLRU
  35. schemaDescriber *schemaDescriber
  36. trace Tracer
  37. hostSource *ringDescriber
  38. ring ring
  39. connCfg *ConnConfig
  40. mu sync.RWMutex
  41. hostFilter HostFilter
  42. control *controlConn
  43. // event handlers
  44. nodeEvents *eventDeouncer
  45. // ring metadata
  46. hosts []HostInfo
  47. cfg ClusterConfig
  48. closeMu sync.RWMutex
  49. isClosed bool
  50. }
  51. // NewSession wraps an existing Node.
  52. func NewSession(cfg ClusterConfig) (*Session, error) {
  53. //Check that hosts in the ClusterConfig is not empty
  54. if len(cfg.Hosts) < 1 {
  55. return nil, ErrNoHosts
  56. }
  57. //Adjust the size of the prepared statements cache to match the latest configuration
  58. stmtsLRU.Lock()
  59. initStmtsLRU(cfg.MaxPreparedStmts)
  60. stmtsLRU.Unlock()
  61. s := &Session{
  62. cons: cfg.Consistency,
  63. prefetch: 0.25,
  64. cfg: cfg,
  65. pageSize: cfg.PageSize,
  66. }
  67. connCfg, err := connConfig(s)
  68. if err != nil {
  69. s.Close()
  70. return nil, fmt.Errorf("gocql: unable to create session: %v", err)
  71. }
  72. s.connCfg = connCfg
  73. s.nodeEvents = newEventDeouncer("NodeEvents", s.handleNodeEvent)
  74. s.routingKeyInfoCache.lru = lru.New(cfg.MaxRoutingKeyInfo)
  75. // I think it might be a good idea to simplify this and make it always discover
  76. // hosts, maybe with more filters.
  77. s.hostSource = &ringDescriber{
  78. session: s,
  79. closeChan: make(chan bool),
  80. }
  81. s.pool = cfg.PoolConfig.buildPool(s)
  82. var hosts []*HostInfo
  83. if !cfg.disableControlConn {
  84. s.control = createControlConn(s)
  85. if err := s.control.connect(cfg.Hosts); err != nil {
  86. s.Close()
  87. return nil, err
  88. }
  89. // need to setup host source to check for broadcast_address in system.local
  90. localHasRPCAddr, _ := checkSystemLocal(s.control)
  91. s.hostSource.localHasRpcAddr = localHasRPCAddr
  92. hosts, _, err = s.hostSource.GetHosts()
  93. if err != nil {
  94. s.Close()
  95. return nil, err
  96. }
  97. } else {
  98. // we dont get host info
  99. hosts = make([]*HostInfo, len(cfg.Hosts))
  100. for i, hostport := range cfg.Hosts {
  101. // TODO: remove duplication
  102. addr, portStr, err := net.SplitHostPort(JoinHostPort(hostport, cfg.Port))
  103. if err != nil {
  104. s.Close()
  105. return nil, fmt.Errorf("NewSession: unable to parse hostport of addr %q: %v", hostport, err)
  106. }
  107. port, err := strconv.Atoi(portStr)
  108. if err != nil {
  109. s.Close()
  110. return nil, fmt.Errorf("NewSession: invalid port for hostport of addr %q: %v", hostport, err)
  111. }
  112. hosts[i] = &HostInfo{peer: addr, port: port, state: NodeUp}
  113. }
  114. }
  115. for _, host := range hosts {
  116. s.handleNodeUp(net.ParseIP(host.Peer()), host.Port())
  117. }
  118. // TODO(zariel): we probably dont need this any more as we verify that we
  119. // can connect to one of the endpoints supplied by using the control conn.
  120. // See if there are any connections in the pool
  121. if s.pool.Size() == 0 {
  122. s.Close()
  123. return nil, ErrNoConnectionsStarted
  124. }
  125. return s, nil
  126. }
  127. // SetConsistency sets the default consistency level for this session. This
  128. // setting can also be changed on a per-query basis and the default value
  129. // is Quorum.
  130. func (s *Session) SetConsistency(cons Consistency) {
  131. s.mu.Lock()
  132. s.cons = cons
  133. s.mu.Unlock()
  134. }
  135. // SetPageSize sets the default page size for this session. A value <= 0 will
  136. // disable paging. This setting can also be changed on a per-query basis.
  137. func (s *Session) SetPageSize(n int) {
  138. s.mu.Lock()
  139. s.pageSize = n
  140. s.mu.Unlock()
  141. }
  142. // SetPrefetch sets the default threshold for pre-fetching new pages. If
  143. // there are only p*pageSize rows remaining, the next page will be requested
  144. // automatically. This value can also be changed on a per-query basis and
  145. // the default value is 0.25.
  146. func (s *Session) SetPrefetch(p float64) {
  147. s.mu.Lock()
  148. s.prefetch = p
  149. s.mu.Unlock()
  150. }
  151. // SetTrace sets the default tracer for this session. This setting can also
  152. // be changed on a per-query basis.
  153. func (s *Session) SetTrace(trace Tracer) {
  154. s.mu.Lock()
  155. s.trace = trace
  156. s.mu.Unlock()
  157. }
  158. // Query generates a new query object for interacting with the database.
  159. // Further details of the query may be tweaked using the resulting query
  160. // value before the query is executed. Query is automatically prepared
  161. // if it has not previously been executed.
  162. func (s *Session) Query(stmt string, values ...interface{}) *Query {
  163. s.mu.RLock()
  164. qry := &Query{stmt: stmt, values: values, cons: s.cons,
  165. session: s, pageSize: s.pageSize, trace: s.trace,
  166. prefetch: s.prefetch, rt: s.cfg.RetryPolicy, serialCons: s.cfg.SerialConsistency,
  167. defaultTimestamp: s.cfg.DefaultTimestamp,
  168. }
  169. s.mu.RUnlock()
  170. return qry
  171. }
  172. type QueryInfo struct {
  173. Id []byte
  174. Args []ColumnInfo
  175. Rval []ColumnInfo
  176. PKeyColumns []int
  177. }
  178. // Bind generates a new query object based on the query statement passed in.
  179. // The query is automatically prepared if it has not previously been executed.
  180. // The binding callback allows the application to define which query argument
  181. // values will be marshalled as part of the query execution.
  182. // During execution, the meta data of the prepared query will be routed to the
  183. // binding callback, which is responsible for producing the query argument values.
  184. func (s *Session) Bind(stmt string, b func(q *QueryInfo) ([]interface{}, error)) *Query {
  185. s.mu.RLock()
  186. qry := &Query{stmt: stmt, binding: b, cons: s.cons,
  187. session: s, pageSize: s.pageSize, trace: s.trace,
  188. prefetch: s.prefetch, rt: s.cfg.RetryPolicy}
  189. s.mu.RUnlock()
  190. return qry
  191. }
  192. // Close closes all connections. The session is unusable after this
  193. // operation.
  194. func (s *Session) Close() {
  195. s.closeMu.Lock()
  196. defer s.closeMu.Unlock()
  197. if s.isClosed {
  198. return
  199. }
  200. s.isClosed = true
  201. s.pool.Close()
  202. if s.hostSource != nil {
  203. close(s.hostSource.closeChan)
  204. }
  205. if s.control != nil {
  206. s.control.close()
  207. }
  208. s.nodeEvents.stop()
  209. }
  210. func (s *Session) Closed() bool {
  211. s.closeMu.RLock()
  212. closed := s.isClosed
  213. s.closeMu.RUnlock()
  214. return closed
  215. }
  216. func (s *Session) executeQuery(qry *Query) *Iter {
  217. // fail fast
  218. if s.Closed() {
  219. return &Iter{err: ErrSessionClosed}
  220. }
  221. var iter *Iter
  222. qry.attempts = 0
  223. qry.totalLatency = 0
  224. for {
  225. host, conn := s.pool.Pick(qry)
  226. //Assign the error unavailable to the iterator
  227. if conn == nil {
  228. if qry.rt == nil || !qry.rt.Attempt(qry) {
  229. iter = &Iter{err: ErrNoConnections}
  230. break
  231. }
  232. continue
  233. }
  234. t := time.Now()
  235. iter = conn.executeQuery(qry)
  236. qry.totalLatency += time.Now().Sub(t).Nanoseconds()
  237. qry.attempts++
  238. //Exit for loop if the query was successful
  239. if iter.err == nil {
  240. host.Mark(iter.err)
  241. break
  242. }
  243. // Mark host as ok
  244. host.Mark(nil)
  245. if qry.rt == nil || !qry.rt.Attempt(qry) {
  246. break
  247. }
  248. }
  249. return iter
  250. }
  251. // KeyspaceMetadata returns the schema metadata for the keyspace specified.
  252. func (s *Session) KeyspaceMetadata(keyspace string) (*KeyspaceMetadata, error) {
  253. // fail fast
  254. if s.Closed() {
  255. return nil, ErrSessionClosed
  256. }
  257. if keyspace == "" {
  258. return nil, ErrNoKeyspace
  259. }
  260. s.mu.Lock()
  261. // lazy-init schemaDescriber
  262. if s.schemaDescriber == nil {
  263. s.schemaDescriber = newSchemaDescriber(s)
  264. }
  265. s.mu.Unlock()
  266. return s.schemaDescriber.getSchema(keyspace)
  267. }
  268. // returns routing key indexes and type info
  269. func (s *Session) routingKeyInfo(stmt string) (*routingKeyInfo, error) {
  270. s.routingKeyInfoCache.mu.Lock()
  271. entry, cached := s.routingKeyInfoCache.lru.Get(stmt)
  272. if cached {
  273. // done accessing the cache
  274. s.routingKeyInfoCache.mu.Unlock()
  275. // the entry is an inflight struct similiar to that used by
  276. // Conn to prepare statements
  277. inflight := entry.(*inflightCachedEntry)
  278. // wait for any inflight work
  279. inflight.wg.Wait()
  280. if inflight.err != nil {
  281. return nil, inflight.err
  282. }
  283. key, _ := inflight.value.(*routingKeyInfo)
  284. return key, nil
  285. }
  286. // create a new inflight entry while the data is created
  287. inflight := new(inflightCachedEntry)
  288. inflight.wg.Add(1)
  289. defer inflight.wg.Done()
  290. s.routingKeyInfoCache.lru.Add(stmt, inflight)
  291. s.routingKeyInfoCache.mu.Unlock()
  292. var (
  293. info *QueryInfo
  294. partitionKey []*ColumnMetadata
  295. )
  296. // get the query info for the statement
  297. host, conn := s.pool.Pick(nil)
  298. if conn == nil {
  299. // no connections
  300. inflight.err = ErrNoConnections
  301. // don't cache this error
  302. s.routingKeyInfoCache.Remove(stmt)
  303. return nil, inflight.err
  304. }
  305. info, inflight.err = conn.prepareStatement(stmt, nil)
  306. if inflight.err != nil {
  307. // don't cache this error
  308. s.routingKeyInfoCache.Remove(stmt)
  309. host.Mark(inflight.err)
  310. return nil, inflight.err
  311. }
  312. // Mark host as OK
  313. host.Mark(nil)
  314. if len(info.Args) == 0 {
  315. // no arguments, no routing key, and no error
  316. return nil, nil
  317. }
  318. // get the table metadata
  319. table := info.Args[0].Table
  320. var keyspaceMetadata *KeyspaceMetadata
  321. keyspaceMetadata, inflight.err = s.KeyspaceMetadata(s.cfg.Keyspace)
  322. if inflight.err != nil {
  323. // don't cache this error
  324. s.routingKeyInfoCache.Remove(stmt)
  325. return nil, inflight.err
  326. }
  327. tableMetadata, found := keyspaceMetadata.Tables[table]
  328. if !found {
  329. // unlikely that the statement could be prepared and the metadata for
  330. // the table couldn't be found, but this may indicate either a bug
  331. // in the metadata code, or that the table was just dropped.
  332. inflight.err = ErrNoMetadata
  333. // don't cache this error
  334. s.routingKeyInfoCache.Remove(stmt)
  335. return nil, inflight.err
  336. }
  337. partitionKey = tableMetadata.PartitionKey
  338. size := len(partitionKey)
  339. routingKeyInfo := &routingKeyInfo{
  340. indexes: make([]int, size),
  341. types: make([]TypeInfo, size),
  342. }
  343. for keyIndex, keyColumn := range partitionKey {
  344. // set an indicator for checking if the mapping is missing
  345. routingKeyInfo.indexes[keyIndex] = -1
  346. // find the column in the query info
  347. for argIndex, boundColumn := range info.Args {
  348. if keyColumn.Name == boundColumn.Name {
  349. // there may be many such bound columns, pick the first
  350. routingKeyInfo.indexes[keyIndex] = argIndex
  351. routingKeyInfo.types[keyIndex] = boundColumn.TypeInfo
  352. break
  353. }
  354. }
  355. if routingKeyInfo.indexes[keyIndex] == -1 {
  356. // missing a routing key column mapping
  357. // no routing key, and no error
  358. return nil, nil
  359. }
  360. }
  361. // cache this result
  362. inflight.value = routingKeyInfo
  363. return routingKeyInfo, nil
  364. }
  365. func (s *Session) executeBatch(batch *Batch) (*Iter, error) {
  366. // fail fast
  367. if s.Closed() {
  368. return nil, ErrSessionClosed
  369. }
  370. // Prevent the execution of the batch if greater than the limit
  371. // Currently batches have a limit of 65536 queries.
  372. // https://datastax-oss.atlassian.net/browse/JAVA-229
  373. if batch.Size() > BatchSizeMaximum {
  374. return nil, ErrTooManyStmts
  375. }
  376. var err error
  377. var iter *Iter
  378. batch.attempts = 0
  379. batch.totalLatency = 0
  380. for {
  381. host, conn := s.pool.Pick(nil)
  382. //Assign the error unavailable and break loop
  383. if conn == nil {
  384. err = ErrNoConnections
  385. break
  386. }
  387. t := time.Now()
  388. iter, err = conn.executeBatch(batch)
  389. batch.totalLatency += time.Now().Sub(t).Nanoseconds()
  390. batch.attempts++
  391. //Exit loop if operation executed correctly
  392. if err == nil {
  393. host.Mark(err)
  394. return iter, err
  395. }
  396. // Mark host as OK
  397. host.Mark(nil)
  398. if batch.rt == nil || !batch.rt.Attempt(batch) {
  399. break
  400. }
  401. }
  402. return nil, err
  403. }
  404. // ExecuteBatch executes a batch operation and returns nil if successful
  405. // otherwise an error is returned describing the failure.
  406. func (s *Session) ExecuteBatch(batch *Batch) error {
  407. _, err := s.executeBatch(batch)
  408. return err
  409. }
  410. // ExecuteBatchCAS executes a batch operation and returns nil if successful and
  411. // an iterator (to scan aditional rows if more than one conditional statement)
  412. // was sent, otherwise an error is returned describing the failure.
  413. // Further scans on the interator must also remember to include
  414. // the applied boolean as the first argument to *Iter.Scan
  415. func (s *Session) ExecuteBatchCAS(batch *Batch, dest ...interface{}) (applied bool, iter *Iter, err error) {
  416. if iter, err := s.executeBatch(batch); err == nil {
  417. if err := iter.checkErrAndNotFound(); err != nil {
  418. return false, nil, err
  419. }
  420. if len(iter.Columns()) > 1 {
  421. dest = append([]interface{}{&applied}, dest...)
  422. iter.Scan(dest...)
  423. } else {
  424. iter.Scan(&applied)
  425. }
  426. return applied, iter, nil
  427. } else {
  428. return false, nil, err
  429. }
  430. }
  431. // MapExecuteBatchCAS executes a batch operation much like ExecuteBatchCAS,
  432. // however it accepts a map rather than a list of arguments for the initial
  433. // scan.
  434. func (s *Session) MapExecuteBatchCAS(batch *Batch, dest map[string]interface{}) (applied bool, iter *Iter, err error) {
  435. if iter, err := s.executeBatch(batch); err == nil {
  436. if err := iter.checkErrAndNotFound(); err != nil {
  437. return false, nil, err
  438. }
  439. iter.MapScan(dest)
  440. applied = dest["[applied]"].(bool)
  441. delete(dest, "[applied]")
  442. // we usually close here, but instead of closing, just returin an error
  443. // if MapScan failed. Although Close just returns err, using Close
  444. // here might be confusing as we are not actually closing the iter
  445. return applied, iter, iter.err
  446. } else {
  447. return false, nil, err
  448. }
  449. }
  450. func (s *Session) connect(addr string, errorHandler ConnErrorHandler) (*Conn, error) {
  451. return Connect(addr, s.connCfg, errorHandler, s)
  452. }
  453. // Query represents a CQL statement that can be executed.
  454. type Query struct {
  455. stmt string
  456. values []interface{}
  457. cons Consistency
  458. pageSize int
  459. routingKey []byte
  460. routingKeyBuffer []byte
  461. pageState []byte
  462. prefetch float64
  463. trace Tracer
  464. session *Session
  465. rt RetryPolicy
  466. binding func(q *QueryInfo) ([]interface{}, error)
  467. attempts int
  468. totalLatency int64
  469. serialCons SerialConsistency
  470. defaultTimestamp bool
  471. disableAutoPage bool
  472. }
  473. // String implements the stringer interface.
  474. func (q Query) String() string {
  475. return fmt.Sprintf("[query statement=%q values=%+v consistency=%s]", q.stmt, q.values, q.cons)
  476. }
  477. //Attempts returns the number of times the query was executed.
  478. func (q *Query) Attempts() int {
  479. return q.attempts
  480. }
  481. //Latency returns the average amount of nanoseconds per attempt of the query.
  482. func (q *Query) Latency() int64 {
  483. if q.attempts > 0 {
  484. return q.totalLatency / int64(q.attempts)
  485. }
  486. return 0
  487. }
  488. // Consistency sets the consistency level for this query. If no consistency
  489. // level have been set, the default consistency level of the cluster
  490. // is used.
  491. func (q *Query) Consistency(c Consistency) *Query {
  492. q.cons = c
  493. return q
  494. }
  495. // GetConsistency returns the currently configured consistency level for
  496. // the query.
  497. func (q *Query) GetConsistency() Consistency {
  498. return q.cons
  499. }
  500. // Trace enables tracing of this query. Look at the documentation of the
  501. // Tracer interface to learn more about tracing.
  502. func (q *Query) Trace(trace Tracer) *Query {
  503. q.trace = trace
  504. return q
  505. }
  506. // PageSize will tell the iterator to fetch the result in pages of size n.
  507. // This is useful for iterating over large result sets, but setting the
  508. // page size to low might decrease the performance. This feature is only
  509. // available in Cassandra 2 and onwards.
  510. func (q *Query) PageSize(n int) *Query {
  511. q.pageSize = n
  512. return q
  513. }
  514. // DefaultTimestamp will enable the with default timestamp flag on the query.
  515. // If enable, this will replace the server side assigned
  516. // timestamp as default timestamp. Note that a timestamp in the query itself
  517. // will still override this timestamp. This is entirely optional.
  518. //
  519. // Only available on protocol >= 3
  520. func (q *Query) DefaultTimestamp(enable bool) *Query {
  521. q.defaultTimestamp = enable
  522. return q
  523. }
  524. // RoutingKey sets the routing key to use when a token aware connection
  525. // pool is used to optimize the routing of this query.
  526. func (q *Query) RoutingKey(routingKey []byte) *Query {
  527. q.routingKey = routingKey
  528. return q
  529. }
  530. // GetRoutingKey gets the routing key to use for routing this query. If
  531. // a routing key has not been explicitly set, then the routing key will
  532. // be constructed if possible using the keyspace's schema and the query
  533. // info for this query statement. If the routing key cannot be determined
  534. // then nil will be returned with no error. On any error condition,
  535. // an error description will be returned.
  536. func (q *Query) GetRoutingKey() ([]byte, error) {
  537. if q.routingKey != nil {
  538. return q.routingKey, nil
  539. }
  540. // try to determine the routing key
  541. routingKeyInfo, err := q.session.routingKeyInfo(q.stmt)
  542. if err != nil {
  543. return nil, err
  544. }
  545. if routingKeyInfo == nil {
  546. return nil, nil
  547. }
  548. if len(routingKeyInfo.indexes) == 1 {
  549. // single column routing key
  550. routingKey, err := Marshal(
  551. routingKeyInfo.types[0],
  552. q.values[routingKeyInfo.indexes[0]],
  553. )
  554. if err != nil {
  555. return nil, err
  556. }
  557. return routingKey, nil
  558. }
  559. // We allocate that buffer only once, so that further re-bind/exec of the
  560. // same query don't allocate more memory.
  561. if q.routingKeyBuffer == nil {
  562. q.routingKeyBuffer = make([]byte, 0, 256)
  563. }
  564. // composite routing key
  565. buf := bytes.NewBuffer(q.routingKeyBuffer)
  566. for i := range routingKeyInfo.indexes {
  567. encoded, err := Marshal(
  568. routingKeyInfo.types[i],
  569. q.values[routingKeyInfo.indexes[i]],
  570. )
  571. if err != nil {
  572. return nil, err
  573. }
  574. lenBuf := []byte{0x00, 0x00}
  575. binary.BigEndian.PutUint16(lenBuf, uint16(len(encoded)))
  576. buf.Write(lenBuf)
  577. buf.Write(encoded)
  578. buf.WriteByte(0x00)
  579. }
  580. routingKey := buf.Bytes()
  581. return routingKey, nil
  582. }
  583. func (q *Query) shouldPrepare() bool {
  584. stmt := strings.TrimLeftFunc(strings.TrimRightFunc(q.stmt, func(r rune) bool {
  585. return unicode.IsSpace(r) || r == ';'
  586. }), unicode.IsSpace)
  587. var stmtType string
  588. if n := strings.IndexFunc(stmt, unicode.IsSpace); n >= 0 {
  589. stmtType = strings.ToLower(stmt[:n])
  590. }
  591. if stmtType == "begin" {
  592. if n := strings.LastIndexFunc(stmt, unicode.IsSpace); n >= 0 {
  593. stmtType = strings.ToLower(stmt[n+1:])
  594. }
  595. }
  596. switch stmtType {
  597. case "select", "insert", "update", "delete", "batch":
  598. return true
  599. }
  600. return false
  601. }
  602. // SetPrefetch sets the default threshold for pre-fetching new pages. If
  603. // there are only p*pageSize rows remaining, the next page will be requested
  604. // automatically.
  605. func (q *Query) Prefetch(p float64) *Query {
  606. q.prefetch = p
  607. return q
  608. }
  609. // RetryPolicy sets the policy to use when retrying the query.
  610. func (q *Query) RetryPolicy(r RetryPolicy) *Query {
  611. q.rt = r
  612. return q
  613. }
  614. // Bind sets query arguments of query. This can also be used to rebind new query arguments
  615. // to an existing query instance.
  616. func (q *Query) Bind(v ...interface{}) *Query {
  617. q.values = v
  618. return q
  619. }
  620. // SerialConsistency sets the consistencyc level for the
  621. // serial phase of conditional updates. That consitency can only be
  622. // either SERIAL or LOCAL_SERIAL and if not present, it defaults to
  623. // SERIAL. This option will be ignored for anything else that a
  624. // conditional update/insert.
  625. func (q *Query) SerialConsistency(cons SerialConsistency) *Query {
  626. q.serialCons = cons
  627. return q
  628. }
  629. // PageState sets the paging state for the query to resume paging from a specific
  630. // point in time. Setting this will disable to query paging for this query, and
  631. // must be used for all subsequent pages.
  632. func (q *Query) PageState(state []byte) *Query {
  633. q.pageState = state
  634. q.disableAutoPage = true
  635. return q
  636. }
  637. // Exec executes the query without returning any rows.
  638. func (q *Query) Exec() error {
  639. iter := q.Iter()
  640. return iter.err
  641. }
  642. func isUseStatement(stmt string) bool {
  643. if len(stmt) < 3 {
  644. return false
  645. }
  646. return strings.ToLower(stmt[0:3]) == "use"
  647. }
  648. // Iter executes the query and returns an iterator capable of iterating
  649. // over all results.
  650. func (q *Query) Iter() *Iter {
  651. if isUseStatement(q.stmt) {
  652. return &Iter{err: ErrUseStmt}
  653. }
  654. return q.session.executeQuery(q)
  655. }
  656. // MapScan executes the query, copies the columns of the first selected
  657. // row into the map pointed at by m and discards the rest. If no rows
  658. // were selected, ErrNotFound is returned.
  659. func (q *Query) MapScan(m map[string]interface{}) error {
  660. iter := q.Iter()
  661. if err := iter.checkErrAndNotFound(); err != nil {
  662. return err
  663. }
  664. iter.MapScan(m)
  665. return iter.Close()
  666. }
  667. // Scan executes the query, copies the columns of the first selected
  668. // row into the values pointed at by dest and discards the rest. If no rows
  669. // were selected, ErrNotFound is returned.
  670. func (q *Query) Scan(dest ...interface{}) error {
  671. iter := q.Iter()
  672. if err := iter.checkErrAndNotFound(); err != nil {
  673. return err
  674. }
  675. iter.Scan(dest...)
  676. return iter.Close()
  677. }
  678. // ScanCAS executes a lightweight transaction (i.e. an UPDATE or INSERT
  679. // statement containing an IF clause). If the transaction fails because
  680. // the existing values did not match, the previous values will be stored
  681. // in dest.
  682. func (q *Query) ScanCAS(dest ...interface{}) (applied bool, err error) {
  683. iter := q.Iter()
  684. if err := iter.checkErrAndNotFound(); err != nil {
  685. return false, err
  686. }
  687. if len(iter.Columns()) > 1 {
  688. dest = append([]interface{}{&applied}, dest...)
  689. iter.Scan(dest...)
  690. } else {
  691. iter.Scan(&applied)
  692. }
  693. return applied, iter.Close()
  694. }
  695. // MapScanCAS executes a lightweight transaction (i.e. an UPDATE or INSERT
  696. // statement containing an IF clause). If the transaction fails because
  697. // the existing values did not match, the previous values will be stored
  698. // in dest map.
  699. //
  700. // As for INSERT .. IF NOT EXISTS, previous values will be returned as if
  701. // SELECT * FROM. So using ScanCAS with INSERT is inherently prone to
  702. // column mismatching. MapScanCAS is added to capture them safely.
  703. func (q *Query) MapScanCAS(dest map[string]interface{}) (applied bool, err error) {
  704. iter := q.Iter()
  705. if err := iter.checkErrAndNotFound(); err != nil {
  706. return false, err
  707. }
  708. iter.MapScan(dest)
  709. applied = dest["[applied]"].(bool)
  710. delete(dest, "[applied]")
  711. return applied, iter.Close()
  712. }
  713. // Iter represents an iterator that can be used to iterate over all rows that
  714. // were returned by a query. The iterator might send additional queries to the
  715. // database during the iteration if paging was enabled.
  716. type Iter struct {
  717. err error
  718. pos int
  719. rows [][][]byte
  720. meta resultMetadata
  721. next *nextIter
  722. framer *framer
  723. once sync.Once
  724. }
  725. // Columns returns the name and type of the selected columns.
  726. func (iter *Iter) Columns() []ColumnInfo {
  727. return iter.meta.columns
  728. }
  729. // Scan consumes the next row of the iterator and copies the columns of the
  730. // current row into the values pointed at by dest. Use nil as a dest value
  731. // to skip the corresponding column. Scan might send additional queries
  732. // to the database to retrieve the next set of rows if paging was enabled.
  733. //
  734. // Scan returns true if the row was successfully unmarshaled or false if the
  735. // end of the result set was reached or if an error occurred. Close should
  736. // be called afterwards to retrieve any potential errors.
  737. func (iter *Iter) Scan(dest ...interface{}) bool {
  738. if iter.err != nil {
  739. return false
  740. }
  741. if iter.pos >= len(iter.rows) {
  742. if iter.next != nil {
  743. *iter = *iter.next.fetch()
  744. return iter.Scan(dest...)
  745. }
  746. return false
  747. }
  748. if iter.next != nil && iter.pos == iter.next.pos {
  749. go iter.next.fetch()
  750. }
  751. // currently only support scanning into an expand tuple, such that its the same
  752. // as scanning in more values from a single column
  753. if len(dest) != iter.meta.actualColCount {
  754. iter.err = errors.New("count mismatch")
  755. return false
  756. }
  757. // i is the current position in dest, could posible replace it and just use
  758. // slices of dest
  759. i := 0
  760. for c, col := range iter.meta.columns {
  761. if dest[i] == nil {
  762. i++
  763. continue
  764. }
  765. switch col.TypeInfo.Type() {
  766. case TypeTuple:
  767. // this will panic, actually a bug, please report
  768. tuple := col.TypeInfo.(TupleTypeInfo)
  769. count := len(tuple.Elems)
  770. // here we pass in a slice of the struct which has the number number of
  771. // values as elements in the tuple
  772. iter.err = Unmarshal(col.TypeInfo, iter.rows[iter.pos][c], dest[i:i+count])
  773. i += count
  774. default:
  775. iter.err = Unmarshal(col.TypeInfo, iter.rows[iter.pos][c], dest[i])
  776. i++
  777. }
  778. if iter.err != nil {
  779. return false
  780. }
  781. }
  782. iter.pos++
  783. return true
  784. }
  785. // Close closes the iterator and returns any errors that happened during
  786. // the query or the iteration.
  787. func (iter *Iter) Close() error {
  788. iter.once.Do(func() {
  789. if iter.framer != nil {
  790. framerPool.Put(iter.framer)
  791. iter.framer = nil
  792. }
  793. })
  794. return iter.err
  795. }
  796. // WillSwitchPage detects if iterator reached end of current page
  797. // and the next page is available.
  798. func (iter *Iter) WillSwitchPage() bool {
  799. return iter.pos >= len(iter.rows) && iter.next != nil
  800. }
  801. // checkErrAndNotFound handle error and NotFound in one method.
  802. func (iter *Iter) checkErrAndNotFound() error {
  803. if iter.err != nil {
  804. return iter.err
  805. } else if len(iter.rows) == 0 {
  806. return ErrNotFound
  807. }
  808. return nil
  809. }
  810. // PageState return the current paging state for a query which can be used for
  811. // subsequent quries to resume paging this point.
  812. func (iter *Iter) PageState() []byte {
  813. return iter.meta.pagingState
  814. }
  815. type nextIter struct {
  816. qry Query
  817. pos int
  818. once sync.Once
  819. next *Iter
  820. }
  821. func (n *nextIter) fetch() *Iter {
  822. n.once.Do(func() {
  823. n.next = n.qry.session.executeQuery(&n.qry)
  824. })
  825. return n.next
  826. }
  827. type Batch struct {
  828. Type BatchType
  829. Entries []BatchEntry
  830. Cons Consistency
  831. rt RetryPolicy
  832. attempts int
  833. totalLatency int64
  834. serialCons SerialConsistency
  835. defaultTimestamp bool
  836. }
  837. // NewBatch creates a new batch operation without defaults from the cluster
  838. func NewBatch(typ BatchType) *Batch {
  839. return &Batch{Type: typ}
  840. }
  841. // NewBatch creates a new batch operation using defaults defined in the cluster
  842. func (s *Session) NewBatch(typ BatchType) *Batch {
  843. s.mu.RLock()
  844. batch := &Batch{Type: typ, rt: s.cfg.RetryPolicy, serialCons: s.cfg.SerialConsistency,
  845. Cons: s.cons, defaultTimestamp: s.cfg.DefaultTimestamp}
  846. s.mu.RUnlock()
  847. return batch
  848. }
  849. // Attempts returns the number of attempts made to execute the batch.
  850. func (b *Batch) Attempts() int {
  851. return b.attempts
  852. }
  853. //Latency returns the average number of nanoseconds to execute a single attempt of the batch.
  854. func (b *Batch) Latency() int64 {
  855. if b.attempts > 0 {
  856. return b.totalLatency / int64(b.attempts)
  857. }
  858. return 0
  859. }
  860. // GetConsistency returns the currently configured consistency level for the batch
  861. // operation.
  862. func (b *Batch) GetConsistency() Consistency {
  863. return b.Cons
  864. }
  865. // Query adds the query to the batch operation
  866. func (b *Batch) Query(stmt string, args ...interface{}) {
  867. b.Entries = append(b.Entries, BatchEntry{Stmt: stmt, Args: args})
  868. }
  869. // Bind adds the query to the batch operation and correlates it with a binding callback
  870. // that will be invoked when the batch is executed. The binding callback allows the application
  871. // to define which query argument values will be marshalled as part of the batch execution.
  872. func (b *Batch) Bind(stmt string, bind func(q *QueryInfo) ([]interface{}, error)) {
  873. b.Entries = append(b.Entries, BatchEntry{Stmt: stmt, binding: bind})
  874. }
  875. // RetryPolicy sets the retry policy to use when executing the batch operation
  876. func (b *Batch) RetryPolicy(r RetryPolicy) *Batch {
  877. b.rt = r
  878. return b
  879. }
  880. // Size returns the number of batch statements to be executed by the batch operation.
  881. func (b *Batch) Size() int {
  882. return len(b.Entries)
  883. }
  884. // SerialConsistency sets the consistencyc level for the
  885. // serial phase of conditional updates. That consitency can only be
  886. // either SERIAL or LOCAL_SERIAL and if not present, it defaults to
  887. // SERIAL. This option will be ignored for anything else that a
  888. // conditional update/insert.
  889. //
  890. // Only available for protocol 3 and above
  891. func (b *Batch) SerialConsistency(cons SerialConsistency) *Batch {
  892. b.serialCons = cons
  893. return b
  894. }
  895. // DefaultTimestamp will enable the with default timestamp flag on the query.
  896. // If enable, this will replace the server side assigned
  897. // timestamp as default timestamp. Note that a timestamp in the query itself
  898. // will still override this timestamp. This is entirely optional.
  899. //
  900. // Only available on protocol >= 3
  901. func (b *Batch) DefaultTimestamp(enable bool) *Batch {
  902. b.defaultTimestamp = enable
  903. return b
  904. }
  905. type BatchType byte
  906. const (
  907. LoggedBatch BatchType = 0
  908. UnloggedBatch BatchType = 1
  909. CounterBatch BatchType = 2
  910. )
  911. type BatchEntry struct {
  912. Stmt string
  913. Args []interface{}
  914. binding func(q *QueryInfo) ([]interface{}, error)
  915. }
  916. type ColumnInfo struct {
  917. Keyspace string
  918. Table string
  919. Name string
  920. TypeInfo TypeInfo
  921. }
  922. func (c ColumnInfo) String() string {
  923. return fmt.Sprintf("[column keyspace=%s table=%s name=%s type=%v]", c.Keyspace, c.Table, c.Name, c.TypeInfo)
  924. }
  925. // routing key indexes LRU cache
  926. type routingKeyInfoLRU struct {
  927. lru *lru.Cache
  928. mu sync.Mutex
  929. }
  930. type routingKeyInfo struct {
  931. indexes []int
  932. types []TypeInfo
  933. }
  934. func (r *routingKeyInfoLRU) Remove(key string) {
  935. r.mu.Lock()
  936. r.lru.Remove(key)
  937. r.mu.Unlock()
  938. }
  939. //Max adjusts the maximum size of the cache and cleans up the oldest records if
  940. //the new max is lower than the previous value. Not concurrency safe.
  941. func (r *routingKeyInfoLRU) Max(max int) {
  942. r.mu.Lock()
  943. for r.lru.Len() > max {
  944. r.lru.RemoveOldest()
  945. }
  946. r.lru.MaxEntries = max
  947. r.mu.Unlock()
  948. }
  949. type inflightCachedEntry struct {
  950. wg sync.WaitGroup
  951. err error
  952. value interface{}
  953. }
  954. // Tracer is the interface implemented by query tracers. Tracers have the
  955. // ability to obtain a detailed event log of all events that happened during
  956. // the execution of a query from Cassandra. Gathering this information might
  957. // be essential for debugging and optimizing queries, but this feature should
  958. // not be used on production systems with very high load.
  959. type Tracer interface {
  960. Trace(traceId []byte)
  961. }
  962. type traceWriter struct {
  963. session *Session
  964. w io.Writer
  965. mu sync.Mutex
  966. }
  967. // NewTraceWriter returns a simple Tracer implementation that outputs
  968. // the event log in a textual format.
  969. func NewTraceWriter(session *Session, w io.Writer) Tracer {
  970. return &traceWriter{session: session, w: w}
  971. }
  972. func (t *traceWriter) Trace(traceId []byte) {
  973. var (
  974. coordinator string
  975. duration int
  976. )
  977. iter := t.session.control.query(`SELECT coordinator, duration
  978. FROM system_traces.sessions
  979. WHERE session_id = ?`, traceId)
  980. iter.Scan(&coordinator, &duration)
  981. if err := iter.Close(); err != nil {
  982. t.mu.Lock()
  983. fmt.Fprintln(t.w, "Error:", err)
  984. t.mu.Unlock()
  985. return
  986. }
  987. var (
  988. timestamp time.Time
  989. activity string
  990. source string
  991. elapsed int
  992. )
  993. fmt.Fprintf(t.w, "Tracing session %016x (coordinator: %s, duration: %v):\n",
  994. traceId, coordinator, time.Duration(duration)*time.Microsecond)
  995. t.mu.Lock()
  996. defer t.mu.Unlock()
  997. iter = t.session.control.query(`SELECT event_id, activity, source, source_elapsed
  998. FROM system_traces.events
  999. WHERE session_id = ?`, traceId)
  1000. for iter.Scan(&timestamp, &activity, &source, &elapsed) {
  1001. fmt.Fprintf(t.w, "%s: %s (source: %s, elapsed: %d)\n",
  1002. timestamp.Format("2006/01/02 15:04:05.999999"), activity, source, elapsed)
  1003. }
  1004. if err := iter.Close(); err != nil {
  1005. fmt.Fprintln(t.w, "Error:", err)
  1006. }
  1007. }
  1008. type Error struct {
  1009. Code int
  1010. Message string
  1011. }
  1012. func (e Error) Error() string {
  1013. return e.Message
  1014. }
  1015. var (
  1016. ErrNotFound = errors.New("not found")
  1017. ErrUnavailable = errors.New("unavailable")
  1018. ErrUnsupported = errors.New("feature not supported")
  1019. ErrTooManyStmts = errors.New("too many statements")
  1020. ErrUseStmt = errors.New("use statements aren't supported. Please see https://github.com/gocql/gocql for explaination.")
  1021. ErrSessionClosed = errors.New("session has been closed")
  1022. ErrNoConnections = errors.New("no connections available")
  1023. ErrNoKeyspace = errors.New("no keyspace provided")
  1024. ErrNoMetadata = errors.New("no metadata available")
  1025. )
  1026. type ErrProtocol struct{ error }
  1027. func NewErrProtocol(format string, args ...interface{}) error {
  1028. return ErrProtocol{fmt.Errorf(format, args...)}
  1029. }
  1030. // BatchSizeMaximum is the maximum number of statements a batch operation can have.
  1031. // This limit is set by cassandra and could change in the future.
  1032. const BatchSizeMaximum = 65535