github.com/google/syzkaller@v0.0.0-20251211124644-a066d2bc4b02/syz-cluster/tools/db-mgmt/main.go (about)

     1  // Copyright 2024 syzkaller project authors. All rights reserved.
     2  // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
     3  
     4  package main
     5  
     6  import (
     7  	"context"
     8  	"fmt"
     9  	"log"
    10  	"os"
    11  
    12  	"cloud.google.com/go/spanner"
    13  	"github.com/google/syzkaller/syz-cluster/pkg/app"
    14  	"github.com/google/syzkaller/syz-cluster/pkg/db"
    15  	"google.golang.org/api/iterator"
    16  )
    17  
    18  func runSQL(ctx context.Context, uri db.ParsedURI, command string) error {
    19  	client, err := spanner.NewClient(ctx, uri.Full)
    20  	if err != nil {
    21  		return err
    22  	}
    23  	defer client.Close()
    24  	stmt := spanner.Statement{SQL: command}
    25  	iter := client.Single().Query(ctx, stmt)
    26  	defer iter.Stop()
    27  
    28  	for {
    29  		row, err := iter.Next()
    30  		if err == iterator.Done {
    31  			break
    32  		}
    33  		if err != nil {
    34  			return err
    35  		}
    36  		cols := row.ColumnNames()
    37  		fmt.Println(cols)
    38  		for i := 0; i < len(cols); i++ {
    39  			fmt.Printf("\t%s", row.ColumnValue(i))
    40  		}
    41  		fmt.Printf("\n")
    42  	}
    43  	return nil
    44  }
    45  
    46  func main() {
    47  	ctx := context.Background()
    48  	uri, err := app.DefaultSpannerURI()
    49  	if err != nil {
    50  		app.Fatalf("failed to get Spanner URI: %v", err)
    51  	}
    52  	if os.Getenv("SPANNER_EMULATOR_HOST") != "" {
    53  		// Should only be done in the test/dev environment.
    54  		log.Printf("check if a Spanner instance is present")
    55  		err = db.CreateSpannerInstance(ctx, uri)
    56  		if err != nil {
    57  			app.Fatalf("failed to create Spanner instance: %v", err)
    58  		}
    59  	}
    60  	log.Printf("check if DB is present")
    61  	err = db.CreateSpannerDB(ctx, uri)
    62  	if err != nil {
    63  		app.Fatalf("failed to create Spanner DB: %v", err)
    64  	}
    65  	if len(os.Args) > 1 {
    66  		switch os.Args[1] {
    67  		case "migrate":
    68  			log.Printf("running schema migrations")
    69  			err = db.RunMigrations(uri.Full)
    70  		case "run":
    71  			if len(os.Args) < 3 {
    72  				app.Fatalf("second argument is the SQL query to run")
    73  			}
    74  			err = runSQL(ctx, uri, os.Args[2])
    75  		default:
    76  			app.Fatalf("unknown command: %s", os.Args[1])
    77  		}
    78  		if err != nil {
    79  			log.Fatal(err)
    80  		}
    81  	}
    82  	log.Printf("finished!")
    83  }