vitess.io/vitess@v0.16.2/go/mysql/streaming_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 mysql
    18  
    19  import (
    20  	"vitess.io/vitess/go/sqltypes"
    21  
    22  	querypb "vitess.io/vitess/go/vt/proto/query"
    23  )
    24  
    25  // This file contains the methods needed to execute streaming queries.
    26  
    27  // ExecuteStreamFetch starts a streaming query.  Fields(), FetchNext() and
    28  // CloseResult() can be called once this is successful.
    29  // Returns a SQLError.
    30  func (c *Conn) ExecuteStreamFetch(query string) (err error) {
    31  	defer func() {
    32  		if err != nil {
    33  			if sqlerr, ok := err.(*SQLError); ok {
    34  				sqlerr.Query = query
    35  			}
    36  		}
    37  	}()
    38  
    39  	// Sanity check.
    40  	if c.fields != nil {
    41  		return NewSQLError(CRCommandsOutOfSync, SSUnknownSQLState, "streaming query already in progress")
    42  	}
    43  
    44  	// Send the query as a COM_QUERY packet.
    45  	if err := c.WriteComQuery(query); err != nil {
    46  		return err
    47  	}
    48  
    49  	// Get the result.
    50  	colNumber, _, err := c.readComQueryResponse()
    51  	if err != nil {
    52  		return err
    53  	}
    54  	if colNumber == 0 {
    55  		// OK packet, means no results. Save an empty Fields array.
    56  		c.fields = make([]*querypb.Field, 0)
    57  		return nil
    58  	}
    59  
    60  	// Read the fields, save them.
    61  	fields := make([]querypb.Field, colNumber)
    62  	fieldsPointers := make([]*querypb.Field, colNumber)
    63  
    64  	// Read column headers. One packet per column.
    65  	// Build the fields.
    66  	for i := 0; i < colNumber; i++ {
    67  		fieldsPointers[i] = &fields[i]
    68  		if err := c.readColumnDefinition(fieldsPointers[i], i); err != nil {
    69  			return err
    70  		}
    71  	}
    72  
    73  	// Read the EOF after the fields if necessary.
    74  	if c.Capabilities&CapabilityClientDeprecateEOF == 0 {
    75  		// EOF is only present here if it's not deprecated.
    76  		data, err := c.readEphemeralPacket()
    77  		if err != nil {
    78  			return NewSQLError(CRServerLost, SSUnknownSQLState, "%v", err)
    79  		}
    80  		defer c.recycleReadPacket()
    81  		if c.isEOFPacket(data) {
    82  			// This is what we expect.
    83  			// Warnings and status flags are ignored.
    84  			// goto: end
    85  		} else if isErrorPacket(data) {
    86  			return ParseErrorPacket(data)
    87  		} else {
    88  			return NewSQLError(CRCommandsOutOfSync, SSUnknownSQLState, "unexpected packet after fields: %v", data)
    89  		}
    90  	}
    91  
    92  	c.fields = fieldsPointers
    93  	return nil
    94  }
    95  
    96  // Fields returns the fields for an ongoing streaming query.
    97  func (c *Conn) Fields() ([]*querypb.Field, error) {
    98  	if c.fields == nil {
    99  		return nil, NewSQLError(CRCommandsOutOfSync, SSUnknownSQLState, "no streaming query in progress")
   100  	}
   101  	if len(c.fields) == 0 {
   102  		// The query returned an empty field list.
   103  		return nil, nil
   104  	}
   105  	return c.fields, nil
   106  }
   107  
   108  // FetchNext returns the next result for an ongoing streaming query.
   109  // It returns (nil, nil) if there is nothing more to read.
   110  func (c *Conn) FetchNext(in []sqltypes.Value) ([]sqltypes.Value, error) {
   111  	if c.fields == nil {
   112  		// We are already done, and the result was closed.
   113  		return nil, NewSQLError(CRCommandsOutOfSync, SSUnknownSQLState, "no streaming query in progress")
   114  	}
   115  
   116  	if len(c.fields) == 0 {
   117  		// We received no fields, so there is no data.
   118  		return nil, nil
   119  	}
   120  
   121  	data, err := c.ReadPacket()
   122  	if err != nil {
   123  		return nil, err
   124  	}
   125  
   126  	if c.isEOFPacket(data) {
   127  		// Warnings and status flags are ignored.
   128  		c.fields = nil
   129  		return nil, nil
   130  	} else if isErrorPacket(data) {
   131  		// Error packet.
   132  		return nil, ParseErrorPacket(data)
   133  	}
   134  
   135  	// Regular row.
   136  	return c.parseRow(data, c.fields, readLenEncStringAsBytes, in)
   137  }
   138  
   139  // CloseResult can be used to terminate a streaming query
   140  // early. It just drains the remaining values.
   141  func (c *Conn) CloseResult() {
   142  	row := make([]sqltypes.Value, 0, len(c.fields))
   143  	for c.fields != nil {
   144  		rows, err := c.FetchNext(row[:0])
   145  		if err != nil || rows == nil {
   146  			// We either got an error, or got the last result.
   147  			c.fields = nil
   148  		}
   149  	}
   150  }