vitess.io/vitess@v0.16.2/go/vt/vttablet/vexec/vexec_test.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 vexec
    18  
    19  import (
    20  	"context"
    21  	"testing"
    22  
    23  	"vitess.io/vitess/go/vt/sqlparser"
    24  
    25  	"github.com/stretchr/testify/assert"
    26  )
    27  
    28  const (
    29  	tWorkflow = "myworkflow"
    30  	tKeyspace = "mykeyspace"
    31  )
    32  
    33  func TestNewTabletVExec(t *testing.T) {
    34  	vx := NewTabletVExec(tWorkflow, tKeyspace)
    35  	assert.NotNil(t, vx)
    36  	assert.Equal(t, vx.Workflow, tWorkflow)
    37  	assert.Equal(t, vx.Keyspace, tKeyspace)
    38  }
    39  
    40  func TestAnalyzeQuerySelect1(t *testing.T) {
    41  	query := `select migration_status, strategy from _vt.schema_migrations where migration_uuid='123'`
    42  	vx := NewTabletVExec(tWorkflow, tKeyspace)
    43  	err := vx.AnalyzeQuery(context.Background(), query)
    44  	assert.NoError(t, err)
    45  
    46  	assert.Equal(t, vx.Query, query)
    47  	assert.Equal(t, vx.TableName, "_vt.schema_migrations")
    48  
    49  	_, ok := vx.WhereCols["migration_uuid"]
    50  	assert.True(t, ok)
    51  	_, ok = vx.WhereCols["strategy"]
    52  	assert.False(t, ok)
    53  
    54  	_, ok = vx.UpdateCols["strategy"]
    55  	assert.False(t, ok)
    56  }
    57  func TestAnalyzeQuerySelect2(t *testing.T) {
    58  	query := `select migration_status, strategy from _vt.schema_migrations where migration_uuid='123' or requested_timestamp<now()`
    59  	vx := NewTabletVExec(tWorkflow, tKeyspace)
    60  	err := vx.AnalyzeQuery(context.Background(), query)
    61  	assert.NoError(t, err)
    62  
    63  	assert.Equal(t, vx.Query, query)
    64  	assert.Equal(t, vx.TableName, "_vt.schema_migrations")
    65  
    66  	_, ok := vx.WhereCols["migration_uuid"]
    67  	assert.False(t, ok)
    68  	_, ok = vx.WhereCols["requested_timestamp"]
    69  	assert.False(t, ok)
    70  	_, ok = vx.WhereCols["strategy"]
    71  	assert.False(t, ok)
    72  
    73  	_, ok = vx.UpdateCols["strategy"]
    74  	assert.False(t, ok)
    75  }
    76  
    77  func TestAnalyzeQueryUpdate1(t *testing.T) {
    78  	query := `update _vt.schema_migrations set migration_status='running', liveness_timestamp=now() where migration_uuid='123' and requested_timestamp<now() and strategy='pt-osc'`
    79  	vx := NewTabletVExec(tWorkflow, tKeyspace)
    80  	err := vx.AnalyzeQuery(context.Background(), query)
    81  	assert.NoError(t, err)
    82  
    83  	assert.Equal(t, vx.Query, query)
    84  	assert.Equal(t, vx.TableName, "_vt.schema_migrations")
    85  
    86  	_, ok := vx.WhereCols["migration_uuid"]
    87  	assert.True(t, ok)
    88  	_, ok = vx.WhereCols["requested_timestamp"]
    89  	assert.False(t, ok)
    90  	_, ok = vx.WhereCols["strategy"]
    91  	assert.True(t, ok)
    92  
    93  	_, ok = vx.UpdateCols["migration_status"]
    94  	assert.True(t, ok)
    95  	_, ok = vx.UpdateCols["strategy"]
    96  	assert.False(t, ok)
    97  	_, ok = vx.UpdateCols["liveness_timestamp"]
    98  	assert.False(t, ok)
    99  }
   100  
   101  func TestAnalyzeQueryInsert1(t *testing.T) {
   102  	query := `insert into _vt.schema_migrations
   103  		(migration_uuid, migration_status, count, liveness_timestamp) values
   104  		('abc123', 'running', 5, now())
   105  		`
   106  	vx := NewTabletVExec(tWorkflow, tKeyspace)
   107  	err := vx.AnalyzeQuery(context.Background(), query)
   108  	assert.NoError(t, err)
   109  
   110  	assert.Equal(t, vx.Query, query)
   111  	assert.Equal(t, vx.TableName, "_vt.schema_migrations")
   112  
   113  	_, ok := vx.InsertCols["migration_uuid"]
   114  	assert.True(t, ok)
   115  	_, ok = vx.InsertCols["count"]
   116  	assert.True(t, ok)
   117  	_, ok = vx.InsertCols["requested_timestamp"]
   118  	assert.False(t, ok) // column does not exist
   119  	_, ok = vx.InsertCols["liveness_timestamp"]
   120  	assert.False(t, ok) // because it's not a literal
   121  
   122  	var val string
   123  	val, err = vx.ColumnStringVal(vx.InsertCols, "migration_uuid")
   124  	assert.NoError(t, err)
   125  	assert.Equal(t, "abc123", val)
   126  
   127  	val, err = vx.ColumnStringVal(vx.InsertCols, "count")
   128  	assert.NoError(t, err)
   129  	assert.Equal(t, "5", val)
   130  
   131  	_, err = vx.ColumnStringVal(vx.InsertCols, "liveness_timestamp")
   132  	assert.Error(t, err)
   133  
   134  	vx.SetColumnStringVal(vx.InsertCols, "migration_uuid", "other")
   135  	val, err = vx.ColumnStringVal(vx.InsertCols, "migration_uuid")
   136  	assert.NoError(t, err)
   137  	assert.Equal(t, "other", val)
   138  
   139  	vx.SetColumnStringVal(vx.InsertCols, "liveness_timestamp", "another")
   140  	val, err = vx.ColumnStringVal(vx.InsertCols, "liveness_timestamp")
   141  	assert.NoError(t, err)
   142  	assert.Equal(t, "another", val)
   143  }
   144  
   145  func TestAnalyzeQueryInsert2(t *testing.T) {
   146  	query := `insert into _vt.schema_migrations
   147  		(migration_uuid, migration_status, count, liveness_timestamp) values
   148  		('abc123', 'running', 5, now())
   149  		`
   150  	vx := NewTabletVExec(tWorkflow, tKeyspace)
   151  	err := vx.AnalyzeQuery(context.Background(), query)
   152  	assert.NoError(t, err)
   153  
   154  	assert.Equal(t, vx.Query, query)
   155  
   156  	newVal := vx.ToStringVal("newval")
   157  	err = vx.ReplaceInsertColumnVal("no_such_column", newVal)
   158  	assert.Error(t, err)
   159  	assert.Equal(t, ErrColumNotFound, err)
   160  	assert.Equal(t, vx.Query, query)
   161  
   162  	err = vx.ReplaceInsertColumnVal("migration_uuid", newVal)
   163  	assert.NoError(t, err)
   164  	assert.NotEqual(t, vx.Query, query)
   165  
   166  	err = vx.ReplaceInsertColumnVal("liveness_timestamp", newVal)
   167  	assert.NoError(t, err)
   168  	assert.NotEqual(t, vx.Query, query)
   169  
   170  	assert.Equal(t, vx.TableName, "_vt.schema_migrations")
   171  
   172  	testVals := func() {
   173  		_, ok := vx.InsertCols["migration_uuid"]
   174  		assert.True(t, ok)
   175  		_, ok = vx.InsertCols["count"]
   176  		assert.True(t, ok)
   177  		_, ok = vx.InsertCols["requested_timestamp"]
   178  		assert.False(t, ok) // column does not exist
   179  		_, ok = vx.InsertCols["liveness_timestamp"]
   180  		assert.True(t, ok) // because it's not a literal
   181  
   182  		var val string
   183  		val, err = vx.ColumnStringVal(vx.InsertCols, "migration_uuid")
   184  		assert.NoError(t, err)
   185  		assert.Equal(t, "newval", val)
   186  
   187  		val, err = vx.ColumnStringVal(vx.InsertCols, "count")
   188  		assert.NoError(t, err)
   189  		assert.Equal(t, "5", val)
   190  
   191  		val, err = vx.ColumnStringVal(vx.InsertCols, "liveness_timestamp")
   192  		assert.NoError(t, err)
   193  		assert.Equal(t, "newval", val)
   194  	}
   195  	testVals()
   196  	rewrittenQuery := sqlparser.String(vx.Stmt)
   197  
   198  	vx = NewTabletVExec(tWorkflow, tKeyspace)
   199  	err = vx.AnalyzeQuery(context.Background(), rewrittenQuery)
   200  	assert.NoError(t, err)
   201  	assert.Equal(t, vx.Query, rewrittenQuery)
   202  	testVals()
   203  }
   204  
   205  func TestAnalyzeQueryInsert3(t *testing.T) {
   206  	query := `insert into _vt.schema_migrations
   207  		(migration_uuid, migration_status, count, liveness_timestamp) values
   208  		('abc123', 'running', 5, now())
   209  		`
   210  	vx := NewTabletVExec(tWorkflow, tKeyspace)
   211  	err := vx.AnalyzeQuery(context.Background(), query)
   212  	assert.NoError(t, err)
   213  
   214  	assert.Equal(t, vx.Query, query)
   215  
   216  	newVal := vx.ToStringVal("newval")
   217  	err = vx.AddOrReplaceInsertColumnVal("a_new_column", newVal)
   218  	assert.NoError(t, err)
   219  	assert.NotEqual(t, vx.Query, query)
   220  	assert.Contains(t, vx.Query, "a_new_column")
   221  	assert.Contains(t, vx.Query, "newval")
   222  
   223  	err = vx.AddOrReplaceInsertColumnVal("migration_uuid", newVal)
   224  	assert.NoError(t, err)
   225  	assert.NotEqual(t, vx.Query, query)
   226  
   227  	err = vx.AddOrReplaceInsertColumnVal("liveness_timestamp", newVal)
   228  	assert.NoError(t, err)
   229  	assert.NotEqual(t, vx.Query, query)
   230  
   231  	assert.Equal(t, vx.TableName, "_vt.schema_migrations")
   232  
   233  	testVals := func() {
   234  		_, ok := vx.InsertCols["migration_uuid"]
   235  		assert.True(t, ok)
   236  		_, ok = vx.InsertCols["count"]
   237  		assert.True(t, ok)
   238  		_, ok = vx.InsertCols["requested_timestamp"]
   239  		assert.False(t, ok) // column does not exist
   240  		_, ok = vx.InsertCols["liveness_timestamp"]
   241  		assert.True(t, ok) // because it's not a literal
   242  		_, ok = vx.InsertCols["a_new_column"]
   243  		assert.True(t, ok)
   244  
   245  		var val string
   246  		val, err = vx.ColumnStringVal(vx.InsertCols, "migration_uuid")
   247  		assert.NoError(t, err)
   248  		assert.Equal(t, "newval", val)
   249  
   250  		val, err = vx.ColumnStringVal(vx.InsertCols, "count")
   251  		assert.NoError(t, err)
   252  		assert.Equal(t, "5", val)
   253  
   254  		val, err = vx.ColumnStringVal(vx.InsertCols, "liveness_timestamp")
   255  		assert.NoError(t, err)
   256  		assert.Equal(t, "newval", val)
   257  
   258  		val, err = vx.ColumnStringVal(vx.InsertCols, "a_new_column")
   259  		assert.NoError(t, err)
   260  		assert.Equal(t, "newval", val)
   261  	}
   262  	testVals()
   263  	rewrittenQuery := sqlparser.String(vx.Stmt)
   264  
   265  	vx = NewTabletVExec(tWorkflow, tKeyspace)
   266  	err = vx.AnalyzeQuery(context.Background(), rewrittenQuery)
   267  	assert.NoError(t, err)
   268  	assert.Equal(t, vx.Query, rewrittenQuery)
   269  	testVals()
   270  }