github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/explain_bundle_test.go (about)

     1  // Copyright 2020 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package sql
    12  
    13  import (
    14  	"archive/zip"
    15  	"bytes"
    16  	"context"
    17  	"fmt"
    18  	"io"
    19  	"regexp"
    20  	"sort"
    21  	"strings"
    22  	"testing"
    23  
    24  	"github.com/cockroachdb/cockroach/pkg/base"
    25  	"github.com/cockroachdb/cockroach/pkg/testutils"
    26  	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
    27  	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
    28  	"github.com/cockroachdb/cockroach/pkg/util/httputil"
    29  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    30  	"github.com/cockroachdb/errors"
    31  	"github.com/lib/pq"
    32  )
    33  
    34  func TestExplainAnalyzeDebug(t *testing.T) {
    35  	defer leaktest.AfterTest(t)()
    36  
    37  	ctx := context.Background()
    38  	srv, godb, _ := serverutils.StartServer(t, base.TestServerArgs{Insecure: true})
    39  	defer srv.Stopper().Stop(ctx)
    40  	r := sqlutils.MakeSQLRunner(godb)
    41  	r.Exec(t, "CREATE TABLE abc (a INT PRIMARY KEY, b INT, c INT UNIQUE)")
    42  
    43  	base := "statement.txt trace.json trace.txt trace-jaeger.json env.sql"
    44  	plans := "schema.sql opt.txt opt-v.txt opt-vv.txt plan.txt"
    45  
    46  	t.Run("basic", func(t *testing.T) {
    47  		rows := r.QueryStr(t, "EXPLAIN ANALYZE (DEBUG) SELECT * FROM abc WHERE c=1")
    48  		checkBundle(
    49  			t, fmt.Sprint(rows),
    50  			base, plans, "stats-defaultdb.public.abc.sql", "distsql.html",
    51  		)
    52  	})
    53  
    54  	// Check that we get separate diagrams for subqueries.
    55  	t.Run("subqueries", func(t *testing.T) {
    56  		rows := r.QueryStr(t, "EXPLAIN ANALYZE (DEBUG) SELECT EXISTS (SELECT * FROM abc WHERE c=1)")
    57  		checkBundle(
    58  			t, fmt.Sprint(rows),
    59  			base, plans, "stats-defaultdb.public.abc.sql", "distsql-1.html distsql-2.html",
    60  		)
    61  	})
    62  
    63  	// Even on query errors we should still get a bundle.
    64  	t.Run("error", func(t *testing.T) {
    65  		_, err := godb.QueryContext(ctx, "EXPLAIN ANALYZE (DEBUG) SELECT * FROM badtable")
    66  		if !testutils.IsError(err, "relation.*does not exist") {
    67  			t.Fatalf("unexpected error %v\n", err)
    68  		}
    69  		// The bundle url is inside the error detail.
    70  		var pqErr *pq.Error
    71  		_ = errors.As(err, &pqErr)
    72  		checkBundle(t, fmt.Sprintf("%+v", pqErr.Detail), base)
    73  	})
    74  
    75  	// Verify that we can issue the statement with prepare (which can happen
    76  	// depending on the client).
    77  	t.Run("prepare", func(t *testing.T) {
    78  		stmt, err := godb.Prepare("EXPLAIN ANALYZE (DEBUG) SELECT * FROM abc WHERE c=1")
    79  		if err != nil {
    80  			t.Fatal(err)
    81  		}
    82  		defer stmt.Close()
    83  		rows, err := stmt.Query()
    84  		if err != nil {
    85  			t.Fatal(err)
    86  		}
    87  		var rowsBuf bytes.Buffer
    88  		for rows.Next() {
    89  			var row string
    90  			if err := rows.Scan(&row); err != nil {
    91  				t.Fatal(err)
    92  			}
    93  			rowsBuf.WriteString(row)
    94  			rowsBuf.WriteByte('\n')
    95  		}
    96  		checkBundle(
    97  			t, rowsBuf.String(),
    98  			base, plans, "stats-defaultdb.public.abc.sql", "distsql.html",
    99  		)
   100  	})
   101  }
   102  
   103  // checkBundle searches text strings for a bundle URL and then verifies that the
   104  // bundle contains the expected files. The expected files are passed as an
   105  // arbitrary number of strings; each string contains one or more filenames
   106  // separated by a space.
   107  func checkBundle(t *testing.T, text string, expectedFiles ...string) {
   108  	t.Helper()
   109  	reg := regexp.MustCompile("http://[a-zA-Z0-9.:]*/_admin/v1/stmtbundle/[0-9]*")
   110  	url := reg.FindString(text)
   111  	if url == "" {
   112  		t.Fatalf("couldn't find URL in response '%s'", text)
   113  	}
   114  	// Download the zip to a BytesBuffer.
   115  	resp, err := httputil.Get(context.Background(), url)
   116  	if err != nil {
   117  		t.Fatal(err)
   118  	}
   119  	defer resp.Body.Close()
   120  	var buf bytes.Buffer
   121  	_, _ = io.Copy(&buf, resp.Body)
   122  
   123  	unzip, err := zip.NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len()))
   124  	if err != nil {
   125  		t.Errorf("%q\n", buf.String())
   126  		t.Fatal(err)
   127  	}
   128  
   129  	// Make sure the bundle contains the expected list of files.
   130  	var files []string
   131  	for _, f := range unzip.File {
   132  		if f.UncompressedSize64 == 0 {
   133  			t.Fatalf("file %s is empty", f.Name)
   134  		}
   135  		files = append(files, f.Name)
   136  	}
   137  
   138  	var expList []string
   139  	for _, s := range expectedFiles {
   140  		expList = append(expList, strings.Split(s, " ")...)
   141  	}
   142  	sort.Strings(files)
   143  	sort.Strings(expList)
   144  	if fmt.Sprint(files) != fmt.Sprint(expList) {
   145  		t.Errorf("unexpected list of files:\n  %v\nexpected:\n  %v", files, expList)
   146  	}
   147  }