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

     1  /*
     2  Copyright 2020 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 godriver
    18  
    19  import (
    20  	"flag"
    21  	"fmt"
    22  	"os"
    23  	"strconv"
    24  	"testing"
    25  	"time"
    26  
    27  	"github.com/stretchr/testify/assert"
    28  	"github.com/stretchr/testify/require"
    29  	"google.golang.org/grpc"
    30  	"google.golang.org/grpc/keepalive"
    31  
    32  	"vitess.io/vitess/go/vt/log"
    33  	"vitess.io/vitess/go/vt/vitessdriver"
    34  
    35  	"vitess.io/vitess/go/test/endtoend/cluster"
    36  )
    37  
    38  var (
    39  	clusterInstance *cluster.LocalProcessCluster
    40  	cell            = "zone1"
    41  	hostname        = "localhost"
    42  	KeyspaceName    = "customer"
    43  	SchemaSQL       = `
    44  create table my_message(
    45  	# required columns
    46  	id bigint NOT NULL COMMENT 'often an event id, can also be auto-increment or a sequence',
    47  	priority tinyint NOT NULL DEFAULT '50' COMMENT 'lower number priorities process first',
    48  	epoch bigint NOT NULL DEFAULT '0' COMMENT 'Vitess increments this each time it sends a message, and is used for incremental backoff doubling',
    49  	time_next bigint DEFAULT 0 COMMENT 'the earliest time the message will be sent in epoch nanoseconds. Must be null if time_acked is set',
    50  	time_acked bigint DEFAULT NULL COMMENT 'the time the message was acked in epoch nanoseconds. Must be null if time_next is set',
    51  
    52  	# add as many custom fields here as required
    53  	# optional - these are suggestions
    54  	tenant_id bigint,
    55  	message json,
    56  
    57  	# required indexes
    58  	primary key(id),
    59  	index poller_idx(time_acked, priority, time_next desc)
    60  
    61  	# add any secondary indexes or foreign keys - no restrictions
    62  ) comment 'vitess_message,vt_min_backoff=30,vt_max_backoff=3600,vt_ack_wait=30,vt_purge_after=86400,vt_batch_size=10,vt_cache_size=10000,vt_poller_interval=30'
    63  `
    64  	VSchema = `
    65  {
    66    "sharded": true,
    67    "vindexes": {
    68      "hash": {
    69        "type": "hash"
    70      }
    71    },
    72    "tables": {
    73      "my_message": {
    74        "column_vindexes": [
    75          {
    76            "column": "id",
    77            "name": "hash"
    78          }
    79        ]
    80      }
    81    }
    82  }
    83  `
    84  
    85  	testMessage = "{\"message\":\"hello world\"}"
    86  )
    87  
    88  func TestMain(m *testing.M) {
    89  	defer cluster.PanicHandler(nil)
    90  	flag.Parse()
    91  
    92  	exitCode := func() int {
    93  		clusterInstance = cluster.NewCluster(cell, hostname)
    94  		defer clusterInstance.Teardown()
    95  
    96  		// Start topo server
    97  		if err := clusterInstance.StartTopo(); err != nil {
    98  			return 1
    99  		}
   100  
   101  		// Start keyspace
   102  		Keyspace := &cluster.Keyspace{
   103  			Name:      KeyspaceName,
   104  			SchemaSQL: SchemaSQL,
   105  			VSchema:   VSchema,
   106  		}
   107  		clusterInstance.VtTabletExtraArgs = []string{
   108  			"--queryserver-config-transaction-timeout", "3",
   109  		}
   110  		if err := clusterInstance.StartKeyspace(*Keyspace, []string{"-80", "80-"}, 1, false); err != nil {
   111  			log.Fatal(err.Error())
   112  			return 1
   113  		}
   114  
   115  		// Start vtgate
   116  		clusterInstance.VtGateExtraArgs = []string{"--warn_sharded_only=true"}
   117  		if err := clusterInstance.StartVtgate(); err != nil {
   118  			log.Fatal(err.Error())
   119  			return 1
   120  		}
   121  
   122  		return m.Run()
   123  	}()
   124  	os.Exit(exitCode)
   125  }
   126  
   127  func TestStreamMessaging(t *testing.T) {
   128  	defer cluster.PanicHandler(t)
   129  
   130  	cnf := vitessdriver.Configuration{
   131  		Protocol: "grpc",
   132  		Address:  clusterInstance.Hostname + ":" + strconv.Itoa(clusterInstance.VtgateGrpcPort),
   133  		GRPCDialOptions: []grpc.DialOption{
   134  			grpc.WithDefaultCallOptions(),
   135  			grpc.WithKeepaliveParams(keepalive.ClientParameters{
   136  				Time:                300 * time.Second,
   137  				Timeout:             600 * time.Second,
   138  				PermitWithoutStream: true,
   139  			}),
   140  		},
   141  	}
   142  
   143  	// for inserting data
   144  	db, err := vitessdriver.OpenWithConfiguration(cnf)
   145  	require.NoError(t, err)
   146  	defer db.Close()
   147  
   148  	// Exec not allowed in streaming
   149  	_, err = db.Exec(fmt.Sprintf("insert into my_message(id, message) values(1, '%s')", testMessage))
   150  	require.NoError(t, err)
   151  
   152  	// for streaming messages
   153  	cnf.Streaming = true
   154  	streamDB, err := vitessdriver.OpenWithConfiguration(cnf)
   155  	require.NoError(t, err)
   156  	defer streamDB.Close()
   157  
   158  	// Exec not allowed in streaming
   159  	_, err = streamDB.Exec("stream * from my_message")
   160  	assert.EqualError(t, err, "Exec not allowed for streaming connections")
   161  
   162  	row := streamDB.QueryRow("stream * from my_message")
   163  	require.NoError(t, row.Err())
   164  }