joao-r-reis commented on code in PR #1868:
URL: 
https://github.com/apache/cassandra-gocql-driver/pull/1868#discussion_r2044793118


##########
session.go:
##########
@@ -743,33 +736,47 @@ func (s *Session) routingKeyInfo(ctx context.Context, 
stmt string, keyspace stri
        return routingKeyInfo, nil
 }
 
-func (b *Batch) execute(ctx context.Context, conn *Conn) *Iter {
-       return conn.executeBatch(ctx, b)
-}
-
 // Exec executes a batch operation and returns nil if successful
 // otherwise an error is returned describing the failure.
 func (b *Batch) Exec() error {
-       iter := b.session.executeBatch(b)
+       iter := b.session.executeBatch(b, nil)
+       return iter.Close()
+}
+
+// ExecContext executes a batch operation with the provided context and 
returns nil if successful
+// otherwise an error is returned describing the failure.
+func (b *Batch) ExecContext(ctx context.Context) error {
+       iter := b.session.executeBatch(b, ctx)
        return iter.Close()
 }
 
-func (s *Session) executeBatch(batch *Batch) *Iter {
+// Iter executes a batch operation and returns an Iter object
+// that can be used to access properties related to the execution like 
Iter.Attempts and Iter.Latency
+func (b *Batch) Iter() *Iter { return b.session.executeBatch(b, nil) }
+
+// Iter executes a batch operation with the provided context and returns an 
Iter object
+// that can be used to access properties related to the execution like 
Iter.Attempts and Iter.Latency
+func (b *Batch) IterContext(ctx context.Context) *Iter {
+       return b.session.executeBatch(b, ctx)
+}

Review Comment:
   done



##########
query_executor.go:
##########
@@ -222,16 +238,452 @@ func (q *queryExecutor) do(ctx context.Context, qry 
ExecutableQuery, hostIter Ne
        }
 
        if lastErr != nil {
-               return &Iter{err: lastErr}
+               return newErrIter(lastErr, qry.getQueryMetrics(), 
qry.Keyspace(), qry.getRoutingInfo(), qry.getKeyspaceFunc())
        }
 
-       return &Iter{err: ErrNoConnections}
+       return newErrIter(ErrNoConnections, qry.getQueryMetrics(), 
qry.Keyspace(), qry.getRoutingInfo(), qry.getKeyspaceFunc())
 }
 
-func (q *queryExecutor) run(ctx context.Context, qry ExecutableQuery, hostIter 
NextHost, results chan<- *Iter) {
+func (q *queryExecutor) run(ctx context.Context, qry internalRequest, hostIter 
NextHost, results chan<- *Iter) {
        select {
        case results <- q.do(ctx, qry, hostIter):
        case <-ctx.Done():
        }
-       qry.releaseAfterExecution()
+}
+
+type queryOptions struct {
+       stmt string
+
+       // Paging
+       pageSize        int
+       disableAutoPage bool
+
+       // Monitoring
+       trace    Tracer
+       observer QueryObserver
+
+       // Parameters
+       values  []interface{}
+       binding func(q *QueryInfo) ([]interface{}, error)
+
+       // Timestamp
+       defaultTimestamp      bool
+       defaultTimestampValue int64
+
+       // Consistency
+       serialCons SerialConsistency
+
+       // Protocol flag
+       disableSkipMetadata bool
+
+       customPayload     map[string][]byte
+       prefetch          float64
+       rt                RetryPolicy
+       spec              SpeculativeExecutionPolicy
+       context           context.Context
+       idempotent        bool
+       keyspace          string
+       skipPrepare       bool
+       routingKey        []byte
+       nowInSecondsValue *int
+       hostID            string
+
+       // getKeyspace is field so that it can be overriden in tests
+       getKeyspace func() string
+}
+
+func newQueryOptions(q *Query, ctx context.Context) *queryOptions {
+       var newPageState, newRoutingKey []byte
+       if q.initialPageState != nil {
+               pageState := q.initialPageState
+               newPageState = make([]byte, len(pageState))
+               copy(newPageState, pageState)
+       }
+       if q.routingKey != nil {
+               routingKey := q.routingKey
+               newRoutingKey = make([]byte, len(routingKey))
+               copy(newRoutingKey, routingKey)
+       }
+       if ctx == nil {
+               ctx = q.Context()
+       }
+       return &queryOptions{
+               stmt:                  q.stmt,
+               values:                q.values,
+               pageSize:              q.pageSize,
+               prefetch:              q.prefetch,
+               trace:                 q.trace,
+               observer:              q.observer,
+               rt:                    q.rt,
+               spec:                  q.spec,
+               binding:               q.binding,
+               serialCons:            q.serialCons,
+               defaultTimestamp:      q.defaultTimestamp,
+               defaultTimestampValue: q.defaultTimestampValue,
+               disableSkipMetadata:   q.disableSkipMetadata,
+               context:               ctx,
+               idempotent:            q.idempotent,
+               customPayload:         q.customPayload,
+               disableAutoPage:       q.disableAutoPage,
+               skipPrepare:           q.skipPrepare,
+               routingKey:            newRoutingKey,
+               getKeyspace:           q.getKeyspace,
+               nowInSecondsValue:     q.nowInSecondsValue,
+               keyspace:              q.keyspace,
+               hostID:                q.hostID,
+       }
+}
+
+type internalQuery struct {
+       originalQuery *Query
+       qryOpts       *queryOptions
+       pageState     []byte
+       metrics       *queryMetrics
+       refCount      uint32
+       conn          *Conn
+       consistency   uint32
+       session       *Session
+       routingInfo   *queryRoutingInfo
+}
+
+func newInternalQuery(q *Query, ctx context.Context) *internalQuery {
+       var newPageState []byte
+       if q.initialPageState != nil {
+               pageState := q.initialPageState
+               newPageState = make([]byte, len(pageState))
+               copy(newPageState, pageState)
+       }
+       return &internalQuery{
+               originalQuery: q,
+               qryOpts:       newQueryOptions(q, ctx),
+               metrics:       &queryMetrics{m: make(map[string]*hostMetrics)},
+               consistency:   uint32(q.initialConsistency),
+               pageState:     newPageState,
+               conn:          nil,
+               session:       q.session,
+               routingInfo:   &queryRoutingInfo{},
+       }
+}
+
+// Attempts returns the number of times the query was executed.
+func (q *internalQuery) Attempts() int {
+       return q.metrics.attempts()
+}
+
+func (q *internalQuery) attempt(keyspace string, end, start time.Time, iter 
*Iter, host *HostInfo) {
+       latency := end.Sub(start)
+       attempt, metricsForHost := q.metrics.attempt(1, latency, host, 
q.qryOpts.observer != nil)
+
+       if q.qryOpts.observer != nil {
+               q.qryOpts.observer.ObserveQuery(q.qryOpts.context, 
ObservedQuery{
+                       Keyspace:  keyspace,
+                       Statement: q.qryOpts.stmt,
+                       Values:    q.qryOpts.values,
+                       Start:     start,
+                       End:       end,
+                       Rows:      iter.numRows,
+                       Host:      host,
+                       Metrics:   metricsForHost,
+                       Err:       iter.err,
+                       Attempt:   attempt,
+                       Query:     q.originalQuery,
+               })
+       }
+}
+
+func (q *internalQuery) execute(ctx context.Context, conn *Conn) *Iter {
+       return conn.executeQuery(ctx, q)
+}
+
+func (q *internalQuery) retryPolicy() RetryPolicy {
+       return q.qryOpts.rt
+}
+
+func (q *internalQuery) speculativeExecutionPolicy() 
SpeculativeExecutionPolicy {
+       return q.qryOpts.spec
+}
+
+func (q *internalQuery) GetRoutingKey() ([]byte, error) {
+       if q.qryOpts.routingKey != nil {
+               return q.qryOpts.routingKey, nil
+       }
+
+       if q.qryOpts.binding != nil && len(q.qryOpts.values) == 0 {
+               // If this query was created using session.Bind we wont have 
the query
+               // values yet, so we have to pass down to the next policy.
+               // TODO: Remove this and handle this case
+               return nil, nil
+       }
+
+       // try to determine the routing key
+       routingKeyInfo, err := q.session.routingKeyInfo(q.Context(), 
q.qryOpts.stmt, q.qryOpts.keyspace)
+       if err != nil {
+               return nil, err
+       }
+
+       if routingKeyInfo != nil {
+               q.routingInfo.mu.Lock()
+               q.routingInfo.keyspace = routingKeyInfo.keyspace
+               q.routingInfo.table = routingKeyInfo.table
+               q.routingInfo.mu.Unlock()
+       }
+       return createRoutingKey(routingKeyInfo, q.qryOpts.values)
+}
+
+func (q *internalQuery) Keyspace() string {
+       if q.qryOpts.getKeyspace != nil {
+               return q.qryOpts.getKeyspace()
+       }
+
+       qrKs := q.routingInfo.getKeyspace()
+       if qrKs != "" {
+               return qrKs
+       }
+       if q.qryOpts.keyspace != "" {
+               return q.qryOpts.keyspace
+       }
+
+       if q.session == nil {
+               return ""
+       }
+       // TODO(chbannis): this should be parsed from the query or we should let
+       // this be set by users.
+       return q.session.cfg.Keyspace
+}
+
+func (q *internalQuery) Table() string {
+       return q.routingInfo.getTable()
+}
+
+func (q *internalQuery) IsIdempotent() bool {
+       return q.qryOpts.idempotent
+}
+
+func (q *internalQuery) getQueryMetrics() *queryMetrics {
+       return q.metrics
+}
+
+func (q *internalQuery) SetConsistency(c Consistency) {
+       atomic.StoreUint32(&q.consistency, uint32(c))
+}
+
+func (q *internalQuery) GetConsistency() Consistency {
+       return Consistency(atomic.LoadUint32(&q.consistency))
+}
+
+func (q *internalQuery) Context() context.Context {
+       return q.qryOpts.context
+}
+
+func (q *internalQuery) Statement() Statement {
+       return q.originalQuery
+}
+
+func (q *internalQuery) GetHostID() string {
+       return q.qryOpts.hostID
+}
+
+func (q *internalQuery) getRoutingInfo() *queryRoutingInfo {
+       return q.routingInfo
+}
+
+func (q *internalQuery) getKeyspaceFunc() func() string {
+       return q.qryOpts.getKeyspace
+}
+
+type batchOptions struct {
+       trace    Tracer
+       observer BatchObserver
+
+       bType   BatchType
+       entries []BatchEntry
+
+       defaultTimestamp      bool
+       defaultTimestampValue int64
+
+       serialCons SerialConsistency
+
+       customPayload map[string][]byte
+       rt            RetryPolicy
+       spec          SpeculativeExecutionPolicy
+       context       context.Context
+       keyspace      string
+       idempotent    bool
+       routingKey    []byte
+       nowInSeconds  *int
+}
+
+func newBatchOptions(b *Batch, ctx context.Context) *batchOptions {
+       // make a new array so if user keeps appending entries on the Batch 
object it doesn't affect this execution
+       newEntries := make([]BatchEntry, len(b.Entries))
+       for i, e := range b.Entries {
+               newEntries[i] = e
+       }
+       var newRoutingKey []byte
+       if b.routingKey != nil {
+               routingKey := b.routingKey
+               newRoutingKey = make([]byte, len(routingKey))
+               copy(newRoutingKey, routingKey)
+       }
+       if ctx == nil {
+               ctx = b.Context()
+       }
+       return &batchOptions{
+               bType:                 b.Type,
+               entries:               newEntries,
+               customPayload:         b.CustomPayload,
+               rt:                    b.rt,
+               spec:                  b.spec,
+               trace:                 b.trace,
+               observer:              b.observer,
+               serialCons:            b.serialCons,
+               defaultTimestamp:      b.defaultTimestamp,
+               defaultTimestampValue: b.defaultTimestampValue,
+               context:               ctx,
+               keyspace:              b.Keyspace(),
+               idempotent:            b.IsIdempotent(),
+               routingKey:            newRoutingKey,
+               nowInSeconds:          b.nowInSeconds,
+       }
+}
+
+type internalBatch struct {
+       originalBatch *Batch
+       batchOpts     *batchOptions
+       metrics       *queryMetrics
+       consistency   uint32
+       routingInfo   *queryRoutingInfo
+       session       *Session
+}
+
+func newInternalBatch(batch *Batch, ctx context.Context) *internalBatch {
+       return &internalBatch{
+               originalBatch: batch,
+               batchOpts:     newBatchOptions(batch, ctx),
+               metrics:       &queryMetrics{m: make(map[string]*hostMetrics)},
+               routingInfo:   &queryRoutingInfo{},
+               session:       batch.session,
+       }
+}

Review Comment:
   nice catch, done



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: pr-unsubscr...@cassandra.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscr...@cassandra.apache.org
For additional commands, e-mail: pr-h...@cassandra.apache.org

Reply via email to