session.go 31 KB

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