session.go 32 KB

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