Skip to content

Commit e995349

Browse files
ankelYuan325
ankel
and
Yuan325
authored
feat(prebuiltconfigs/alloydb-omni): Implement Alloydb omni dataplane tools (#2340)
## Description Implement Dataplane tools for alloydb omni using Postgres source and postgres tools. I also add more omni-specific tool (columnar engine). --------- Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com>
1 parent 4c9765f commit e995349

8 files changed

Lines changed: 657 additions & 48 deletions

File tree

‎.ci/integration.cloudbuild.yaml‎

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -171,6 +171,23 @@ steps:
171171
alloydbainl \
172172
alloydbainl
173173
174+
- id: "alloydb-omni"
175+
name: golang:1
176+
waitFor: ["compile-test-binary"]
177+
entrypoint: /bin/bash
178+
env:
179+
- "GOPATH=/gopath"
180+
volumes:
181+
- name: "go"
182+
path: "/gopath"
183+
args:
184+
- -c
185+
- |
186+
.ci/test_with_coverage.sh \
187+
"AlloyDB Omni" \
188+
alloydbomni \
189+
postgres
190+
174191
- id: "bigtable"
175192
name: golang:1
176193
waitFor: ["compile-test-binary"]
@@ -676,7 +693,7 @@ steps:
676693
"Looker" \
677694
looker \
678695
looker
679-
696+
680697
- id: "mindsdb"
681698
name: golang:1
682699
waitFor: ["compile-test-binary"]
@@ -864,7 +881,7 @@ steps:
864881
"Snowflake" \
865882
snowflake \
866883
snowflake
867-
884+
868885
- id: "cassandra"
869886
name: golang:1
870887
waitFor: ["compile-test-binary"]
@@ -909,7 +926,7 @@ steps:
909926
go test -v ./tests/oracle/... \
910927
-coverprofile=oracle_coverage.out \
911928
-coverpkg=./internal/sources/oracle/...,./internal/tools/oracle/...
912-
929+
913930
# Coverage check
914931
total_coverage=$(go tool cover -func=oracle_coverage.out | grep "total:" | awk '{print $3}')
915932
echo "Oracle total coverage: $total_coverage"

‎cmd/root_test.go‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1624,7 +1624,7 @@ func TestEnvVarReplacement(t *testing.T) {
16241624
${toolset_name}:
16251625
- example_tool
16261626
1627-
1627+
16281628
prompts:
16291629
${prompt_name}:
16301630
description: A test prompt for {{.name}}.
@@ -2054,6 +2054,7 @@ func TestSingleEdit(t *testing.T) {
20542054

20552055
func TestPrebuiltTools(t *testing.T) {
20562056
// Get prebuilt configs
2057+
alloydb_omni_config, _ := prebuiltconfigs.Get("alloydb-omni")
20572058
alloydb_admin_config, _ := prebuiltconfigs.Get("alloydb-postgres-admin")
20582059
alloydb_config, _ := prebuiltconfigs.Get("alloydb-postgres")
20592060
bigquery_config, _ := prebuiltconfigs.Get("bigquery")
@@ -2104,6 +2105,12 @@ func TestPrebuiltTools(t *testing.T) {
21042105
t.Setenv("ALLOYDB_POSTGRES_USER", "your_alloydb_user")
21052106
t.Setenv("ALLOYDB_POSTGRES_PASSWORD", "your_alloydb_password")
21062107

2108+
t.Setenv("ALLOYDB_OMNI_HOST", "localhost")
2109+
t.Setenv("ALLOYDB_OMNI_PORT", "5432")
2110+
t.Setenv("ALLOYDB_OMNI_DATABASE", "your_alloydb_db")
2111+
t.Setenv("ALLOYDB_OMNI_USER", "your_alloydb_user")
2112+
t.Setenv("ALLOYDB_OMNI_PASSWORD", "your_alloydb_password")
2113+
21072114
t.Setenv("CLICKHOUSE_PROTOCOL", "your_clickhouse_protocol")
21082115
t.Setenv("CLICKHOUSE_DATABASE", "your_clickhouse_database")
21092116
t.Setenv("CLICKHOUSE_PASSWORD", "your_clickhouse_password")
@@ -2197,6 +2204,16 @@ func TestPrebuiltTools(t *testing.T) {
21972204
in []byte
21982205
wantToolset server.ToolsetConfigs
21992206
}{
2207+
{
2208+
name: "alloydb omni prebuilt tools",
2209+
in: alloydb_omni_config,
2210+
wantToolset: server.ToolsetConfigs{
2211+
"alloydb_omni_database_tools": tools.ToolsetConfig{
2212+
Name: "alloydb_omni_database_tools",
2213+
ToolNames: []string{"execute_sql", "list_tables", "list_active_queries", "list_available_extensions", "list_installed_extensions", "list_autovacuum_configurations", "list_columnar_configurations", "list_columnar_recommended_columns", "list_memory_configurations", "list_top_bloated_tables", "list_replication_slots", "list_invalid_indexes", "get_query_plan", "list_views", "list_schemas", "database_overview", "list_triggers", "list_indexes", "list_sequences", "long_running_transactions", "list_locks", "replication_stats", "list_query_stats", "get_column_cardinality", "list_publication_tables", "list_tablespaces", "list_pg_settings", "list_database_stats", "list_roles", "list_table_stats", "list_stored_procedure"},
2214+
},
2215+
},
2216+
},
22002217
{
22012218
name: "alloydb postgres admin prebuilt tools",
22022219
in: alloydb_admin_config,

‎docs/en/reference/prebuilt-tools.md‎

Lines changed: 49 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -54,17 +54,17 @@ See [Usage Examples](../reference/cli.md#examples).
5454
* `get_query_plan`: Generate the execution plan of a statement.
5555
* `list_views`: Lists views in the database from pg_views with a default
5656
limit of 50 rows. Returns schemaname, viewname and the ownername.
57-
* `list_schemas`: Lists schemas in the database.
58-
* `database_overview`: Fetches the current state of the PostgreSQL server.
59-
* `list_triggers`: Lists triggers in the database.
57+
* `list_schemas`: Lists schemas in the database.
58+
* `database_overview`: Fetches the current state of the PostgreSQL server.
59+
* `list_triggers`: Lists triggers in the database.
6060
* `list_indexes`: List available user indexes in a PostgreSQL database.
6161
* `list_sequences`: List sequences in a PostgreSQL database.
6262
* `list_publication_tables`: List publication tables in a PostgreSQL database.
6363
* `list_tablespaces`: Lists tablespaces in the database.
6464
* `list_pg_settings`: List configuration parameters for the PostgreSQL server.
6565
* `list_database_stats`: Lists the key performance and activity statistics for
6666
each database in the AlloyDB instance.
67-
* `list_roles`: Lists all the user-created roles in PostgreSQL database.
67+
* `list_roles`: Lists all the user-created roles in PostgreSQL database.
6868

6969
## AlloyDB Postgres Admin
7070

@@ -100,6 +100,43 @@ See [Usage Examples](../reference/cli.md#examples).
100100
(timeseries metrics) for queries running in an AlloyDB instance using a
101101
PromQL query.
102102

103+
## AlloyDB Omni
104+
105+
* `--prebuilt` value: `alloydb-omni`
106+
* **Environment Variables:**
107+
* `ALLOYDB_OMNI_HOST`: (Optional) The hostname or IP address (Default: localhost).
108+
* `ALLOYDB_OMNI_PORT`: (Optional) The port number (Default: 5432).
109+
* `ALLOYDB_OMNI_DATABASE`: The name of the database to connect to.
110+
* `ALLOYDB_OMNI_USER`: The database username.
111+
* `ALLOYDB_OMNI_PASSWORD`: (Optional) The password for the database user.
112+
* `ALLOYDB_OMNI_QUERY_PARAMS`: (Optional) Connection query parameters.
113+
* **Tools:**
114+
* `execute_sql`: Executes a SQL query.
115+
* `list_tables`: Lists tables in the database.
116+
* `list_autovacuum_configurations`: Lists autovacuum configurations in the
117+
database.
118+
* `list_columnar_configurations`: List AlloyDB Omni columnar-related configurations.
119+
* `list_columnar_recommended_columns`: Lists columns that AlloyDB Omni recommends adding to the columnar engine.
120+
* `list_memory_configurations`: Lists memory-related configurations in the
121+
database.
122+
* `list_top_bloated_tables`: List top bloated tables in the database.
123+
* `list_replication_slots`: Lists replication slots in the database.
124+
* `list_invalid_indexes`: Lists invalid indexes in the database.
125+
* `get_query_plan`: Generate the execution plan of a statement.
126+
* `list_views`: Lists views in the database from pg_views with a default
127+
limit of 50 rows. Returns schemaname, viewname and the ownername.
128+
* `list_schemas`: Lists schemas in the database.
129+
* `database_overview`: Fetches the current state of the PostgreSQL server.
130+
* `list_triggers`: Lists triggers in the database.
131+
* `list_indexes`: List available user indexes in a PostgreSQL database.
132+
* `list_sequences`: List sequences in a PostgreSQL database.
133+
* `list_publication_tables`: List publication tables in a PostgreSQL database.
134+
* `list_tablespaces`: Lists tablespaces in the database.
135+
* `list_pg_settings`: List configuration parameters for the PostgreSQL server.
136+
* `list_database_stats`: Lists the key performance and activity statistics for
137+
each database in the AlloyDB instance.
138+
* `list_roles`: Lists all the user-created roles in PostgreSQL database.
139+
103140
## BigQuery
104141

105142
* `--prebuilt` value: `bigquery`
@@ -243,17 +280,17 @@ See [Usage Examples](../reference/cli.md#examples).
243280
* `get_query_plan`: Generate the execution plan of a statement.
244281
* `list_views`: Lists views in the database from pg_views with a default
245282
limit of 50 rows. Returns schemaname, viewname and the ownername.
246-
* `list_schemas`: Lists schemas in the database.
247-
* `database_overview`: Fetches the current state of the PostgreSQL server.
248-
* `list_triggers`: Lists triggers in the database.
283+
* `list_schemas`: Lists schemas in the database.
284+
* `database_overview`: Fetches the current state of the PostgreSQL server.
285+
* `list_triggers`: Lists triggers in the database.
249286
* `list_indexes`: List available user indexes in a PostgreSQL database.
250287
* `list_sequences`: List sequences in a PostgreSQL database.
251288
* `list_publication_tables`: List publication tables in a PostgreSQL database.
252289
* `list_tablespaces`: Lists tablespaces in the database.
253290
* `list_pg_settings`: List configuration parameters for the PostgreSQL server.
254291
* `list_database_stats`: Lists the key performance and activity statistics for
255292
each database in the postgreSQL instance.
256-
* `list_roles`: Lists all the user-created roles in PostgreSQL database.
293+
* `list_roles`: Lists all the user-created roles in PostgreSQL database.
257294

258295
## Cloud SQL for PostgreSQL Observability
259296

@@ -564,17 +601,17 @@ See [Usage Examples](../reference/cli.md#examples).
564601
* `get_query_plan`: Generate the execution plan of a statement.
565602
* `list_views`: Lists views in the database from pg_views with a default
566603
limit of 50 rows. Returns schemaname, viewname and the ownername.
567-
* `list_schemas`: Lists schemas in the database.
568-
* `database_overview`: Fetches the current state of the PostgreSQL server.
569-
* `list_triggers`: Lists triggers in the database.
604+
* `list_schemas`: Lists schemas in the database.
605+
* `database_overview`: Fetches the current state of the PostgreSQL server.
606+
* `list_triggers`: Lists triggers in the database.
570607
* `list_indexes`: List available user indexes in a PostgreSQL database.
571608
* `list_sequences`: List sequences in a PostgreSQL database.
572609
* `list_publication_tables`: List publication tables in a PostgreSQL database.
573610
* `list_tablespaces`: Lists tablespaces in the database.
574611
* `list_pg_settings`: List configuration parameters for the PostgreSQL server.
575612
* `list_database_stats`: Lists the key performance and activity statistics for
576613
each database in the PostgreSQL server.
577-
* `list_roles`: Lists all the user-created roles in PostgreSQL database.
614+
* `list_roles`: Lists all the user-created roles in PostgreSQL database.
578615

579616
## Google Cloud Serverless for Apache Spark
580617

‎go.mod‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ require (
4747
github.com/sijms/go-ora/v2 v2.9.0
4848
github.com/snowflakedb/gosnowflake v1.18.1
4949
github.com/spf13/cobra v1.10.1
50+
github.com/testcontainers/testcontainers-go v0.40.0
5051
github.com/thlib/go-timezone-local v0.0.7
5152
github.com/trinodb/trino-go-client v0.330.0
5253
github.com/valkey-io/valkey-go v1.0.68
@@ -91,16 +92,19 @@ require (
9192
cloud.google.com/go/iam v1.5.3 // indirect
9293
cloud.google.com/go/monitoring v1.24.3 // indirect
9394
cloud.google.com/go/trace v1.11.7 // indirect
95+
dario.cat/mergo v1.0.2 // indirect
9496
filippo.io/edwards25519 v1.1.0 // indirect
9597
github.com/99designs/go-keychain v0.0.0-20191008050251-8e49817e8af4 // indirect
9698
github.com/99designs/keyring v1.2.2 // indirect
9799
github.com/Azure/azure-sdk-for-go/sdk/azcore v1.18.0 // indirect
98100
github.com/Azure/azure-sdk-for-go/sdk/internal v1.11.1 // indirect
99101
github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v1.0.0 // indirect
102+
github.com/Azure/go-ansiterm v0.0.0-20250102033503-faa5f7b0171c // indirect
100103
github.com/BurntSushi/toml v1.4.0 // indirect
101104
github.com/GoogleCloudPlatform/grpc-gcp-go/grpcgcp v1.5.3 // indirect
102105
github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.29.0 // indirect
103106
github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.54.0 // indirect
107+
github.com/Microsoft/go-winio v0.6.2 // indirect
104108
github.com/PuerkitoBio/goquery v1.10.3 // indirect
105109
github.com/VictoriaMetrics/easyproto v0.1.4 // indirect
106110
github.com/ajg/form v1.5.1 // indirect
@@ -126,17 +130,29 @@ require (
126130
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.34.4 // indirect
127131
github.com/aws/aws-sdk-go-v2/service/sts v1.38.4 // indirect
128132
github.com/aws/smithy-go v1.23.0 // indirect
133+
github.com/cenkalti/backoff/v4 v4.3.0 // indirect
129134
github.com/cespare/xxhash/v2 v2.3.0 // indirect
130135
github.com/cncf/xds/go v0.0.0-20250501225837-2ac532fd4443 // indirect
136+
github.com/containerd/errdefs v1.0.0 // indirect
137+
github.com/containerd/errdefs/pkg v0.3.0 // indirect
138+
github.com/containerd/log v0.1.0 // indirect
139+
github.com/containerd/platforms v0.2.1 // indirect
131140
github.com/couchbase/gocbcore/v10 v10.8.1 // indirect
132141
github.com/couchbase/gocbcoreps v0.1.4 // indirect
133142
github.com/couchbase/goprotostellar v1.0.2 // indirect
134143
github.com/couchbase/tools-common/errors v1.0.0 // indirect
135144
github.com/couchbaselabs/gocbconnstr/v2 v2.0.0 // indirect
145+
github.com/cpuguy83/dockercfg v0.3.2 // indirect
136146
github.com/danieljoos/wincred v1.2.2 // indirect
147+
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
137148
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
149+
github.com/distribution/reference v0.6.0 // indirect
150+
github.com/docker/docker v28.5.1+incompatible // indirect
151+
github.com/docker/go-connections v0.6.0 // indirect
152+
github.com/docker/go-units v0.5.0 // indirect
138153
github.com/dustin/go-humanize v1.0.1 // indirect
139154
github.com/dvsekhvalnov/jose2go v1.7.0 // indirect
155+
github.com/ebitengine/purego v0.8.4 // indirect
140156
github.com/envoyproxy/go-control-plane/envoy v1.32.4 // indirect
141157
github.com/envoyproxy/protoc-gen-validate v1.2.1 // indirect
142158
github.com/felixge/httpsnoop v1.0.4 // indirect
@@ -145,6 +161,7 @@ require (
145161
github.com/go-logfmt/logfmt v0.6.0 // indirect
146162
github.com/go-logr/logr v1.4.3 // indirect
147163
github.com/go-logr/stdr v1.2.2 // indirect
164+
github.com/go-ole/go-ole v1.2.6 // indirect
148165
github.com/go-playground/locales v0.14.1 // indirect
149166
github.com/go-playground/universal-translator v0.18.1 // indirect
150167
github.com/goccy/go-json v0.10.5 // indirect
@@ -180,27 +197,47 @@ require (
180197
github.com/klauspost/compress v1.18.0 // indirect
181198
github.com/klauspost/cpuid/v2 v2.2.11 // indirect
182199
github.com/leodido/go-urn v1.4.0 // indirect
200+
github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect
201+
github.com/magiconair/properties v1.8.10 // indirect
183202
github.com/mattn/go-isatty v0.0.20 // indirect
184203
github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 // indirect
185204
github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3 // indirect
205+
github.com/moby/docker-image-spec v1.3.1 // indirect
206+
github.com/moby/go-archive v0.1.0 // indirect
207+
github.com/moby/patternmatcher v0.6.0 // indirect
208+
github.com/moby/sys/sequential v0.6.0 // indirect
209+
github.com/moby/sys/user v0.4.0 // indirect
210+
github.com/moby/sys/userns v0.1.0 // indirect
211+
github.com/moby/term v0.5.2 // indirect
186212
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
187213
github.com/modern-go/reflect2 v1.0.2 // indirect
188214
github.com/montanaflynn/stats v0.7.1 // indirect
215+
github.com/morikuni/aec v1.0.0 // indirect
189216
github.com/mtibben/percent v0.2.1 // indirect
190217
github.com/nakagami/chacha20 v0.1.0 // indirect
191218
github.com/ncruces/go-strftime v0.1.9 // indirect
219+
github.com/opencontainers/go-digest v1.0.0 // indirect
220+
github.com/opencontainers/image-spec v1.1.1 // indirect
192221
github.com/pierrec/lz4 v2.6.1+incompatible // indirect
193222
github.com/pierrec/lz4/v4 v4.1.22 // indirect
194223
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c // indirect
224+
github.com/pkg/errors v0.9.1 // indirect
195225
github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 // indirect
226+
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
227+
github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c // indirect
196228
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
229+
github.com/shirou/gopsutil/v4 v4.25.6 // indirect
197230
github.com/sirupsen/logrus v1.9.3 // indirect
198231
github.com/spf13/pflag v1.0.9 // indirect
199232
github.com/spiffe/go-spiffe/v2 v2.5.0 // indirect
233+
github.com/stretchr/testify v1.11.1 // indirect
234+
github.com/tklauser/go-sysconf v0.3.12 // indirect
235+
github.com/tklauser/numcpus v0.6.1 // indirect
200236
github.com/xdg-go/pbkdf2 v1.0.0 // indirect
201237
github.com/xdg-go/scram v1.1.2 // indirect
202238
github.com/xdg-go/stringprep v1.0.4 // indirect
203239
github.com/youmark/pkcs8 v0.0.0-20240726163527-a2c0da244d78 // indirect
240+
github.com/yusufpapurcu/wmi v1.2.4 // indirect
204241
github.com/zeebo/errs v1.4.0 // indirect
205242
github.com/zeebo/xxh3 v1.0.2 // indirect
206243
gitlab.com/nyarla/go-crypt v0.0.0-20160106005555-d9a5dc2b789b // indirect
@@ -233,6 +270,7 @@ require (
233270
google.golang.org/genproto/googleapis/rpc v0.0.0-20251103181224-f26f9409b101 // indirect
234271
gopkg.in/inf.v0 v0.9.1 // indirect
235272
gopkg.in/ini.v1 v1.67.0 // indirect
273+
gopkg.in/yaml.v3 v3.0.1 // indirect
236274
modernc.org/libc v1.66.10 // indirect
237275
modernc.org/mathutil v1.7.1 // indirect
238276
modernc.org/memory v1.11.0 // indirect

0 commit comments

Comments
 (0)