From c27c3fc15ebf314a543116adf6db3a6f128e4a80 Mon Sep 17 00:00:00 2001 From: Lukasz Zajaczkowski Date: Mon, 3 Aug 2026 11:44:35 +0200 Subject: [PATCH] add ability to list schemas for multiple tables in cloud query --- go/cloud-query/api/proto/cloudquery.proto | 3 ++ go/cloud-query/docs/api-reference.md | 22 +++++++++++++-- .../internal/connection/connection.go | 2 +- go/cloud-query/internal/connection/schema.go | 28 +++++++++++++------ .../proto/cloudquery/cloudquery.pb.go | 23 +++++++++++---- .../internal/proto/toolquery/toolquery.pb.go | 2 +- .../internal/service/cloudquery_schema.go | 6 ++-- lib/cloud_query/cloudquery.pb.ex | 1 + 8 files changed, 66 insertions(+), 21 deletions(-) diff --git a/go/cloud-query/api/proto/cloudquery.proto b/go/cloud-query/api/proto/cloudquery.proto index bcebb24b9d..ca7b70797e 100644 --- a/go/cloud-query/api/proto/cloudquery.proto +++ b/go/cloud-query/api/proto/cloudquery.proto @@ -51,6 +51,9 @@ message QueryInput { message SchemaInput { Connection connection = 1; optional string table = 2; + // Exact table names to return schemas for. When non-empty, takes + // precedence over `table` and matches with table_name = ANY(tables). + repeated string tables = 3; } message ExtractInput { diff --git a/go/cloud-query/docs/api-reference.md b/go/cloud-query/docs/api-reference.md index fd4da9019e..e38c1fddc7 100644 --- a/go/cloud-query/docs/api-reference.md +++ b/go/cloud-query/docs/api-reference.md @@ -186,9 +186,15 @@ The Schema method retrieves the schema information for cloud resources. message SchemaInput { Connection connection = 1; optional string table = 2; + // Exact table names to return schemas for. When non-empty, takes + // precedence over `table` and matches with table_name = ANY(tables). + repeated string tables = 3; } ``` +- `table`: optional substring / LIKE filter (`%table%`). When omitted and `tables` is empty, returns schemas for all provider tables. +- `tables`: optional list of exact table names. When non-empty, takes precedence over `table`. + #### Response: SchemaOutput ```protobuf @@ -212,7 +218,7 @@ message SchemaResult { Using curl (with grpcurl): ```bash -# Using grpcurl to make a Schema request +# Using grpcurl to make a Schema request (LIKE filter) grpcurl -d '{ "connection": { "provider": "aws", @@ -223,6 +229,18 @@ grpcurl -d '{ }, "table": "aws_ec2_%" }' -plaintext localhost:9192 cloudquery.CloudQuery/Schema + +# Exact schemas for multiple tables +grpcurl -d '{ + "connection": { + "provider": "aws", + "aws": { + "access_key_id": "YOUR_ACCESS_KEY", + "secret_access_key": "YOUR_SECRET_KEY" + } + }, + "tables": ["aws_vpc", "aws_ec2_instance", "aws_eks_cluster"] +}' -plaintext localhost:9192 cloudquery.CloudQuery/Schema ``` Using Postman: @@ -242,7 +260,7 @@ Using Postman: "secret_access_key": "YOUR_SECRET_KEY" } }, - "table": "aws_ec2_%" + "tables": ["aws_vpc", "aws_ec2_instance"] } ``` diff --git a/go/cloud-query/internal/connection/connection.go b/go/cloud-query/internal/connection/connection.go index 43cff59351..b85b9c53d6 100644 --- a/go/cloud-query/internal/connection/connection.go +++ b/go/cloud-query/internal/connection/connection.go @@ -20,7 +20,7 @@ var defaultDataSource = common.DataSource(args.DatabaseHost(), args.DatabasePort type Connection interface { Configure() error - Schema(table string) ([]cloudquery.SchemaResult, error) + Schema(table string, tables []string) ([]cloudquery.SchemaResult, error) Tables(table string) ([]string, error) Query(q string, args ...any) (columns []string, rows [][]any, err error) Exec(q string, args ...any) (sql.Result, error) diff --git a/go/cloud-query/internal/connection/schema.go b/go/cloud-query/internal/connection/schema.go index 213874c13b..c93b7c2293 100644 --- a/go/cloud-query/internal/connection/schema.go +++ b/go/cloud-query/internal/connection/schema.go @@ -1,8 +1,10 @@ package connection import ( + "database/sql" "fmt" + "github.com/lib/pq" "github.com/samber/lo" "k8s.io/klog/v2" @@ -10,15 +12,10 @@ import ( "github.com/pluralsh/console/go/cloud-query/internal/proto/cloudquery" ) -func (in *connection) Schema(table string) ([]cloudquery.SchemaResult, error) { - klog.V(log.LogLevelDebug).InfoS("running schema query", "table", table) +func (in *connection) Schema(table string, tables []string) ([]cloudquery.SchemaResult, error) { + klog.V(log.LogLevelDebug).InfoS("running schema query", "table", table, "tables", tables) - prefix := fmt.Sprintf("%s_", in.provider()) - - qResponse, err := in.db.Query(` - SELECT table_name, column_name, data_type - FROM information_schema.columns - WHERE table_name LIKE $1;`, lo.Ternary(lo.IsEmpty(table), prefix+"%", "%"+table+"%")) + qResponse, err := in.querySchema(table, tables) if err != nil { return nil, err } @@ -55,6 +52,21 @@ func (in *connection) Schema(table string) ([]cloudquery.SchemaResult, error) { return result, nil } +func (in *connection) querySchema(table string, tables []string) (*sql.Rows, error) { + if len(tables) > 0 { + return in.db.Query(` + SELECT table_name, column_name, data_type + FROM information_schema.columns + WHERE table_name = ANY($1);`, pq.Array(tables)) + } + + prefix := fmt.Sprintf("%s_", in.provider()) + return in.db.Query(` + SELECT table_name, column_name, data_type + FROM information_schema.columns + WHERE table_name LIKE $1;`, lo.Ternary(lo.IsEmpty(table), prefix+"%", "%"+table+"%")) +} + func (in *connection) Tables(table string) ([]string, error) { klog.V(log.LogLevelDebug).InfoS("running tables query", "table", table) diff --git a/go/cloud-query/internal/proto/cloudquery/cloudquery.pb.go b/go/cloud-query/internal/proto/cloudquery/cloudquery.pb.go index 5ce41767df..b90882c5a0 100644 --- a/go/cloud-query/internal/proto/cloudquery/cloudquery.pb.go +++ b/go/cloud-query/internal/proto/cloudquery/cloudquery.pb.go @@ -1,6 +1,6 @@ // Code generated by protoc-gen-go. DO NOT EDIT. // versions: -// protoc-gen-go v1.36.11 +// protoc-gen-go v1.36.11-devel // protoc v6.31.1 // source: cloudquery.proto @@ -462,9 +462,12 @@ func (x *QueryInput) GetQuery() string { } type SchemaInput struct { - state protoimpl.MessageState `protogen:"open.v1"` - Connection *Connection `protobuf:"bytes,1,opt,name=connection,proto3" json:"connection,omitempty"` - Table *string `protobuf:"bytes,2,opt,name=table,proto3,oneof" json:"table,omitempty"` + state protoimpl.MessageState `protogen:"open.v1"` + Connection *Connection `protobuf:"bytes,1,opt,name=connection,proto3" json:"connection,omitempty"` + Table *string `protobuf:"bytes,2,opt,name=table,proto3,oneof" json:"table,omitempty"` + // Exact table names to return schemas for. When non-empty, takes + // precedence over `table` and matches with table_name = ANY(tables). + Tables []string `protobuf:"bytes,3,rep,name=tables,proto3" json:"tables,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -513,6 +516,13 @@ func (x *SchemaInput) GetTable() string { return "" } +func (x *SchemaInput) GetTables() []string { + if x != nil { + return x.Tables + } + return nil +} + type ExtractInput struct { state protoimpl.MessageState `protogen:"open.v1"` Connection *Connection `protobuf:"bytes,1,opt,name=connection,proto3" json:"connection,omitempty"` @@ -956,12 +966,13 @@ const file_cloudquery_proto_rawDesc = "" + "\n" + "connection\x18\x01 \x01(\v2\x16.cloudquery.ConnectionR\n" + "connection\x12\x14\n" + - "\x05query\x18\x02 \x01(\tR\x05query\"j\n" + + "\x05query\x18\x02 \x01(\tR\x05query\"\x82\x01\n" + "\vSchemaInput\x126\n" + "\n" + "connection\x18\x01 \x01(\v2\x16.cloudquery.ConnectionR\n" + "connection\x12\x19\n" + - "\x05table\x18\x02 \x01(\tH\x00R\x05table\x88\x01\x01B\b\n" + + "\x05table\x18\x02 \x01(\tH\x00R\x05table\x88\x01\x01\x12\x16\n" + + "\x06tables\x18\x03 \x03(\tR\x06tablesB\b\n" + "\x06_table\"F\n" + "\fExtractInput\x126\n" + "\n" + diff --git a/go/cloud-query/internal/proto/toolquery/toolquery.pb.go b/go/cloud-query/internal/proto/toolquery/toolquery.pb.go index d004253d7e..31b4b4be87 100644 --- a/go/cloud-query/internal/proto/toolquery/toolquery.pb.go +++ b/go/cloud-query/internal/proto/toolquery/toolquery.pb.go @@ -1,6 +1,6 @@ // Code generated by protoc-gen-go. DO NOT EDIT. // versions: -// protoc-gen-go v1.36.11 +// protoc-gen-go v1.36.11-devel // protoc v6.31.1 // source: toolquery.proto diff --git a/go/cloud-query/internal/service/cloudquery_schema.go b/go/cloud-query/internal/service/cloudquery_schema.go index e3dd9fadc8..8cfbebd9db 100644 --- a/go/cloud-query/internal/service/cloudquery_schema.go +++ b/go/cloud-query/internal/service/cloudquery_schema.go @@ -18,11 +18,11 @@ func (in *CloudQueryService) Schema(_ context.Context, input *cloudquery.SchemaI return nil, err } - return in.handleSchema(c, input.GetTable()) + return in.handleSchema(c, input.GetTable(), input.GetTables()) } -func (in *CloudQueryService) handleSchema(c connection.Connection, table string) (*cloudquery.SchemaOutput, error) { - result, err := c.Schema(table) +func (in *CloudQueryService) handleSchema(c connection.Connection, table string, tables []string) (*cloudquery.SchemaOutput, error) { + result, err := c.Schema(table, tables) if err != nil { return nil, status.Errorf(codes.Internal, "failed to execute schema query '%s': %v", table, err) } diff --git a/lib/cloud_query/cloudquery.pb.ex b/lib/cloud_query/cloudquery.pb.ex index 74e337959b..478a8512c9 100644 --- a/lib/cloud_query/cloudquery.pb.ex +++ b/lib/cloud_query/cloudquery.pb.ex @@ -96,6 +96,7 @@ defmodule Cloudquery.SchemaInput do field :connection, 1, type: Cloudquery.Connection field :table, 2, proto3_optional: true, type: :string + field :tables, 3, repeated: true, type: :string end defmodule Cloudquery.ExtractInput do