diff --git a/DEPS.bzl b/DEPS.bzl index a2769bcb8744..374be1d5e32f 100644 --- a/DEPS.bzl +++ b/DEPS.bzl @@ -1445,14 +1445,17 @@ def go_deps(): name = "com_github_cenkalti_backoff_v4", build_file_proto_mode = "disable_global", importpath = "github.com/cenkalti/backoff/v4", - sha256 = "73ff572a901c0307aa1c16db43812da7ca2555aa403cfdd9d3a239ecbdad2274", - strip_prefix = "github.com/cenkalti/backoff/v4@v4.1.3", + sha256 = "0b1d9cedebb1b814f4fbc03a47fdd2c2bb91d8cf14dbb1a71d3bc1482600cd2a", + strip_prefix = "github.com/cenkalti/backoff/v4@v4.2.1", urls = [ - "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cenkalti/backoff/v4/com_github_cenkalti_backoff_v4-v4.1.3.zip", + "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cenkalti/backoff/v4/com_github_cenkalti_backoff_v4-v4.2.1.zip", ], ) go_repository( name = "com_github_census_instrumentation_opencensus_proto", + build_directives = [ + "gazelle:resolve go go github.com/census-instrumentation/opencensus-proto/gen-go/resource/v1 @com_github_census_instrumentation_opencensus_proto//gen-go/resource/v1:go_default_library", + ], build_file_proto_mode = "disable_global", importpath = "github.com/census-instrumentation/opencensus-proto", sha256 = "6fce66b7dcd2cba031ed9d73d77d6b21c2fe749c5de27cbb416a2d2cc1c68719", @@ -2902,9 +2905,18 @@ def go_deps(): ) go_repository( name = "com_github_envoyproxy_protoc_gen_validate", + build_directives = [ + "gazelle:resolve go go github.com/golang/protobuf/protoc-gen-go/descriptor @com_github_golang_protobuf//protoc-gen-go/descriptor:go_default_library", + "gazelle:resolve go go github.com/golang/protobuf/ptypes/duration @com_github_golang_protobuf//ptypes/duration:go_default_library", + "gazelle:resolve go go github.com/golang/protobuf/ptypes/timestamp @com_github_golang_protobuf//ptypes/timestamp:go_default_library", + ], build_file_proto_mode = "disable_global", build_naming_convention = "go_default_library", importpath = "github.com/envoyproxy/protoc-gen-validate", + patch_args = ["-p1"], + patches = [ + "@com_github_cockroachdb_cockroach//build/patches:com_github_envoyproxy_protoc_gen_validate.patch", + ], sha256 = "d9803ca09d732a15d994f16ef297f901496e8861eabfbe6bdfa00d280f1d608f", strip_prefix = "github.com/envoyproxy/protoc-gen-validate@v0.10.1", urls = [ @@ -4440,6 +4452,20 @@ def go_deps(): "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/grpc-ecosystem/grpc-gateway/com_github_grpc_ecosystem_grpc_gateway-v1.16.0.zip", ], ) + go_repository( + name = "com_github_grpc_ecosystem_grpc_gateway_v2", + build_file_proto_mode = "disable_global", + importpath = "github.com/grpc-ecosystem/grpc-gateway/v2", + patch_args = ["-p1"], + patches = [ + "@com_github_cockroachdb_cockroach//build/patches:com_github_grpc_ecosystem_grpc_gateway_v2.patch", + ], + sha256 = "12995bc9a8d7a18ff3fb7fe3575dd2d9b7447d2f2749f0dbd99ef2efbaf9a6e2", + strip_prefix = "github.com/grpc-ecosystem/grpc-gateway/v2@v2.16.0", + urls = [ + "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/grpc-ecosystem/grpc-gateway/v2/com_github_grpc_ecosystem_grpc_gateway_v2-v2.16.0.zip", + ], + ) go_repository( name = "com_github_gsterjov_go_libsecret", build_file_proto_mode = "disable_global", @@ -10927,20 +10953,20 @@ def go_deps(): name = "io_opentelemetry_go_otel_exporters_otlp_otlptrace", build_file_proto_mode = "disable_global", importpath = "go.opentelemetry.io/otel/exporters/otlp/otlptrace", - sha256 = "1ff3c17fc607e5fd94c6e88127138fc2a2489fe31862b2b969e25caeb3d57db1", - strip_prefix = "go.opentelemetry.io/otel/exporters/otlp/otlptrace@v1.3.0", + sha256 = "8c315880f1b7f58438a2fc7c1ed7edbafea6104724984eb7e564a730a249110c", + strip_prefix = "go.opentelemetry.io/otel/exporters/otlp/otlptrace@v1.17.0", urls = [ - "https://storage.googleapis.com/cockroach-godeps/gomod/go.opentelemetry.io/otel/exporters/otlp/otlptrace/io_opentelemetry_go_otel_exporters_otlp_otlptrace-v1.3.0.zip", + "https://storage.googleapis.com/cockroach-godeps/gomod/go.opentelemetry.io/otel/exporters/otlp/otlptrace/io_opentelemetry_go_otel_exporters_otlp_otlptrace-v1.17.0.zip", ], ) go_repository( name = "io_opentelemetry_go_otel_exporters_otlp_otlptrace_otlptracegrpc", build_file_proto_mode = "disable_global", importpath = "go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc", - sha256 = "f898ba889e06d82790ea6819243c3217bcfe1f01bcce4dbb497946843f93f83f", - strip_prefix = "go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc@v1.3.0", + sha256 = "0965d989c01bad8217e2f937c54d82d200896f7e2df6888ce54813c6219b4bcc", + strip_prefix = "go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc@v1.17.0", urls = [ - "https://storage.googleapis.com/cockroach-godeps/gomod/go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc/io_opentelemetry_go_otel_exporters_otlp_otlptrace_otlptracegrpc-v1.3.0.zip", + "https://storage.googleapis.com/cockroach-godeps/gomod/go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc/io_opentelemetry_go_otel_exporters_otlp_otlptrace_otlptracegrpc-v1.17.0.zip", ], ) go_repository( @@ -11010,10 +11036,10 @@ def go_deps(): ], build_file_proto_mode = "disable_global", importpath = "go.opentelemetry.io/proto/otlp", - sha256 = "f22d677bc272c65f45ca31b1ca80a28d1bdb922858e86fbd1579e1852fdb51d8", - strip_prefix = "go.opentelemetry.io/proto/otlp@v0.11.0", + sha256 = "264e0db82a8c572ebc52fbe3da3f4a2a60d7e5fcf81238d1389390f66e2da246", + strip_prefix = "go.opentelemetry.io/proto/otlp@v1.0.0", urls = [ - "https://storage.googleapis.com/cockroach-godeps/gomod/go.opentelemetry.io/proto/otlp/io_opentelemetry_go_proto_otlp-v0.11.0.zip", + "https://storage.googleapis.com/cockroach-godeps/gomod/go.opentelemetry.io/proto/otlp/io_opentelemetry_go_proto_otlp-v1.0.0.zip", ], ) go_repository( diff --git a/build/bazelutil/distdir_files.bzl b/build/bazelutil/distdir_files.bzl index 5cc404254c61..5abdcbb78477 100644 --- a/build/bazelutil/distdir_files.bzl +++ b/build/bazelutil/distdir_files.bzl @@ -316,7 +316,7 @@ DISTDIR_FILES = { "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/campoy/embedmd/com_github_campoy_embedmd-v1.0.0.zip": "a0e0daed0e40d30dfaf7ba58bc8057450f5c1964d5672c49d3b4817a82f9a512", "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/casbin/casbin/v2/com_github_casbin_casbin_v2-v2.1.2.zip": "753df5c3fa5de68592e95fd55427f264dc7590a0bf781a77eb56ae721d6d3351", "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cenkalti/backoff/com_github_cenkalti_backoff-v2.2.1+incompatible.zip": "f8196815a1b4d25e5b8158029d5264801fc8aa5ff128ccf30752fd169693d43b", - "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cenkalti/backoff/v4/com_github_cenkalti_backoff_v4-v4.1.3.zip": "73ff572a901c0307aa1c16db43812da7ca2555aa403cfdd9d3a239ecbdad2274", + "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cenkalti/backoff/v4/com_github_cenkalti_backoff_v4-v4.2.1.zip": "0b1d9cedebb1b814f4fbc03a47fdd2c2bb91d8cf14dbb1a71d3bc1482600cd2a", "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/census-instrumentation/opencensus-proto/com_github_census_instrumentation_opencensus_proto-v0.4.1.zip": "6fce66b7dcd2cba031ed9d73d77d6b21c2fe749c5de27cbb416a2d2cc1c68719", "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cespare/xxhash/com_github_cespare_xxhash-v1.1.0.zip": "fe98c56670b21631f7fd3305a29a3b17e86a6cce3876a2119460717a18538e2e", "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cespare/xxhash/v2/com_github_cespare_xxhash_v2-v2.3.0.zip": "145a26cdc7c49db566017b807c4989ee7f7ddeb569423e9cb99f995fac3621d3", @@ -609,6 +609,7 @@ DISTDIR_FILES = { "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/grpc-ecosystem/go-grpc-middleware/com_github_grpc_ecosystem_go_grpc_middleware-v1.3.0.zip": "081d63238be37f9f7fd2688642dc0f2c9c37374f99e7ac1d42c1f9184521723a", "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/grpc-ecosystem/go-grpc-prometheus/com_github_grpc_ecosystem_go_grpc_prometheus-v1.2.0.zip": "124dfc63aa52611a2882417e685c0452d4d99d64c13836a6a6747675e911fc17", "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/grpc-ecosystem/grpc-gateway/com_github_grpc_ecosystem_grpc_gateway-v1.16.0.zip": "377b03aef288b34ed894449d3ddba40d525dd7fb55de6e79045cdf499e7fe565", + "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/grpc-ecosystem/grpc-gateway/v2/com_github_grpc_ecosystem_grpc_gateway_v2-v2.16.0.zip": "12995bc9a8d7a18ff3fb7fe3575dd2d9b7447d2f2749f0dbd99ef2efbaf9a6e2", "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/gsterjov/go-libsecret/com_github_gsterjov_go_libsecret-v0.0.0-20161001094733-a6f4afe4910c.zip": "cffe0a452fd3f00e4d07730caeb254417a720d907294b5b4a3428322655fb130", "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/guptarohit/asciigraph/com_github_guptarohit_asciigraph-v0.7.3.zip": "ec30034bd6d082f3242a5410ae1d02d9a4d164504e735f8448766461207be5a5", "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/hailocab/go-hostpool/com_github_hailocab_go_hostpool-v0.0.0-20160125115350-e80d13ce29ed.zip": "faf2b985681cda77ab928976b620b790585e364b6aff351483227d474db85e9a", @@ -1044,8 +1045,8 @@ DISTDIR_FILES = { "https://storage.googleapis.com/cockroach-godeps/gomod/go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc/io_opentelemetry_go_contrib_instrumentation_google_golang_org_grpc_otelgrpc-v0.28.0.zip": "056ab084a48e44c5fb5ba4e1970880570de71eac6133d23e4b55d009e71be2c3", "https://storage.googleapis.com/cockroach-godeps/gomod/go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp/io_opentelemetry_go_contrib_instrumentation_net_http_otelhttp-v0.65.0.zip": "f68a6e3f8f93b5857390583f3e855801ce93244ea18bf56ab45c221a4ea32a4c", "https://storage.googleapis.com/cockroach-godeps/gomod/go.opentelemetry.io/otel/exporters/otlp/internal/retry/io_opentelemetry_go_otel_exporters_otlp_internal_retry-v1.3.0.zip": "3ac72c80a4ef44c5df534587d5228891922f6f1978a10c736fac5f90b4050c73", - "https://storage.googleapis.com/cockroach-godeps/gomod/go.opentelemetry.io/otel/exporters/otlp/otlptrace/io_opentelemetry_go_otel_exporters_otlp_otlptrace-v1.3.0.zip": "1ff3c17fc607e5fd94c6e88127138fc2a2489fe31862b2b969e25caeb3d57db1", - "https://storage.googleapis.com/cockroach-godeps/gomod/go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc/io_opentelemetry_go_otel_exporters_otlp_otlptrace_otlptracegrpc-v1.3.0.zip": "f898ba889e06d82790ea6819243c3217bcfe1f01bcce4dbb497946843f93f83f", + "https://storage.googleapis.com/cockroach-godeps/gomod/go.opentelemetry.io/otel/exporters/otlp/otlptrace/io_opentelemetry_go_otel_exporters_otlp_otlptrace-v1.17.0.zip": "8c315880f1b7f58438a2fc7c1ed7edbafea6104724984eb7e564a730a249110c", + "https://storage.googleapis.com/cockroach-godeps/gomod/go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc/io_opentelemetry_go_otel_exporters_otlp_otlptrace_otlptracegrpc-v1.17.0.zip": "0965d989c01bad8217e2f937c54d82d200896f7e2df6888ce54813c6219b4bcc", "https://storage.googleapis.com/cockroach-godeps/gomod/go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp/io_opentelemetry_go_otel_exporters_otlp_otlptrace_otlptracehttp-v1.3.0.zip": "4b1bdf3c1d7f535acecb51f46be28636484bb18867c6e72ccf7743f689650994", "https://storage.googleapis.com/cockroach-godeps/gomod/go.opentelemetry.io/otel/exporters/zipkin/io_opentelemetry_go_otel_exporters_zipkin-v1.0.0-RC3.zip": "4e4074dc5fa0ae55cddbba06d4c266c23b0461bd5056ab69055f203b506fa64b", "https://storage.googleapis.com/cockroach-godeps/gomod/go.opentelemetry.io/otel/io_opentelemetry_go_otel-v1.40.0.zip": "8a284921d3c8f46b65a5c42105243e2d38a12edd592aa103b5e20a02bc96fa57", @@ -1053,7 +1054,7 @@ DISTDIR_FILES = { "https://storage.googleapis.com/cockroach-godeps/gomod/go.opentelemetry.io/otel/sdk/io_opentelemetry_go_otel_sdk-v1.40.0.zip": "f9efa73532d2644edda9a9f4a92c6941aed00bccf4d3801bc5771e250edad805", "https://storage.googleapis.com/cockroach-godeps/gomod/go.opentelemetry.io/otel/sdk/metric/io_opentelemetry_go_otel_sdk_metric-v1.40.0.zip": "25f84a504723feec88a7313e5c088c967cceb2b432e5038ba2cb22ce01e32517", "https://storage.googleapis.com/cockroach-godeps/gomod/go.opentelemetry.io/otel/trace/io_opentelemetry_go_otel_trace-v1.40.0.zip": "74c0bcedeb905aee6ef6b753141b7b5b38f817cd2dc6c7c5ced80929ea0408a2", - "https://storage.googleapis.com/cockroach-godeps/gomod/go.opentelemetry.io/proto/otlp/io_opentelemetry_go_proto_otlp-v0.11.0.zip": "f22d677bc272c65f45ca31b1ca80a28d1bdb922858e86fbd1579e1852fdb51d8", + "https://storage.googleapis.com/cockroach-godeps/gomod/go.opentelemetry.io/proto/otlp/io_opentelemetry_go_proto_otlp-v1.0.0.zip": "264e0db82a8c572ebc52fbe3da3f4a2a60d7e5fcf81238d1389390f66e2da246", "https://storage.googleapis.com/cockroach-godeps/gomod/go.uber.org/atomic/org_uber_go_atomic-v1.10.0.zip": "1a3a7303a5d7372db8184404b09f3142bf206e3a0001be468be2fc2540893d7a", "https://storage.googleapis.com/cockroach-godeps/gomod/go.uber.org/automaxprocs/org_uber_go_automaxprocs-v1.3.0.zip": "8f3ac8ce408b75928367ef26bbcb40dc98bbd197e2e9c51129859b2e6073542b", "https://storage.googleapis.com/cockroach-godeps/gomod/go.uber.org/goleak/org_uber_go_goleak-v1.3.0.zip": "70edef0ce7d830d992f024e527fd3452069b884f94a27787a718bd68dd620702", diff --git a/build/patches/com_github_envoyproxy_protoc_gen_validate.patch b/build/patches/com_github_envoyproxy_protoc_gen_validate.patch new file mode 100644 index 000000000000..a3c6e7689047 --- /dev/null +++ b/build/patches/com_github_envoyproxy_protoc_gen_validate.patch @@ -0,0 +1,14 @@ +diff --git a/validate/BUILD b/validate/BUILD +--- a/validate/BUILD ++++ b/validate/BUILD +@@ -34,5 +34,10 @@ + go_proto_library( + name = "validate_go_proto", + importpath = "github.com/envoyproxy/protoc-gen-validate/validate", + proto = ":validate_proto", ++ deps = [ ++ "@com_github_golang_protobuf//protoc-gen-go/descriptor:go_default_library", ++ "@com_github_golang_protobuf//ptypes/duration:go_default_library", ++ "@com_github_golang_protobuf//ptypes/timestamp:go_default_library", ++ ], + ) diff --git a/build/patches/com_github_grpc_ecosystem_grpc_gateway_v2.patch b/build/patches/com_github_grpc_ecosystem_grpc_gateway_v2.patch new file mode 100644 index 000000000000..474ee61fdf6a --- /dev/null +++ b/build/patches/com_github_grpc_ecosystem_grpc_gateway_v2.patch @@ -0,0 +1,26 @@ +diff --git a/runtime/BUILD.bazel b/runtime/BUILD.bazel +--- a/runtime/BUILD.bazel ++++ b/runtime/BUILD.bazel +@@ -26,7 +26,7 @@ go_library( + deps = [ + "//internal/httprule", + "//utilities", +- "@go_googleapis//google/api:httpbody_go_proto", ++ "@org_golang_google_genproto_googleapis_api//httpbody", + "@org_golang_google_grpc//codes", + "@org_golang_google_grpc//grpclog", + "@org_golang_google_grpc//health/grpc_health_v1", +@@ -71,9 +71,9 @@ go_test( + "//runtime/internal/examplepb", + "//utilities", + "@com_github_google_go_cmp//cmp", + "@com_github_google_go_cmp//cmp/cmpopts", +- "@go_googleapis//google/api:httpbody_go_proto", +- "@go_googleapis//google/rpc:errdetails_go_proto", +- "@go_googleapis//google/rpc:status_go_proto", ++ "@org_golang_google_genproto_googleapis_api//httpbody", ++ "@org_golang_google_genproto_googleapis_rpc//errdetails", ++ "@org_golang_google_genproto_googleapis_rpc//status", + "@org_golang_google_grpc//:go_default_library", + "@org_golang_google_grpc//codes", + "@org_golang_google_grpc//health/grpc_health_v1", diff --git a/go.mod b/go.mod index 208f24010a7c..82d6b145fbbc 100644 --- a/go.mod +++ b/go.mod @@ -90,6 +90,7 @@ require ( require ( cloud.google.com/go/compute v1.19.3 cloud.google.com/go/profiler v0.3.1 + cloud.google.com/go/spanner v1.45.0 github.com/Azure/azure-sdk-for-go/sdk/azcore v1.11.1 github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.7.0 github.com/Azure/azure-sdk-for-go/sdk/keyvault/azkeys v0.9.0 @@ -246,11 +247,11 @@ require ( github.com/zabawaba99/go-gitignore v0.0.0-20200117185801-39e6bddfb292 gitlab.com/golang-commonmark/markdown v0.0.0-20211110145824-bf3e522c626a go.opentelemetry.io/otel v1.40.0 - go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.3.0 + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.17.0 go.opentelemetry.io/otel/exporters/zipkin v1.0.0-RC3 go.opentelemetry.io/otel/sdk v1.40.0 go.opentelemetry.io/otel/trace v1.40.0 - go.opentelemetry.io/proto/otlp v0.11.0 + go.opentelemetry.io/proto/otlp v1.0.0 go.yaml.in/yaml/v4 v4.0.0-rc.3 golang.org/x/perf v0.0.0-20251112180420-cfbd823f7301 golang.org/x/term v0.37.0 @@ -314,11 +315,14 @@ require ( github.com/bufbuild/buf v0.56.0 // indirect github.com/bytedance/sonic v1.9.1 // indirect github.com/campoy/embedmd v1.0.0 // indirect - github.com/cenkalti/backoff/v4 v4.1.3 // indirect + github.com/cenkalti/backoff/v4 v4.2.1 // indirect + github.com/census-instrumentation/opencensus-proto v0.4.1 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/charmbracelet/bubbletea v0.23.1 // indirect github.com/charmbracelet/lipgloss v0.6.0 // indirect github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 // indirect + github.com/cncf/udpa/go v0.0.0-20220112060539-c52dc94e7fbe // indirect + github.com/cncf/xds/go v0.0.0-20230607035331-e9ce68804cb4 // indirect github.com/cockroachdb/crlfmt v0.0.0-20221214225007-b2fc5c302548 // indirect github.com/cockroachdb/go-test-teamcity v0.0.0-20191211140407-cff980ad0a55 // indirect github.com/cockroachdb/stress v0.0.0-20220803192808-1806698b1b7b // indirect @@ -336,6 +340,7 @@ require ( github.com/eapache/go-resiliency v1.6.0 // indirect github.com/eapache/go-xerial-snappy v0.0.0-20230731223053-c322873962e3 // indirect github.com/eapache/queue v1.1.0 // indirect + github.com/envoyproxy/go-control-plane v0.11.1-0.20230524094728-9239064ad72f // indirect github.com/felixge/httpsnoop v1.0.4 // indirect github.com/form3tech-oss/jwt-go v3.2.5+incompatible // indirect github.com/gabriel-vasile/mimetype v1.4.8 // indirect @@ -377,6 +382,7 @@ require ( github.com/gorilla/websocket v1.4.2 // indirect github.com/goware/modvendor v0.5.0 // indirect github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0 // indirect + github.com/grpc-ecosystem/grpc-gateway/v2 v2.16.0 // indirect github.com/gsterjov/go-libsecret v0.0.0-20161001094733-a6f4afe4910c // indirect github.com/hashicorp/errwrap v1.1.0 // indirect github.com/hashicorp/go-cleanhttp v0.5.2 // indirect @@ -484,7 +490,6 @@ require ( go.opencensus.io v0.24.0 // indirect go.opentelemetry.io/auto/sdk v1.2.1 // indirect go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.65.0 // indirect - go.opentelemetry.io/otel/exporters/otlp/internal/retry v1.3.0 // indirect go.opentelemetry.io/otel/metric v1.40.0 // indirect go.uber.org/atomic v1.10.0 // indirect go.uber.org/multierr v1.7.0 // indirect @@ -514,9 +519,7 @@ require ( github.com/muesli/ansi v0.0.0-20211031195517-c9f0611b6c70 // indirect github.com/muesli/cancelreader v0.2.2 // indirect github.com/muesli/reflow v0.3.0 // indirect - // The indicated commit is required on top of v1.0.0-RC3 because - // it fixes an import comment that otherwise breaks our prereqs tool. - go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.3.0 + go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.17.0 google.golang.org/grpc/examples v0.0.0-20210324172016-702608ffae4d // indirect gopkg.in/ini.v1 v1.67.0 // indirect ) diff --git a/go.sum b/go.sum index a29ddec2ad0a..a1f7de6ade46 100644 --- a/go.sum +++ b/go.sum @@ -61,6 +61,8 @@ cloud.google.com/go/pubsub v1.2.0/go.mod h1:jhfEVHT8odbXTkndysNHCcx0awwzvfOlguIA cloud.google.com/go/pubsub v1.3.1/go.mod h1:i+ucay31+CNRpDW4Lu78I4xXG+O1r/MAHgjpRVR+TSU= cloud.google.com/go/pubsub v1.30.0 h1:vCge8m7aUKBJYOgrZp7EsNDf6QMd2CAlXZqWTn3yq6s= cloud.google.com/go/pubsub v1.30.0/go.mod h1:qWi1OPS0B+b5L+Sg6Gmc9zD1Y+HaM0MdUr7LsupY1P4= +cloud.google.com/go/spanner v1.45.0 h1:7VdjZ8zj4sHbDw55atp5dfY6kn1j9sam9DRNpPQhqR4= +cloud.google.com/go/spanner v1.45.0/go.mod h1:FIws5LowYz8YAE1J8fOS7DJup8ff7xJeetWEo5REA2M= cloud.google.com/go/storage v1.0.0/go.mod h1:IhtSnM/ZTZV8YYJWCY8RULGVqBDmpoyjwiyrjsg+URw= cloud.google.com/go/storage v1.5.0/go.mod h1:tpKbwo567HUNpVclU5sGELwQWBDZ8gh0ZeosJ0Rtdos= cloud.google.com/go/storage v1.6.0/go.mod h1:N7U0C8pVQ/+NIKOBQyamJIeKQKkZ+mxpohlUTyfDhBk= @@ -497,10 +499,12 @@ github.com/casbin/casbin/v2 v2.1.2/go.mod h1:YcPU1XXisHhLzuxH9coDNf2FbKpjGlbCg3n github.com/cenkalti/backoff v0.0.0-20181003080854-62661b46c409/go.mod h1:90ReRw6GdpyfrHakVjL/QHaoyV4aDUVVkXQJJJ3NXXM= github.com/cenkalti/backoff v2.2.1+incompatible/go.mod h1:90ReRw6GdpyfrHakVjL/QHaoyV4aDUVVkXQJJJ3NXXM= github.com/cenkalti/backoff/v4 v4.1.1/go.mod h1:scbssz8iZGpm3xbr14ovlUdkxfGXNInqkPWOWmG2CLw= -github.com/cenkalti/backoff/v4 v4.1.2/go.mod h1:scbssz8iZGpm3xbr14ovlUdkxfGXNInqkPWOWmG2CLw= -github.com/cenkalti/backoff/v4 v4.1.3 h1:cFAlzYUlVYDysBEH2T5hyJZMh3+5+WCBvSnK6Q8UtC4= github.com/cenkalti/backoff/v4 v4.1.3/go.mod h1:scbssz8iZGpm3xbr14ovlUdkxfGXNInqkPWOWmG2CLw= +github.com/cenkalti/backoff/v4 v4.2.1 h1:y4OZtCnogmCPw98Zjyt5a6+QwPLGkiQsYW5oUqylYbM= +github.com/cenkalti/backoff/v4 v4.2.1/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE= github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU= +github.com/census-instrumentation/opencensus-proto v0.4.1 h1:iKLQ0xPNFxR/2hzXZMrBo8f1j86j5WHzznCCQxV/b8g= +github.com/census-instrumentation/opencensus-proto v0.4.1/go.mod h1:4T9NM4+4Vw91VeyqjLS6ao50K5bOcLKN6Q42XnYaRYw= github.com/cespare/xxhash v1.1.0/go.mod h1:XrSqR1VqqWfGrhpAt58auRo0WTKS1nRRg3ghfAqPWnc= github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cespare/xxhash/v2 v2.1.2/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= @@ -544,11 +548,15 @@ github.com/cncf/udpa/go v0.0.0-20191209042840-269d4d468f6f/go.mod h1:M8M6+tZqaGX github.com/cncf/udpa/go v0.0.0-20200629203442-efcf912fb354/go.mod h1:WmhPx2Nbnhtbo57+VJT5O0JRkEi1Wbu0z5j0R8u5Hbk= github.com/cncf/udpa/go v0.0.0-20201120205902-5459f2c99403/go.mod h1:WmhPx2Nbnhtbo57+VJT5O0JRkEi1Wbu0z5j0R8u5Hbk= github.com/cncf/udpa/go v0.0.0-20210930031921-04548b0d99d4/go.mod h1:6pvJx4me5XPnfI9Z40ddWsdw2W/uZgQLFXToKeRcDiI= +github.com/cncf/udpa/go v0.0.0-20220112060539-c52dc94e7fbe h1:QQ3GSy+MqSHxm/d8nCtnAiZdYFd45cYZPs8vOOIYKfk= +github.com/cncf/udpa/go v0.0.0-20220112060539-c52dc94e7fbe/go.mod h1:6pvJx4me5XPnfI9Z40ddWsdw2W/uZgQLFXToKeRcDiI= github.com/cncf/xds/go v0.0.0-20210312221358-fbca930ec8ed/go.mod h1:eXthEFrGJvWHgFFCl3hGmgk+/aYT6PnTQLykKQRLhEs= github.com/cncf/xds/go v0.0.0-20210805033703-aa0b78936158/go.mod h1:eXthEFrGJvWHgFFCl3hGmgk+/aYT6PnTQLykKQRLhEs= github.com/cncf/xds/go v0.0.0-20210922020428-25de7278fc84/go.mod h1:eXthEFrGJvWHgFFCl3hGmgk+/aYT6PnTQLykKQRLhEs= github.com/cncf/xds/go v0.0.0-20211001041855-01bcc9b48dfe/go.mod h1:eXthEFrGJvWHgFFCl3hGmgk+/aYT6PnTQLykKQRLhEs= github.com/cncf/xds/go v0.0.0-20211011173535-cb28da3451f1/go.mod h1:eXthEFrGJvWHgFFCl3hGmgk+/aYT6PnTQLykKQRLhEs= +github.com/cncf/xds/go v0.0.0-20230607035331-e9ce68804cb4 h1:/inchEIKaYC1Akx+H+gqO04wryn5h75LSazbRlnya1k= +github.com/cncf/xds/go v0.0.0-20230607035331-e9ce68804cb4/go.mod h1:eXthEFrGJvWHgFFCl3hGmgk+/aYT6PnTQLykKQRLhEs= github.com/cockroachdb/apd v1.1.0 h1:3LFP3629v+1aKXU5Q37mxmRxX/pIu1nijXydLShEq5I= github.com/cockroachdb/apd v1.1.0/go.mod h1:8Sl8LxpKi29FqWXR16WEFZRNSz3SoPzUzeMeY4+DwBQ= github.com/cockroachdb/apd/v3 v3.2.1 h1:U+8j7t0axsIgvQUqthuNm82HIrYXodOV2iWLWtEaIwg= @@ -840,6 +848,8 @@ github.com/envoyproxy/go-control-plane v0.9.9-0.20210512163311-63b5d3c536b0/go.m github.com/envoyproxy/go-control-plane v0.9.9/go.mod h1:hliV/p42l8fGbc6Y9bQ70uLwIvmJyVE5k4iMKlh8wCQ= github.com/envoyproxy/go-control-plane v0.9.10-0.20210907150352-cf90f659a021/go.mod h1:AFq3mo9L8Lqqiid3OhADV3RfLJnjiw63cSpi+fDTRC0= github.com/envoyproxy/go-control-plane v0.10.2-0.20220325020618-49ff273808a1/go.mod h1:KJwIaB5Mv44NWtYuAOFCVOjcI94vtpEz2JU/D2v6IjE= +github.com/envoyproxy/go-control-plane v0.11.1-0.20230524094728-9239064ad72f h1:7T++XKzy4xg7PKy+bM+Sa9/oe1OC88yz2hXQUISoXfA= +github.com/envoyproxy/go-control-plane v0.11.1-0.20230524094728-9239064ad72f/go.mod h1:sfYdkwUW4BA3PbKjySwjJy+O4Pu0h62rlqCMHNk+K+Q= github.com/envoyproxy/protoc-gen-validate v0.0.14/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c= github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c= github.com/envoyproxy/protoc-gen-validate v0.6.1/go.mod h1:txg5va2Qkip90uYoSKH+nkAAmXrb2j3iq4FLwdrCbXQ= @@ -944,12 +954,9 @@ github.com/go-logfmt/logfmt v0.5.1/go.mod h1:WYhtIu8zTZfxdn5+rREduYbwxfcBr/Vr6KE github.com/go-logr/logr v0.1.0/go.mod h1:ixOQHD9gLJUVQQ2ZOR7zLEifBX6tGkNJF4QyIY7sIas= github.com/go-logr/logr v0.2.0/go.mod h1:z6/tIYblkpsD+a4lm/fGIIU9mZ+XfAiaFtq7xTgseGU= github.com/go-logr/logr v0.4.0/go.mod h1:z6/tIYblkpsD+a4lm/fGIIU9mZ+XfAiaFtq7xTgseGU= -github.com/go-logr/logr v1.2.0/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= -github.com/go-logr/logr v1.2.1/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= -github.com/go-logr/stdr v1.2.0/go.mod h1:YkVgnZu1ZjjL7xTxrfm/LLZBfkhTqSR1ydtm6jTKKwI= github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= github.com/go-ole/go-ole v1.2.6 h1:/Fpf6oFPoeFik9ty7siob0G6Ke8QvQEuVcuChpwXzpY= @@ -1353,6 +1360,8 @@ github.com/grpc-ecosystem/grpc-gateway v1.13.0/go.mod h1:8XEsbTttt/W+VvjtQhLACqC github.com/grpc-ecosystem/grpc-gateway v1.14.4/go.mod h1:6CwZWGDSPRJidgKAtJVvND6soZe6fT7iteq8wDPdhb0= github.com/grpc-ecosystem/grpc-gateway v1.16.0 h1:gmcG1KaJ57LophUzW0Hy8NmPhnMZb4M0+kPpLofRdBo= github.com/grpc-ecosystem/grpc-gateway v1.16.0/go.mod h1:BDjrQk3hbvj6Nolgz8mAMFbcEtjT1g+wF4CSlocrBnw= +github.com/grpc-ecosystem/grpc-gateway/v2 v2.16.0 h1:YBftPWNWd4WwGqtY2yeZL2ef8rHAxPBD8KFhJpmcqms= +github.com/grpc-ecosystem/grpc-gateway/v2 v2.16.0/go.mod h1:YN5jB8ie0yfIUg6VvR9Kz84aCaG7AsGZnLjhHbUqwPg= github.com/gsterjov/go-libsecret v0.0.0-20161001094733-a6f4afe4910c h1:6rhixN/i8ZofjG1Y75iExal34USq5p+wiN1tpie8IrU= github.com/gsterjov/go-libsecret v0.0.0-20161001094733-a6f4afe4910c/go.mod h1:NMPJylDgVpX0MLRlPy15sqSwOFv/U1GZ2m21JhFfek0= github.com/guptarohit/asciigraph v0.7.3 h1:p05XDDn7cBTWiBqWb30mrwxd6oU0claAjqeytllnsPY= @@ -2406,15 +2415,14 @@ go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbE go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.65.0 h1:7iP2uCb7sGddAr30RRS6xjKy7AZ2JtTOPA3oolgVSw8= go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.65.0/go.mod h1:c7hN3ddxs/z6q9xwvfLPk+UHlWRQyaeR1LdgfL/66l0= go.opentelemetry.io/otel v1.0.0-RC3/go.mod h1:Ka5j3ua8tZs4Rkq4Ex3hwgBgOchyPVq5S6P2lz//nKQ= -go.opentelemetry.io/otel v1.3.0/go.mod h1:PWIKzi6JCp7sM0k9yZ43VX+T345uNbAkDKwHVjb2PTs= go.opentelemetry.io/otel v1.40.0 h1:oA5YeOcpRTXq6NN7frwmwFR0Cn3RhTVZvXsP4duvCms= go.opentelemetry.io/otel v1.40.0/go.mod h1:IMb+uXZUKkMXdPddhwAHm6UfOwJyh4ct1ybIlV14J0g= go.opentelemetry.io/otel/exporters/otlp/internal/retry v1.3.0 h1:R/OBkMoGgfy2fLhs2QhkCI1w4HLEQX92GCcJB6SSdNk= go.opentelemetry.io/otel/exporters/otlp/internal/retry v1.3.0/go.mod h1:VpP4/RMn8bv8gNo9uK7/IMY4mtWLELsS+JIP0inH0h4= -go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.3.0 h1:giGm8w67Ja7amYNfYMdme7xSp2pIxThWopw8+QP51Yk= -go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.3.0/go.mod h1:hO1KLR7jcKaDDKDkvI9dP/FIhpmna5lkqPUQdEjFAM8= -go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.3.0 h1:VQbUHoJqytHHSJ1OZodPH9tvZZSVzUHjPHpkO85sT6k= -go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.3.0/go.mod h1:keUU7UfnwWTWpJ+FWnyqmogPa82nuU5VUANFq49hlMY= +go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.17.0 h1:U5GYackKpVKlPrd/5gKMlrTlP2dCESAAFU682VCpieY= +go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.17.0/go.mod h1:aFsJfCEnLzEu9vRRAcUiB/cpRTbVsNdF3OHSPpdjxZQ= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.17.0 h1:iGeIsSYwpYSvh5UGzWrJfTDJvPjrXtxl3GUppj6IXQU= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.17.0/go.mod h1:1j3H3G1SBYpZFti6OI4P0uRQCW20MXkG5v4UWXppLLE= go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.3.0 h1:Ydage/P0fRrSPpZeCVxzjqGcI6iVmG2xb43+IR8cjqM= go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.3.0/go.mod h1:QNX1aly8ehqqX1LEa6YniTU7VY9I6R3X/oPxhGdTceE= go.opentelemetry.io/otel/exporters/zipkin v1.0.0-RC3 h1:NcjZQwPFW0V2OvasQQJWSOTGA6R+4SLLgcPGIiDC3fk= @@ -2422,18 +2430,16 @@ go.opentelemetry.io/otel/exporters/zipkin v1.0.0-RC3/go.mod h1:3ASq7NNNoDb2P6dIB go.opentelemetry.io/otel/metric v1.40.0 h1:rcZe317KPftE2rstWIBitCdVp89A2HqjkxR3c11+p9g= go.opentelemetry.io/otel/metric v1.40.0/go.mod h1:ib/crwQH7N3r5kfiBZQbwrTge743UDc7DTFVZrrXnqc= go.opentelemetry.io/otel/sdk v1.0.0-RC3/go.mod h1:78H6hyg2fka0NYT9fqGuFLvly2yCxiBXDJAgLKo/2Us= -go.opentelemetry.io/otel/sdk v1.3.0/go.mod h1:rIo4suHNhQwBIPg9axF8V9CA72Wz2mKF1teNrup8yzs= go.opentelemetry.io/otel/sdk v1.40.0 h1:KHW/jUzgo6wsPh9At46+h4upjtccTmuZCFAc9OJ71f8= go.opentelemetry.io/otel/sdk v1.40.0/go.mod h1:Ph7EFdYvxq72Y8Li9q8KebuYUr2KoeyHx0DRMKrYBUE= go.opentelemetry.io/otel/sdk/metric v1.40.0 h1:mtmdVqgQkeRxHgRv4qhyJduP3fYJRMX4AtAlbuWdCYw= go.opentelemetry.io/otel/sdk/metric v1.40.0/go.mod h1:4Z2bGMf0KSK3uRjlczMOeMhKU2rhUqdWNoKcYrtcBPg= go.opentelemetry.io/otel/trace v1.0.0-RC3/go.mod h1:VUt2TUYd8S2/ZRX09ZDFZQwn2RqfMB5MzO17jBojGxo= -go.opentelemetry.io/otel/trace v1.3.0/go.mod h1:c/VDhno8888bvQYmbYLqe41/Ldmr/KKunbvWM4/fEjk= go.opentelemetry.io/otel/trace v1.40.0 h1:WA4etStDttCSYuhwvEa8OP8I5EWu24lkOzp+ZYblVjw= go.opentelemetry.io/otel/trace v1.40.0/go.mod h1:zeAhriXecNGP/s2SEG3+Y8X9ujcJOTqQ5RgdEJcawiA= go.opentelemetry.io/proto/otlp v0.7.0/go.mod h1:PqfVotwruBrMGOCsRd/89rSnXhoiJIqeYNgFYFoEGnI= -go.opentelemetry.io/proto/otlp v0.11.0 h1:cLDgIBTf4lLOlztkhzAEdQsJ4Lj+i5Wc9k6Nn0K1VyU= -go.opentelemetry.io/proto/otlp v0.11.0/go.mod h1:QpEjXPrNQzrFDZgoTo49dgHR9RYRSrg3NAKnUGl9YpQ= +go.opentelemetry.io/proto/otlp v1.0.0 h1:T0TX0tmXU8a3CbNXzEKGeU5mIVOdf0oykP+u2lIVU/I= +go.opentelemetry.io/proto/otlp v1.0.0/go.mod h1:Sy6pihPLfYHkr3NkUbEhGHFhINUSI/v80hjKIs5JXpM= go.uber.org/atomic v1.3.2/go.mod h1:gD2HeocX3+yG+ygLZcrzQJaqmWj9AIm7n08wl/qW/PE= go.uber.org/atomic v1.4.0/go.mod h1:gD2HeocX3+yG+ygLZcrzQJaqmWj9AIm7n08wl/qW/PE= go.uber.org/atomic v1.5.0/go.mod h1:sABNBOSYdrvTF6hTgEIbc7YasKWGhgEQZyfxyTvoXHQ= @@ -2445,7 +2451,6 @@ go.uber.org/atomic v1.10.0 h1:9qC72Qh0+3MqyJbAn8YU5xVq1frD8bn3JtD2oXtafVQ= go.uber.org/atomic v1.10.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0= go.uber.org/automaxprocs v1.3.0/go.mod h1:9CWT6lKIep8U41DDaPiH6eFscnTyjfTANNQNx6LrIcA= go.uber.org/goleak v1.1.10/go.mod h1:8a7PlsEVH3e/a/GLqe5IIrQx6GzcnRmZEufDUTk4A7A= -go.uber.org/goleak v1.1.12/go.mod h1:cwTWslyiVhfpKIDGSZEM2HlOvcqm+tG4zioyIeLoqMQ= go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= go.uber.org/multierr v1.1.0/go.mod h1:wR5kodmAFQ0UK8QlbwjlSNy0Z68gJhDJUG5sjR94q/0= @@ -3192,7 +3197,6 @@ google.golang.org/grpc v1.39.0/go.mod h1:PImNr+rS9TWYb2O4/emRugxiyHZ5JyHW5F+RPnD google.golang.org/grpc v1.39.1/go.mod h1:PImNr+rS9TWYb2O4/emRugxiyHZ5JyHW5F+RPnDzfrE= google.golang.org/grpc v1.40.0/go.mod h1:ogyxbiOoUXAkP+4+xa6PZSE9DZgIHtSpzjDTB9KAK34= google.golang.org/grpc v1.41.0-dev.0.20210907181116-2f3355d2244e/go.mod h1:U3l9uK9J0sini8mHphKoXyaqDA/8VyGnDee1zzIUK6k= -google.golang.org/grpc v1.42.0/go.mod h1:k+4IHHFw41K8+bbowsex27ge2rCb65oeWqe4jJ590SU= google.golang.org/grpc v1.45.0/go.mod h1:lN7owxKUQEqMfSyQikvvk5tf/6zMPsrK+ONuO11+0rQ= google.golang.org/grpc v1.49.0/go.mod h1:ZgQEeidpAuNRZ8iRrlBKXZQP1ghovWIVhdJRyCDK+GI= google.golang.org/grpc v1.57.2 h1:uw37EN34aMFFXB2QPW7Tq6tdTbind1GpRxw5aOX3a5k= diff --git a/pkg/BUILD.bazel b/pkg/BUILD.bazel index 4fbb52042ffc..316320de2503 100644 --- a/pkg/BUILD.bazel +++ b/pkg/BUILD.bazel @@ -887,6 +887,7 @@ ALL_TESTS = [ "//pkg/workload/bank:bank_test", "//pkg/workload/cli:cli_test", "//pkg/workload/conflict:conflict_test", + "//pkg/workload/dbexec:dbexec_test", "//pkg/workload/faker:faker_test", "//pkg/workload/histogram/exporter:exporter_test", "//pkg/workload/histogram:histogram_test", @@ -3056,6 +3057,8 @@ GO_TARGETS = [ "//pkg/workload/conflict:conflict", "//pkg/workload/conflict:conflict_test", "//pkg/workload/connectionlatency:connectionlatency", + "//pkg/workload/dbexec:dbexec", + "//pkg/workload/dbexec:dbexec_test", "//pkg/workload/debug:debug", "//pkg/workload/examples:examples", "//pkg/workload/faker:faker", diff --git a/pkg/cmd/mirror/go/mirror.go b/pkg/cmd/mirror/go/mirror.go index 74bc6ad8ffe0..72f931dc24ae 100644 --- a/pkg/cmd/mirror/go/mirror.go +++ b/pkg/cmd/mirror/go/mirror.go @@ -317,6 +317,14 @@ func dumpBuildDirectivesForRepo(repoName string) { } else if repoName == "io_opentelemetry_go_proto_otlp" { directives = append(directives, "gazelle:resolve go go github.com/golang/protobuf/descriptor @com_github_golang_protobuf//descriptor:go_default_library_gen") + } else if repoName == "com_github_census_instrumentation_opencensus_proto" { + directives = append(directives, + "gazelle:resolve go go github.com/census-instrumentation/opencensus-proto/gen-go/resource/v1 @com_github_census_instrumentation_opencensus_proto//gen-go/resource/v1:go_default_library") + } else if repoName == "com_github_envoyproxy_protoc_gen_validate" { + directives = append(directives, + "gazelle:resolve go go github.com/golang/protobuf/protoc-gen-go/descriptor @com_github_golang_protobuf//protoc-gen-go/descriptor:go_default_library", + "gazelle:resolve go go github.com/golang/protobuf/ptypes/duration @com_github_golang_protobuf//ptypes/duration:go_default_library", + "gazelle:resolve go go github.com/golang/protobuf/ptypes/timestamp @com_github_golang_protobuf//ptypes/timestamp:go_default_library") } if len(directives) > 0 { diff --git a/pkg/util/log/otlp_client.go b/pkg/util/log/otlp_client.go index d594e6b419d8..2cef5ab1ea03 100644 --- a/pkg/util/log/otlp_client.go +++ b/pkg/util/log/otlp_client.go @@ -81,9 +81,9 @@ func newOTLPSink(config logconfig.OTLPSinkConfig) (*otlpSink, error) { }, }, }, - InstrumentationLibraryLogs: []*lpb.InstrumentationLibraryLogs{ + ScopeLogs: []*lpb.ScopeLogs{ { - Logs: nil, + LogRecords: nil, }, }, }, @@ -147,7 +147,7 @@ func (sink *otlpSink) output(b []byte, opts sinkOutputOptions) error { ctx := context.Background() records := otlpExtractRecords(b) - sink.requestObject.ResourceLogs[0].InstrumentationLibraryLogs[0].Logs = records + sink.requestObject.ResourceLogs[0].ScopeLogs[0].LogRecords = records // transmit the log over the network _, err := sink.client.Export(ctx, sink.requestObject) diff --git a/pkg/util/log/otlp_client_test.go b/pkg/util/log/otlp_client_test.go index d3f3576facc0..987dfa8464cf 100644 --- a/pkg/util/log/otlp_client_test.go +++ b/pkg/util/log/otlp_client_test.go @@ -155,9 +155,9 @@ func TestOTLPSink(t *testing.T) { check: func(t *testing.T, reqs []*collpb.ExportLogsServiceRequest) { require.Len(t, reqs, 1) require.Len(t, reqs[0].ResourceLogs, 1) - require.Len(t, reqs[0].ResourceLogs[0].InstrumentationLibraryLogs, 1) - require.Len(t, reqs[0].ResourceLogs[0].InstrumentationLibraryLogs[0].Logs, 1) - logRecord := reqs[0].ResourceLogs[0].InstrumentationLibraryLogs[0].Logs[0] + require.Len(t, reqs[0].ResourceLogs[0].ScopeLogs, 1) + require.Len(t, reqs[0].ResourceLogs[0].ScopeLogs[0].LogRecords, 1) + logRecord := reqs[0].ResourceLogs[0].ScopeLogs[0].LogRecords[0] data, err := parseJSON(logRecord.Body.GetStringValue()) require.NoError(t, err) @@ -189,7 +189,7 @@ func TestOTLPSink(t *testing.T) { check: func(t *testing.T, reqs []*collpb.ExportLogsServiceRequest) { severities := []string{"INFO", "WARNING", "ERROR"} for i, sev := range severities { - logRecord := reqs[i].ResourceLogs[0].InstrumentationLibraryLogs[0].Logs[0] + logRecord := reqs[i].ResourceLogs[0].ScopeLogs[0].LogRecords[0] data, err := parseJSON(logRecord.Body.GetStringValue()) require.NoError(t, err) diff --git a/pkg/workload/cli/BUILD.bazel b/pkg/workload/cli/BUILD.bazel index 34a2649d64dd..11e29d4d6859 100644 --- a/pkg/workload/cli/BUILD.bazel +++ b/pkg/workload/cli/BUILD.bazel @@ -11,6 +11,8 @@ go_library( "run.go", "run_unix.go", "run_windows.go", + "spanner_init.go", + "sql_init.go", ], importpath = "github.com/cockroachdb/cockroach/pkg/workload/cli", visibility = ["//visibility:public"], @@ -30,10 +32,15 @@ go_library( "//pkg/workload/workloadsql", "@com_github_cockroachdb_errors//:errors", "@com_github_datadog_datadog_api_client_go_v2//api/datadogV1", + "@com_github_jackc_pgx_v5//stdlib", "@com_github_prometheus_client_golang//prometheus/collectors", "@com_github_prometheus_client_golang//prometheus/promhttp", "@com_github_spf13_cobra//:cobra", "@com_github_spf13_pflag//:pflag", + "@com_google_cloud_go_spanner//:spanner", + "@com_google_cloud_go_spanner//admin/database/apiv1", + "@com_google_cloud_go_spanner//admin/database/apiv1/databasepb", + "@org_golang_x_sync//errgroup", "@org_golang_x_time//rate", ] + select({ "@io_bazel_rules_go//go/platform:aix": [ diff --git a/pkg/workload/cli/run.go b/pkg/workload/cli/run.go index b1735838cb9e..cab9cbdbfba7 100644 --- a/pkg/workload/cli/run.go +++ b/pkg/workload/cli/run.go @@ -57,6 +57,8 @@ var drop = initFlags.Bool("drop", false, "Drop the existing database, if it exis var sharedFlags = pflag.NewFlagSet(`shared`, pflag.ContinueOnError) var pprofport = sharedFlags.Int("pprofport", 33333, "Port for pprof endpoint.") +var dialect = sharedFlags.String("dialect", "crdb", + "Database dialect: crdb (default), postgres, aurora, dsql, spanner") var dataLoader = sharedFlags.String("data-loader", `AUTO`, "How to load initial table data. All workloads support INSERT; some support IMPORT, which AUTO prefers if available.") var initConns = sharedFlags.Int("init-conns", 16, @@ -217,6 +219,18 @@ func CmdHelper( } } + dialectName := strings.ToLower(*dialect) + if dialectName == "" { + dialectName = "crdb" + } + if ds, ok := gen.(workload.DialectSetter); ok { + if err := ds.SetDialect(dialectName); err != nil { + return errors.Wrapf(err, "could not set dialect") + } + } else if dialectName != "crdb" { + return errors.Errorf("workload %s does not support --dialect=%s", gen.Meta().Name, dialectName) + } + if h, ok := gen.(workload.Hookser); ok { if h.Hooks().Validate != nil { if err := h.Hooks().Validate(); err != nil { @@ -330,13 +344,30 @@ func workerRun( func runInit(gen workload.Generator, urls []string, dbName string) error { ctx := context.Background() + startPProfEndPoint(ctx) + maybeLogRandomSeed(ctx, gen) + + dialectName := strings.ToLower(*dialect) + if dialectName == "" { + dialectName = "crdb" + } + switch dialectName { + case "spanner": + if len(urls) != 1 { + return errors.New("spanner init requires exactly one database URL") + } + return runSpannerInit(ctx, gen, urls[0], *drop) + case "postgres", "aurora", "dsql": + return runSQLInitNonCRDB(ctx, gen, urls) + case "crdb": + default: + return errors.Errorf("unknown dialect: %s", dialectName) + } + initDB, err := gosql.Open(`cockroach`, strings.Join(urls, ` `)) if err != nil { return err } - - startPProfEndPoint(ctx) - maybeLogRandomSeed(ctx, gen) return runInitImpl(ctx, gen, initDB, dbName) } diff --git a/pkg/workload/cli/spanner_init.go b/pkg/workload/cli/spanner_init.go new file mode 100644 index 000000000000..721486f96f64 --- /dev/null +++ b/pkg/workload/cli/spanner_init.go @@ -0,0 +1,323 @@ +// Copyright 2024 The Cockroach Authors. +// +// Use of this software is governed by the CockroachDB Software License +// included in the /LICENSE file. + +package cli + +import ( + "context" + "fmt" + "strings" + + "cloud.google.com/go/spanner" + database "cloud.google.com/go/spanner/admin/database/apiv1" + databasepb "cloud.google.com/go/spanner/admin/database/apiv1/databasepb" + "github.com/cockroachdb/cockroach/pkg/util/log" + "github.com/cockroachdb/cockroach/pkg/workload" + "github.com/cockroachdb/errors" + "golang.org/x/sync/errgroup" +) + +func runSpannerInit(ctx context.Context, gen workload.Generator, dbPath string, drop bool) error { + lc := strings.ToLower(*dataLoader) + if lc == "auto" { + lc = "insert" + } + if lc != "insert" { + return errors.Errorf("data loader %q is not supported for dialect %q", *dataLoader, *dialect) + } + + admin, err := database.NewDatabaseAdminClient(ctx) + if err != nil { + return errors.Wrap(err, "creating spanner database admin client") + } + defer func() { + err := admin.Close() + if err != nil { + log.Dev.Errorf(ctx, "error closing spanner admin client: %v", err) + } + }() + + existingTables, existingIndexes, err := fetchSpannerSchema(ctx, admin, dbPath) + if err != nil { + return err + } + + tables := gen.Tables() + if drop { + dropStmts := make([]string, 0, len(tables)) + for _, table := range tables { + if existingTables[table.Name] { + dropStmts = append(dropStmts, fmt.Sprintf("DROP TABLE %s", table.Name)) + } + } + if len(dropStmts) > 0 { + if err := applySpannerDDL(ctx, admin, dbPath, dropStmts); err != nil { + return errors.Wrap(err, "dropping Spanner tables") + } + for _, table := range tables { + delete(existingTables, table.Name) + } + existingIndexes = map[string]bool{} + } + } + + createStmts := make([]string, 0, len(tables)) + indexStmts := make([]string, 0, len(tables)) + for _, table := range tables { + if existingTables[table.Name] { + continue + } + createStmts = append(createStmts, fmt.Sprintf("CREATE TABLE %s %s", table.Name, table.Schema)) + } + for _, table := range tables { + for _, stmt := range table.Indexes { + name, ok := spannerDDLIndexName(stmt) + if !ok { + return errors.Errorf("invalid Spanner index statement for table %q: %s", table.Name, stmt) + } + if existingIndexes[name] { + continue + } + indexStmts = append(indexStmts, stmt) + } + } + if len(createStmts)+len(indexStmts) > 0 { + statements := append(createStmts, indexStmts...) + if err := applySpannerDDL(ctx, admin, dbPath, statements); err != nil { + return errors.Wrap(err, "creating Spanner tables") + } + } + + if err := loadSpannerInitialData(ctx, dbPath, tables); err != nil { + return err + } + return nil +} + +const spannerInitMaxMutations = 200 + +func loadSpannerInitialData(ctx context.Context, dbPath string, tables []workload.Table) error { + needsData := false + for _, table := range tables { + if table.InitialRows.NumBatches > 0 { + needsData = true + break + } + } + if !needsData { + return nil + } + + client, err := spanner.NewClient(ctx, dbPath) + if err != nil { + return errors.Wrap(err, "creating spanner data client") + } + defer client.Close() + + for _, table := range tables { + if table.InitialRows.NumBatches == 0 { + continue + } + if table.InitialRows.FillBatch == nil { + return errors.Errorf("initial data is not supported for workload %s", table.Name) + } + if table.ObjectPrefix != nil && table.ObjectPrefix.ExplicitCatalog { + return errors.Errorf("multi-database table %q is not supported for Spanner init", table.Name) + } + columns, err := spannerColumnsFromSchema(table.Schema) + if err != nil { + return errors.Wrapf(err, "parsing columns for table %q", table.Name) + } + + workers := *initConns + if workers < 1 { + workers = 1 + } + batchesPerWorker := table.InitialRows.NumBatches / workers + g, gCtx := errgroup.WithContext(ctx) + for i := 0; i < workers; i++ { + startIdx := i * batchesPerWorker + endIdx := startIdx + batchesPerWorker + if i == workers-1 { + endIdx = table.InitialRows.NumBatches + } + table := table + g.Go(func() error { + mutations := make([]*spanner.Mutation, 0, spannerInitMaxMutations) + flush := func() error { + if len(mutations) == 0 { + return nil + } + if _, err := client.Apply(gCtx, mutations); err != nil { + return err + } + mutations = mutations[:0] + return nil + } + for batchIdx := startIdx; batchIdx < endIdx; batchIdx++ { + for _, row := range table.InitialRows.BatchRows(batchIdx) { + if len(row) != len(columns) { + return errors.Errorf("row/column mismatch for table %q: %d values, %d columns", + table.Name, len(row), len(columns)) + } + mutations = append(mutations, spanner.InsertOrUpdate(table.Name, columns, row)) + if len(mutations) >= spannerInitMaxMutations { + if err := flush(); err != nil { + return err + } + } + } + } + return flush() + }) + } + if err := g.Wait(); err != nil { + return errors.Wrapf(err, "loading initial data for table %q", table.Name) + } + } + return nil +} + +func spannerColumnsFromSchema(schema string) ([]string, error) { + start := strings.Index(schema, "(") + if start == -1 { + return nil, errors.New("schema missing opening parenthesis") + } + depth := 0 + segmentStart := -1 + var segments []string +scan: + for i := start; i < len(schema); i++ { + switch schema[i] { + case '(': + if depth == 0 { + segmentStart = i + 1 + } + depth++ + case ')': + depth-- + if depth == 0 { + if segmentStart != -1 && segmentStart < i { + segments = append(segments, schema[segmentStart:i]) + } + break scan + } + case ',': + if depth == 1 && segmentStart != -1 { + segments = append(segments, schema[segmentStart:i]) + segmentStart = i + 1 + } + } + if depth < 0 { + return nil, errors.New("schema has mismatched parentheses") + } + } + if depth != 0 { + return nil, errors.New("schema has mismatched parentheses") + } + if len(segments) == 0 { + return nil, errors.New("schema has no columns") + } + cols := make([]string, 0, len(segments)) + for _, seg := range segments { + seg = strings.TrimSpace(seg) + if seg == "" { + continue + } + parts := strings.Fields(seg) + if len(parts) == 0 { + continue + } + col := strings.Trim(parts[0], "`") + switch strings.ToUpper(col) { + case "PRIMARY", "UNIQUE", "INDEX", "CONSTRAINT", "FOREIGN": + continue + } + cols = append(cols, col) + } + if len(cols) == 0 { + return nil, errors.New("schema has no column definitions") + } + return cols, nil +} + +func applySpannerDDL( + ctx context.Context, admin *database.DatabaseAdminClient, dbPath string, statements []string, +) error { + op, err := admin.UpdateDatabaseDdl(ctx, &databasepb.UpdateDatabaseDdlRequest{ + Database: dbPath, + Statements: statements, + }) + if err != nil { + return err + } + return op.Wait(ctx) +} + +func fetchSpannerSchema( + ctx context.Context, admin *database.DatabaseAdminClient, dbPath string, +) (map[string]bool, map[string]bool, error) { + resp, err := admin.GetDatabaseDdl(ctx, &databasepb.GetDatabaseDdlRequest{ + Database: dbPath, + }) + if err != nil { + return nil, nil, errors.Wrap(err, "fetching Spanner database DDL") + } + tables := make(map[string]bool, len(resp.Statements)) + indexes := make(map[string]bool, len(resp.Statements)) + for _, stmt := range resp.Statements { + if name, ok := spannerDDLTableName(stmt); ok { + tables[name] = true + continue + } + if name, ok := spannerDDLIndexName(stmt); ok { + indexes[name] = true + } + } + return tables, indexes, nil +} + +func spannerDDLTableName(stmt string) (string, bool) { + trimmed := strings.TrimSpace(stmt) + const prefix = "CREATE TABLE " + if !strings.HasPrefix(strings.ToUpper(trimmed), prefix) { + return "", false + } + rest := strings.TrimSpace(trimmed[len(prefix):]) + end := strings.IndexAny(rest, " (") + if end == -1 { + return "", false + } + name := strings.TrimSpace(rest[:end]) + name = strings.Trim(name, "`") + if name == "" { + return "", false + } + return name, true +} + +func spannerDDLIndexName(stmt string) (string, bool) { + trimmed := strings.TrimSpace(stmt) + upper := strings.ToUpper(trimmed) + var prefix string + switch { + case strings.HasPrefix(upper, "CREATE INDEX "): + prefix = "CREATE INDEX " + case strings.HasPrefix(upper, "CREATE UNIQUE INDEX "): + prefix = "CREATE UNIQUE INDEX " + default: + return "", false + } + rest := strings.TrimSpace(trimmed[len(prefix):]) + parts := strings.Fields(rest) + if len(parts) == 0 { + return "", false + } + name := strings.Trim(parts[0], "`") + if name == "" { + return "", false + } + return name, true +} diff --git a/pkg/workload/cli/sql_init.go b/pkg/workload/cli/sql_init.go new file mode 100644 index 000000000000..88e7cfadf019 --- /dev/null +++ b/pkg/workload/cli/sql_init.go @@ -0,0 +1,76 @@ +// Copyright 2024 The Cockroach Authors. +// +// Use of this software is governed by the CockroachDB Software License +// included in the /LICENSE file. + +package cli + +import ( + "context" + gosql "database/sql" + "fmt" + "strings" + + _ "github.com/jackc/pgx/v5/stdlib" + + "github.com/cockroachdb/cockroach/pkg/workload" + "github.com/cockroachdb/cockroach/pkg/workload/workloadsql" + "github.com/cockroachdb/errors" +) + +func runSQLInitNonCRDB(ctx context.Context, gen workload.Generator, urls []string) error { + lc := strings.ToLower(*dataLoader) + if lc == "auto" { + lc = "insert" + } + if lc != "insert" { + return errors.Errorf("data loader %q is not supported for dialect %q", *dataLoader, *dialect) + } + + initDB, err := gosql.Open(`pgx`, strings.Join(urls, ` `)) + if err != nil { + return err + } + defer initDB.Close() + initDB.SetMaxOpenConns(*initConns) + initDB.SetMaxIdleConns(*initConns) + + if err := dropSQLTables(ctx, initDB, gen); err != nil { + return err + } + + l := workloadsql.InsertsDataLoader{ + BatchSize: 0, + Concurrency: *initConns, + DisableCRDBDDL: true, + } + if _, err := l.InitialDataLoad(ctx, initDB, gen); err != nil { + return err + } + + if h, ok := gen.(workload.Hookser); ok { + if h.Hooks().PostLoad != nil { + if err := h.Hooks().PostLoad(ctx, initDB); err != nil { + return errors.Wrapf(err, "could not postload") + } + } + } + return nil +} + +func dropSQLTables(ctx context.Context, db *gosql.DB, gen workload.Generator) error { + if !*drop { + return nil + } + for _, table := range gen.Tables() { + if table.ObjectPrefix != nil && table.ObjectPrefix.ExplicitCatalog { + return errors.Errorf("drop not supported for multi-database table %q", table.Name) + } + tableName := table.GetResolvedName() + stmt := fmt.Sprintf(`DROP TABLE IF EXISTS %s`, tableName.String()) + if _, err := db.ExecContext(ctx, stmt); err != nil { + return errors.Wrapf(err, "dropping table %q", table.Name) + } + } + return nil +} diff --git a/pkg/workload/connection.go b/pkg/workload/connection.go index ed5c81ad481a..edf9fa4e3949 100644 --- a/pkg/workload/connection.go +++ b/pkg/workload/connection.go @@ -83,14 +83,38 @@ func NewConnFlags(genFlags *Flags) *ConnFlags { return c } +func dialectForGen(gen Generator) string { + if dp, ok := gen.(DialectProvider); ok { + if dialect := strings.ToLower(dp.Dialect()); dialect != "" { + return dialect + } + } + return "crdb" +} + // SanitizeUrls verifies that the give SQL connection strings have the correct // SQL database set, rewriting them in place if necessary. This database name is // returned. +// +// For Spanner URLs (projects/PROJECT/instances/INSTANCE/databases/DATABASE), +// the URL is passed through unchanged and the database name from the path is returned. func SanitizeUrls(gen Generator, connFlags *ConnFlags, urls []string) (string, error) { dbName := gen.Meta().Name if connFlags != nil && connFlags.DBOverride != `` { dbName = connFlags.DBOverride } + + if dialectForGen(gen) == "spanner" { + // For Spanner, accept database paths and skip URL validation. + if len(urls) > 0 { + parts := strings.Split(urls[0], "/") + if len(parts) >= 6 && parts[4] == "databases" { + dbName = parts[5] + } + } + return dbName, nil + } + for i := range urls { parsed, err := url.Parse(urls[i]) if err != nil { @@ -115,10 +139,27 @@ func SanitizeUrls(gen Generator, connFlags *ConnFlags, urls []string) (string, e // SetUrlConnVars augments the provided URLs with additional query parameters // which are used by the SQL server during connection establishment to configure // default session variables. +// +// For Spanner URLs, this function is a no-op since Spanner doesn't use +// PostgreSQL-style connection variables. +// +// For non-CockroachDB PostgreSQL-compatible dialects (postgres, aurora, dsql), +// CockroachDB-specific variables like allow_unsafe_internals are omitted. func SetUrlConnVars(gen Generator, connFlags *ConnFlags, urls []string) error { + // Skip for Spanner - it doesn't use PostgreSQL connection variables. + if dialectForGen(gen) == "spanner" { + return nil + } + + // Check if the generator uses a non-CRDB dialect. + isCRDB := dialectForGen(gen) == "crdb" + vars := make(map[string]string) vars["application_name"] = gen.Meta().Name - vars["allow_unsafe_internals"] = "true" + // Only set CRDB-specific variables for CockroachDB. + if isCRDB { + vars["allow_unsafe_internals"] = "true" + } if connFlags != nil { if connFlags.IsoLevel != "" { // As a convenience, replace underscores with spaces. This allows users of diff --git a/pkg/workload/dbexec/BUILD.bazel b/pkg/workload/dbexec/BUILD.bazel new file mode 100644 index 000000000000..fe12fff05756 --- /dev/null +++ b/pkg/workload/dbexec/BUILD.bazel @@ -0,0 +1,33 @@ +load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "dbexec", + srcs = [ + "crdb_executor.go", + "dialect.go", + "executor.go", + "pgx_executor.go", + "spanner_executor.go", + ], + importpath = "github.com/cockroachdb/cockroach/pkg/workload/dbexec", + visibility = ["//visibility:public"], + deps = [ + "//pkg/workload", + "@com_github_cockroachdb_errors//:errors", + "@com_github_jackc_pgx_v5//:pgx", + "@com_github_jackc_pgx_v5//pgxpool", + "@com_google_cloud_go_spanner//:spanner", + "@org_golang_google_api//iterator", + "@org_golang_google_grpc//codes", + ], +) + +go_test( + name = "dbexec_test", + srcs = [ + "crdb_executor_test.go", + "dialect_test.go", + ], + embed = [":dbexec"], + deps = ["@com_github_stretchr_testify//require"], +) diff --git a/pkg/workload/dbexec/crdb_executor.go b/pkg/workload/dbexec/crdb_executor.go new file mode 100644 index 000000000000..f9c58b200f91 --- /dev/null +++ b/pkg/workload/dbexec/crdb_executor.go @@ -0,0 +1,290 @@ +// Copyright 2024 The Cockroach Authors. +// +// Use of this software is governed by the CockroachDB Software License +// included in the /LICENSE file. + +package dbexec + +import ( + "context" + "fmt" + "strings" + + "github.com/cockroachdb/cockroach/pkg/workload" + "github.com/cockroachdb/errors" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" +) + +// CRDBExecutor implements Executor for CockroachDB with full feature support. +// Unlike PGXExecutor which is designed for generic PostgreSQL-compatible databases, +// CRDBExecutor provides access to CockroachDB-specific features like transaction +// QoS, SELECT FOR UPDATE, and follower reads. +// +// The executor is safe for concurrent use from multiple goroutines. It maintains +// a connection pool and prepared statements that are shared across workers. +type CRDBExecutor struct { + dialect CockroachDialect + cfg Config + + mcp *workload.MultiConnPool + sr workload.SQLRunner + + // Statement handles for prepared statements. + readStmt workload.StmtHandle + followerReadStmt workload.StmtHandle + writeStmt workload.StmtHandle + deleteStmt workload.StmtHandle + spanStmt workload.StmtHandle + sfuStmt workload.StmtHandle // SELECT ... FOR UPDATE + sel1Stmt workload.StmtHandle // SELECT 1 +} + +var _ Executor = (*CRDBExecutor)(nil) + +// NewCRDBExecutor creates a new CRDBExecutor. Init must be called before using +// the executor. +func NewCRDBExecutor() *CRDBExecutor { + return &CRDBExecutor{} +} + +// Init prepares the executor by creating connection pools and defining +// statements via SQLRunner. This must be called before any other methods. +func (e *CRDBExecutor) Init(ctx context.Context, cfg Config, connFlags *workload.ConnFlags) error { + e.cfg = cfg + + // Create multi-connection pool. + poolCfg := workload.NewMultiConnPoolCfgFromFlags(connFlags) + poolCfg.MaxTotalConnections = connFlags.Concurrency + 1 + mcp, err := workload.NewMultiConnPool(ctx, poolCfg, cfg.URLs...) + if err != nil { + return errors.Wrap(err, "creating connection pool") + } + e.mcp = mcp + + // Define statements using the CockroachDB dialect. + e.readStmt = e.sr.Define(e.dialect.ReadStmt(cfg.Table, cfg.BatchSize, cfg.Enum)) + e.followerReadStmt = e.sr.Define(e.dialect.FollowerReadStmt(cfg.Table, cfg.BatchSize, cfg.Enum)) + e.writeStmt = e.sr.Define(e.dialect.UpsertStmt(cfg.Table, cfg.BatchSize, cfg.Enum)) + e.deleteStmt = e.sr.Define(e.dialect.DeleteStmt(cfg.Table, cfg.BatchSize)) + e.spanStmt = e.sr.Define(e.dialect.SpanStmt(cfg.Table, cfg.SpanLimit)) + e.sfuStmt = e.sr.Define(e.dialect.SelectForUpdateStmt(cfg.Table, cfg.BatchSize)) + e.sel1Stmt = e.sr.Define("SELECT 1") + + // Initialize the SQLRunner with prepared statements. + if err := e.sr.Init(ctx, e.dialect.Name(), e.mcp); err != nil { + e.mcp.Close() + return errors.Wrap(err, "initializing SQL runner") + } + + return nil +} + +// SetupWorker performs per-worker initialization. For CRDBExecutor, QoS is set +// per-transaction in BeginTx rather than per-worker, so this is a no-op. +func (e *CRDBExecutor) SetupWorker(ctx context.Context) error { + return nil +} + +// Read executes a batch read for the given keys. +func (e *CRDBExecutor) Read(ctx context.Context, keys []interface{}) (Rows, error) { + rows, err := e.readStmt.Query(ctx, keys...) + if err != nil { + return nil, err + } + return &pgxRowsWrapper{rows: rows}, nil +} + +// FollowerRead executes a stale/replica read using CockroachDB's +// AS OF SYSTEM TIME follower_read_timestamp() feature. +func (e *CRDBExecutor) FollowerRead(ctx context.Context, keys []interface{}) (Rows, error) { + rows, err := e.followerReadStmt.Query(ctx, keys...) + if err != nil { + return nil, err + } + return &pgxRowsWrapper{rows: rows}, nil +} + +// Write upserts a batch of key-value pairs. The enum column is not included +// because it is a computed stored column in CockroachDB (AS ('v') STORED). +func (e *CRDBExecutor) Write(ctx context.Context, rows []Row) error { + args := flattenRowArgs(rows, false) + _, err := e.writeStmt.Exec(ctx, args...) + return err +} + +// Delete removes a batch of keys. +func (e *CRDBExecutor) Delete(ctx context.Context, keys []interface{}) error { + _, err := e.deleteStmt.Exec(ctx, keys...) + return err +} + +// Span executes a spanning query starting at startKey with the configured limit. +// Returns the count of values scanned. +func (e *CRDBExecutor) Span(ctx context.Context, startKey interface{}, limit int) (int64, error) { + // The limit is baked into the prepared statement via SpanStmt. When + // SpanLimit == 0, the statement is a full table scan with no arguments. + // When SpanLimit > 0, it takes a startKey argument for the WHERE clause. + var row pgx.Row + if e.cfg.SpanLimit == 0 { + row = e.spanStmt.QueryRow(ctx) + } else { + row = e.spanStmt.QueryRow(ctx, startKey) + } + var count int64 + if err := row.Scan(&count); err != nil { + return 0, err + } + return count, nil +} + +// BeginTx starts a transaction with the given options. For CRDBExecutor, this +// returns a crdbTxWrapper that implements both Tx and ExtendedTx interfaces, +// providing access to CockroachDB-specific transaction operations. +func (e *CRDBExecutor) BeginTx(ctx context.Context, opts TxOptions) (Tx, error) { + pool := e.mcp.Get() + conn, err := pool.Acquire(ctx) + if err != nil { + return nil, errors.Wrap(err, "acquiring connection") + } + + txOpts := pgx.TxOptions{} + + // Set transaction priority if specified. + if opts.Priority != "" { + p := strings.ToLower(opts.Priority) + switch p { + case "low", "normal", "high": + txOpts.BeginQuery = fmt.Sprintf("BEGIN PRIORITY %s", p) + default: + conn.Release() + return nil, errors.Newf("invalid transaction priority: %q", opts.Priority) + } + } + + tx, err := conn.BeginTx(ctx, txOpts) + if err != nil { + conn.Release() + return nil, errors.Wrap(err, "beginning transaction") + } + + // Set transaction QoS if configured. + if e.cfg.TxnQoS != "" && e.cfg.TxnQoS != "regular" { + switch e.cfg.TxnQoS { + case "background", "critical": + // valid QoS values + default: + _ = tx.Rollback(ctx) + conn.Release() + return nil, errors.Newf("invalid transaction QoS: %q (must be background, regular, or critical)", e.cfg.TxnQoS) + } + _, err = tx.Exec(ctx, fmt.Sprintf("SET default_transaction_quality_of_service = '%s'", e.cfg.TxnQoS)) + if err != nil { + _ = tx.Rollback(ctx) + conn.Release() + return nil, errors.Wrapf(err, "setting transaction QoS to %q", e.cfg.TxnQoS) + } + } + + return &crdbTxWrapper{ + tx: tx, + conn: conn, + executor: e, + }, nil +} + +// Close releases all resources including connection pools. +func (e *CRDBExecutor) Close() error { + if e.mcp != nil { + e.mcp.Close() + } + return nil +} + +// Capabilities returns CockroachDB's full feature set. All capabilities are +// supported by CRDBExecutor. +func (e *CRDBExecutor) Capabilities() Capabilities { + return Capabilities{ + FollowerReads: true, + Scatter: true, + Splits: true, + HashShardedPK: true, + TransactionQoS: true, + SelectForUpdate: true, + TransactionPriority: true, + SerializationRetry: true, + Select1: true, + EnumColumn: true, + } +} + +// crdbTxWrapper wraps a pgx transaction and provides access to prepared +// statement operations from the executor. It implements both Tx and ExtendedTx. +// +// The wrapper holds a reference to the executor to access prepared statements for +// operations like SELECT FOR UPDATE and Write within transactions. +type crdbTxWrapper struct { + tx pgx.Tx + conn *pgxpool.Conn + executor *CRDBExecutor +} + +var _ Tx = (*crdbTxWrapper)(nil) +var _ ExtendedTx = (*crdbTxWrapper)(nil) + +// Exec executes a statement that does not return rows within the transaction. +func (t *crdbTxWrapper) Exec(ctx context.Context, sql string, args ...interface{}) error { + _, err := t.tx.Exec(ctx, sql, args...) + return err +} + +// Query executes a query that returns rows within the transaction. +func (t *crdbTxWrapper) Query(ctx context.Context, sql string, args ...interface{}) (Rows, error) { + rows, err := t.tx.Query(ctx, sql, args...) + if err != nil { + return nil, err + } + return &pgxRowsWrapper{rows: rows}, nil +} + +// Commit commits the transaction and releases the connection back to the pool. +func (t *crdbTxWrapper) Commit(ctx context.Context) error { + err := t.tx.Commit(ctx) + t.conn.Release() + return err +} + +// Rollback aborts the transaction and releases the connection back to the pool. +func (t *crdbTxWrapper) Rollback(ctx context.Context) error { + err := t.tx.Rollback(ctx) + t.conn.Release() + return err +} + +// SelectForUpdate executes a SELECT ... FOR UPDATE query within the transaction, +// locking the selected rows until the transaction completes. +func (t *crdbTxWrapper) SelectForUpdate(ctx context.Context, keys []interface{}) (Rows, error) { + rows, err := t.executor.sfuStmt.QueryTx(ctx, t.tx, keys...) + if err != nil { + return nil, err + } + return &pgxRowsWrapper{rows: rows}, nil +} + +// Select1 executes a SELECT 1 query within the transaction. This is useful +// for keeping transactions alive or as a lightweight operation in transaction +// retry loops. +func (t *crdbTxWrapper) Select1(ctx context.Context) error { + row := t.executor.sel1Stmt.QueryRowTx(ctx, t.tx) + var dummy int + return row.Scan(&dummy) +} + +// Write executes an upsert within the transaction using prepared statements. +// The enum column is not included because it is a computed stored column in +// CockroachDB. +func (t *crdbTxWrapper) Write(ctx context.Context, rows []Row) error { + args := flattenRowArgs(rows, false) + _, err := t.executor.writeStmt.ExecTx(ctx, t.tx, args...) + return err +} diff --git a/pkg/workload/dbexec/crdb_executor_test.go b/pkg/workload/dbexec/crdb_executor_test.go new file mode 100644 index 000000000000..f62b925827dd --- /dev/null +++ b/pkg/workload/dbexec/crdb_executor_test.go @@ -0,0 +1,34 @@ +// Copyright 2024 The Cockroach Authors. +// +// Use of this software is governed by the CockroachDB Software License +// included in the /LICENSE file. + +package dbexec + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestCRDBExecutor_Capabilities(t *testing.T) { + exec := NewCRDBExecutor() + caps := exec.Capabilities() + + // CRDBExecutor should support all CockroachDB features. + require.True(t, caps.FollowerReads, "CRDBExecutor should support follower reads") + require.True(t, caps.Scatter, "CRDBExecutor should support scatter") + require.True(t, caps.Splits, "CRDBExecutor should support splits") + require.True(t, caps.HashShardedPK, "CRDBExecutor should support hash-sharded primary keys") + require.True(t, caps.TransactionQoS, "CRDBExecutor should support transaction QoS") + require.True(t, caps.SelectForUpdate, "CRDBExecutor should support SELECT FOR UPDATE") + require.True(t, caps.TransactionPriority, "CRDBExecutor should support transaction priority") + require.True(t, caps.SerializationRetry, "CRDBExecutor should support serialization retry") + require.True(t, caps.Select1, "CRDBExecutor should support SELECT 1 warmup") + require.True(t, caps.EnumColumn, "CRDBExecutor should support enum columns") +} + +func TestNewCRDBExecutor(t *testing.T) { + exec := NewCRDBExecutor() + require.NotNil(t, exec, "NewCRDBExecutor should return non-nil executor") +} diff --git a/pkg/workload/dbexec/dialect.go b/pkg/workload/dbexec/dialect.go new file mode 100644 index 000000000000..baf17e95a8e5 --- /dev/null +++ b/pkg/workload/dbexec/dialect.go @@ -0,0 +1,517 @@ +// Copyright 2024 The Cockroach Authors. +// +// Use of this software is governed by the CockroachDB Software License +// included in the /LICENSE file. + +package dbexec + +import ( + "fmt" + "strings" +) + +// SchemaDialect generates DDL for table creation during workload init. +// All Dialect implementations satisfy this interface. Standalone +// implementations exist for databases (like Spanner) that don't use +// the full Dialect for queries. +type SchemaDialect interface { + // SchemaFragment returns the CREATE TABLE body (columns, constraints) + // for use as workload.Table.Schema. The result must NOT include the + // "CREATE TABLE " prefix -- the workload framework adds that. + // + // keySize controls the primary key type: + // - keySize == 0: integer key (INT/BIGINT/INT64 depending on dialect) + // - keySize > 0: string key of that width (STRING(N)/VARCHAR(N)) + // + // If secondaryIndex is true, the schema includes a secondary index on v + // (inline for CRDB, separate statement via SecondaryIndexStmts for others). + // + // numShards specifies hash bucket count for hash-sharded primary keys. + // Ignored by dialects that don't support hash sharding. + SchemaFragment(keySize int, secondaryIndex bool, numShards int) string + + // SecondaryIndexStmts returns DDL to create secondary indexes as separate + // statements (outside CREATE TABLE), or nil if the index is included + // inline in SchemaFragment. Each statement must be valid for the target + // dialect. Only called when secondaryIndex is true. + SecondaryIndexStmts(table string) []string + + // PostLoadStmts returns DDL to execute after table creation and data + // loading. Used for operations that must happen after the table exists, + // such as adding computed columns or scattering ranges. + // Returns nil if no post-load DDL is needed. + PostLoadStmts(table string, enum, scatter bool) []string + + // Capabilities returns what features this dialect supports. + Capabilities() Capabilities +} + +// Dialect defines SQL syntax variations between databases. Each database +// family has different syntax for upserts, follower reads, schema features, +// and other operations. Implementations generate the appropriate SQL for +// their target database. +type Dialect interface { + SchemaDialect + + // Name returns a short identifier for this dialect (e.g., "crdb", "postgres"). + Name() string + + // UpsertStmt returns an upsert statement for batchSize rows. + // If hasEnum is true, the statement includes the enum column. + UpsertStmt(table string, batchSize int, hasEnum bool) string + + // ReadStmt returns a batch read statement for batchSize keys. + // If hasEnum is true, the statement selects the enum column. + ReadStmt(table string, batchSize int, hasEnum bool) string + + // FollowerReadStmt returns a stale/replica read statement for batchSize keys. + // If hasEnum is true, the statement selects the enum column. + // Returns the same as ReadStmt if follower reads are not supported. + FollowerReadStmt(table string, batchSize int, hasEnum bool) string + + // DeleteStmt returns a batch delete statement for batchSize keys. + DeleteStmt(table string, batchSize int) string + + // SpanStmt returns a spanning query statement with the given limit. + // The statement counts values starting from a given key. + SpanStmt(table string, limit int) string + + // SelectForUpdateStmt returns a SELECT ... FOR UPDATE statement for + // batchSize keys. + SelectForUpdateStmt(table string, batchSize int) string +} + +// placeholders generates "$1, $2, ..., $n" for parameterized queries. +func placeholders(n int) string { + if n <= 0 { + return "" + } + parts := make([]string, n) + for i := range parts { + parts[i] = fmt.Sprintf("$%d", i+1) + } + return strings.Join(parts, ", ") +} + +// valueTuples generates "($1, $2), ($3, $4), ..." for batch inserts. +// Each tuple has tupleSize placeholders, and there are count tuples. +func valueTuples(count, tupleSize int) string { + if count <= 0 || tupleSize <= 0 { + return "" + } + tuples := make([]string, count) + for i := range tuples { + parts := make([]string, tupleSize) + for j := range parts { + parts[j] = fmt.Sprintf("$%d", i*tupleSize+j+1) + } + tuples[i] = "(" + strings.Join(parts, ", ") + ")" + } + return strings.Join(tuples, ", ") +} + +// CockroachDialect implements Dialect for CockroachDB. It supports all +// CockroachDB-specific features including follower reads, hash-sharded +// primary keys, and transaction QoS. +type CockroachDialect struct{} + +var _ Dialect = CockroachDialect{} + +// Name returns "crdb" as the dialect identifier. +func (CockroachDialect) Name() string { + return "crdb" +} + +// UpsertStmt generates a CockroachDB UPSERT statement. The hasEnum parameter +// is ignored because CockroachDB's enum column is a computed stored column +// (AS ('v') STORED) that cannot be written to directly. +// +// Example (batchSize=2): +// +// UPSERT INTO kv (k, v) VALUES ($1, $2), ($3, $4) +func (CockroachDialect) UpsertStmt(table string, batchSize int, hasEnum bool) string { + // hasEnum is intentionally ignored: the e column is a computed stored column + // in CockroachDB and cannot be included in UPSERT statements. + return fmt.Sprintf("UPSERT INTO %s (k, v) VALUES %s", table, valueTuples(batchSize, 2)) +} + +// ReadStmt generates a batch SELECT statement. +// +// Example without enum (batchSize=3): +// +// SELECT k, v FROM kv WHERE k IN ($1, $2, $3) +// +// Example with enum (batchSize=3): +// +// SELECT k, v, e FROM kv WHERE k IN ($1, $2, $3) +func (CockroachDialect) ReadStmt(table string, batchSize int, hasEnum bool) string { + cols := "k, v" + if hasEnum { + cols = "k, v, e" + } + return fmt.Sprintf("SELECT %s FROM %s WHERE k IN (%s)", cols, table, placeholders(batchSize)) +} + +// FollowerReadStmt generates a CockroachDB follower read statement using +// AS OF SYSTEM TIME follower_read_timestamp(). +// +// Example (batchSize=2): +// +// SELECT k, v FROM kv AS OF SYSTEM TIME follower_read_timestamp() WHERE k IN ($1, $2) +func (CockroachDialect) FollowerReadStmt(table string, batchSize int, hasEnum bool) string { + cols := "k, v" + if hasEnum { + cols = "k, v, e" + } + return fmt.Sprintf( + "SELECT %s FROM %s AS OF SYSTEM TIME follower_read_timestamp() WHERE k IN (%s)", + cols, table, placeholders(batchSize), + ) +} + +// DeleteStmt generates a batch DELETE statement. +// +// Example (batchSize=3): +// +// DELETE FROM kv WHERE k IN ($1, $2, $3) +func (CockroachDialect) DeleteStmt(table string, batchSize int) string { + return fmt.Sprintf("DELETE FROM %s WHERE k IN (%s)", table, placeholders(batchSize)) +} + +// SpanStmt generates a spanning query using CockroachDB's scalar subquery syntax. +// When limit > 0, the statement counts values starting from key $1 with the +// given limit. When limit == 0, it performs a full table scan with no arguments. +// +// Example (limit=100): +// +// SELECT count(v) FROM [SELECT v FROM kv WHERE k >= $1 ORDER BY k LIMIT 100] +// +// Example (limit=0): +// +// SELECT count(v) FROM [SELECT v FROM kv] +func (CockroachDialect) SpanStmt(table string, limit int) string { + if limit == 0 { + return fmt.Sprintf("SELECT count(v) FROM [SELECT v FROM %s]", table) + } + return fmt.Sprintf( + "SELECT count(v) FROM [SELECT v FROM %s WHERE k >= $1 ORDER BY k LIMIT %d]", + table, limit, + ) +} + +// SelectForUpdateStmt generates a SELECT ... FOR UPDATE statement. +// +// Example (batchSize=2): +// +// SELECT k, v FROM kv WHERE k IN ($1, $2) FOR UPDATE +func (CockroachDialect) SelectForUpdateStmt(table string, batchSize int) string { + return fmt.Sprintf("SELECT k, v FROM %s WHERE k IN (%s) FOR UPDATE", table, placeholders(batchSize)) +} + +// SchemaFragment generates the CockroachDB CREATE TABLE body for use as +// workload.Table.Schema. The output matches the production schema constants +// in kv.go (kvSchema, kvSchemaWithIndex, shardedKvSchema, etc.). +// +// Example (keySize=0, no shards, no index): +// +// ( +// k BIGINT NOT NULL PRIMARY KEY, +// v BYTES NOT NULL +// ) +// +// Example (keySize=0, numShards=8, secondaryIndex=true): +// +// ( +// k BIGINT NOT NULL PRIMARY KEY USING HASH WITH (bucket_count = 8), +// v BYTES NOT NULL, +// INDEX (v) +// ) +func (CockroachDialect) SchemaFragment(keySize int, secondaryIndex bool, numShards int) string { + keyType := "BIGINT" + if keySize > 0 { + keyType = "STRING" + } + + var sb strings.Builder + sb.WriteString("(\n") + + if numShards > 0 { + sb.WriteString(fmt.Sprintf( + "\t\tk %s NOT NULL PRIMARY KEY USING HASH WITH (bucket_count = %d),\n", + keyType, numShards)) + } else { + sb.WriteString(fmt.Sprintf("\t\tk %s NOT NULL PRIMARY KEY,\n", keyType)) + } + + sb.WriteString("\t\tv BYTES NOT NULL") + if secondaryIndex { + sb.WriteString(",\n\t\tINDEX (v)") + } + sb.WriteString("\n\t)") + return sb.String() +} + +// SecondaryIndexStmts returns nil because CockroachDB includes secondary +// indexes inline in the CREATE TABLE statement via SchemaFragment. +func (CockroachDialect) SecondaryIndexStmts(table string) []string { + return nil +} + +// PostLoadStmts returns DDL to execute after table creation and data loading. +// For CockroachDB, this includes adding computed enum columns and scattering +// ranges across the cluster. +// +// The enum DDL combines CREATE TYPE and ALTER TABLE into a single string +// because they are executed together as a multi-statement batch via db.Exec(). +// This matches the original behavior in kv.go's PostLoad hook. +func (CockroachDialect) PostLoadStmts(table string, enum, scatter bool) []string { + var stmts []string + if enum { + stmts = append(stmts, fmt.Sprintf( + "CREATE TYPE enum_type AS ENUM ('v');\n"+ + "ALTER TABLE %s ADD COLUMN e enum_type NOT NULL AS ('v') STORED;", + table)) + } + if scatter { + stmts = append(stmts, fmt.Sprintf("ALTER TABLE %s SCATTER", table)) + } + return stmts +} + +// Capabilities returns CockroachDB's full feature set. +func (CockroachDialect) Capabilities() Capabilities { + return Capabilities{ + FollowerReads: true, + Scatter: true, + Splits: true, + HashShardedPK: true, + TransactionQoS: true, + SelectForUpdate: true, + TransactionPriority: true, + SerializationRetry: true, + Select1: true, + EnumColumn: true, + } +} + +// PostgresDialect implements Dialect for PostgreSQL-compatible databases. +// This includes PostgreSQL, Aurora PostgreSQL, and other compatible databases. +// It does not support CockroachDB-specific features like follower reads or +// hash-sharded primary keys. +type PostgresDialect struct{} + +var _ Dialect = PostgresDialect{} + +// Name returns "postgres" as the dialect identifier. +func (PostgresDialect) Name() string { + return "postgres" +} + +// UpsertStmt generates a PostgreSQL INSERT ... ON CONFLICT statement. +// +// Example without enum (batchSize=2): +// +// INSERT INTO kv (k, v) VALUES ($1, $2), ($3, $4) +// ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v +// +// Example with enum (batchSize=2): +// +// INSERT INTO kv (k, v, e) VALUES ($1, $2, $3), ($4, $5, $6) +// ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v, e = EXCLUDED.e +func (PostgresDialect) UpsertStmt(table string, batchSize int, hasEnum bool) string { + if hasEnum { + return fmt.Sprintf( + "INSERT INTO %s (k, v, e) VALUES %s ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v, e = EXCLUDED.e", + table, valueTuples(batchSize, 3), + ) + } + return fmt.Sprintf( + "INSERT INTO %s (k, v) VALUES %s ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v", + table, valueTuples(batchSize, 2), + ) +} + +// ReadStmt generates a batch SELECT statement. +// +// Example without enum (batchSize=3): +// +// SELECT k, v FROM kv WHERE k IN ($1, $2, $3) +// +// Example with enum (batchSize=3): +// +// SELECT k, v, e FROM kv WHERE k IN ($1, $2, $3) +func (PostgresDialect) ReadStmt(table string, batchSize int, hasEnum bool) string { + cols := "k, v" + if hasEnum { + cols = "k, v, e" + } + return fmt.Sprintf("SELECT %s FROM %s WHERE k IN (%s)", cols, table, placeholders(batchSize)) +} + +// FollowerReadStmt returns the same as ReadStmt since PostgreSQL does not +// support follower reads. +func (d PostgresDialect) FollowerReadStmt(table string, batchSize int, hasEnum bool) string { + return d.ReadStmt(table, batchSize, hasEnum) +} + +// DeleteStmt generates a batch DELETE statement. +// +// Example (batchSize=3): +// +// DELETE FROM kv WHERE k IN ($1, $2, $3) +func (PostgresDialect) DeleteStmt(table string, batchSize int) string { + return fmt.Sprintf("DELETE FROM %s WHERE k IN (%s)", table, placeholders(batchSize)) +} + +// SpanStmt generates a spanning query using standard SQL derived table syntax. +// When limit > 0, the statement counts values starting from key $1 with the +// given limit. When limit == 0, it performs a full table scan with no arguments. +// +// Example (limit=100): +// +// SELECT count(v) FROM (SELECT v FROM kv WHERE k >= $1 ORDER BY k LIMIT 100) AS t +// +// Example (limit=0): +// +// SELECT count(v) FROM (SELECT v FROM kv) AS t +func (PostgresDialect) SpanStmt(table string, limit int) string { + if limit == 0 { + return fmt.Sprintf("SELECT count(v) FROM (SELECT v FROM %s) AS t", table) + } + return fmt.Sprintf( + "SELECT count(v) FROM (SELECT v FROM %s WHERE k >= $1 ORDER BY k LIMIT %d) AS t", + table, limit, + ) +} + +// SelectForUpdateStmt generates a SELECT ... FOR UPDATE statement. +// +// Example (batchSize=2): +// +// SELECT k, v FROM kv WHERE k IN ($1, $2) FOR UPDATE +func (PostgresDialect) SelectForUpdateStmt(table string, batchSize int) string { + return fmt.Sprintf("SELECT k, v FROM %s WHERE k IN (%s) FOR UPDATE", table, placeholders(batchSize)) +} + +// SchemaFragment generates the PostgreSQL CREATE TABLE body for use as +// workload.Table.Schema. Uses BYTEA for binary data and standard SQL +// types for keys. Hash sharding (numShards) is ignored because PostgreSQL +// does not support hash-sharded primary keys. +// +// Example (keySize=0): +// +// (k BIGINT NOT NULL PRIMARY KEY, v BYTEA NOT NULL) +// +// Example (keySize=100): +// +// (k VARCHAR(100) NOT NULL PRIMARY KEY, v BYTEA NOT NULL) +func (PostgresDialect) SchemaFragment(keySize int, secondaryIndex bool, numShards int) string { + // numShards and secondaryIndex are ignored: PostgreSQL doesn't support + // hash-sharded PKs, and secondary indexes are created via + // SecondaryIndexStmts as separate DDL statements. + keyType := "BIGINT" + if keySize > 0 { + keyType = fmt.Sprintf("VARCHAR(%d)", keySize) + } + return fmt.Sprintf("(k %s NOT NULL PRIMARY KEY, v BYTEA NOT NULL)", keyType) +} + +// SecondaryIndexStmts returns a CREATE INDEX statement for the secondary +// index on column v. Uses IF NOT EXISTS for idempotency. +func (PostgresDialect) SecondaryIndexStmts(table string) []string { + return []string{ + fmt.Sprintf("CREATE INDEX IF NOT EXISTS %s_v_idx ON %s (v)", table, table), + } +} + +// PostLoadStmts returns nil because PostgreSQL does not require any +// post-load DDL for the kv workload. +func (PostgresDialect) PostLoadStmts(table string, enum, scatter bool) []string { + return nil +} + +// Capabilities returns PostgreSQL's feature set. SelectForUpdate and Select1 +// are supported because PGXExecutor implements ExtendedTx with prepared +// statement access for transactional write workflows. +func (PostgresDialect) Capabilities() Capabilities { + return Capabilities{ + FollowerReads: false, + Scatter: false, + Splits: false, + HashShardedPK: false, + TransactionQoS: false, + SelectForUpdate: true, + TransactionPriority: false, + SerializationRetry: false, + Select1: true, + EnumColumn: false, + } +} + +// SpannerSchemaDialect implements SchemaDialect for Google Cloud Spanner. +// Unlike CockroachDialect and PostgresDialect, SpannerSchemaDialect only +// implements SchemaDialect (not the full Dialect) because SpannerExecutor +// uses the Spanner mutations API for queries instead of SQL. +type SpannerSchemaDialect struct{} + +var _ SchemaDialect = SpannerSchemaDialect{} + +// SchemaFragment generates the Spanner CREATE TABLE body for use as +// workload.Table.Schema. Spanner uses INT64 for integer keys and +// BYTES(MAX) for binary data. The PRIMARY KEY clause appears outside the +// column definition parentheses, matching Spanner's DDL syntax. +// +// Example (keySize=0): +// +// (k INT64 NOT NULL, v BYTES(MAX) NOT NULL) PRIMARY KEY (k) +// +// Example (keySize=100): +// +// (k STRING(100) NOT NULL, v BYTES(MAX) NOT NULL) PRIMARY KEY (k) +func (SpannerSchemaDialect) SchemaFragment(keySize int, secondaryIndex bool, numShards int) string { + // numShards and secondaryIndex are ignored: Spanner does not support + // hash-sharded PKs, and secondary indexes are created via + // SecondaryIndexStmts as separate DDL statements. + if keySize > 0 { + return fmt.Sprintf( + "(k STRING(%d) NOT NULL, v BYTES(MAX) NOT NULL) PRIMARY KEY (k)", + keySize) + } + return "(k INT64 NOT NULL, v BYTES(MAX) NOT NULL) PRIMARY KEY (k)" +} + +// SecondaryIndexStmts returns a CREATE INDEX statement for the secondary +// index on column v. Spanner does not support IF NOT EXISTS on indexes. +func (SpannerSchemaDialect) SecondaryIndexStmts(table string) []string { + return []string{ + fmt.Sprintf("CREATE INDEX %s_v_idx ON %s (v)", table, table), + } +} + +// PostLoadStmts returns nil because Spanner does not require any post-load DDL +// for the kv workload. Although Spanner supports enum columns (EnumColumn +// capability is true), enum values are stored as regular strings by +// SpannerExecutor and don't require DDL changes like CockroachDB's computed +// stored columns. +func (SpannerSchemaDialect) PostLoadStmts(table string, enum, scatter bool) []string { + return nil +} + +// Capabilities returns Spanner's feature set. Spanner supports stale reads +// (follower reads) and enum columns (stored as strings) but does not support +// CockroachDB-specific features like scatter, splits, or hash-sharded +// primary keys. +func (SpannerSchemaDialect) Capabilities() Capabilities { + return Capabilities{ + FollowerReads: true, // via stale reads + Scatter: false, + Splits: false, + HashShardedPK: false, + TransactionQoS: false, + SelectForUpdate: false, + TransactionPriority: false, + SerializationRetry: false, + Select1: false, + EnumColumn: true, // stored as strings + } +} diff --git a/pkg/workload/dbexec/dialect_test.go b/pkg/workload/dbexec/dialect_test.go new file mode 100644 index 000000000000..0690c5bcfc71 --- /dev/null +++ b/pkg/workload/dbexec/dialect_test.go @@ -0,0 +1,566 @@ +// Copyright 2024 The Cockroach Authors. +// +// Use of this software is governed by the CockroachDB Software License +// included in the /LICENSE file. + +package dbexec + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestPlaceholders(t *testing.T) { + tests := []struct { + n int + want string + }{{ + n: 0, + want: "", + }, { + n: 1, + want: "$1", + }, { + n: 3, + want: "$1, $2, $3", + }, { + n: 5, + want: "$1, $2, $3, $4, $5", + }} + + for _, tt := range tests { + t.Run("", func(t *testing.T) { + got := placeholders(tt.n) + require.Equal(t, tt.want, got) + }) + } +} + +func TestValueTuples(t *testing.T) { + tests := []struct { + count int + tupleSize int + want string + }{{ + count: 0, + tupleSize: 2, + want: "", + }, { + count: 1, + tupleSize: 2, + want: "($1, $2)", + }, { + count: 2, + tupleSize: 2, + want: "($1, $2), ($3, $4)", + }, { + count: 2, + tupleSize: 3, + want: "($1, $2, $3), ($4, $5, $6)", + }, { + count: 3, + tupleSize: 1, + want: "($1), ($2), ($3)", + }} + + for _, tt := range tests { + t.Run("", func(t *testing.T) { + got := valueTuples(tt.count, tt.tupleSize) + require.Equal(t, tt.want, got) + }) + } +} + +func TestCockroachDialect_Name(t *testing.T) { + d := CockroachDialect{} + require.Equal(t, "crdb", d.Name()) +} + +func TestPostgresDialect_Name(t *testing.T) { + d := PostgresDialect{} + require.Equal(t, "postgres", d.Name()) +} + +func TestCockroachDialect_UpsertStmt(t *testing.T) { + d := CockroachDialect{} + + t.Run("without enum", func(t *testing.T) { + stmt := d.UpsertStmt("kv", 2, false) + require.Equal(t, "UPSERT INTO kv (k, v) VALUES ($1, $2), ($3, $4)", stmt) + }) + + t.Run("with enum ignored", func(t *testing.T) { + // CockroachDB's enum column is a computed stored column, so it is never + // included in UPSERT statements. + stmt := d.UpsertStmt("kv", 2, true) + require.Equal(t, "UPSERT INTO kv (k, v) VALUES ($1, $2), ($3, $4)", stmt) + }) + + t.Run("single row", func(t *testing.T) { + stmt := d.UpsertStmt("kv", 1, false) + require.Equal(t, "UPSERT INTO kv (k, v) VALUES ($1, $2)", stmt) + }) +} + +func TestPostgresDialect_UpsertStmt(t *testing.T) { + d := PostgresDialect{} + + t.Run("without enum", func(t *testing.T) { + stmt := d.UpsertStmt("kv", 2, false) + expected := "INSERT INTO kv (k, v) VALUES ($1, $2), ($3, $4) ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v" + require.Equal(t, expected, stmt) + }) + + t.Run("with enum", func(t *testing.T) { + stmt := d.UpsertStmt("kv", 2, true) + expected := "INSERT INTO kv (k, v, e) VALUES ($1, $2, $3), ($4, $5, $6) ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v, e = EXCLUDED.e" + require.Equal(t, expected, stmt) + }) + + t.Run("single row", func(t *testing.T) { + stmt := d.UpsertStmt("kv", 1, false) + expected := "INSERT INTO kv (k, v) VALUES ($1, $2) ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v" + require.Equal(t, expected, stmt) + }) +} + +func TestCockroachDialect_ReadStmt(t *testing.T) { + d := CockroachDialect{} + + t.Run("without enum", func(t *testing.T) { + stmt := d.ReadStmt("kv", 3, false) + require.Equal(t, "SELECT k, v FROM kv WHERE k IN ($1, $2, $3)", stmt) + }) + + t.Run("with enum", func(t *testing.T) { + stmt := d.ReadStmt("kv", 3, true) + require.Equal(t, "SELECT k, v, e FROM kv WHERE k IN ($1, $2, $3)", stmt) + }) + + t.Run("single key", func(t *testing.T) { + stmt := d.ReadStmt("kv", 1, false) + require.Equal(t, "SELECT k, v FROM kv WHERE k IN ($1)", stmt) + }) +} + +func TestPostgresDialect_ReadStmt(t *testing.T) { + d := PostgresDialect{} + + t.Run("without enum", func(t *testing.T) { + stmt := d.ReadStmt("kv", 3, false) + require.Equal(t, "SELECT k, v FROM kv WHERE k IN ($1, $2, $3)", stmt) + }) + + t.Run("with enum", func(t *testing.T) { + stmt := d.ReadStmt("kv", 3, true) + require.Equal(t, "SELECT k, v, e FROM kv WHERE k IN ($1, $2, $3)", stmt) + }) +} + +func TestCockroachDialect_FollowerReadStmt(t *testing.T) { + d := CockroachDialect{} + + t.Run("contains AS OF SYSTEM TIME", func(t *testing.T) { + stmt := d.FollowerReadStmt("kv", 2, false) + require.Contains(t, stmt, "AS OF SYSTEM TIME follower_read_timestamp()") + require.Contains(t, stmt, "WHERE k IN ($1, $2)") + }) + + t.Run("without enum", func(t *testing.T) { + stmt := d.FollowerReadStmt("kv", 2, false) + expected := "SELECT k, v FROM kv AS OF SYSTEM TIME follower_read_timestamp() WHERE k IN ($1, $2)" + require.Equal(t, expected, stmt) + }) + + t.Run("with enum", func(t *testing.T) { + stmt := d.FollowerReadStmt("kv", 2, true) + expected := "SELECT k, v, e FROM kv AS OF SYSTEM TIME follower_read_timestamp() WHERE k IN ($1, $2)" + require.Equal(t, expected, stmt) + }) +} + +func TestPostgresDialect_FollowerReadStmt(t *testing.T) { + d := PostgresDialect{} + + t.Run("equals ReadStmt", func(t *testing.T) { + followerStmt := d.FollowerReadStmt("kv", 3, false) + readStmt := d.ReadStmt("kv", 3, false) + require.Equal(t, readStmt, followerStmt) + }) + + t.Run("does not contain AS OF SYSTEM TIME", func(t *testing.T) { + stmt := d.FollowerReadStmt("kv", 2, false) + require.NotContains(t, stmt, "AS OF SYSTEM TIME") + }) + + t.Run("with enum equals ReadStmt", func(t *testing.T) { + followerStmt := d.FollowerReadStmt("kv", 2, true) + readStmt := d.ReadStmt("kv", 2, true) + require.Equal(t, readStmt, followerStmt) + }) +} + +func TestCockroachDialect_DeleteStmt(t *testing.T) { + d := CockroachDialect{} + + stmt := d.DeleteStmt("kv", 3) + require.Equal(t, "DELETE FROM kv WHERE k IN ($1, $2, $3)", stmt) +} + +func TestPostgresDialect_DeleteStmt(t *testing.T) { + d := PostgresDialect{} + + stmt := d.DeleteStmt("kv", 3) + require.Equal(t, "DELETE FROM kv WHERE k IN ($1, $2, $3)", stmt) +} + +func TestCockroachDialect_SpanStmt(t *testing.T) { + d := CockroachDialect{} + + t.Run("uses scalar subquery syntax", func(t *testing.T) { + stmt := d.SpanStmt("kv", 100) + // CockroachDB uses [SELECT ...] scalar subquery syntax. + require.Contains(t, stmt, "[SELECT v FROM kv") + require.Contains(t, stmt, "LIMIT 100") + }) + + t.Run("full statement", func(t *testing.T) { + stmt := d.SpanStmt("kv", 50) + expected := "SELECT count(v) FROM [SELECT v FROM kv WHERE k >= $1 ORDER BY k LIMIT 50]" + require.Equal(t, expected, stmt) + }) + + t.Run("limit zero scans full table", func(t *testing.T) { + stmt := d.SpanStmt("kv", 0) + expected := "SELECT count(v) FROM [SELECT v FROM kv]" + require.Equal(t, expected, stmt) + }) +} + +func TestPostgresDialect_SpanStmt(t *testing.T) { + d := PostgresDialect{} + + t.Run("uses derived table syntax", func(t *testing.T) { + stmt := d.SpanStmt("kv", 100) + // PostgreSQL uses (SELECT ...) AS t derived table syntax. + require.Contains(t, stmt, "(SELECT v FROM kv") + require.Contains(t, stmt, ") AS t") + require.Contains(t, stmt, "LIMIT 100") + }) + + t.Run("full statement", func(t *testing.T) { + stmt := d.SpanStmt("kv", 50) + expected := "SELECT count(v) FROM (SELECT v FROM kv WHERE k >= $1 ORDER BY k LIMIT 50) AS t" + require.Equal(t, expected, stmt) + }) + + t.Run("limit zero scans full table", func(t *testing.T) { + stmt := d.SpanStmt("kv", 0) + expected := "SELECT count(v) FROM (SELECT v FROM kv) AS t" + require.Equal(t, expected, stmt) + }) +} + +func TestCockroachDialect_SelectForUpdateStmt(t *testing.T) { + d := CockroachDialect{} + + stmt := d.SelectForUpdateStmt("kv", 2) + require.Equal(t, "SELECT k, v FROM kv WHERE k IN ($1, $2) FOR UPDATE", stmt) +} + +func TestPostgresDialect_SelectForUpdateStmt(t *testing.T) { + d := PostgresDialect{} + + stmt := d.SelectForUpdateStmt("kv", 2) + require.Equal(t, "SELECT k, v FROM kv WHERE k IN ($1, $2) FOR UPDATE", stmt) +} + +func TestCockroachDialect_SchemaFragment(t *testing.T) { + d := CockroachDialect{} + + t.Run("basic table", func(t *testing.T) { + frag := d.SchemaFragment(0, false, 0) + require.Equal(t, "(\n\t\tk BIGINT NOT NULL PRIMARY KEY,\n\t\tv BYTES NOT NULL\n\t)", frag) + }) + + t.Run("with hash sharding", func(t *testing.T) { + frag := d.SchemaFragment(0, false, 8) + require.Contains(t, frag, "k BIGINT NOT NULL") + require.Contains(t, frag, "USING HASH WITH (bucket_count = 8)") + }) + + t.Run("with secondary index", func(t *testing.T) { + frag := d.SchemaFragment(0, true, 0) + require.Contains(t, frag, "INDEX (v)") + }) + + t.Run("with string key", func(t *testing.T) { + frag := d.SchemaFragment(20, false, 0) + require.Contains(t, frag, "k STRING NOT NULL") + }) + + t.Run("hash sharding and secondary index", func(t *testing.T) { + frag := d.SchemaFragment(0, true, 8) + require.Contains(t, frag, "USING HASH WITH (bucket_count = 8)") + require.Contains(t, frag, "INDEX (v)") + }) +} + +func TestPostgresDialect_SchemaFragment(t *testing.T) { + d := PostgresDialect{} + + t.Run("basic table", func(t *testing.T) { + frag := d.SchemaFragment(0, false, 0) + require.Equal(t, "(k BIGINT NOT NULL PRIMARY KEY, v BYTEA NOT NULL)", frag) + }) + + t.Run("with string key", func(t *testing.T) { + frag := d.SchemaFragment(100, false, 0) + require.Equal(t, "(k VARCHAR(100) NOT NULL PRIMARY KEY, v BYTEA NOT NULL)", frag) + }) + + t.Run("ignores hash sharding", func(t *testing.T) { + withShards := d.SchemaFragment(0, false, 8) + withoutShards := d.SchemaFragment(0, false, 0) + require.Equal(t, withoutShards, withShards) + }) + + t.Run("ignores secondaryIndex", func(t *testing.T) { + withIndex := d.SchemaFragment(0, true, 0) + withoutIndex := d.SchemaFragment(0, false, 0) + require.Equal(t, withoutIndex, withIndex) + }) + + t.Run("uses BYTEA not BYTES", func(t *testing.T) { + frag := d.SchemaFragment(0, false, 0) + require.Contains(t, frag, "BYTEA") + }) +} + +func TestCockroachDialect_Capabilities(t *testing.T) { + d := CockroachDialect{} + caps := d.Capabilities() + + require.True(t, caps.FollowerReads, "CockroachDB supports follower reads") + require.True(t, caps.Scatter, "CockroachDB supports scatter") + require.True(t, caps.Splits, "CockroachDB supports splits") + require.True(t, caps.HashShardedPK, "CockroachDB supports hash-sharded PKs") + require.True(t, caps.TransactionQoS, "CockroachDB supports transaction QoS") + require.True(t, caps.SelectForUpdate, "CockroachDB supports SELECT FOR UPDATE") + require.True(t, caps.TransactionPriority, "CockroachDB supports transaction priority") + require.True(t, caps.SerializationRetry, "CockroachDB supports serialization retry") + require.True(t, caps.Select1, "CockroachDB supports SELECT 1 warmup") + require.True(t, caps.EnumColumn, "CockroachDB supports enum columns") +} + +func TestPostgresDialect_Capabilities(t *testing.T) { + d := PostgresDialect{} + caps := d.Capabilities() + + require.False(t, caps.FollowerReads, "PostgreSQL does not support follower reads") + require.False(t, caps.Scatter, "PostgreSQL does not support scatter") + require.False(t, caps.Splits, "PostgreSQL does not support splits") + require.False(t, caps.HashShardedPK, "PostgreSQL does not support hash-sharded PKs") + require.False(t, caps.TransactionQoS, "PostgreSQL does not support transaction QoS") + require.True(t, caps.SelectForUpdate, "PostgreSQL supports SELECT FOR UPDATE") + require.False(t, caps.TransactionPriority, "PostgreSQL does not support transaction priority") + require.False(t, caps.SerializationRetry, "PostgreSQL does not support serialization retry") + require.True(t, caps.Select1, "PostgreSQL supports SELECT 1") + require.False(t, caps.EnumColumn, "PostgreSQL does not support enum columns") +} + +func TestSpannerSchemaDialect_SchemaFragment(t *testing.T) { + d := SpannerSchemaDialect{} + + t.Run("integer key", func(t *testing.T) { + frag := d.SchemaFragment(0, false, 0) + require.Equal(t, "(k INT64 NOT NULL, v BYTES(MAX) NOT NULL) PRIMARY KEY (k)", frag) + }) + + t.Run("string key", func(t *testing.T) { + frag := d.SchemaFragment(100, false, 0) + require.Equal(t, "(k STRING(100) NOT NULL, v BYTES(MAX) NOT NULL) PRIMARY KEY (k)", frag) + }) + + t.Run("ignores hash sharding", func(t *testing.T) { + withShards := d.SchemaFragment(0, false, 8) + withoutShards := d.SchemaFragment(0, false, 0) + require.Equal(t, withoutShards, withShards) + }) + + t.Run("ignores secondaryIndex", func(t *testing.T) { + withIndex := d.SchemaFragment(0, true, 0) + withoutIndex := d.SchemaFragment(0, false, 0) + require.Equal(t, withoutIndex, withIndex) + }) + + t.Run("uses BYTES(MAX)", func(t *testing.T) { + frag := d.SchemaFragment(0, false, 0) + require.Contains(t, frag, "BYTES(MAX)") + }) +} + +func TestSecondaryIndexStmts(t *testing.T) { + t.Run("CockroachDialect returns nil", func(t *testing.T) { + d := CockroachDialect{} + require.Nil(t, d.SecondaryIndexStmts("kv")) + }) + + t.Run("PostgresDialect returns CREATE INDEX IF NOT EXISTS", func(t *testing.T) { + d := PostgresDialect{} + stmts := d.SecondaryIndexStmts("kv") + require.Equal(t, []string{"CREATE INDEX IF NOT EXISTS kv_v_idx ON kv (v)"}, stmts) + }) + + t.Run("SpannerSchemaDialect returns CREATE INDEX without IF NOT EXISTS", func(t *testing.T) { + d := SpannerSchemaDialect{} + stmts := d.SecondaryIndexStmts("kv") + require.Equal(t, []string{"CREATE INDEX kv_v_idx ON kv (v)"}, stmts) + }) +} + +func TestPostLoadStmts(t *testing.T) { + t.Run("CockroachDialect", func(t *testing.T) { + d := CockroachDialect{} + + t.Run("enum only", func(t *testing.T) { + stmts := d.PostLoadStmts("kv", true, false) + require.NotNil(t, stmts) + require.Len(t, stmts, 1) + require.Contains(t, stmts[0], "CREATE TYPE") + require.Contains(t, stmts[0], "ALTER TABLE") + }) + + t.Run("scatter only", func(t *testing.T) { + stmts := d.PostLoadStmts("kv", false, true) + require.NotNil(t, stmts) + require.Len(t, stmts, 1) + require.Contains(t, stmts[0], "ALTER TABLE kv SCATTER") + }) + + t.Run("enum and scatter", func(t *testing.T) { + stmts := d.PostLoadStmts("kv", true, true) + require.NotNil(t, stmts) + require.Len(t, stmts, 2) + }) + + t.Run("neither", func(t *testing.T) { + stmts := d.PostLoadStmts("kv", false, false) + require.Nil(t, stmts) + }) + }) + + t.Run("PostgresDialect always returns nil", func(t *testing.T) { + d := PostgresDialect{} + require.Nil(t, d.PostLoadStmts("kv", true, true)) + require.Nil(t, d.PostLoadStmts("kv", false, false)) + }) + + t.Run("SpannerSchemaDialect always returns nil", func(t *testing.T) { + d := SpannerSchemaDialect{} + require.Nil(t, d.PostLoadStmts("kv", true, true)) + require.Nil(t, d.PostLoadStmts("kv", false, false)) + }) +} + +func TestSpannerSchemaDialect_Capabilities(t *testing.T) { + d := SpannerSchemaDialect{} + caps := d.Capabilities() + + require.True(t, caps.FollowerReads, "Spanner supports follower reads") + require.True(t, caps.EnumColumn, "Spanner supports enum columns") + require.False(t, caps.Select1, "Spanner does not support SELECT 1") + require.False(t, caps.SelectForUpdate, "Spanner does not support SELECT FOR UPDATE") + require.False(t, caps.Scatter, "Spanner does not support scatter") + require.False(t, caps.Splits, "Spanner does not support splits") + require.False(t, caps.HashShardedPK, "Spanner does not support hash-sharded PKs") + require.False(t, caps.TransactionQoS, "Spanner does not support transaction QoS") + require.False(t, caps.TransactionPriority, "Spanner does not support transaction priority") + require.False(t, caps.SerializationRetry, "Spanner does not support serialization retry") +} + +func TestDialectInterface(t *testing.T) { + // Verify both dialects implement the Dialect interface. + dialects := []Dialect{ + CockroachDialect{}, + PostgresDialect{}, + } + + for _, d := range dialects { + t.Run(d.Name(), func(t *testing.T) { + // All Dialect methods should be callable without panicking. + require.NotEmpty(t, d.Name()) + require.NotEmpty(t, d.UpsertStmt("kv", 1, false)) + require.NotEmpty(t, d.ReadStmt("kv", 1, false)) + require.NotEmpty(t, d.FollowerReadStmt("kv", 1, false)) + require.NotEmpty(t, d.DeleteStmt("kv", 1)) + require.NotEmpty(t, d.SpanStmt("kv", 10)) + require.NotEmpty(t, d.SelectForUpdateStmt("kv", 1)) + + // SchemaDialect methods should be callable. + require.NotEmpty(t, d.SchemaFragment(0, false, 0)) + _ = d.SecondaryIndexStmts("kv") + _ = d.Capabilities() + }) + } + + // SpannerSchemaDialect only implements SchemaDialect, not full Dialect. + t.Run("SpannerSchemaDialect", func(t *testing.T) { + var d SchemaDialect = SpannerSchemaDialect{} + require.NotEmpty(t, d.SchemaFragment(0, false, 0)) + require.NotNil(t, d.SecondaryIndexStmts("kv")) + _ = d.Capabilities() + }) +} + +func TestCustomTableName(t *testing.T) { + dialects := []Dialect{ + CockroachDialect{}, + PostgresDialect{}, + } + + for _, d := range dialects { + t.Run(d.Name(), func(t *testing.T) { + customTable := "my_custom_table" + + stmt := d.UpsertStmt(customTable, 1, false) + require.True(t, strings.Contains(stmt, customTable)) + + stmt = d.ReadStmt(customTable, 1, false) + require.True(t, strings.Contains(stmt, customTable)) + + stmt = d.DeleteStmt(customTable, 1) + require.True(t, strings.Contains(stmt, customTable)) + + stmt = d.SpanStmt(customTable, 10) + require.True(t, strings.Contains(stmt, customTable)) + + // SecondaryIndexStmts should include the custom table name. + stmts := d.SecondaryIndexStmts(customTable) + if stmts != nil { + for _, s := range stmts { + require.True(t, strings.Contains(s, customTable)) + } + } + + // PostLoadStmts should include the custom table name. + postStmts := d.PostLoadStmts(customTable, true, true) + for _, s := range postStmts { + require.Contains(t, s, customTable) + } + }) + } + + // SpannerSchemaDialect SecondaryIndexStmts with custom table name. + t.Run("SpannerSchemaDialect", func(t *testing.T) { + d := SpannerSchemaDialect{} + customTable := "my_custom_table" + stmts := d.SecondaryIndexStmts(customTable) + require.NotNil(t, stmts) + for _, s := range stmts { + require.True(t, strings.Contains(s, customTable)) + } + }) +} diff --git a/pkg/workload/dbexec/executor.go b/pkg/workload/dbexec/executor.go new file mode 100644 index 000000000000..5616a59e5da3 --- /dev/null +++ b/pkg/workload/dbexec/executor.go @@ -0,0 +1,231 @@ +// Copyright 2024 The Cockroach Authors. +// +// Use of this software is governed by the CockroachDB Software License +// included in the /LICENSE file. + +// Package dbexec provides database-agnostic execution interfaces for workloads. +// +// This package abstracts database operations for workloads like kv to support +// multiple databases (CockroachDB, PostgreSQL, Aurora, DSQL, Spanner). Each +// database has different SQL syntax, capabilities, and connection handling, +// which are encapsulated by the Executor and Dialect interfaces. +package dbexec + +import ( + "context" + + "github.com/cockroachdb/cockroach/pkg/workload" +) + +// Row represents a key-value pair for the kv workload. +type Row struct { + K interface{} // int64 or string depending on key type + V []byte + E *string // optional enum value, nil if not using enum +} + +// Rows is an iterator over query results. Implementations must be safe +// to use from a single goroutine and must be closed after use. +type Rows interface { + // Next advances to the next row, returning false when no more rows exist + // or an error occurred. + Next() bool + + // Scan copies the columns from the current row into dest. The number of + // dest values must match the number of columns in the result set. + Scan(dest ...interface{}) error + + // Err returns any error encountered during iteration. + Err() error + + // Close releases resources associated with the result set. + Close() +} + +// Capabilities indicates which features the executor supports. This allows +// workloads to adapt their behavior based on database capabilities. +type Capabilities struct { + // FollowerReads indicates support for stale/replica reads (e.g., CockroachDB's + // AS OF SYSTEM TIME follower_read_timestamp()). + FollowerReads bool + + // Scatter indicates support for the ALTER TABLE ... SCATTER command to + // distribute ranges across the cluster. + Scatter bool + + // Splits indicates support for the ALTER TABLE ... SPLIT AT command to + // manually split ranges at specific keys. + Splits bool + + // HashShardedPK indicates support for hash-sharded primary keys using + // PRIMARY KEY USING HASH WITH (bucket_count = N). + HashShardedPK bool + + // TransactionQoS indicates support for transaction priority/QoS settings + // (e.g., CockroachDB's SET TRANSACTION PRIORITY). + TransactionQoS bool + + // SelectForUpdate indicates support for SELECT ... FOR UPDATE locking. + SelectForUpdate bool + + // TransactionPriority indicates support for BEGIN PRIORITY {low|normal|high}. + TransactionPriority bool + + // SerializationRetry indicates the executor handles serialization failures gracefully. + SerializationRetry bool + + // Select1 indicates the executor supports SELECT 1 as a transaction warmup + // to prevent automatic transaction-level retries. This is a CockroachDB- + // specific optimization used by --sel1-writes. + Select1 bool + + // EnumColumn indicates support for enum types in schema. + EnumColumn bool +} + +// Config holds executor configuration. These values are typically derived +// from workload flags and control table schema and query generation. +type Config struct { + // URLs contains database connection strings. Multiple URLs enable + // connection distribution across nodes. + URLs []string + + // Table is the name of the kv table to operate on. + Table string + + // BatchSize is the number of rows per batch operation (read, write, delete). + BatchSize int + + // KeyType specifies the primary key column type, either "BIGINT" or + // "STRING(N)" where N is the string length. + KeyType string + + // SecondaryIndex controls whether to create a secondary index on the + // value column. + SecondaryIndex bool + + // NumShards specifies the number of hash buckets for hash-sharded primary + // keys. Only used when the database supports HashShardedPK. + NumShards int + + // Concurrency is the number of concurrent workers accessing the database. + Concurrency int + + // Enum controls whether to include an enum column in the schema. + Enum bool + + // SpanLimit is the default LIMIT for spanning queries. + SpanLimit int + + // TxnQoS sets transaction quality of service (background/regular/critical). + // Only supported by CockroachDB. + TxnQoS string +} + +// TxOptions configures transaction behavior. +type TxOptions struct { + // Priority sets the transaction priority level. Valid values are "low", + // "normal", and "high". Not all databases support transaction priorities. + Priority string +} + +// Tx represents a database transaction. All operations within a transaction +// see a consistent snapshot and are committed atomically. +type Tx interface { + // Exec executes a statement that does not return rows. + Exec(ctx context.Context, sql string, args ...interface{}) error + + // Query executes a query that returns rows. + Query(ctx context.Context, sql string, args ...interface{}) (Rows, error) + + // Commit commits the transaction. After Commit returns, the transaction + // is no longer valid. + Commit(ctx context.Context) error + + // Rollback aborts the transaction. After Rollback returns, the transaction + // is no longer valid. + Rollback(ctx context.Context) error +} + +// ExtendedTx extends Tx with transactional operations that use prepared +// statements. Both CRDBExecutor and PGXExecutor implement this interface, +// enabling transactional write workflows like SELECT FOR UPDATE → delay → +// write across multiple database backends. +type ExtendedTx interface { + Tx + + // SelectForUpdate executes SELECT ... FOR UPDATE within the transaction. + SelectForUpdate(ctx context.Context, keys []interface{}) (Rows, error) + + // Select1 executes SELECT 1 within the transaction. On CockroachDB, this + // prevents automatic transaction-level retries. On other databases, it is + // a harmless no-op. + Select1(ctx context.Context) error + + // Write executes the prepared write statement within the transaction. + Write(ctx context.Context, rows []Row) error +} + +// flattenRowArgs converts a slice of Rows into a flat argument list for +// parameterized SQL statements. If includeEnum is true, each row contributes +// (K, V, E); otherwise each row contributes (K, V). +func flattenRowArgs(rows []Row, includeEnum bool) []interface{} { + if includeEnum { + args := make([]interface{}, 0, len(rows)*3) + for _, row := range rows { + args = append(args, row.K, row.V, row.E) + } + return args + } + args := make([]interface{}, 0, len(rows)*2) + for _, row := range rows { + args = append(args, row.K, row.V) + } + return args +} + +// Executor abstracts database operations for the kv workload. Each +// implementation handles connection pooling, SQL generation, and execution +// for a specific database or family of databases. +// +// Implementations must be safe for concurrent use from multiple goroutines. +type Executor interface { + // Init prepares the executor by creating connection pools and optionally + // preparing statements. This must be called before any other methods. + Init(ctx context.Context, cfg Config, connFlags *workload.ConnFlags) error + + // Read executes a batch read for the given keys. The returned Rows + // iterator yields (k, v) pairs for keys that exist. + Read(ctx context.Context, keys []interface{}) (Rows, error) + + // FollowerRead executes a stale/replica read if supported by the database. + // Falls back to a regular read if follower reads are not supported. + FollowerRead(ctx context.Context, keys []interface{}) (Rows, error) + + // Write upserts a batch of key-value pairs. Existing keys are updated; + // new keys are inserted. + Write(ctx context.Context, rows []Row) error + + // Delete removes a batch of keys. Non-existent keys are silently ignored. + Delete(ctx context.Context, keys []interface{}) error + + // Span executes a spanning query starting at startKey with the given limit. + // Returns the count of values scanned. + Span(ctx context.Context, startKey interface{}, limit int) (int64, error) + + // BeginTx starts a transaction with the given options. The returned Tx + // must be committed or rolled back. + BeginTx(ctx context.Context, opts TxOptions) (Tx, error) + + // Close releases all resources including connection pools. After Close + // returns, the executor must not be used. + Close() error + + // Capabilities returns what features this executor supports. + Capabilities() Capabilities + + // SetupWorker performs per-worker initialization. Called once when each + // worker starts, before any operations. Currently a no-op for all + // executors; retained as an extension point for future executor types. + SetupWorker(ctx context.Context) error +} diff --git a/pkg/workload/dbexec/pgx_executor.go b/pkg/workload/dbexec/pgx_executor.go new file mode 100644 index 000000000000..45e5ece2bce1 --- /dev/null +++ b/pkg/workload/dbexec/pgx_executor.go @@ -0,0 +1,265 @@ +// Copyright 2024 The Cockroach Authors. +// +// Use of this software is governed by the CockroachDB Software License +// included in the /LICENSE file. + +package dbexec + +import ( + "context" + "fmt" + + "github.com/cockroachdb/cockroach/pkg/workload" + "github.com/cockroachdb/errors" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" +) + +// PGXExecutor implements Executor for PostgreSQL-compatible databases using pgx. +// This includes CockroachDB, PostgreSQL, Aurora PostgreSQL, and other compatible +// databases. The executor uses a Dialect to generate database-specific SQL. +type PGXExecutor struct { + dialect Dialect + cfg Config + + mcp *workload.MultiConnPool + sr workload.SQLRunner + + // Statement handles for prepared statements. + readStmt workload.StmtHandle + followerReadStmt workload.StmtHandle + writeStmt workload.StmtHandle + deleteStmt workload.StmtHandle + spanStmt workload.StmtHandle + sfuStmt workload.StmtHandle // SELECT ... FOR UPDATE + sel1Stmt workload.StmtHandle // SELECT 1 +} + +var _ Executor = (*PGXExecutor)(nil) + +// NewPGXExecutor creates a new PGXExecutor with the given dialect. +func NewPGXExecutor(dialect Dialect) *PGXExecutor { + return &PGXExecutor{dialect: dialect} +} + +// Init prepares the executor by creating connection pools and defining +// statements via SQLRunner. +func (e *PGXExecutor) Init(ctx context.Context, cfg Config, connFlags *workload.ConnFlags) error { + e.cfg = cfg + + // Create multi-connection pool. + poolCfg := workload.NewMultiConnPoolCfgFromFlags(connFlags) + poolCfg.MaxTotalConnections = connFlags.Concurrency + 1 + mcp, err := workload.NewMultiConnPool(ctx, poolCfg, cfg.URLs...) + if err != nil { + return errors.Wrap(err, "creating connection pool") + } + e.mcp = mcp + + // Define statements using the dialect. + e.readStmt = e.sr.Define(e.dialect.ReadStmt(cfg.Table, cfg.BatchSize, cfg.Enum)) + e.followerReadStmt = e.sr.Define(e.dialect.FollowerReadStmt(cfg.Table, cfg.BatchSize, cfg.Enum)) + e.writeStmt = e.sr.Define(e.dialect.UpsertStmt(cfg.Table, cfg.BatchSize, cfg.Enum)) + e.deleteStmt = e.sr.Define(e.dialect.DeleteStmt(cfg.Table, cfg.BatchSize)) + e.spanStmt = e.sr.Define(e.dialect.SpanStmt(cfg.Table, cfg.SpanLimit)) + e.sfuStmt = e.sr.Define(e.dialect.SelectForUpdateStmt(cfg.Table, cfg.BatchSize)) + e.sel1Stmt = e.sr.Define("SELECT 1") + + // Initialize the SQLRunner with prepared statements. + if err := e.sr.Init(ctx, e.dialect.Name(), e.mcp); err != nil { + e.mcp.Close() + return errors.Wrap(err, "initializing SQL runner") + } + + return nil +} + +// Read executes a batch read for the given keys. +func (e *PGXExecutor) Read(ctx context.Context, keys []interface{}) (Rows, error) { + rows, err := e.readStmt.Query(ctx, keys...) + if err != nil { + return nil, err + } + return &pgxRowsWrapper{rows: rows}, nil +} + +// FollowerRead executes a stale/replica read if supported by the dialect. +func (e *PGXExecutor) FollowerRead(ctx context.Context, keys []interface{}) (Rows, error) { + rows, err := e.followerReadStmt.Query(ctx, keys...) + if err != nil { + return nil, err + } + return &pgxRowsWrapper{rows: rows}, nil +} + +// Write upserts a batch of key-value pairs. +func (e *PGXExecutor) Write(ctx context.Context, rows []Row) error { + args := flattenRowArgs(rows, e.cfg.Enum) + _, err := e.writeStmt.Exec(ctx, args...) + return err +} + +// Delete removes a batch of keys. +func (e *PGXExecutor) Delete(ctx context.Context, keys []interface{}) error { + _, err := e.deleteStmt.Exec(ctx, keys...) + return err +} + +// Span executes a spanning query starting at startKey with the given limit. +func (e *PGXExecutor) Span(ctx context.Context, startKey interface{}, limit int) (int64, error) { + // The limit is baked into the prepared statement via SpanStmt. When + // SpanLimit == 0, the statement is a full table scan with no arguments. + // When SpanLimit > 0, it takes a startKey argument for the WHERE clause. + var row pgx.Row + if e.cfg.SpanLimit == 0 { + row = e.spanStmt.QueryRow(ctx) + } else { + row = e.spanStmt.QueryRow(ctx, startKey) + } + var count int64 + if err := row.Scan(&count); err != nil { + return 0, err + } + return count, nil +} + +// BeginTx starts a transaction with the given options. +func (e *PGXExecutor) BeginTx(ctx context.Context, opts TxOptions) (Tx, error) { + pool := e.mcp.Get() + conn, err := pool.Acquire(ctx) + if err != nil { + return nil, errors.Wrap(err, "acquiring connection") + } + + txOpts := pgx.TxOptions{} + + // Set transaction priority if supported and specified. + if opts.Priority != "" && e.dialect.Capabilities().TransactionQoS { + txOpts.BeginQuery = fmt.Sprintf("BEGIN PRIORITY %s", opts.Priority) + } + + tx, err := conn.BeginTx(ctx, txOpts) + if err != nil { + conn.Release() + return nil, errors.Wrap(err, "beginning transaction") + } + + return &pgxTxWrapper{ + tx: tx, + conn: conn, + executor: e, + }, nil +} + +// Close releases all resources including connection pools. +func (e *PGXExecutor) Close() error { + if e.mcp != nil { + e.mcp.Close() + } + return nil +} + +// Capabilities returns what features this executor supports. +func (e *PGXExecutor) Capabilities() Capabilities { + return e.dialect.Capabilities() +} + +// SetupWorker is a no-op for PGXExecutor. CRDB-specific per-worker setup +// (like transaction QoS) is not supported for generic PostgreSQL-compatible databases. +func (e *PGXExecutor) SetupWorker(ctx context.Context) error { + return nil +} + +// pgxRowsWrapper wraps pgx.Rows to implement our Rows interface. +type pgxRowsWrapper struct { + rows pgx.Rows +} + +var _ Rows = (*pgxRowsWrapper)(nil) + +// Next advances to the next row. +func (r *pgxRowsWrapper) Next() bool { + return r.rows.Next() +} + +// Scan copies the columns from the current row into dest. +func (r *pgxRowsWrapper) Scan(dest ...interface{}) error { + return r.rows.Scan(dest...) +} + +// Err returns any error encountered during iteration. +func (r *pgxRowsWrapper) Err() error { + return r.rows.Err() +} + +// Close releases resources associated with the result set. +func (r *pgxRowsWrapper) Close() { + r.rows.Close() +} + +// pgxTxWrapper wraps pgx.Tx and the connection to implement Tx and ExtendedTx. +// It holds the connection so it can be released when the transaction completes, +// and a reference to the executor for prepared statement access. +type pgxTxWrapper struct { + tx pgx.Tx + conn *pgxpool.Conn + executor *PGXExecutor +} + +var _ Tx = (*pgxTxWrapper)(nil) +var _ ExtendedTx = (*pgxTxWrapper)(nil) + +// Exec executes a statement that does not return rows. +func (t *pgxTxWrapper) Exec(ctx context.Context, sql string, args ...interface{}) error { + _, err := t.tx.Exec(ctx, sql, args...) + return err +} + +// Query executes a query that returns rows. +func (t *pgxTxWrapper) Query(ctx context.Context, sql string, args ...interface{}) (Rows, error) { + rows, err := t.tx.Query(ctx, sql, args...) + if err != nil { + return nil, err + } + return &pgxRowsWrapper{rows: rows}, nil +} + +// Commit commits the transaction. +func (t *pgxTxWrapper) Commit(ctx context.Context) error { + err := t.tx.Commit(ctx) + t.conn.Release() + return err +} + +// Rollback aborts the transaction. +func (t *pgxTxWrapper) Rollback(ctx context.Context) error { + err := t.tx.Rollback(ctx) + t.conn.Release() + return err +} + +// SelectForUpdate executes a SELECT ... FOR UPDATE query within the transaction, +// locking the selected rows until the transaction completes. +func (t *pgxTxWrapper) SelectForUpdate(ctx context.Context, keys []interface{}) (Rows, error) { + rows, err := t.executor.sfuStmt.QueryTx(ctx, t.tx, keys...) + if err != nil { + return nil, err + } + return &pgxRowsWrapper{rows: rows}, nil +} + +// Select1 executes a SELECT 1 query within the transaction. On CockroachDB, +// this prevents automatic transaction-level retries. On PostgreSQL, it is a +// harmless no-op that keeps the transaction active. +func (t *pgxTxWrapper) Select1(ctx context.Context) error { + row := t.executor.sel1Stmt.QueryRowTx(ctx, t.tx) + var dummy int + return row.Scan(&dummy) +} + +// Write executes an upsert within the transaction using prepared statements. +func (t *pgxTxWrapper) Write(ctx context.Context, rows []Row) error { + args := flattenRowArgs(rows, t.executor.cfg.Enum) + _, err := t.executor.writeStmt.ExecTx(ctx, t.tx, args...) + return err +} diff --git a/pkg/workload/dbexec/spanner_executor.go b/pkg/workload/dbexec/spanner_executor.go new file mode 100644 index 000000000000..e87ec611b899 --- /dev/null +++ b/pkg/workload/dbexec/spanner_executor.go @@ -0,0 +1,410 @@ +// Copyright 2024 The Cockroach Authors. +// +// Use of this software is governed by the CockroachDB Software License +// included in the /LICENSE file. + +package dbexec + +import ( + "context" + "fmt" + "log" + "strings" + "time" + + "cloud.google.com/go/spanner" + "github.com/cockroachdb/cockroach/pkg/workload" + "github.com/cockroachdb/errors" + "google.golang.org/api/iterator" + "google.golang.org/grpc/codes" +) + +// SpannerExecutor implements Executor for Google Cloud Spanner. Unlike SQL-based +// executors, SpannerExecutor uses the Spanner mutations API for writes, which +// provides better performance and is the idiomatic way to interact with Spanner. +// +// Connection format: projects/PROJECT/instances/INSTANCE/databases/DATABASE +type SpannerExecutor struct { + client *spanner.Client + cfg Config +} + +// NewSpannerExecutor creates a new SpannerExecutor. Init must be called before +// using the executor. +func NewSpannerExecutor() *SpannerExecutor { + return &SpannerExecutor{} +} + +var _ Executor = (*SpannerExecutor)(nil) + +// Init creates the Spanner client from the database path in cfg.URLs[0]. +// The connFlags parameter is not used for Spanner as it has its own connection +// handling. +func (e *SpannerExecutor) Init( + ctx context.Context, cfg Config, connFlags *workload.ConnFlags, +) error { + if len(cfg.URLs) == 0 { + return errors.New("spanner executor requires at least one URL") + } + + // The URL should be in the format: + // projects/PROJECT/instances/INSTANCE/databases/DATABASE + dbPath := cfg.URLs[0] + log.Printf("[spanner] connecting to database: %s", dbPath) + log.Printf("[spanner] target table: %s", cfg.Table) + client, err := spanner.NewClient(ctx, dbPath) + if err != nil { + return errors.Wrap(err, "creating spanner client") + } + + e.client = client + e.cfg = cfg + + // Validate connection by executing a simple query to check if the table exists. + // This provides a clearer error message than failing on first write. + if err := e.validateTable(ctx); err != nil { + client.Close() + return err + } + + return nil +} + +// validateTable checks that the target table exists and is accessible. +// This uses a simple read operation that will trigger session pool creation +// and validate the table exists. +func (e *SpannerExecutor) validateTable(ctx context.Context) error { + // Don't add a timeout here - Spanner's session pool initialization can take + // 60+ seconds on first connection. The Spanner client has its own timeouts. + stmt := spanner.Statement{ + SQL: fmt.Sprintf("SELECT k FROM %s LIMIT 1", e.cfg.Table), + } + log.Printf("[spanner] validating table with query: %s", stmt.SQL) + iter := e.client.Single().Query(ctx, stmt) + defer iter.Stop() + + // Try to get one row - we don't care about the result, just that the query works. + log.Printf("[spanner] waiting for query result (session pool init may take 60+ seconds)...") + _, err := iter.Next() + if err != nil && !errors.Is(err, iterator.Done) { + // Check if this is a "table not found" error. + if spanner.ErrCode(err) == codes.NotFound || + strings.Contains(err.Error(), "not found") || + strings.Contains(err.Error(), "Table not found") || + strings.Contains(err.Error(), "does not exist") { + return errors.Newf( + "table %q does not exist in Spanner database. "+ + "Please create it manually before running the workload:\n\n"+ + " CREATE TABLE %s (\n"+ + " k INT64 NOT NULL,\n"+ + " v BYTES(MAX),\n"+ + " ) PRIMARY KEY (k)", + e.cfg.Table, e.cfg.Table) + } + // For timeout errors during session creation, provide helpful context. + if spanner.ErrCode(err) == codes.DeadlineExceeded { + return errors.Wrapf(err, + "timeout connecting to Spanner (session pool initialization can take 60+ seconds). "+ + "Try running again, or check network connectivity and authentication") + } + return errors.Wrapf(err, "validating Spanner table %q", e.cfg.Table) + } + log.Printf("[spanner] table %q validated successfully", e.cfg.Table) + return nil +} + +// Read executes a batch read for the given keys using Spanner's Read API. +func (e *SpannerExecutor) Read(ctx context.Context, keys []interface{}) (Rows, error) { + keySet := spanner.KeySetFromKeys(toSpannerKeys(keys)...) + iter := e.client.Single().Read(ctx, e.cfg.Table, keySet, []string{"k", "v"}) + return &spannerRowsWrapper{iter: iter}, nil +} + +// FollowerRead executes a stale read using Spanner's exact staleness bound. +// This allows reads from any replica, reducing latency for read-heavy workloads. +func (e *SpannerExecutor) FollowerRead(ctx context.Context, keys []interface{}) (Rows, error) { + keySet := spanner.KeySetFromKeys(toSpannerKeys(keys)...) + // 15 seconds staleness is a reasonable default for follower reads. + ro := e.client.ReadOnlyTransaction().WithTimestampBound(spanner.ExactStaleness(15 * time.Second)) + + iter := ro.Read(ctx, e.cfg.Table, keySet, []string{"k", "v"}) + return &spannerStaleRowsWrapper{ + spannerRowsWrapper: spannerRowsWrapper{iter: iter}, + txn: ro, + }, nil +} + +// maxMutationsPerBatch is the recommended limit for Spanner mutations per Apply call. +// Larger batches can cause timeout issues. +const maxMutationsPerBatch = 100 + +// Write upserts a batch of key-value pairs using Spanner mutations. +// Uses InsertOrUpdate which inserts new rows or updates existing ones. +// For large batches, mutations are split into smaller batches to avoid timeouts. +func (e *SpannerExecutor) Write(ctx context.Context, rows []Row) error { + // For small batches, apply directly. + if len(rows) <= maxMutationsPerBatch { + return e.applyWriteMutations(ctx, rows) + } + + // For large batches, split into smaller chunks. + for i := 0; i < len(rows); i += maxMutationsPerBatch { + end := i + maxMutationsPerBatch + if end > len(rows) { + end = len(rows) + } + if err := e.applyWriteMutations(ctx, rows[i:end]); err != nil { + return errors.Wrapf(err, "batch [%d:%d]", i, end) + } + } + return nil +} + +// applyWriteMutations applies a batch of write mutations. +func (e *SpannerExecutor) applyWriteMutations(ctx context.Context, rows []Row) error { + mutations := make([]*spanner.Mutation, len(rows)) + for i, row := range rows { + if e.cfg.Enum { + var enumVal interface{} + if row.E != nil { + enumVal = *row.E + } + mutations[i] = spanner.InsertOrUpdate( + e.cfg.Table, + []string{"k", "v", "e"}, + []interface{}{row.K, row.V, enumVal}, + ) + } else { + mutations[i] = spanner.InsertOrUpdate( + e.cfg.Table, + []string{"k", "v"}, + []interface{}{row.K, row.V}, + ) + } + } + _, err := e.client.Apply(ctx, mutations) + if err != nil { + return errors.Wrap(err, "applying write mutations") + } + return nil +} + +// Delete removes a batch of keys using Spanner mutations. +func (e *SpannerExecutor) Delete(ctx context.Context, keys []interface{}) error { + mutations := make([]*spanner.Mutation, len(keys)) + for i, key := range keys { + mutations[i] = spanner.Delete(e.cfg.Table, spanner.Key{key}) + } + _, err := e.client.Apply(ctx, mutations) + if err != nil { + return errors.Wrap(err, "applying delete mutations") + } + return nil +} + +// Span executes a spanning query starting at startKey with the given limit. +// Uses SQL via Spanner's Query API since this requires ordered iteration. +func (e *SpannerExecutor) Span( + ctx context.Context, startKey interface{}, limit int, +) (int64, error) { + var stmt spanner.Statement + if limit == 0 { + // Full table scan -- no WHERE clause, no parameters. + stmt = spanner.Statement{ + SQL: fmt.Sprintf("SELECT COUNT(v) FROM %s", e.cfg.Table), + } + } else { + stmt = spanner.Statement{ + SQL: fmt.Sprintf( + "SELECT COUNT(v) FROM (SELECT v FROM %s WHERE k >= @startKey ORDER BY k LIMIT %d)", + e.cfg.Table, limit, + ), + Params: map[string]interface{}{ + "startKey": startKey, + }, + } + } + + iter := e.client.Single().Query(ctx, stmt) + defer iter.Stop() + + row, err := iter.Next() + if err != nil { + return 0, errors.Wrap(err, "executing span query") + } + + var count int64 + if err := row.Columns(&count); err != nil { + return 0, errors.Wrap(err, "scanning count") + } + + return count, nil +} + +// BeginTx starts a transaction that accumulates mutations. +// Spanner transactions work differently from SQL databases - mutations are +// accumulated and applied atomically on commit. +func (e *SpannerExecutor) BeginTx(ctx context.Context, opts TxOptions) (Tx, error) { + return &spannerTxWrapper{ + client: e.client, + table: e.cfg.Table, + }, nil +} + +// Close releases all resources including the Spanner client. +func (e *SpannerExecutor) Close() error { + if e.client != nil { + e.client.Close() + } + return nil +} + +// Capabilities returns Spanner's feature set. Spanner supports stale reads +// (follower reads) but does not support CockroachDB-specific features like +// scatter, splits, or hash-sharded primary keys. Enum columns are supported +// by storing enum values as strings. +func (e *SpannerExecutor) Capabilities() Capabilities { + return Capabilities{ + FollowerReads: true, // via stale reads + Scatter: false, + Splits: false, // Spanner auto-splits + HashShardedPK: false, + TransactionQoS: false, + SelectForUpdate: false, + TransactionPriority: false, + SerializationRetry: false, + EnumColumn: true, // stored as strings + } +} + +// SetupWorker is a no-op for SpannerExecutor. CRDB-specific per-worker setup +// (like transaction QoS) is not supported for Spanner. +func (e *SpannerExecutor) SetupWorker(ctx context.Context) error { + return nil +} + +// toSpannerKeys converts a slice of interface{} keys to spanner.Key values. +func toSpannerKeys(keys []interface{}) []spanner.Key { + result := make([]spanner.Key, len(keys)) + for i, k := range keys { + result[i] = spanner.Key{k} + } + return result +} + +// spannerRowsWrapper wraps a Spanner RowIterator to implement the Rows interface. +type spannerRowsWrapper struct { + iter *spanner.RowIterator + row *spanner.Row + err error +} + +var _ Rows = (*spannerRowsWrapper)(nil) + +// Next advances to the next row. +func (w *spannerRowsWrapper) Next() bool { + w.row, w.err = w.iter.Next() + return w.err == nil +} + +// Scan copies the columns from the current row into dest. +func (w *spannerRowsWrapper) Scan(dest ...interface{}) error { + if w.row == nil { + return errors.New("no current row") + } + return w.row.Columns(dest...) +} + +// Err returns any error encountered during iteration. +// Returns nil if the error is iterator.Done (normal end of iteration). +func (w *spannerRowsWrapper) Err() error { + if errors.Is(w.err, iterator.Done) { + return nil + } + return w.err +} + +// Close releases resources associated with the iterator. +func (w *spannerRowsWrapper) Close() { + w.iter.Stop() +} + +// spannerStaleRowsWrapper extends spannerRowsWrapper to also close the +// read-only transaction when the rows are closed. +type spannerStaleRowsWrapper struct { + spannerRowsWrapper + txn *spanner.ReadOnlyTransaction +} + +var _ Rows = (*spannerStaleRowsWrapper)(nil) + +// Close releases resources associated with the iterator and transaction. +func (w *spannerStaleRowsWrapper) Close() { + w.iter.Stop() + w.txn.Close() +} + +// spannerTxWrapper implements Tx for Spanner by accumulating mutations. +// Unlike SQL transactions, Spanner mutations are buffered and applied +// atomically when Commit is called. +type spannerTxWrapper struct { + client *spanner.Client + table string + mutations []*spanner.Mutation +} + +var _ Tx = (*spannerTxWrapper)(nil) + +// Exec executes a statement by converting it to mutations. +// This is a simplified implementation that only supports basic INSERT/UPDATE/DELETE. +// For complex operations, use Query or direct mutation methods. +func (t *spannerTxWrapper) Exec(ctx context.Context, sql string, args ...interface{}) error { + // For Spanner, we expect callers to use the mutation-based Write/Delete methods + // rather than SQL. This Exec implementation is provided for compatibility but + // has limited functionality. + return errors.New("Exec not fully supported in Spanner transactions; use Write/Delete methods") +} + +// Query executes a query within the transaction context. +func (t *spannerTxWrapper) Query( + ctx context.Context, sql string, args ...interface{}, +) (Rows, error) { + // Build parameters map from positional args. + params := make(map[string]interface{}) + for i, arg := range args { + params[fmt.Sprintf("p%d", i+1)] = arg + } + + stmt := spanner.Statement{ + SQL: sql, + Params: params, + } + + iter := t.client.Single().Query(ctx, stmt) + return &spannerRowsWrapper{iter: iter}, nil +} + +// Commit applies all accumulated mutations atomically. +func (t *spannerTxWrapper) Commit(ctx context.Context) error { + if len(t.mutations) == 0 { + return nil + } + _, err := t.client.Apply(ctx, t.mutations) + if err != nil { + return errors.Wrap(err, "committing spanner transaction") + } + t.mutations = nil + return nil +} + +// Rollback discards all accumulated mutations. +func (t *spannerTxWrapper) Rollback(ctx context.Context) error { + t.mutations = nil + return nil +} + +// AddMutation adds a mutation to be applied on commit. +// This is a Spanner-specific method for building up transaction mutations. +func (t *spannerTxWrapper) AddMutation(m *spanner.Mutation) { + t.mutations = append(t.mutations, m) +} diff --git a/pkg/workload/kv/BUILD.bazel b/pkg/workload/kv/BUILD.bazel index 6cc3d7b3732a..554b91ca7b77 100644 --- a/pkg/workload/kv/BUILD.bazel +++ b/pkg/workload/kv/BUILD.bazel @@ -16,6 +16,7 @@ go_library( "//pkg/util/randutil", "//pkg/util/timeutil", "//pkg/workload", + "//pkg/workload/dbexec", "//pkg/workload/histogram", "@com_github_cockroachdb_errors//:errors", "@com_github_jackc_pgx_v5//:pgx", diff --git a/pkg/workload/kv/kv.go b/pkg/workload/kv/kv.go index ffbfed2a633e..5d458e74b433 100644 --- a/pkg/workload/kv/kv.go +++ b/pkg/workload/kv/kv.go @@ -27,6 +27,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/util/randutil" "github.com/cockroachdb/cockroach/pkg/util/timeutil" "github.com/cockroachdb/cockroach/pkg/workload" + "github.com/cockroachdb/cockroach/pkg/workload/dbexec" "github.com/cockroachdb/cockroach/pkg/workload/histogram" "github.com/cockroachdb/errors" "github.com/jackc/pgx/v5" @@ -34,26 +35,9 @@ import ( "github.com/spf13/pflag" ) -const ( - kvSchema = `( - k %s NOT NULL PRIMARY KEY, - v BYTES NOT NULL - )` - kvSchemaWithIndex = `( - k %s NOT NULL PRIMARY KEY, - v BYTES NOT NULL, - INDEX (v) - )` - shardedKvSchema = `( - k %s NOT NULL PRIMARY KEY USING HASH WITH (bucket_count = %d), - v BYTES NOT NULL - )` - shardedKvSchemaWithIndex = `( - k %s NOT NULL PRIMARY KEY USING HASH WITH (bucket_count = %d), - v BYTES NOT NULL, - INDEX (v) - )` -) +// Ensure kv implements the DialectProvider interface. +var _ workload.DialectProvider = (*kv)(nil) +var _ workload.DialectSetter = (*kv)(nil) var RandomSeed = workload.NewInt64RandomSeed() @@ -98,6 +82,7 @@ type kv struct { prepareReadOnly bool writesUseSelect1 bool alwaysIncKeySeq bool + dialect string } func init() { @@ -212,6 +197,18 @@ var kvMeta = workload.Meta{ // Meta implements the Generator interface. func (*kv) Meta() workload.Meta { return kvMeta } +// Dialect implements the DialectProvider interface. +func (w *kv) Dialect() string { return w.dialect } + +// SetDialect implements the DialectSetter interface. +func (w *kv) SetDialect(dialect string) error { + w.dialect = strings.ToLower(dialect) + if w.dialect == "" { + w.dialect = "crdb" + } + return nil +} + // Flags implements the Flagser interface. func (w *kv) Flags() workload.Flags { return w.flags } @@ -222,16 +219,9 @@ func (w *kv) ConnFlags() *workload.ConnFlags { return w.connFlags } func (w *kv) Hooks() workload.Hooks { return workload.Hooks{ PostLoad: func(_ context.Context, db *gosql.DB) error { - if w.enum { - _, err := db.Exec(` -CREATE TYPE enum_type AS ENUM ('v'); -ALTER TABLE kv ADD COLUMN e enum_type NOT NULL AS ('v') STORED;`) - if err != nil { - return err - } - } - if w.scatter { - if _, err := db.Exec(`ALTER TABLE kv SCATTER`); err != nil { + sd := w.createSchemaDialect() + for _, stmt := range sd.PostLoadStmts("kv", w.enum, w.scatter) { + if _, err := db.Exec(stmt); err != nil { return err } } @@ -314,6 +304,55 @@ func (w *kv) validateConfig() (err error) { "`--insert-count` (%d) is greater than the number of unique keys that could be possibly generated [%d,%d)", w.insertCount, kg.kr.min, kg.kr.max) } + + // Validate dialect. + validDialects := map[string]bool{ + "crdb": true, "": true, "postgres": true, "aurora": true, "dsql": true, "spanner": true, + } + if !validDialects[w.dialect] { + return errors.Errorf("unknown dialect: %s", w.dialect) + } + + // Validate dialect-specific features. + if w.dialect != "crdb" && w.dialect != "" { + var unsupported []string + if w.shards > 0 { + unsupported = append(unsupported, "--num-shards") + } + if w.txnQoS != "regular" && w.txnQoS != "" { + unsupported = append(unsupported, "--txn-qos") + } + if w.splits > 0 { + unsupported = append(unsupported, "--splits") + } + if w.scatter { + unsupported = append(unsupported, "--scatter") + } + if w.enum { + unsupported = append(unsupported, "--enum") + } + // --sfu-writes and --sel1-writes are handled by validateCapabilities + // since they are supported by both CRDBExecutor and PGXExecutor. + if w.longRunningTxn { + unsupported = append(unsupported, "--long-running-txn") + } + if w.prepareReadOnly { + unsupported = append(unsupported, "--prepare-read-only") + } + // For postgres dialects, follower reads are not supported. + if (w.dialect == "postgres" || w.dialect == "aurora" || w.dialect == "dsql") && + w.followerReadPercent > 0 { + unsupported = append(unsupported, "--follower-read-percent") + } + if len(unsupported) > 0 { + return errors.Errorf( + "flags %s are not supported for dialect %s", + strings.Join(unsupported, ", "), + w.dialect, + ) + } + } + return nil } @@ -382,6 +421,21 @@ func (w *kv) createKeyGenerator() keyGeneratorConfig { return kg } +// createSchemaDialect returns the appropriate SchemaDialect for the configured +// dialect. Used by Tables() and Hooks().PostLoad for schema generation. +func (w *kv) createSchemaDialect() dbexec.SchemaDialect { + switch w.dialect { + case "crdb", "": + return dbexec.CockroachDialect{} + case "postgres", "aurora", "dsql": + return dbexec.PostgresDialect{} + case "spanner": + return dbexec.SpannerSchemaDialect{} + default: + return dbexec.CockroachDialect{} // validated elsewhere + } +} + func splitFinder(i, splits int, r keyRange, k keyTransformer) interface{} { if splits < 0 || i >= splits { panic(fmt.Sprintf("programming error: split index (%d) cannot be less than 0, "+ @@ -403,230 +457,224 @@ func insertCountKey(idx, count int64, kr keyRange) int64 { // Tables implements the Generator interface. func (w *kv) Tables() []workload.Table { - // Tables should only run on initialized workload, safe to call create without - // having a panic. We don't need to defer this to the actual table callbacks - // like Splits or InitialRows. + sd := w.createSchemaDialect() kg := w.createKeyGenerator() - table := workload.Table{Name: `kv`} - table.Splits = workload.Tuples( - w.splits, - func(splitIdx int) []interface{} { - return []interface{}{splitFinder(splitIdx, w.splits, kg.kr, kg.transformer)} - }, - ) + table := workload.Table{ + Name: "kv", + Schema: sd.SchemaFragment(w.keySize, w.secondaryIndex, w.shards), + } - if w.shards > 0 { - schema := shardedKvSchema - if w.secondaryIndex { - schema = shardedKvSchemaWithIndex - } - table.Schema = fmt.Sprintf(schema, kg.transformer.keySQLType(), w.shards) - } else { - schema := kvSchema - if w.secondaryIndex { - schema = kvSchemaWithIndex - } - table.Schema = fmt.Sprintf(schema, kg.transformer.keySQLType()) + // Add secondary index DDL if requested. For CRDB, the index is inline + // in the schema; for other dialects, it's a separate CREATE INDEX statement. + if w.secondaryIndex { + table.Indexes = sd.SecondaryIndexStmts("kv") } - if w.insertCount > 0 { - const batchSize = 1000 - table.InitialRows = workload.BatchedTuples{ - NumBatches: (w.insertCount + batchSize - 1) / batchSize, - // If the key sequence is not sequential, duplicates are possible. - // The zipfian distribution produces duplicates by design, and the - // hash key mapper can also produce duplicates at larger insert - // counts (it's at least inevitable at ~1b rows). Marking that the - // keys may contain duplicates will cause the data loader to use - // INSERT ... ON CONFLICT DO NOTHING statements. - MayContainDuplicates: !w.sequential, - FillBatch: func(batchIdx int, cb coldata.Batch, a *bufalloc.ByteAllocator) { - rowBegin, rowEnd := batchIdx*batchSize, (batchIdx+1)*batchSize - if rowEnd > w.insertCount { - rowEnd = w.insertCount - } + // Add split points if configured. Only effective for CRDB (other dialects + // don't support splits, and validateConfig rejects --splits for them). + if w.splits > 0 { + table.Splits = workload.Tuples( + w.splits, + func(splitIdx int) []interface{} { + return []interface{}{splitFinder(splitIdx, w.splits, kg.kr, kg.transformer)} + }, + ) + } - var kvtableTypes = []*types.T{ - kg.transformer.getColumnType(), - types.Bytes, - } + if w.insertCount > 0 { + w.addInitialRows(&table, kg) + } - cb.Reset(kvtableTypes, rowEnd-rowBegin, coldata.StandardColumnFactory) - - { - seq := rowBegin - kg.transformer.fillColumnBatch(cb, a, func() (s int64, ok bool) { - if seq < rowEnd { - seq++ - return insertCountKey(int64(seq-1), int64(w.insertCount), kg.kr), true - } - return 0, false - }) - } + return []workload.Table{table} +} - valCol := cb.ColVec(1).Bytes() - // coldata.Bytes only allows appends so we have to reset it. - valCol.Reset() - rndBlock := rand.New(rand.NewSource(RandomSeed.Seed())) - - for rowIdx := rowBegin; rowIdx < rowEnd; rowIdx++ { - rowOffset := rowIdx - rowBegin - var payload []byte - blockSize, uniqueSize := w.randBlockSize(rndBlock) - *a, payload = a.Alloc(blockSize) - w.randFillBlock(rndBlock, payload, uniqueSize) - valCol.Set(rowOffset, payload) - } - }, - } +func (w *kv) addInitialRows(table *workload.Table, kg keyGeneratorConfig) { + if w.insertCount <= 0 { + return } + const batchSize = 1000 + table.InitialRows = workload.BatchedTuples{ + NumBatches: (w.insertCount + batchSize - 1) / batchSize, + // If the key sequence is not sequential, duplicates are possible. + // The zipfian distribution produces duplicates by design, and the + // hash key mapper can also produce duplicates at larger insert + // counts (it's at least inevitable at ~1b rows). Marking that the + // keys may contain duplicates will cause the data loader to use + // INSERT ... ON CONFLICT DO NOTHING statements. + MayContainDuplicates: !w.sequential, + FillBatch: func(batchIdx int, cb coldata.Batch, a *bufalloc.ByteAllocator) { + rowBegin, rowEnd := batchIdx*batchSize, (batchIdx+1)*batchSize + if rowEnd > w.insertCount { + rowEnd = w.insertCount + } - return []workload.Table{table} + var kvtableTypes = []*types.T{ + kg.transformer.getColumnType(), + types.Bytes, + } + + cb.Reset(kvtableTypes, rowEnd-rowBegin, coldata.StandardColumnFactory) + + { + seq := rowBegin + kg.transformer.fillColumnBatch(cb, a, func() (s int64, ok bool) { + if seq < rowEnd { + seq++ + return insertCountKey(int64(seq-1), int64(w.insertCount), kg.kr), true + } + return 0, false + }) + } + + valCol := cb.ColVec(1).Bytes() + // coldata.Bytes only allows appends so we have to reset it. + valCol.Reset() + rndBlock := rand.New(rand.NewSource(RandomSeed.Seed())) + + for rowIdx := rowBegin; rowIdx < rowEnd; rowIdx++ { + rowOffset := rowIdx - rowBegin + var payload []byte + blockSize, uniqueSize := w.randBlockSize(rndBlock) + *a, payload = a.Alloc(blockSize) + w.randFillBlock(rndBlock, payload, uniqueSize) + valCol.Set(rowOffset, payload) + } + }, + } } // Ops implements the Opser interface. func (w *kv) Ops( ctx context.Context, urls []string, reg *histogram.Registry, ) (workload.QueryLoad, error) { - cfg := workload.NewMultiConnPoolCfgFromFlags(w.connFlags) - cfg.MaxTotalConnections = w.connFlags.Concurrency + 1 - mcp, err := workload.NewMultiConnPool(ctx, cfg, urls...) + exec, err := w.createExecutor(ctx, urls) if err != nil { return workload.QueryLoad{}, err } - // Read statement - var buf strings.Builder - var folBuf strings.Builder - if w.enum { - buf.WriteString(`SELECT k, v, e FROM kv WHERE k IN (`) - folBuf.WriteString(`SELECT k, v, e FROM kv AS OF SYSTEM TIME follower_read_timestamp() WHERE k IN (`) - } else { - buf.WriteString(`SELECT k, v FROM kv WHERE k IN (`) - folBuf.WriteString(`SELECT k, v FROM kv AS OF SYSTEM TIME follower_read_timestamp() WHERE k IN (`) - } - for i := 0; i < w.batchSize; i++ { - if i > 0 { - buf.WriteString(", ") - folBuf.WriteString(", ") - } - fmt.Fprintf(&buf, `$%d`, i+1) - fmt.Fprintf(&folBuf, `$%d`, i+1) - } - buf.WriteString(`)`) - folBuf.WriteString(`)`) - - readStmtStr := buf.String() - followerReadStmtStr := folBuf.String() - - // Write statement - buf.Reset() - buf.WriteString(`UPSERT INTO kv (k, v) VALUES`) - for i := 0; i < w.batchSize; i++ { - j := i * 2 - if i > 0 { - buf.WriteString(", ") - } - fmt.Fprintf(&buf, ` ($%d, $%d)`, j+1, j+2) - } - writeStmtStr := buf.String() - - // Select for update statement - var sfuStmtStr string - if w.writesUseSelectForUpdate { - buf.Reset() - buf.WriteString(`SELECT k, v FROM kv WHERE k IN (`) - for i := 0; i < w.batchSize; i++ { - if i > 0 { - buf.WriteString(", ") - } - fmt.Fprintf(&buf, `$%d`, i+1) - } - buf.WriteString(`) FOR UPDATE`) - sfuStmtStr = buf.String() - } - - // Span statement - buf.Reset() - buf.WriteString(`SELECT count(v) FROM [SELECT v FROM kv`) - if w.spanLimit > 0 { - // Span statements without a limit query all ranges. However, if there's - // a span limit specified, we want to randomly choose the range from which - // the limited scan starts at. We do this by introducing the k >= $1 - // predicate. - fmt.Fprintf(&buf, ` WHERE k >= $1 ORDER BY k LIMIT %d`, w.spanLimit) - } - buf.WriteString(`]`) - spanStmtStr := buf.String() - - // Del statement - buf.Reset() - buf.WriteString(`DELETE FROM kv WHERE k IN (`) - for i := 0; i < w.batchSize; i++ { - if i > 0 { - buf.WriteString(", ") - } - fmt.Fprintf(&buf, `$%d`, i+1) + if err := w.validateCapabilities(exec.Capabilities()); err != nil { + _ = exec.Close() + return workload.QueryLoad{}, err } - buf.WriteString(`)`) - delStmtStr := buf.String() kg := w.createKeyGenerator() ql := workload.QueryLoad{} var numEmptyResults atomic.Int64 + for i := 0; i < w.connFlags.Concurrency; i++ { op := &kvOp{ config: w, hists: reg.GetHandle(), + exec: exec, + kg: kg, + ks: kg.newState(), numEmptyResults: &numEmptyResults, } - op.readStmt = op.sr.Define(readStmtStr) - op.followerReadStmt = op.sr.Define(followerReadStmtStr) - if !op.config.prepareReadOnly { - op.writeStmt = op.sr.Define(writeStmtStr) - } - if len(sfuStmtStr) > 0 && !op.config.prepareReadOnly { - op.sfuStmt = op.sr.Define(sfuStmtStr) - } - op.sel1Stmt = op.sr.Define("SELECT 1") - op.spanStmt = op.sr.Define(spanStmtStr) - if w.txnQoS != `regular` { - stmt := op.sr.Define(fmt.Sprintf( - " SET default_transaction_quality_of_service = %s", w.txnQoS)) - op.qosStmt = &stmt - } - if !op.config.prepareReadOnly { - op.delStmt = op.sr.Define(delStmtStr) - } - if err := op.sr.Init(ctx, "kv", mcp); err != nil { - return workload.QueryLoad{}, err - } - op.mcp = mcp - op.kg = kg - op.ks = kg.newState() ql.WorkerFns = append(ql.WorkerFns, op.run) - ql.Close = op.close } + + ql.Close = func(ctx context.Context) error { + if empty := numEmptyResults.Load(); empty != 0 { + fmt.Printf("Number of reads that didn't return any results: %d.\n", empty) + } + fmt.Printf("Write sequence could be resumed by passing --write-seq=%s to the next run.\n", + kg.cursor()) + return exec.Close() + } + return ql, nil } +// keyType returns the SQL key type string for the dbexec.Config. +func (w *kv) keyType() string { + if w.keySize > 0 { + return fmt.Sprintf("VARCHAR(%d)", w.keySize) + } + return "BIGINT" +} + +// createExecutor creates and initializes the appropriate Executor based on the +// configured dialect. +func (w *kv) createExecutor(ctx context.Context, urls []string) (dbexec.Executor, error) { + var exec dbexec.Executor + switch w.dialect { + case "crdb", "": + exec = dbexec.NewCRDBExecutor() + case "postgres", "aurora", "dsql": + exec = dbexec.NewPGXExecutor(dbexec.PostgresDialect{}) + case "spanner": + exec = dbexec.NewSpannerExecutor() + default: + return nil, errors.Errorf("unknown dialect: %s", w.dialect) + } + + cfg := dbexec.Config{ + URLs: urls, + Table: "kv", + BatchSize: w.batchSize, + KeyType: w.keyType(), + SecondaryIndex: w.secondaryIndex, + NumShards: w.shards, + Concurrency: w.connFlags.Concurrency, + Enum: w.enum, + SpanLimit: w.spanLimit, + TxnQoS: w.txnQoS, + } + + if err := exec.Init(ctx, cfg, w.connFlags); err != nil { + return nil, err + } + return exec, nil +} + +// validateCapabilities checks that all requested features are supported by the +// executor. Returns an error listing any unsupported flags. +func (w *kv) validateCapabilities(caps dbexec.Capabilities) error { + var unsupported []string + + if w.followerReadPercent > 0 && !caps.FollowerReads { + unsupported = append(unsupported, "--follower-read-percent") + } + if w.writesUseSelectForUpdate && !caps.SelectForUpdate { + unsupported = append(unsupported, "--sfu-writes") + } + if w.longRunningTxn && !caps.TransactionPriority { + unsupported = append(unsupported, "--long-running-txn") + } + if w.txnQoS != "regular" && w.txnQoS != "" && !caps.TransactionQoS { + unsupported = append(unsupported, "--txn-qos") + } + if w.shards > 0 && !caps.HashShardedPK { + unsupported = append(unsupported, "--num-shards") + } + if w.splits > 0 && !caps.Splits { + unsupported = append(unsupported, "--splits") + } + if w.scatter && !caps.Scatter { + unsupported = append(unsupported, "--scatter") + } + if w.enum && !caps.EnumColumn { + unsupported = append(unsupported, "--enum") + } + if w.writesUseSelect1 && !caps.Select1 { + unsupported = append(unsupported, "--sel1-writes") + } + + if len(unsupported) > 0 { + return errors.Errorf("flags %s are not supported for dialect %q", + strings.Join(unsupported, ", "), w.dialect) + } + return nil +} + type kvOp struct { - config *kv - hists *histogram.Histograms - sr workload.SQLRunner - mcp *workload.MultiConnPool - qosStmt *workload.StmtHandle - readStmt workload.StmtHandle - followerReadStmt workload.StmtHandle - writeStmt workload.StmtHandle - spanStmt workload.StmtHandle - sfuStmt workload.StmtHandle - sel1Stmt workload.StmtHandle - delStmt workload.StmtHandle - kg keyGeneratorConfig - ks *keyGeneratorState - numEmptyResults *atomic.Int64 + config *kv + hists *histogram.Histograms + exec dbexec.Executor + kg keyGeneratorConfig + ks *keyGeneratorState + numEmptyResults *atomic.Int64 } func (o *kvOp) run(ctx context.Context) (retErr error) { @@ -636,172 +684,223 @@ func (o *kvOp) run(ctx context.Context) (retErr error) { defer cancel() } - if o.qosStmt != nil { - _, err := o.qosStmt.Exec(ctx) - if err != nil { - return err - } - } statementProbability := o.ks.rand.Intn(100) // Determines what statement is executed. + + // Read operation. if statementProbability < o.config.readPercent { - args := make([]interface{}, o.config.batchSize) - for i := 0; i < o.config.batchSize; i++ { - key := o.kg.readKey(o.ks) - if o.config.alwaysIncKeySeq { - key = o.kg.writeKey(o.ks) - } - args[i] = o.kg.transformer.getKey(key) - } - start := timeutil.Now() - readStmt := o.readStmt - opName := `read` + return o.doRead(ctx) + } + statementProbability -= o.config.readPercent - if o.ks.rand.Intn(100) < o.config.followerReadPercent { - readStmt = o.followerReadStmt - opName = `follower-read` - } - rows, err := readStmt.Query(ctx, args...) - if err != nil { - return err + // Delete operation. + if statementProbability < o.config.delPercent { + return o.doDelete(ctx) + } + statementProbability -= o.config.delPercent + + // Span operation. + if statementProbability < o.config.spanPercent { + return o.doSpan(ctx) + } + + // Write operation (default). + return o.doWrite(ctx) +} + +// doRead executes a read operation using the executor. +func (o *kvOp) doRead(ctx context.Context) error { + keys := make([]interface{}, o.config.batchSize) + for i := range keys { + key := o.kg.readKey(o.ks) + if o.config.alwaysIncKeySeq { + key = o.kg.writeKey(o.ks) } - empty := true - for rows.Next() { - empty = false + keys[i] = o.kg.transformer.getKey(key) + } + + start := timeutil.Now() + var rows dbexec.Rows + var err error + opName := "read" + + if o.ks.rand.Intn(100) < o.config.followerReadPercent { + rows, err = o.exec.FollowerRead(ctx, keys) + opName = "follower-read" + } else { + rows, err = o.exec.Read(ctx, keys) + } + + if err != nil { + return err + } + defer rows.Close() + + empty := true + for rows.Next() { + empty = false + } + if err := rows.Err(); err != nil { + return err + } + if empty { + o.numEmptyResults.Add(1) + } + + elapsed := timeutil.Since(start) + o.hists.Get(opName).Record(elapsed) + return nil +} + +// doDelete executes a delete operation using the executor. +func (o *kvOp) doDelete(ctx context.Context) error { + keys := make([]interface{}, o.config.batchSize) + for i := range keys { + key := o.kg.readKey(o.ks) + if o.config.alwaysIncKeySeq { + key = o.kg.writeKey(o.ks) } - if empty { - o.numEmptyResults.Add(1) + keys[i] = o.kg.transformer.getKey(key) + } + + start := timeutil.Now() + err := o.exec.Delete(ctx, keys) + elapsed := timeutil.Since(start) + o.hists.Get("del").Record(elapsed) + return err +} + +// doSpan executes a span operation using the executor. +func (o *kvOp) doSpan(ctx context.Context) error { + start := timeutil.Now() + var startKey interface{} + if o.config.spanLimit > 0 { + key := o.kg.readKey(o.ks) + if o.config.alwaysIncKeySeq { + key = o.kg.writeKey(o.ks) } - elapsed := timeutil.Since(start) - o.hists.Get(opName).Record(elapsed) - return rows.Err() + startKey = o.kg.transformer.getKey(key) } - // Since we know the statement is not a read, we recalibrate - // statementProbability to only consider the other statements. - statementProbability -= o.config.readPercent - if statementProbability < o.config.delPercent { - start := timeutil.Now() - args := make([]interface{}, o.config.batchSize) - for i := 0; i < o.config.batchSize; i++ { - args[i] = o.kg.readKey(o.ks) - if o.config.alwaysIncKeySeq { - args[i] = o.kg.writeKey(o.ks) - } + _, err := o.exec.Span(ctx, startKey, o.config.spanLimit) + elapsed := timeutil.Since(start) + o.hists.Get("span").Record(elapsed) + return err +} + +// doWrite executes a write operation. For simple writes, it uses the executor's +// Write method. For transactional writes (SFU, SELECT 1, long-running txn), it +// uses BeginTx and type-asserts to ExtendedTx for prepared statement access. +func (o *kvOp) doWrite(ctx context.Context) (retErr error) { + // Helper to generate write rows. + generateWriteRows := func() []dbexec.Row { + rows := make([]dbexec.Row, o.config.batchSize) + for i := range rows { + rows[i].K = o.kg.transformer.getKey(o.kg.writeKey(o.ks)) + rows[i].V = o.config.randBlock(o.ks.rand) } - _, err := o.delStmt.Exec(ctx, args...) + return rows + } + + start := timeutil.Now() + + // Check if we need transactional writes (CRDB-specific features). + needsTxn := o.config.writesUseSelect1 || o.config.writesUseSelectForUpdate || o.config.longRunningTxn + if !needsTxn { + // Simple non-transactional write. + rows := generateWriteRows() + err := o.exec.Write(ctx, rows) if err != nil { return err } elapsed := timeutil.Since(start) - o.hists.Get(`del`).Record(elapsed) + o.hists.Get("write").Record(elapsed) return nil } - statementProbability -= o.config.delPercent - if statementProbability < o.config.spanPercent { - start := timeutil.Now() - var err error - if o.config.spanLimit > 0 { - arg := o.kg.readKey(o.ks) - if o.config.alwaysIncKeySeq { - arg = o.kg.writeKey(o.ks) - } - _, err = o.spanStmt.Exec(ctx, arg) - } else { - _, err = o.spanStmt.Exec(ctx) - } - if err != nil { - return err - } - elapsed := timeutil.Since(start) - o.hists.Get(`span`).Record(elapsed) + + // Transactional write path (CRDB-specific). + txOpts := dbexec.TxOptions{} + if o.config.longRunningTxn { + txOpts.Priority = o.config.longRunningTxnPriority + } + + tx, err := o.exec.BeginTx(ctx, txOpts) + if err != nil { return err } - makeWriteBatchArgs := func() ([]interface{}, []interface{}) { - const argCount = 2 - writeArgs := make([]interface{}, argCount*o.config.batchSize) - var sfuArgs []interface{} - if o.config.writesUseSelectForUpdate { - sfuArgs = make([]interface{}, o.config.batchSize) + defer func() { + rollbackErr := tx.Rollback(ctx) + if rollbackErr != nil && !errors.Is(rollbackErr, pgx.ErrTxClosed) { + retErr = errors.CombineErrors(retErr, rollbackErr) } - for i := 0; i < o.config.batchSize; i++ { - j := i * argCount - writeArgs[j+0] = o.kg.transformer.getKey(o.kg.writeKey(o.ks)) - if sfuArgs != nil { - sfuArgs[i] = writeArgs[j] - } - writeArgs[j+1] = o.config.randBlock(o.ks.rand) - } - return writeArgs, sfuArgs + }() + + // Type assert to ExtendedTx for transactional operations using prepared + // statements (SFU, SELECT 1, Write within transaction). + extTx, ok := tx.(dbexec.ExtendedTx) + if !ok { + // This shouldn't happen since validateCapabilities should have rejected + // these flags for executors that don't support ExtendedTx. + return errors.New("transactional writes with SFU/SELECT1/long-running-txn require ExtendedTx support") } - start := timeutil.Now() - var err error - if o.config.writesUseSelect1 || o.config.writesUseSelectForUpdate || o.config.longRunningTxn { - // We could use crdb.ExecuteTx, but we avoid retries in this workload so - // that each run call makes 1 attempt, so that rate limiting in workerRun - // behaves as expected. - var tx pgx.Tx - txnOptions := pgx.TxOptions{} - txnOptions.BeginQuery = fmt.Sprintf("BEGIN PRIORITY %s", o.config.longRunningTxnPriority) - tx, err := o.mcp.Get().BeginTx(ctx, txnOptions) - if err != nil { - return err + + iterations := 1 + if o.config.longRunningTxn { + iterations = o.config.longRunningTxnNumWrites + } + + for i := 0; i < iterations; i++ { + rows := generateWriteRows() + keys := make([]interface{}, len(rows)) + for j, r := range rows { + keys[j] = r.K } - defer func() { - rollbackErr := tx.Rollback(ctx) - if !errors.Is(rollbackErr, pgx.ErrTxClosed) { - retErr = errors.CombineErrors(retErr, rollbackErr) + + // SELECT 1 warmup. + if o.config.writesUseSelect1 { + if err := extTx.Select1(ctx); err != nil { + return err } - }() - iterations := 1 - if o.config.longRunningTxn { - iterations = o.config.longRunningTxnNumWrites } - for i := 0; i < iterations; i++ { - writeArgs, sfuArgs := makeWriteBatchArgs() - if o.config.writesUseSelect1 { - rows, err := o.sel1Stmt.QueryTx(ctx, tx) - if err != nil { - return err - } - rows.Close() - if err = rows.Err(); err != nil { - return err - } - } - if o.config.writesUseSelectForUpdate { - rows, err := o.sfuStmt.QueryTx(ctx, tx, sfuArgs...) - if err != nil { - return err - } - rows.Close() - if err = rows.Err(); err != nil { - // The transaction may have experienced an error in the meantime. - return o.tryHandleWriteErr("write-write-err", start, err) - } + + // SELECT FOR UPDATE. + if o.config.writesUseSelectForUpdate { + sfuRows, err := extTx.SelectForUpdate(ctx, keys) + if err != nil { + return o.tryHandleWriteErr("write-write-err", start, err) } - // Simulate a transaction that does other work between the sel1 / SFU and write. - time.Sleep(o.config.sfuDelay) - if _, err = o.writeStmt.ExecTx(ctx, tx, writeArgs...); err != nil { - // Multiple write transactions can contend and encounter - // a serialization failure. We swallow such an error. + sfuRows.Close() + if err = sfuRows.Err(); err != nil { return o.tryHandleWriteErr("write-write-err", start, err) } } - if err = tx.Commit(ctx); err != nil { - return o.tryHandleWriteErr("write-commit-err", start, err) + + // Simulate a transaction that does other work between the + // sel1 / SFU and write. + time.Sleep(o.config.sfuDelay) + + // Write within the transaction. + if err = extTx.Write(ctx, rows); err != nil { + return o.tryHandleWriteErr("write-write-err", start, err) } - } else { - writeArgs, _ := makeWriteBatchArgs() - _, err = o.writeStmt.Exec(ctx, writeArgs...) } - if err != nil { - return err + + if err = tx.Commit(ctx); err != nil { + return o.tryHandleWriteErr("write-commit-err", start, err) } + elapsed := timeutil.Since(start) - o.hists.Get(`write`).Record(elapsed) - return err + o.hists.Get("write").Record(elapsed) + return nil } +// tryHandleWriteErr handles write errors by checking if they are serialization +// failures. For CRDB executors that support serialization retry, these errors +// are recorded as metrics rather than returned as failures. func (o *kvOp) tryHandleWriteErr(name string, start time.Time, err error) error { + // Only handle serialization failures if the executor supports it. + if !o.exec.Capabilities().SerializationRetry { + return err + } + // If the error is not an instance of pgconn.PgError, then it is unexpected. pgErr := new(pgconn.PgError) if !errors.As(err, &pgErr) { @@ -817,15 +916,6 @@ func (o *kvOp) tryHandleWriteErr(name string, start time.Time, err error) error return err } -func (o *kvOp) close(context.Context) error { - if empty := o.numEmptyResults.Load(); empty != 0 { - fmt.Printf("Number of reads that didn't return any results: %d.\n", empty) - } - fmt.Printf("Write sequence could be resumed by passing --write-seq=%s to the next run.\n", - o.kg.cursor()) - return nil -} - type sequence struct { val *atomic.Int64 max int64 diff --git a/pkg/workload/kv/kv_test.go b/pkg/workload/kv/kv_test.go index 2e7c8c00299b..815f92dc20d3 100644 --- a/pkg/workload/kv/kv_test.go +++ b/pkg/workload/kv/kv_test.go @@ -190,3 +190,81 @@ func TestInitialSeqValidation(t *testing.T) { }) } } + +func TestDialectValidation(t *testing.T) { + testCases := []struct { + name string + dialect string + setup func(*kv) + wantErr string + }{ + { + name: "valid crdb dialect", + dialect: "crdb", + wantErr: "", + }, + { + name: "valid postgres dialect", + dialect: "postgres", + wantErr: "", + }, + { + name: "valid spanner dialect", + dialect: "spanner", + wantErr: "", + }, + { + name: "invalid dialect", + dialect: "mysql", + wantErr: "unknown dialect", + }, + { + name: "postgres with hash sharding", + dialect: "postgres", + setup: func(k *kv) { + k.shards = 4 + k.sequential = true // Required for shards to pass basic validation + }, + wantErr: "--num-shards", + }, + { + name: "postgres with follower reads", + dialect: "postgres", + setup: func(k *kv) { k.followerReadPercent = 50 }, + wantErr: "--follower-read-percent", + }, + { + name: "postgres with txn qos", + dialect: "postgres", + setup: func(k *kv) { k.txnQoS = "critical" }, + wantErr: "--txn-qos", + }, + { + name: "spanner allows follower reads", + dialect: "spanner", + setup: func(k *kv) { k.followerReadPercent = 50 }, + wantErr: "", // Spanner supports stale reads + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + k := &kv{ + dialect: tc.dialect, + cycleLength: 1000, + txnQoS: "regular", + targetCompressionRatio: 1.0, + } + if tc.setup != nil { + tc.setup(k) + } + err := k.validateConfig() + if tc.wantErr == "" { + require.NoError(t, err) + } else { + require.Error(t, err) + require.Contains(t, err.Error(), tc.wantErr) + } + }) + } +} diff --git a/pkg/workload/workload.go b/pkg/workload/workload.go index bf1e0999454e..f36301ffbe88 100644 --- a/pkg/workload/workload.go +++ b/pkg/workload/workload.go @@ -107,6 +107,24 @@ type Hookser interface { Hooks() Hooks } +// DialectProvider is optionally implemented by generators that support multiple +// database dialects. This allows the workload framework to adjust behavior +// (e.g., skip CockroachDB-specific connection variables) based on the target database. +type DialectProvider interface { + Generator + // Dialect returns the target database dialect (e.g., "crdb", "postgres", "spanner"). + // An empty string means CockroachDB (the default). + Dialect() string +} + +// DialectSetter is optionally implemented by generators that support multiple +// database dialects. The workload CLI uses this to apply the global --dialect +// flag to the generator before validation and execution. +type DialectSetter interface { + Generator + SetDialect(string) error +} + // Hooks stores functions to be called at points in the workload lifecycle. type Hooks struct { // Validate is called after workload flags are parsed. It should return an @@ -176,6 +194,9 @@ type Table struct { // Schema is the SQL formatted schema for this table, with the `CREATE TABLE // ` prefix omitted. Schema string + // Indexes holds optional index DDL statements to run after table creation. + // These statements must be valid for the target dialect. + Indexes []string // InitialRows is the initial rows that will be present in the table after // setup is completed. Note that the default value of NumBatches (zero) is // special - such a Table will be skipped during `init`; non-zero NumBatches diff --git a/pkg/workload/workloadsql/dataload.go b/pkg/workload/workloadsql/dataload.go index fab331636ae2..277383b64e1d 100644 --- a/pkg/workload/workloadsql/dataload.go +++ b/pkg/workload/workloadsql/dataload.go @@ -27,6 +27,12 @@ import ( type InsertsDataLoader struct { BatchSize int Concurrency int + // DisableCRDBDDL skips CockroachDB-specific DDL optimizations when creating tables. + DisableCRDBDDL bool +} + +type sqlExecer interface { + ExecContext(context.Context, string, ...interface{}) (gosql.Result, error) } // InitialDataLoad implements the InitialDataLoader interface. @@ -65,7 +71,22 @@ func (l InsertsDataLoader) InitialDataLoad( for currentTable < len(tables) { batchEnd := min(currentTable+maxTableBatchSize, len(tables)) nextBatch := tables[currentTable:batchEnd] - if err := crdb.ExecuteTx(ctx, db, &gosql.TxOptions{}, func(tx *gosql.Tx) error { + if l.DisableCRDBDDL { + for _, table := range nextBatch { + if table.ObjectPrefix != nil && table.ObjectPrefix.ExplicitCatalog { + return 0, errors.Errorf("multi-database table %q is not supported for non-CRDB init", table.Name) + } + tableName := table.GetResolvedName() + createStmt := fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s %s`, tableName.String(), table.Schema) + if _, err := db.ExecContext(ctx, createStmt); err != nil { + return 0, errors.WithDetailf(errors.Wrapf(err, "could not create table: %q", table.Name), + "SQL: %s", createStmt) + } + if err := execIndexStatements(ctx, db, table); err != nil { + return 0, err + } + } + } else if err := crdb.ExecuteTx(ctx, db, &gosql.TxOptions{}, func(tx *gosql.Tx) error { // Run the operations in a single txn so they complete more quickly. if _, err := tx.Exec("SET LOCAL autocommit_before_ddl = false"); err != nil { return err @@ -89,6 +110,9 @@ func (l InsertsDataLoader) InitialDataLoad( return errors.WithDetailf(errors.Wrapf(err, "could not create table: %q", table.Name), "SQL: %s", createStmt) } + if err := execIndexStatements(ctx, tx, table); err != nil { + return err + } } return nil }); err != nil { @@ -183,3 +207,13 @@ func (l InsertsDataLoader) InitialDataLoad( } return bytesAtomic.Load(), nil } + +func execIndexStatements(ctx context.Context, execer sqlExecer, table workload.Table) error { + for _, stmt := range table.Indexes { + if _, err := execer.ExecContext(ctx, stmt); err != nil { + return errors.WithDetailf(errors.Wrapf(err, "could not create index for table: %q", table.Name), + "SQL: %s", stmt) + } + } + return nil +} diff --git a/pkg/workload/workloadsql/workloadsql.go b/pkg/workload/workloadsql/workloadsql.go index ca48db372537..4258fe67529b 100644 --- a/pkg/workload/workloadsql/workloadsql.go +++ b/pkg/workload/workloadsql/workloadsql.go @@ -93,15 +93,17 @@ func maybeDisableMergeQueue(db *gosql.DB) error { // Split creates the range splits defined by the given table. func Split(ctx context.Context, db *gosql.DB, table workload.Table, concurrency int) error { + // Early exit if there are no splits to create. This also allows non-CockroachDB + // databases to skip the CRDB-specific merge queue logic below. + if table.Splits.NumBatches <= 0 { + return nil + } + // Prevent the merge queue from immediately discarding our splits. if err := maybeDisableMergeQueue(db); err != nil { return err } - if table.Splits.NumBatches <= 0 { - return nil - } - splitPoints := make([][]interface{}, 0, table.Splits.NumBatches) for splitIdx := 0; splitIdx < table.Splits.NumBatches; splitIdx++ { splitPoints = append(splitPoints, table.Splits.BatchRows(splitIdx)...)