vitess.io/vitess@v0.16.2/go/vt/vttablet/tabletmanager/rpc_query.go (about)

     1  /*
     2  Copyright 2019 The Vitess Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package tabletmanager
    18  
    19  import (
    20  	"context"
    21  
    22  	"vitess.io/vitess/go/sqlescape"
    23  	"vitess.io/vitess/go/sqltypes"
    24  	"vitess.io/vitess/go/vt/log"
    25  	"vitess.io/vitess/go/vt/sqlparser"
    26  
    27  	querypb "vitess.io/vitess/go/vt/proto/query"
    28  	tabletmanagerdatapb "vitess.io/vitess/go/vt/proto/tabletmanagerdata"
    29  )
    30  
    31  // ExecuteFetchAsDba will execute the given query, possibly disabling binlogs and reload schema.
    32  func (tm *TabletManager) ExecuteFetchAsDba(ctx context.Context, req *tabletmanagerdatapb.ExecuteFetchAsDbaRequest) (*querypb.QueryResult, error) {
    33  	// get a connection
    34  	conn, err := tm.MysqlDaemon.GetDbaConnection(ctx)
    35  	if err != nil {
    36  		return nil, err
    37  	}
    38  	defer conn.Close()
    39  
    40  	// disable binlogs if necessary
    41  	if req.DisableBinlogs {
    42  		_, err := conn.ExecuteFetch("SET sql_log_bin = OFF", 0, false)
    43  		if err != nil {
    44  			return nil, err
    45  		}
    46  	}
    47  
    48  	if req.DbName != "" {
    49  		// This execute might fail if db does not exist.
    50  		// Error is ignored because given query might create this database.
    51  		_, _ = conn.ExecuteFetch("USE "+sqlescape.EscapeID(req.DbName), 1, false)
    52  	}
    53  
    54  	// Handle special possible directives
    55  	var directives *sqlparser.CommentDirectives
    56  	if stmt, err := sqlparser.Parse(string(req.Query)); err == nil {
    57  		if cmnt, ok := stmt.(sqlparser.Commented); ok {
    58  			directives = cmnt.GetParsedComments().Directives()
    59  		}
    60  	}
    61  	if directives.IsSet("allowZeroInDate") {
    62  		if _, err := conn.ExecuteFetch("set @@session.sql_mode=REPLACE(REPLACE(@@session.sql_mode, 'NO_ZERO_DATE', ''), 'NO_ZERO_IN_DATE', '')", 1, false); err != nil {
    63  			return nil, err
    64  		}
    65  	}
    66  	// run the query
    67  	result, err := conn.ExecuteFetch(string(req.Query), int(req.MaxRows), true /*wantFields*/)
    68  
    69  	// re-enable binlogs if necessary
    70  	if req.DisableBinlogs && !conn.IsClosed() {
    71  		_, err := conn.ExecuteFetch("SET sql_log_bin = ON", 0, false)
    72  		if err != nil {
    73  			// if we can't reset the sql_log_bin flag,
    74  			// let's just close the connection.
    75  			conn.Close()
    76  		}
    77  	}
    78  
    79  	if err == nil && req.ReloadSchema {
    80  		reloadErr := tm.QueryServiceControl.ReloadSchema(ctx)
    81  		if reloadErr != nil {
    82  			log.Errorf("failed to reload the schema %v", reloadErr)
    83  		}
    84  	}
    85  	return sqltypes.ResultToProto3(result), err
    86  }
    87  
    88  // ExecuteFetchAsAllPrivs will execute the given query, possibly reloading schema.
    89  func (tm *TabletManager) ExecuteFetchAsAllPrivs(ctx context.Context, req *tabletmanagerdatapb.ExecuteFetchAsAllPrivsRequest) (*querypb.QueryResult, error) {
    90  	// get a connection
    91  	conn, err := tm.MysqlDaemon.GetAllPrivsConnection(ctx)
    92  	if err != nil {
    93  		return nil, err
    94  	}
    95  	defer conn.Close()
    96  
    97  	if req.DbName != "" {
    98  		// This execute might fail if db does not exist.
    99  		// Error is ignored because given query might create this database.
   100  		_, _ = conn.ExecuteFetch("USE "+sqlescape.EscapeID(req.DbName), 1, false)
   101  	}
   102  
   103  	// run the query
   104  	result, err := conn.ExecuteFetch(string(req.Query), int(req.MaxRows), true /*wantFields*/)
   105  
   106  	if err == nil && req.ReloadSchema {
   107  		reloadErr := tm.QueryServiceControl.ReloadSchema(ctx)
   108  		if reloadErr != nil {
   109  			log.Errorf("failed to reload the schema %v", reloadErr)
   110  		}
   111  	}
   112  	return sqltypes.ResultToProto3(result), err
   113  }
   114  
   115  // ExecuteFetchAsApp will execute the given query.
   116  func (tm *TabletManager) ExecuteFetchAsApp(ctx context.Context, req *tabletmanagerdatapb.ExecuteFetchAsAppRequest) (*querypb.QueryResult, error) {
   117  	// get a connection
   118  	conn, err := tm.MysqlDaemon.GetAppConnection(ctx)
   119  	if err != nil {
   120  		return nil, err
   121  	}
   122  	defer conn.Recycle()
   123  	result, err := conn.ExecuteFetch(string(req.Query), int(req.MaxRows), true /*wantFields*/)
   124  	return sqltypes.ResultToProto3(result), err
   125  }
   126  
   127  // ExecuteQuery submits a new online DDL request
   128  func (tm *TabletManager) ExecuteQuery(ctx context.Context, req *tabletmanagerdatapb.ExecuteQueryRequest) (*querypb.QueryResult, error) {
   129  	// get the db name from the tablet
   130  	tablet := tm.Tablet()
   131  	target := &querypb.Target{Keyspace: tablet.Keyspace, Shard: tablet.Shard, TabletType: tablet.Type}
   132  	result, err := tm.QueryServiceControl.QueryService().Execute(ctx, target, string(req.Query), nil, 0, 0, nil)
   133  	return sqltypes.ResultToProto3(result), err
   134  }