session.go 32 KB

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