WIP: Add initial schema and storage provider for CockroachDB

.gitignore

@@ -22,3 +22,4 @@
 /trillian_log_signer
 /trillian_map_server
 default.etcd
+cockroach-data/

cmd/trillian_log_server/main.go

@@ -47,6 +47,7 @@ import (
 
 	// Register supported storage providers.
 	_ "github.com/google/trillian/storage/cloudspanner"
+	_ "github.com/google/trillian/storage/crdb"
 	_ "github.com/google/trillian/storage/mysql"
 
 	// Load MySQL quota provider

cmd/trillian_log_signer/main.go

@@ -53,6 +53,7 @@ import (
 
 	// Register supported storage providers.
 	_ "github.com/google/trillian/storage/cloudspanner"
+	_ "github.com/google/trillian/storage/crdb"
 	_ "github.com/google/trillian/storage/mysql"
 
 	// Load MySQL quota provider

go.mod

@@ -7,6 +7,7 @@ require (
 	cloud.google.com/go/spanner v1.40.0
 	contrib.go.opencensus.io/exporter/stackdriver v0.13.12
 	github.com/apache/beam/sdks/v2 v2.0.0-20211012030016-ef4364519c94
+	github.com/cockroachdb/cockroach-go/v2 v2.2.16
 	github.com/fullstorydev/grpcurl v1.8.7
 	github.com/go-redis/redis v6.15.9+incompatible
 	github.com/go-sql-driver/mysql v1.6.0
@@ -16,6 +17,7 @@ require (
 	github.com/google/go-licenses v0.0.0-20210329231322-ce1d9163b77d
 	github.com/grpc-ecosystem/go-grpc-middleware v1.3.0
 	github.com/letsencrypt/pkcs11key/v4 v4.0.0
+	github.com/lib/pq v1.10.7
 	github.com/prometheus/client_golang v1.13.1
 	github.com/prometheus/client_model v0.3.0
 	github.com/pseudomuto/protoc-gen-doc v1.5.1
@@ -67,6 +69,7 @@ require (
 	github.com/fatih/color v1.9.0 // indirect
 	github.com/form3tech-oss/jwt-go v3.2.5+incompatible // indirect
 	github.com/go-logr/logr v1.2.0 // indirect
+	github.com/gofrs/flock v0.8.1 // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
 	github.com/golang/glog v1.0.0 // indirect
 	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
@@ -82,13 +85,20 @@ require (
 	github.com/huandu/xstrings v1.2.0 // indirect
 	github.com/imdario/mergo v0.3.9 // indirect
 	github.com/inconshreveable/mousetrap v1.0.0 // indirect
+	github.com/jackc/chunkreader/v2 v2.0.1 // indirect
+	github.com/jackc/pgconn v1.12.1 // indirect
+	github.com/jackc/pgio v1.0.0 // indirect
+	github.com/jackc/pgpassfile v1.0.0 // indirect
+	github.com/jackc/pgproto3/v2 v2.3.0 // indirect
+	github.com/jackc/pgservicefile v0.0.0-20200714003250-2b9c44734f2b // indirect
+	github.com/jackc/pgtype v1.11.0 // indirect
+	github.com/jackc/pgx/v4 v4.16.1 // indirect
 	github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99 // indirect
 	github.com/jhump/protoreflect v1.12.0 // indirect
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
 	github.com/jonboulle/clockwork v0.3.0 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
 	github.com/kevinburke/ssh_config v0.0.0-20190725054713-01f96b0aa0cd // indirect
-	github.com/mattn/go-colorable v0.1.4 // indirect
 	github.com/mattn/go-runewidth v0.0.13 // indirect
 	github.com/matttproud/golang_protobuf_extensions v1.0.1 // indirect
 	github.com/miekg/pkcs11 v1.0.3 // indirect
@@ -102,6 +112,7 @@ require (
 	github.com/onsi/ginkgo v1.10.3 // indirect
 	github.com/onsi/gomega v1.7.1 // indirect
 	github.com/otiai10/copy v1.2.0 // indirect
+	github.com/pkg/errors v0.9.1 // indirect
 	github.com/prometheus/common v0.37.0 // indirect
 	github.com/prometheus/procfs v0.8.0 // indirect
 	github.com/prometheus/prometheus v2.5.0+incompatible // indirect

integration/admin/admin_integration_test.go

@@ -582,7 +582,8 @@ func setupAdminServer(ctx context.Context, t *testing.T) (*testServer, error) {
 		return nil, err
 	}
 
-	registry, done, err := integration.NewRegistryForTests(ctx)
+	// TODO(jaosorior): Make this configurable for Cockroach or MySQL
+	registry, done, err := integration.NewRegistryForTests(ctx, testdb.DriverMySQL)
 	if err != nil {
 		ts.closeAll()
 		return nil, err

integration/quota/quota_test.go

@@ -51,7 +51,8 @@ func TestEtcdRateLimiting(t *testing.T) {
 	testdb.SkipIfNoMySQL(t)
 	ctx := context.Background()
 
-	registry, done, err := integration.NewRegistryForTests(ctx)
+	// TODO(jaosorior): Make this configurable for Cockroach or MySQL
+	registry, done, err := integration.NewRegistryForTests(ctx, testdb.DriverMySQL)
 	if err != nil {
 		t.Fatalf("NewRegistryForTests() returned err = %v", err)
 	}
@@ -100,7 +101,7 @@ func TestEtcdRateLimiting(t *testing.T) {
 func TestMySQLRateLimiting(t *testing.T) {
 	testdb.SkipIfNoMySQL(t)
 	ctx := context.Background()
-	db, done, err := testdb.NewTrillianDB(ctx)
+	db, done, err := testdb.NewTrillianDB(ctx, testdb.DriverMySQL)
 	if err != nil {
 		t.Fatalf("GetTestDB() returned err = %v", err)
 	}

quota/mysqlqm/mysql_quota_test.go

@@ -37,7 +37,7 @@ func TestQuotaManager_GetTokens(t *testing.T) {
 	testdb.SkipIfNoMySQL(t)
 	ctx := context.Background()
 
-	db, done, err := testdb.NewTrillianDB(ctx)
+	db, done, err := testdb.NewTrillianDB(ctx, testdb.DriverMySQL)
 	if err != nil {
 		t.Fatalf("GetTestDB() returned err = %v", err)
 	}
@@ -131,7 +131,7 @@ func TestQuotaManager_GetTokens_InformationSchema(t *testing.T) {
 	for _, test := range tests {
 		desc := fmt.Sprintf("useSelectCount = %v", test.useSelectCount)
 		t.Run(desc, func(t *testing.T) {
-			db, done, err := testdb.NewTrillianDB(ctx)
+			db, done, err := testdb.NewTrillianDB(ctx, testdb.DriverMySQL)
 			if err != nil {
 				t.Fatalf("NewTrillianDB() returned err = %v", err)
 			}
@@ -182,7 +182,7 @@ func TestQuotaManager_Noops(t *testing.T) {
 	testdb.SkipIfNoMySQL(t)
 	ctx := context.Background()
 
-	db, done, err := testdb.NewTrillianDB(ctx)
+	db, done, err := testdb.NewTrillianDB(ctx, testdb.DriverMySQL)
 	if err != nil {
 		t.Fatalf("GetTestDB() returned err = %v", err)
 	}

scripts/resetcrdb.sh

@@ -0,0 +1,106 @@
+#!/bin/bash
+
+set -e
+
+usage() {
+  cat <<EOF
+$(basename $0) [--force] [--verbose] ...
+All unrecognised arguments will be passed through to the 'cockroach' command.
+Accepts environment variables:
+- CRDB_ROOT_USER: A user with sufficient rights to create/reset the Trillian
+  database (default: root).
+- CRDB_ROOT_PASSWORD: The password for \$CRDB_ROOT_USER (default: none).
+- CRDB_HOST: The hostname of the MySQL server (default: localhost).
+- CRDB_PORT: The port the MySQL server is listening on (default: 3306).
+- CRDB_DATABASE: The name to give to the new Trillian user and database
+  (default: test).
+- CRDB_USER: The name to give to the new Trillian user (default: test).
+- CRDB_PASSWORD: The password to use for the new Trillian user
+  (default: zaphod).
+- CRDB_USER_HOST: The host that the Trillian user will connect from; use '%' as
+  a wildcard (default: localhost).
+EOF
+}
+
+die() {
+  echo "$*" > /dev/stderr
+  exit 1
+}
+
+collect_vars() {
+  # set unset environment variables to defaults
+  [ -z ${CRDB_ROOT_USER+x} ] && CRDB_ROOT_USER="root"
+  [ -z ${CRDB_HOST+x} ] && CRDB_HOST="localhost"
+  [ -z ${CRDB_PORT+x} ] && CRDB_PORT="26257"
+  [ -z ${CRDB_DATABASE+x} ] && CRDB_DATABASE="defaultdb"
+  [ -z ${CRDB_USER+x} ] && CRDB_USER="test"
+  [ -z ${CRDB_PASSWORD+x} ] && CRDB_PASSWORD="zaphod"
+  [ -z ${CRDB_USER_HOST+x} ] && CRDB_USER_HOST="localhost"
+  [ -z ${CRDB_INSECURE+x} ] && CRDB_INSECURE="true" 
+  FLAGS=()
+
+  # handle flags
+  FORCE=false
+  VERBOSE=false
+  while [[ $# -gt 0 ]]; do
+    case "$1" in
+      --force) FORCE=true ;;
+      --verbose) VERBOSE=true ;;
+      --help) usage; exit ;;
+      *) FLAGS+=("$1")
+    esac
+    shift 1
+  done
+
+  FLAGS+=(-u "${CRDB_ROOT_USER}")
+  FLAGS+=(--host "${CRDB_HOST}")
+  FLAGS+=(--port "${CRDB_PORT}")
+
+  # Useful for debugging
+  FLAGS+=(--echo-sql)
+
+  if [[ ${CRDB_INSECURE} = 'true' ]]; then
+    FLAGS+=(--insecure)
+  fi
+
+  # Optionally print flags (before appending password)
+  [[ ${VERBOSE} = 'true' ]] && echo "- Using MySQL Flags: ${FLAGS[@]}"
+
+  # append password if supplied
+  [ -z ${CRDB_ROOT_PASSWORD+x} ] || FLAGS+=(-p"${CRDB_ROOT_PASSWORD}")
+}
+
+main() {
+  collect_vars "$@"
+
+  readonly TRILLIAN_PATH=$(go list -f '{{.Dir}}' github.com/google/trillian)
+
+  echo "Warning: about to destroy and reset database '${CRDB_DATABASE}'"
+
+  [[ ${FORCE} = true ]] || read -p "Are you sure? [Y/N]: " -n 1 -r
+  echo # Print newline following the above prompt
+
+  if [ -z ${REPLY+x} ] || [[ $REPLY =~ ^[Yy]$ ]]
+  then
+      echo "Resetting DB..."
+      set -eux
+      cockroach sql "${FLAGS[@]}" -e "DROP DATABASE IF EXISTS ${CRDB_DATABASE};" || \
+        die "Error: Failed to drop database '${CRDB_DATABASE}'."
+      cockroach sql "${FLAGS[@]}" -e "CREATE DATABASE ${CRDB_DATABASE};" || \
+        die "Error: Failed to create database '${CRDB_DATABASE}'."
+      if [[ ${CRDB_INSECURE} = 'true' ]]; then
+        cockroach sql "${FLAGS[@]}" -e "CREATE USER IF NOT EXISTS ${CRDB_USER};" || \
+          die "Error: Failed to create user '${CRDB_USER}'."
+      else
+        cockroach sql "${FLAGS[@]}" -e "CREATE USER IF NOT EXISTS ${CRDB_USER} WITH PASSWORD '${CRDB_PASSWORD}';" || \
+          die "Error: Failed to create user '${CRDB_USER}'."
+      fi
+      cockroach sql "${FLAGS[@]}" -e "GRANT ALL PRIVILEGES ON DATABASE ${CRDB_DATABASE} TO ${CRDB_USER} WITH GRANT OPTION" || \
+        die "Error: Failed to grant '${CRDB_USER}' user all privileges on '${CRDB_DATABASE}'."
+      cockroach sql "${FLAGS[@]}" -d ${CRDB_DATABASE} < ${TRILLIAN_PATH}/storage/crdb/schema/storage.sql || \
+        die "Error: Failed to create tables in '${CRDB_DATABASE}' database."
+      echo "Reset Complete"
+  fi
+}
+
+main "$@"

storage/crdb/common_test.go

@@ -0,0 +1,107 @@
+// Copyright 2022 <TBD>
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package crdb
+
+import (
+	"context"
+	"database/sql"
+	"flag"
+	"os"
+	"sync"
+	"testing"
+
+	"github.com/cockroachdb/cockroach-go/v2/testserver"
+	"k8s.io/klog/v2"
+
+	"github.com/google/trillian/storage/testdb"
+)
+
+// testDBs holds a set of test databases, one per test.
+var testDBs sync.Map
+
+type testDBHandle struct {
+	db   *sql.DB
+	done func(context.Context)
+}
+
+func (db *testDBHandle) GetDB() *sql.DB {
+	return db.db
+}
+
+func TestMain(m *testing.M) {
+	flag.Parse()
+
+	ts, err := testserver.NewTestServer()
+	if err != nil {
+		klog.Errorf("Failed to start test server: %v", err)
+		os.Exit(1)
+	}
+	defer ts.Stop()
+
+	// reset the test server URL path. By default cockroach sets it
+	// to point to a default database, we don't want that.
+	dburl := ts.PGURL()
+	dburl.Path = "/"
+
+	// Set the environment variable for the test server
+	os.Setenv(testdb.CockroachDBURIEnv, dburl.String())
+
+	if !testdb.CockroachDBAvailable() {
+		klog.Errorf("CockroachDB not available, skipping all CockroachDB storage tests")
+		return
+	}
+
+	status := m.Run()
+
+	// Clean up databases
+	testDBs.Range(func(key, value interface{}) bool {
+		testName := key.(string)
+		klog.Infof("Cleaning up database for test %s", testName)
+
+		db := value.(*testDBHandle)
+
+		// TODO(jaosorior): Set a timeout instead of using Background
+		db.done(context.Background())
+
+		return true
+	})
+
+	os.Exit(status)
+}
+
+// This is used to identify a database from the map
+func getDBID(t *testing.T) string {
+	t.Helper()
+	return t.Name()
+}
+
+func openTestDBOrDie(t *testing.T) *testDBHandle {
+	t.Helper()
+
+	// TODO(jaosorior): Check for Cockroach
+	db, done, err := testdb.NewTrillianDB(context.TODO(), testdb.DriverCockroachDB)
+	if err != nil {
+		panic(err)
+	}
+
+	handle := &testDBHandle{
+		db:   db,
+		done: done,
+	}
+
+	testDBs.Store(getDBID(t), handle)
+
+	return handle
+}

storage/crdb/drop_storage.sql

@@ -0,0 +1,9 @@
+-- Caution - this removes all tables in our schema
+
+DROP TABLE IF EXISTS Unsequenced;
+DROP TABLE IF EXISTS Subtree;
+DROP TABLE IF EXISTS SequencedLeafData;
+DROP TABLE IF EXISTS TreeHead;
+DROP TABLE IF EXISTS LeafData;
+DROP TABLE IF EXISTS TreeControl;
+DROP TABLE IF EXISTS Trees;

storage/crdb/errors.go

@@ -0,0 +1,42 @@
+// Copyright 2022 <TBD>
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package crdb
+
+import (
+	"github.com/lib/pq"
+)
+
+var uniqueViolationErrorCode = pq.ErrorCode("23505")
+
+// crdbToGRPC converts some types of CockroachDB errors to GRPC errors. This gives
+// clients more signal when the operation can be retried.
+func crdbToGRPC(err error) error {
+	_, ok := err.(*pq.Error)
+	if !ok {
+		return err
+	}
+	// TODO(jaosorior): Do we have a crdb equivalent for a deadlock
+	// error code?
+	return err
+}
+
+func isDuplicateErr(err error) bool {
+	switch err := err.(type) {
+	case *pq.Error:
+		return err.Code == uniqueViolationErrorCode
+	default:
+		return false
+	}
+}