vitess.io/vitess@v0.16.2/go/test/endtoend/vtgate/queries/reference/main_test.go (about)

     1  /*
     2  Copyright 2022 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 reference
    18  
    19  import (
    20  	"context"
    21  	"flag"
    22  	"fmt"
    23  	"os"
    24  	"testing"
    25  	"time"
    26  
    27  	"vitess.io/vitess/go/mysql"
    28  
    29  	querypb "vitess.io/vitess/go/vt/proto/query"
    30  	"vitess.io/vitess/go/vt/vtgate/vtgateconn"
    31  
    32  	"vitess.io/vitess/go/test/endtoend/cluster"
    33  )
    34  
    35  var (
    36  	clusterInstance *cluster.LocalProcessCluster
    37  	cell            = "zone1"
    38  	hostname        = "localhost"
    39  	vtParams        mysql.ConnParams
    40  
    41  	unshardedKeyspaceName = "uks"
    42  	unshardedSQLSchema    = `
    43  		CREATE TABLE IF NOT EXISTS zip(
    44  			id BIGINT NOT NULL AUTO_INCREMENT,
    45  			code5 INT(5) NOT NULL,
    46  			PRIMARY KEY(id)
    47  		) ENGINE=InnoDB;
    48  
    49  		INSERT INTO zip(id, code5)
    50  		VALUES (1, 47107),
    51  			   (2, 82845),
    52  			   (3, 11237);
    53  
    54  		CREATE TABLE IF NOT EXISTS zip_detail(
    55  			id BIGINT NOT NULL AUTO_INCREMENT,
    56  			zip_id BIGINT NOT NULL,
    57  			discontinued_at DATE,
    58  			PRIMARY KEY(id)
    59  		) ENGINE=InnoDB;
    60  
    61  	`
    62  	unshardedVSchema = `
    63  		{
    64  			"sharded":false,
    65  			"tables": {
    66  				"zip": {},
    67  				"zip_detail": {}
    68  			}
    69  		}
    70  	`
    71  	shardedKeyspaceName = "sks"
    72  	shardedSQLSchema    = `
    73  		CREATE TABLE IF NOT EXISTS delivery_failure (
    74  			id BIGINT NOT NULL,
    75  			zip_detail_id BIGINT NOT NULL,
    76  			reason VARCHAR(255),
    77  			PRIMARY KEY(id)
    78  		) ENGINE=InnoDB;
    79  	`
    80  	shardedVSchema = `
    81  		{
    82  			"sharded": true,
    83  			"vindexes": {
    84  				"hash": {
    85  					"type": "hash"
    86  				}
    87  			},
    88  			"tables": {
    89  				"delivery_failure": {
    90  					"columnVindexes": [
    91  						{
    92  							"column": "id",
    93  							"name": "hash"
    94  						}
    95  					]
    96  				},
    97  				"zip_detail": {
    98  					"type": "reference",
    99  					"source": "` + unshardedKeyspaceName + `.zip_detail"
   100  				}
   101  			}
   102  		}
   103  	`
   104  )
   105  
   106  func TestMain(m *testing.M) {
   107  	defer cluster.PanicHandler(nil)
   108  	flag.Parse()
   109  
   110  	exitCode := func() int {
   111  		clusterInstance = cluster.NewCluster(cell, hostname)
   112  		defer clusterInstance.Teardown()
   113  
   114  		// Start topo server
   115  		if err := clusterInstance.StartTopo(); err != nil {
   116  			return 1
   117  		}
   118  
   119  		// Start keyspace
   120  		uKeyspace := &cluster.Keyspace{
   121  			Name:      unshardedKeyspaceName,
   122  			SchemaSQL: unshardedSQLSchema,
   123  			VSchema:   unshardedVSchema,
   124  		}
   125  		if err := clusterInstance.StartUnshardedKeyspace(*uKeyspace, 0, false); err != nil {
   126  			return 1
   127  		}
   128  
   129  		sKeyspace := &cluster.Keyspace{
   130  			Name:      shardedKeyspaceName,
   131  			SchemaSQL: shardedSQLSchema,
   132  			VSchema:   shardedVSchema,
   133  		}
   134  		if err := clusterInstance.StartKeyspace(*sKeyspace, []string{"-80", "80-"}, 0, false); err != nil {
   135  			return 1
   136  		}
   137  
   138  		// Start vtgate
   139  		if err := clusterInstance.StartVtgate(); err != nil {
   140  			return 1
   141  		}
   142  
   143  		if err := clusterInstance.WaitForTabletsToHealthyInVtgate(); err != nil {
   144  			return 1
   145  		}
   146  
   147  		vtParams = mysql.ConnParams{
   148  			Host: "localhost",
   149  			Port: clusterInstance.VtgateMySQLPort,
   150  		}
   151  
   152  		// TODO(maxeng) remove when we have a proper way to check
   153  		// materialization lag and cutover.
   154  		done := make(chan bool, 1)
   155  		expectRows := 2
   156  		go func() {
   157  			ctx := context.Background()
   158  			vtgateAddr := fmt.Sprintf("%s:%d", clusterInstance.Hostname, clusterInstance.VtgateProcess.GrpcPort)
   159  			vtgateConn, err := vtgateconn.Dial(ctx, vtgateAddr)
   160  			if err != nil {
   161  				done <- false
   162  				return
   163  			}
   164  			defer vtgateConn.Close()
   165  
   166  			maxWait := time.After(300 * time.Second)
   167  			for _, ks := range clusterInstance.Keyspaces {
   168  				if ks.Name != shardedKeyspaceName {
   169  					continue
   170  				}
   171  				for _, s := range ks.Shards {
   172  					var ok bool
   173  					for !ok {
   174  						select {
   175  						case <-maxWait:
   176  							fmt.Println("Waited too long for materialization, cancelling.")
   177  							done <- false
   178  							return
   179  						default:
   180  						}
   181  						shard := fmt.Sprintf("%s/%s@primary", ks.Name, s.Name)
   182  						session := vtgateConn.Session(shard, nil)
   183  						_, err := session.Execute(ctx, "SHOW CREATE TABLE zip_detail", map[string]*querypb.BindVariable{})
   184  						if err != nil {
   185  							fmt.Fprintf(os.Stderr, "Failed to SHOW CREATE TABLE zip_detail; might not exist yet: %v\n", err)
   186  							time.Sleep(1 * time.Second)
   187  							continue
   188  						}
   189  						qr, err := session.Execute(ctx, "SELECT * FROM zip_detail", map[string]*querypb.BindVariable{})
   190  						if err != nil {
   191  							fmt.Fprintf(os.Stderr, "Failed to query sharded keyspace for zip_detail rows: %v\n", err)
   192  							done <- false
   193  							return
   194  						}
   195  						if len(qr.Rows) != expectRows {
   196  							fmt.Fprintf(os.Stderr, "Shard %s doesn't yet have expected number of zip_detail rows\n", shard)
   197  							time.Sleep(10 * time.Second)
   198  							continue
   199  						}
   200  						fmt.Fprintf(os.Stdout, "Shard %s has expected number of zip_detail rows.\n", shard)
   201  						ok = true
   202  					}
   203  				}
   204  				fmt.Println("All shards have expected number of zip_detail rows.")
   205  				done <- true
   206  			}
   207  		}()
   208  
   209  		// Materialize zip_detail to sharded keyspace.
   210  		output, err := clusterInstance.VtctlProcess.ExecuteCommandWithOutput(
   211  			"Materialize",
   212  			"--",
   213  			"--tablet_types",
   214  			"PRIMARY",
   215  			`{
   216  				"workflow": "copy_zip_detail",
   217  				"source_keyspace": "`+unshardedKeyspaceName+`",
   218  				"target_keyspace": "`+shardedKeyspaceName+`",
   219  				"tablet_types": "PRIMARY",
   220  				"table_settings": [
   221  					{
   222  						"target_table": "zip_detail",
   223  						"source_expression": "select * from zip_detail",
   224  						"create_ddl": "copy"
   225  					}
   226  				]
   227  			}`,
   228  		)
   229  		fmt.Fprintf(os.Stderr, "Output from materialize: %s\n", output)
   230  		if err != nil {
   231  			fmt.Fprintf(os.Stderr, "Got error trying to start materialize zip_detail: %v\n", err)
   232  			return 1
   233  		}
   234  
   235  		ctx := context.Background()
   236  		vtgateAddr := fmt.Sprintf("%s:%d", clusterInstance.Hostname, clusterInstance.VtgateProcess.GrpcPort)
   237  		vtgateConn, err := vtgateconn.Dial(ctx, vtgateAddr)
   238  		if err != nil {
   239  			return 1
   240  		}
   241  		defer vtgateConn.Close()
   242  
   243  		session := vtgateConn.Session("@primary", nil)
   244  		// INSERT some zip_detail rows.
   245  		if _, err := session.Execute(ctx, `
   246  			INSERT INTO zip_detail(id, zip_id, discontinued_at)
   247  			VALUES (1, 1, '2022-05-13'),
   248  				   (2, 2, '2022-08-15')
   249  		`, map[string]*querypb.BindVariable{}); err != nil {
   250  			return 1
   251  		}
   252  
   253  		// INSERT some delivery_failure rows.
   254  		if _, err := session.Execute(ctx, `
   255  			INSERT INTO delivery_failure(id, zip_detail_id, reason)
   256  			VALUES (1, 1, 'Failed delivery due to discontinued zipcode.'),
   257  			       (2, 2, 'Failed delivery due to discontinued zipcode.'),
   258  			       (3, 3, 'Failed delivery due to unknown reason.');
   259  		`, map[string]*querypb.BindVariable{}); err != nil {
   260  			return 1
   261  		}
   262  
   263  		if ok := <-done; !ok {
   264  			fmt.Fprintf(os.Stderr, "Materialize did not succeed.\n")
   265  			return 1
   266  		}
   267  
   268  		// Stop materialize zip_detail to sharded keyspace.
   269  		err = clusterInstance.VtctlProcess.ExecuteCommand(
   270  			"Workflow",
   271  			"--",
   272  			shardedKeyspaceName+".copy_zip_detail",
   273  			"delete",
   274  		)
   275  		if err != nil {
   276  			fmt.Fprintf(os.Stderr, "Failed to stop materialization workflow: %v", err)
   277  			return 1
   278  		}
   279  
   280  		return m.Run()
   281  	}()
   282  	os.Exit(exitCode)
   283  }