Ejemplo n.º 1
0
// the first QueryResult will have Fields set (and Rows nil)
// the subsequent QueryResult will have Rows set (and Fields nil)
func (qe *QueryEngine) StreamExecute(logStats *sqlQueryStats, query *proto.Query, sendReply func(*mproto.QueryResult) error) {
	qe.mu.RLock()
	defer qe.mu.RUnlock()

	if query.BindVariables == nil { // will help us avoid repeated nil checks
		query.BindVariables = make(map[string]interface{})
	}
	logStats.BindVariables = query.BindVariables
	// cheap hack: strip trailing comment into a special bind var
	stripTrailing(query)

	plan := qe.schemaInfo.GetStreamPlan(query.Sql)
	logStats.PlanType = "SELECT_STREAM"
	logStats.OriginalSql = plan.DisplayQuery
	defer queryStats.Record("SELECT_STREAM", time.Now())

	// does the real work: first get a connection
	waitingForConnectionStart := time.Now()
	conn := qe.streamConnPool.Get()
	logStats.WaitingForConnection += time.Now().Sub(waitingForConnectionStart)
	defer conn.Recycle()

	// then let's stream!
	qe.fullStreamFetch(logStats, conn, plan.FullQuery, query.BindVariables, nil, nil, sendReply)
}
Ejemplo n.º 2
0
// StreamExecute executes the query and streams the result.
// The first QueryResult will have Fields set (and Rows nil).
// The subsequent QueryResult will have Rows set (and Fields nil).
func (sq *SqlQuery) StreamExecute(ctx context.Context, query *proto.Query, sendReply func(*mproto.QueryResult) error) (err error) {
	// check cases we don't handle yet
	if query.TransactionId != 0 {
		return NewTabletError(ErrFail, "Transactions not supported with streaming")
	}

	logStats := newSqlQueryStats("StreamExecute", ctx)
	defer sq.handleExecError(query, &err, logStats)

	if err = sq.startRequest(query.SessionId, false, false); err != nil {
		return err
	}
	defer sq.endRequest()

	if query.BindVariables == nil {
		query.BindVariables = make(map[string]interface{})
	}
	stripTrailing(query)
	qre := &QueryExecutor{
		query:         query.Sql,
		bindVars:      query.BindVariables,
		transactionID: query.TransactionId,
		plan:          sq.qe.schemaInfo.GetStreamPlan(query.Sql),
		ctx:           ctx,
		logStats:      logStats,
		qe:            sq.qe,
	}
	qre.Stream(sendReply)
	return nil
}
Ejemplo n.º 3
0
// StreamExecute executes the query and streams the result.
// The first QueryResult will have Fields set (and Rows nil).
// The subsequent QueryResult will have Rows set (and Fields nil).
func (tsv *TabletServer) StreamExecute(ctx context.Context, target *pb.Target, query *proto.Query, sendReply func(*mproto.QueryResult) error) (err error) {
	// check cases we don't handle yet
	if query.TransactionId != 0 {
		return NewTabletError(ErrFail, vtrpc.ErrorCode_BAD_INPUT, "Transactions not supported with streaming")
	}

	logStats := newLogStats("StreamExecute", ctx)
	defer tsv.handleExecError(query, &err, logStats)

	if err = tsv.startRequest(target, query.SessionId, false); err != nil {
		return err
	}
	defer tsv.endRequest()

	if query.BindVariables == nil {
		query.BindVariables = make(map[string]interface{})
	}
	stripTrailing(query)
	qre := &QueryExecutor{
		query:         query.Sql,
		bindVars:      query.BindVariables,
		transactionID: query.TransactionId,
		plan:          tsv.qe.schemaInfo.GetStreamPlan(query.Sql),
		ctx:           ctx,
		logStats:      logStats,
		qe:            tsv.qe,
	}
	err = qre.Stream(sendReply)
	if err != nil {
		return tsv.handleExecErrorNoPanic(query, err, logStats)
	}
	return nil
}
Ejemplo n.º 4
0
// Execute executes the query and returns the result as response.
func (sq *SqlQuery) Execute(ctx context.Context, query *proto.Query, reply *mproto.QueryResult) (err error) {
	logStats := newSqlQueryStats("Execute", ctx)
	defer sq.handleExecError(query, &err, logStats)

	allowShutdown := (query.TransactionId != 0)
	if err = sq.startRequest(query.SessionId, false, allowShutdown); err != nil {
		return err
	}
	ctx, cancel := withTimeout(ctx, sq.qe.queryTimeout.Get())
	defer func() {
		cancel()
		sq.endRequest()
	}()

	if query.BindVariables == nil {
		query.BindVariables = make(map[string]interface{})
	}
	stripTrailing(query)
	qre := &QueryExecutor{
		query:         query.Sql,
		bindVars:      query.BindVariables,
		transactionID: query.TransactionId,
		plan:          sq.qe.schemaInfo.GetPlan(ctx, logStats, query.Sql),
		ctx:           ctx,
		logStats:      logStats,
		qe:            sq.qe,
	}
	*reply = *qre.Execute()
	return nil
}
Ejemplo n.º 5
0
func TestComments(t *testing.T) {
	for _, testCase := range testCases {
		query := proto.Query{
			Sql:           testCase.input,
			BindVariables: make(map[string]interface{}),
		}

		stripTrailing(&query)

		want := proto.Query{
			Sql: testCase.outSQL,
		}
		want.BindVariables = make(map[string]interface{})
		if testCase.outVar != "" {
			want.BindVariables[trailingComment] = testCase.outVar
		}
		if !reflect.DeepEqual(query, want) {
			t.Errorf("test input: '%s', got\n%+v, want\n%+v", testCase.input, query, want)
		}
		sql := string(restoreTrailing([]byte(testCase.outSQL), want.BindVariables))
		if !reflect.DeepEqual(testCase.input, sql) {
			t.Fatalf("failed to restore to original sql, got: %s, want: %s", sql, testCase.input)
		}
	}
}
Ejemplo n.º 6
0
// Execute executes the query and returns the result as response.
func (tsv *TabletServer) Execute(ctx context.Context, target *pb.Target, query *proto.Query, reply *mproto.QueryResult) (err error) {
	logStats := newLogStats("Execute", ctx)
	defer tsv.handleExecError(query, &err, logStats)

	allowShutdown := (query.TransactionId != 0)
	if err = tsv.startRequest(target, query.SessionId, allowShutdown); err != nil {
		return err
	}
	ctx, cancel := withTimeout(ctx, tsv.QueryTimeout.Get())
	defer func() {
		cancel()
		tsv.endRequest()
	}()

	if query.BindVariables == nil {
		query.BindVariables = make(map[string]interface{})
	}
	stripTrailing(query)
	qre := &QueryExecutor{
		query:         query.Sql,
		bindVars:      query.BindVariables,
		transactionID: query.TransactionId,
		plan:          tsv.qe.schemaInfo.GetPlan(ctx, logStats, query.Sql),
		ctx:           ctx,
		logStats:      logStats,
		qe:            tsv.qe,
	}
	result, err := qre.Execute()
	if err != nil {
		return tsv.handleExecErrorNoPanic(query, err, logStats)
	}
	*reply = *result
	return nil
}
Ejemplo n.º 7
0
// StreamExecute executes the query and streams its result.
// The first QueryResult will have Fields set (and Rows nil)
// The subsequent QueryResult will have Rows set (and Fields nil)
func (qe *QueryEngine) StreamExecute(logStats *SQLQueryStats, query *proto.Query, sendReply func(*mproto.QueryResult) error) {
	if query.BindVariables == nil { // will help us avoid repeated nil checks
		query.BindVariables = make(map[string]interface{})
	}
	logStats.BindVariables = query.BindVariables
	// cheap hack: strip trailing comment into a special bind var
	stripTrailing(query)

	plan := qe.schemaInfo.GetStreamPlan(query.Sql)
	logStats.PlanType = "SELECT_STREAM"
	logStats.OriginalSql = query.Sql
	defer queryStats.Record("SELECT_STREAM", time.Now())

	authorized := tableacl.Authorized(plan.TableName, plan.PlanId.MinRole())
	qe.checkTableAcl(plan.TableName, plan.PlanId, authorized, logStats.context.GetUsername())

	// does the real work: first get a connection
	waitingForConnectionStart := time.Now()
	conn := getOrPanic(qe.streamConnPool)
	logStats.WaitingForConnection += time.Now().Sub(waitingForConnectionStart)
	defer conn.Recycle()

	qd := NewQueryDetail(query, logStats.context, conn.Id())
	qe.streamQList.Add(qd)
	defer qe.streamQList.Remove(qd)

	// then let's stream! Wrap callback function to return an
	// error on query termination to stop further streaming
	qe.fullStreamFetch(logStats, conn, plan.FullQuery, query.BindVariables, nil, nil, sendReply)
}
Ejemplo n.º 8
0
// Execute executes the query and returns the result as response.
func (sq *SqlQuery) Execute(context context.Context, query *proto.Query, reply *mproto.QueryResult) (err error) {
	logStats := newSqlQueryStats("Execute", context)
	allowShutdown := (query.TransactionId != 0)
	if err = sq.startRequest(query.SessionId, allowShutdown); err != nil {
		return err
	}
	defer sq.endRequest()
	defer handleExecError(query, &err, logStats)

	// TODO(sougou): Change usage such that we don't have to do this.
	if query.BindVariables == nil {
		query.BindVariables = make(map[string]interface{})
	}
	stripTrailing(query)
	qre := &QueryExecutor{
		query:         query.Sql,
		bindVars:      query.BindVariables,
		transactionID: query.TransactionId,
		plan:          sq.qe.schemaInfo.GetPlan(logStats, query.Sql),
		RequestContext: RequestContext{
			ctx:      context,
			logStats: logStats,
			qe:       sq.qe,
		},
	}
	*reply = *qre.Execute()
	return nil
}
Ejemplo n.º 9
0
// StreamExecute executes the query and streams the result.
// The first QueryResult will have Fields set (and Rows nil).
// The subsequent QueryResult will have Rows set (and Fields nil).
func (sq *SqlQuery) StreamExecute(context context.Context, query *proto.Query, sendReply func(*mproto.QueryResult) error) (err error) {
	// check cases we don't handle yet
	if query.TransactionId != 0 {
		return NewTabletError(FAIL, "Transactions not supported with streaming")
	}

	logStats := newSqlQueryStats("StreamExecute", context)
	if err = sq.startRequest(query.SessionId, false); err != nil {
		return err
	}
	defer sq.endRequest()
	defer handleExecError(query, &err, logStats)

	// TODO(sougou): Change usage such that we don't have to do this.
	if query.BindVariables == nil {
		query.BindVariables = make(map[string]interface{})
	}
	stripTrailing(query)
	qre := &QueryExecutor{
		query:         query.Sql,
		bindVars:      query.BindVariables,
		transactionID: query.TransactionId,
		plan:          sq.qe.schemaInfo.GetStreamPlan(query.Sql),
		RequestContext: RequestContext{
			ctx:      context,
			logStats: logStats,
			qe:       sq.qe,
			deadline: NewDeadline(sq.qe.queryTimeout.Get()),
		},
	}
	qre.Stream(sendReply)
	return nil
}
Ejemplo n.º 10
0
func TestComments(t *testing.T) {
	for _, testCase := range testCases {
		query := proto.Query{
			Sql:           testCase.input,
			BindVariables: make(map[string]interface{}),
		}

		stripTrailing(&query)

		want := proto.Query{
			Sql: testCase.outSql,
		}
		want.BindVariables = make(map[string]interface{})
		if testCase.outVar != "" {
			want.BindVariables[TRAILING_COMMENT] = testCase.outVar
		}
		if !reflect.DeepEqual(query, want) {
			t.Errorf("test input: '%s', got\n%+v, want\n%+v", testCase.input, query, want)
		}
	}
}
Ejemplo n.º 11
0
func (qe *QueryEngine) Execute(logStats *sqlQueryStats, query *proto.Query) (reply *mproto.QueryResult) {
	qe.mu.RLock()
	defer qe.mu.RUnlock()

	if query.BindVariables == nil { // will help us avoid repeated nil checks
		query.BindVariables = make(map[string]interface{})
	}
	logStats.BindVariables = query.BindVariables
	// cheap hack: strip trailing comment into a special bind var
	stripTrailing(query)
	basePlan := qe.schemaInfo.GetPlan(logStats, query.Sql)
	planName := basePlan.PlanId.String()
	logStats.PlanType = planName
	logStats.OriginalSql = basePlan.DisplayQuery
	defer func(start time.Time) {
		duration := time.Now().Sub(start)
		queryStats.Add(planName, duration)
		if reply == nil {
			basePlan.AddStats(1, duration, 0, 1)
		} else {
			basePlan.AddStats(1, duration, int64(len(reply.Rows)), 0)
		}
	}(time.Now())

	// Run it by the rules engine
	action, desc := basePlan.Rules.getAction(logStats.RemoteAddr(), logStats.Username(), query.BindVariables)
	if action == QR_FAIL_QUERY {
		panic(NewTabletError(FAIL, "Query disallowed due to rule: %s", desc))
	}

	if basePlan.PlanId == sqlparser.PLAN_DDL {
		return qe.execDDL(logStats, query.Sql)
	}

	plan := &CompiledPlan{
		Query:         query.Sql,
		ExecPlan:      basePlan,
		BindVars:      query.BindVariables,
		TransactionId: query.TransactionId,
	}
	if query.TransactionId != 0 {
		// Need upfront connection for DMLs and transactions
		conn := qe.activeTxPool.Get(query.TransactionId)
		defer conn.Recycle()
		conn.RecordQuery(plan.DisplayQuery)
		var invalidator CacheInvalidator
		if plan.TableInfo != nil && plan.TableInfo.CacheType != schema.CACHE_NONE {
			invalidator = conn.DirtyKeys(plan.TableName)
		}
		switch plan.PlanId {
		case sqlparser.PLAN_PASS_DML:
			// TODO(sougou): Delete code path that leads here.
			// We need to permanently disallow this plan.
			panic(NewTabletError(FAIL, "DML too complex"))
		case sqlparser.PLAN_INSERT_PK:
			reply = qe.execInsertPK(logStats, conn, plan, invalidator)
		case sqlparser.PLAN_INSERT_SUBQUERY:
			reply = qe.execInsertSubquery(logStats, conn, plan, invalidator)
		case sqlparser.PLAN_DML_PK:
			reply = qe.execDMLPK(logStats, conn, plan, invalidator)
		case sqlparser.PLAN_DML_SUBQUERY:
			reply = qe.execDMLSubquery(logStats, conn, plan, invalidator)
		default: // select or set in a transaction, just count as select
			reply = qe.execDirect(logStats, plan, conn)
		}
	} else {
		switch plan.PlanId {
		case sqlparser.PLAN_PASS_SELECT:
			if plan.Reason == sqlparser.REASON_FOR_UPDATE {
				panic(NewTabletError(FAIL, "Disallowed outside transaction"))
			}
			reply = qe.execSelect(logStats, plan)
		case sqlparser.PLAN_PK_EQUAL:
			reply = qe.execPKEqual(logStats, plan)
		case sqlparser.PLAN_PK_IN:
			reply = qe.execPKIN(logStats, plan)
		case sqlparser.PLAN_SELECT_SUBQUERY:
			reply = qe.execSubquery(logStats, plan)
		case sqlparser.PLAN_SET:
			waitingForConnectionStart := time.Now()
			conn := qe.connPool.Get()
			logStats.WaitingForConnection += time.Now().Sub(waitingForConnectionStart)
			defer conn.Recycle()
			reply = qe.execSet(logStats, conn, plan)
		default:
			panic(NewTabletError(NOT_IN_TX, "DMLs not allowed outside of transactions"))
		}
	}
	if plan.PlanId.IsSelect() {
		logStats.RowsAffected = int(reply.RowsAffected)
		resultStats.Add(int64(reply.RowsAffected))
		logStats.Rows = reply.Rows
	}

	return reply
}
Ejemplo n.º 12
0
// Execute executes the specified query and returns its result.
func (qe *QueryEngine) Execute(logStats *SQLQueryStats, query *proto.Query) (reply *mproto.QueryResult) {
	if query.BindVariables == nil { // will help us avoid repeated nil checks
		query.BindVariables = make(map[string]interface{})
	}
	logStats.BindVariables = query.BindVariables
	// cheap hack: strip trailing comment into a special bind var
	stripTrailing(query)
	basePlan := qe.schemaInfo.GetPlan(logStats, query.Sql)
	planName := basePlan.PlanId.String()
	logStats.PlanType = planName
	logStats.OriginalSql = query.Sql
	defer func(start time.Time) {
		duration := time.Now().Sub(start)
		queryStats.Add(planName, duration)
		if reply == nil {
			basePlan.AddStats(1, duration, 0, 1)
		} else {
			basePlan.AddStats(1, duration, int64(len(reply.Rows)), 0)
		}
	}(time.Now())

	// Run it by the rules engine
	action, desc := basePlan.Rules.getAction(logStats.RemoteAddr(), logStats.Username(), query.BindVariables)
	switch action {
	case QR_FAIL:
		panic(NewTabletError(FAIL, "Query disallowed due to rule: %s", desc))
	case QR_FAIL_RETRY:
		panic(NewTabletError(RETRY, "Query disallowed due to rule: %s", desc))
	}

	qe.checkTableAcl(basePlan.TableName, basePlan.PlanId, basePlan.Authorized, logStats.context.GetUsername())

	if basePlan.PlanId == planbuilder.PLAN_DDL {
		return qe.execDDL(logStats, query.Sql)
	}

	plan := &compiledPlan{
		Query:         query.Sql,
		ExecPlan:      basePlan,
		BindVars:      query.BindVariables,
		TransactionID: query.TransactionId,
	}
	if query.TransactionId != 0 {
		// Need upfront connection for DMLs and transactions
		conn := qe.activeTxPool.Get(query.TransactionId)
		defer conn.Recycle()
		conn.RecordQuery(plan.Query)
		switch plan.PlanId {
		case planbuilder.PLAN_PASS_DML:
			if qe.strictMode.Get() != 0 {
				panic(NewTabletError(FAIL, "DML too complex"))
			}
			reply = qe.directFetch(logStats, conn, plan.FullQuery, plan.BindVars, nil, nil)
		case planbuilder.PLAN_INSERT_PK:
			reply = qe.execInsertPK(logStats, conn, plan)
		case planbuilder.PLAN_INSERT_SUBQUERY:
			reply = qe.execInsertSubquery(logStats, conn, plan)
		case planbuilder.PLAN_DML_PK:
			reply = qe.execDMLPK(logStats, conn, plan)
		case planbuilder.PLAN_DML_SUBQUERY:
			reply = qe.execDMLSubquery(logStats, conn, plan)
		case planbuilder.PLAN_OTHER:
			reply = qe.executeSqlString(logStats, conn, query.Sql, true)
		default: // select or set in a transaction, just count as select
			reply = qe.execDirect(logStats, plan, conn)
		}
	} else {
		switch plan.PlanId {
		case planbuilder.PLAN_PASS_SELECT:
			if plan.Reason == planbuilder.REASON_LOCK {
				panic(NewTabletError(FAIL, "Disallowed outside transaction"))
			}
			reply = qe.execSelect(logStats, plan)
		case planbuilder.PLAN_PK_EQUAL:
			reply = qe.execPKEqual(logStats, plan)
		case planbuilder.PLAN_PK_IN:
			reply = qe.execPKIN(logStats, plan)
		case planbuilder.PLAN_SELECT_SUBQUERY:
			reply = qe.execSubquery(logStats, plan)
		case planbuilder.PLAN_SET:
			reply = qe.execSet(logStats, plan)
		case planbuilder.PLAN_OTHER:
			waitingForConnectionStart := time.Now()
			conn := getOrPanic(qe.connPool)
			logStats.WaitingForConnection += time.Now().Sub(waitingForConnectionStart)
			defer conn.Recycle()
			reply = qe.executeSqlString(logStats, conn, query.Sql, true)
		default:
			panic(NewTabletError(NOT_IN_TX, "DMLs not allowed outside of transactions"))
		}
	}
	if plan.PlanId.IsSelect() {
		logStats.RowsAffected = int(reply.RowsAffected)
		resultStats.Add(int64(reply.RowsAffected))
		logStats.Rows = reply.Rows
	}

	return reply
}