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)
}
}
}
func TestTabletServerStreamExecute(t *testing.T) {
db := setUpTabletServerTest()
testUtils := newTestUtils()
// sql that will be executed in this test
executeSQL := "select * from test_table limit 1000"
executeSQLResult := &sqltypes.Result{
RowsAffected: 1,
Rows: [][]sqltypes.Value{
[]sqltypes.Value{sqltypes.MakeString([]byte("row01"))},
},
}
db.AddQuery(executeSQL, executeSQLResult)
config := testUtils.newQueryServiceConfig()
tsv := NewTabletServer(config)
dbconfigs := testUtils.newDBConfigs(db)
target := querypb.Target{TabletType: topodatapb.TabletType_MASTER}
err := tsv.StartService(target, dbconfigs, []SchemaOverride{}, testUtils.newMysqld(&dbconfigs))
if err != nil {
t.Fatalf("StartService failed: %v", err)
}
defer tsv.StopService()
ctx := context.Background()
session := proto.Session{
SessionId: tsv.sessionID,
TransactionId: 0,
}
txInfo := proto.TransactionInfo{TransactionId: 0}
if err = tsv.Begin(ctx, nil, &session, &txInfo); err != nil {
t.Fatalf("call TabletServer.Begin failed")
}
session.TransactionId = txInfo.TransactionId
query := proto.Query{
Sql: executeSQL,
BindVariables: nil,
SessionId: session.SessionId,
TransactionId: session.TransactionId,
}
sendReply := func(*sqltypes.Result) error { return nil }
if err := tsv.StreamExecute(ctx, nil, &query, sendReply); err == nil {
t.Fatalf("TabletServer.StreamExecute should fail: %s", query.Sql)
}
if err := tsv.Rollback(ctx, nil, &session); err != nil {
t.Fatalf("call TabletServer.Rollback failed")
}
query.TransactionId = 0
if err := tsv.StreamExecute(ctx, nil, &query, sendReply); err != nil {
t.Fatalf("TabletServer.StreamExecute should success: %s, but get error: %v",
query.Sql, err)
}
}
func TestSqlQueryStreamExecute(t *testing.T) {
db := setUpSQLQueryTest()
testUtils := newTestUtils()
// sql that will be executed in this test
executeSQL := "select * from test_table limit 1000"
executeSQLResult := &mproto.QueryResult{
RowsAffected: 1,
Rows: [][]sqltypes.Value{
[]sqltypes.Value{sqltypes.MakeString([]byte("row01"))},
},
}
db.AddQuery(executeSQL, executeSQLResult)
config := testUtils.newQueryServiceConfig()
sqlQuery := NewSqlQuery(config)
dbconfigs := testUtils.newDBConfigs()
err := sqlQuery.allowQueries(nil, &dbconfigs, []SchemaOverride{}, testUtils.newMysqld(&dbconfigs))
if err != nil {
t.Fatalf("allowQueries failed: %v", err)
}
defer sqlQuery.disallowQueries()
ctx := context.Background()
session := proto.Session{
SessionId: sqlQuery.sessionID,
TransactionId: 0,
}
txInfo := proto.TransactionInfo{TransactionId: 0}
if err = sqlQuery.Begin(ctx, nil, &session, &txInfo); err != nil {
t.Fatalf("call SqlQuery.Begin failed")
}
session.TransactionId = txInfo.TransactionId
query := proto.Query{
Sql: executeSQL,
BindVariables: nil,
SessionId: session.SessionId,
TransactionId: session.TransactionId,
}
sendReply := func(*mproto.QueryResult) error { return nil }
if err := sqlQuery.StreamExecute(ctx, nil, &query, sendReply); err == nil {
t.Fatalf("SqlQuery.StreamExecute should fail: %s", query.Sql)
}
if err := sqlQuery.Rollback(ctx, nil, &session); err != nil {
t.Fatalf("call SqlQuery.Rollback failed")
}
query.TransactionId = 0
if err := sqlQuery.StreamExecute(ctx, nil, &query, sendReply); err != nil {
t.Fatalf("SqlQuery.StreamExecute should success: %s, but get error: %v",
query.Sql, err)
}
}
// 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)
}
// 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
}
// 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
}
// 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
}
// 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
}
// 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)
}
// 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
}
// 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
}
// stripTrailing strips out trailing comments if any and puts them in a bind variable.
// This code is a hack. Will need cleaning if it evolves beyond this.
func stripTrailing(query *proto.Query) {
tracker := matchtracker{query.Sql, len(query.Sql)}
pos := tracker.matchComments()
if pos >= 0 {
query.Sql = tracker.query[:pos]
query.BindVariables[TRAILING_COMMENT] = tracker.query[pos:]
}
}
// stripTrailing strips out trailing comments if any and puts them in a bind variable.
// This code is a hack. Will need cleaning if it evolves beyond this.
func stripTrailing(query *proto.Query) {
tracker := matchtracker{
query: query.Sql,
index: len(query.Sql),
}
pos := tracker.matchComments()
if pos >= 0 {
query.Sql = tracker.query[:pos]
query.BindVariables[trailingComment] = tracker.query[pos:]
}
}
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)
}
}
}
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
}
// 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
}
