connection_test.go

Documentation: github.com/GoogleCloudPlatform/cloudsql-proxy/tests

     1  // Copyright 2020 Google LLC
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // connection_test.go provides some helpers for basic connectivity tests to Cloud SQL instances.
    16  package tests
    17  
    18  import (
    19  	"context"
    20  	"database/sql"
    21  	"fmt"
    22  	"sync"
    23  	"testing"
    24  )
    25  
    26  // proxyConnTest is a test helper to verify the proxy works with a basic connectivity test.
    27  func proxyConnTest(t *testing.T, connName, driver, dsn string, port int, dir string) {
    28  	ctx := context.Background()
    29  
    30  	var args []string
    31  	if dir != "" { // unix port
    32  		args = append(args, fmt.Sprintf("-dir=%s", dir), fmt.Sprintf("-instances=%s", connName))
    33  	} else { // tcp socket
    34  		args = append(args, fmt.Sprintf("-instances=%s=tcp:%d", connName, port))
    35  	}
    36  
    37  	// Start the proxy
    38  	p, err := StartProxy(ctx, args...)
    39  	if err != nil {
    40  		t.Fatalf("unable to start proxy: %v", err)
    41  	}
    42  	defer p.Close()
    43  	output, err := p.WaitForServe(ctx)
    44  	if err != nil {
    45  		t.Fatalf("unable to verify proxy was serving: %s \n %s", err, output)
    46  	}
    47  
    48  	// Connect to the instance
    49  	db, err := sql.Open(driver, dsn)
    50  	if err != nil {
    51  		t.Fatalf("unable to connect to db: %s", err)
    52  	}
    53  	defer db.Close()
    54  	_, err = db.Exec("SELECT 1;")
    55  	if err != nil {
    56  
    57  		t.Fatalf("unable to exec on db: %s", err)
    58  	}
    59  }
    60  
    61  func proxyConnLimitTest(t *testing.T, connName, driver, dsn string, port int) {
    62  	ctx := context.Background()
    63  
    64  	maxConn, totConn := 5, 10
    65  
    66  	// Start the proxy
    67  	p, err := StartProxy(ctx, fmt.Sprintf("-instances=%s=tcp:%d", connName, port), fmt.Sprintf("-max_connections=%d", maxConn))
    68  	if err != nil {
    69  		t.Fatalf("unable to start proxy: %v", err)
    70  	}
    71  	defer p.Close()
    72  	output, err := p.WaitForServe(ctx)
    73  	if err != nil {
    74  		t.Fatalf("unable to verify proxy was serving: %s \n %s", err, output)
    75  	}
    76  
    77  	// Create connection pool
    78  	var stmt string
    79  	switch driver {
    80  	case "mysql":
    81  		stmt = "SELECT sleep(2);"
    82  	case "postgres":
    83  		stmt = "SELECT pg_sleep(2);"
    84  	case "sqlserver":
    85  		stmt = "WAITFOR DELAY '00:00:02'"
    86  	default:
    87  		t.Fatalf("unsupported driver: no sleep query found")
    88  	}
    89  	db, err := sql.Open(driver, dsn)
    90  	if err != nil {
    91  		t.Fatalf("unable to connect to db: %s", err)
    92  	}
    93  	db.SetMaxIdleConns(0)
    94  	defer db.Close()
    95  
    96  	// Connect with up to totConn and count errors
    97  	var wg sync.WaitGroup
    98  	c := make(chan error, totConn)
    99  	for i := 0; i < totConn; i++ {
   100  		wg.Add(1)
   101  		go func() {
   102  			defer wg.Done()
   103  			_, err := db.ExecContext(ctx, stmt)
   104  			if err != nil {
   105  				c <- err
   106  			}
   107  		}()
   108  	}
   109  	wg.Wait()
   110  	close(c)
   111  
   112  	var errs []error
   113  	for e := range c {
   114  		errs = append(errs, e)
   115  	}
   116  	want, got := totConn-maxConn, len(errs)
   117  	if want != got {
   118  		t.Errorf("wrong errCt - want: %d, got %d", want, got)
   119  		for _, e := range errs {
   120  			t.Errorf("%s\n", e)
   121  		}
   122  		t.Fail()
   123  	}
   124  }
   125
View as plain text