From 016f1e5e39c8692b135a40a41e5bfdcbb632e94d Mon Sep 17 00:00:00 2001 From: crisfiorini-oss <301773486+crisfiorini-oss@users.noreply.github.com> Date: Thu, 8 Oct 2026 16:00:49 -0400 Subject: [PATCH 01/14] feat(proto): give RowError a portable refusal reason InsertRows reported a refused row as an index and free text, so a caller that must set aside exactly the rows a backend permanently refused, and retry the rest, could only read the text. That is the question codefly-dev/service-warehouse#12 asks: can per-row refusal detail survive the error normalization? RowError now carries a RowRefusal: unspecified, invalid value, row too large, or schema mismatch. The names are the ones every warehouse can mean the same thing by, and schema mismatch is the one that says the table, not the row, is what disagrees, so a caller does not set a row aside for a table's fault. A failure that says nothing against the row (throttle, quota, timeout) stays a status code and is never a row refusal. The Go-native backend.RowError gets the same field and the server maps it in one place; a reason the server does not know stays unspecified instead of being guessed. mem still returns Unsupported for InsertRows, now pinned by a test, and the wire mapping is covered with a backend that refuses one row per reason. gen/ was regenerated through `codefly generate proto`; the generator version banners are unchanged. Co-Authored-By: Claude Sonnet 5.5 --- gen/codefly/warehouse/v0/warehouse.pb.go | 305 +++++++++++++-------- internal/backend/backend.go | 29 +- internal/backend/mem/mem_test.go | 10 + internal/server/convert.go | 20 ++ internal/server/server.go | 2 +- internal/server/server_test.go | 73 ++++- proto/codefly/warehouse/v0/warehouse.proto | 28 +- 7 files changed, 354 insertions(+), 113 deletions(-) diff --git a/gen/codefly/warehouse/v0/warehouse.pb.go b/gen/codefly/warehouse/v0/warehouse.pb.go index d6109bf..c67bb8f 100644 --- a/gen/codefly/warehouse/v0/warehouse.pb.go +++ b/gen/codefly/warehouse/v0/warehouse.pb.go @@ -353,6 +353,73 @@ func (WriteDisposition) EnumDescriptor() ([]byte, []int) { return file_codefly_warehouse_v0_warehouse_proto_rawDescGZIP(), []int{4} } +// RowRefusal is the portable reason a backend permanently refused one row of an +// InsertRows. It answers what a caller may do next: retrying the same row against +// the same table gets the same answer, unless the reason says the table, not the +// row, is what disagrees. A failure that says nothing against the row (a +// throttle, a quota, a timeout, a dead backend) is never a row refusal; it fails +// the call with a status code instead. +type RowRefusal int32 + +const ( + // The backend refused the row without saying which of the reasons below + // applies, or none of them does. Treat it as a refusal of the row. + RowRefusal_ROW_REFUSAL_UNSPECIFIED RowRefusal = 0 + // A value cannot be stored in its column: the wrong type, out of range, + // malformed, not valid text, or null in a column that is required. + RowRefusal_ROW_REFUSAL_INVALID_VALUE RowRefusal = 1 + // The row alone is bigger than the backend accepts for one row. + RowRefusal_ROW_REFUSAL_ROW_TOO_LARGE RowRefusal = 2 + // The row's columns disagree with the table's: a column the table does not + // have, or a required column the row does not carry. The table would have to + // change for the same row to be accepted, so a caller must not conclude the + // row itself is bad. + RowRefusal_ROW_REFUSAL_SCHEMA_MISMATCH RowRefusal = 3 +) + +// Enum value maps for RowRefusal. +var ( + RowRefusal_name = map[int32]string{ + 0: "ROW_REFUSAL_UNSPECIFIED", + 1: "ROW_REFUSAL_INVALID_VALUE", + 2: "ROW_REFUSAL_ROW_TOO_LARGE", + 3: "ROW_REFUSAL_SCHEMA_MISMATCH", + } + RowRefusal_value = map[string]int32{ + "ROW_REFUSAL_UNSPECIFIED": 0, + "ROW_REFUSAL_INVALID_VALUE": 1, + "ROW_REFUSAL_ROW_TOO_LARGE": 2, + "ROW_REFUSAL_SCHEMA_MISMATCH": 3, + } +) + +func (x RowRefusal) Enum() *RowRefusal { + p := new(RowRefusal) + *p = x + return p +} + +func (x RowRefusal) String() string { + return protoimpl.X.EnumStringOf(x.Descriptor(), protoreflect.EnumNumber(x)) +} + +func (RowRefusal) Descriptor() protoreflect.EnumDescriptor { + return file_codefly_warehouse_v0_warehouse_proto_enumTypes[5].Descriptor() +} + +func (RowRefusal) Type() protoreflect.EnumType { + return &file_codefly_warehouse_v0_warehouse_proto_enumTypes[5] +} + +func (x RowRefusal) Number() protoreflect.EnumNumber { + return protoreflect.EnumNumber(x) +} + +// Deprecated: Use RowRefusal.Descriptor instead. +func (RowRefusal) EnumDescriptor() ([]byte, []int) { + return file_codefly_warehouse_v0_warehouse_proto_rawDescGZIP(), []int{5} +} + // Column describes one field of a table or a result set. native_type is the // exact backend spelling (e.g. "NUMERIC(38,9)", "LowCardinality(String)") for // clients that need it; type is the portable bucket. @@ -2321,11 +2388,15 @@ func (x *InsertRowsResult) GetErrors() []*RowError { } // RowError reports a per-row rejection — streaming insert is not atomic, so -// partial success is normal and surfaced per offset. +// partial success is normal and surfaced per offset. row_index counts rows from +// 0 across every batch of the call, in the order they were sent. error is a +// description written by this server, never a backend's own text; branch on +// reason. type RowError struct { state protoimpl.MessageState `protogen:"open.v1"` RowIndex int64 `protobuf:"varint,1,opt,name=row_index,json=rowIndex,proto3" json:"row_index,omitempty"` Error string `protobuf:"bytes,2,opt,name=error,proto3" json:"error,omitempty"` + Reason RowRefusal `protobuf:"varint,3,opt,name=reason,proto3,enum=codefly.warehouse.v0.RowRefusal" json:"reason,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -2374,6 +2445,13 @@ func (x *RowError) GetError() string { return "" } +func (x *RowError) GetReason() RowRefusal { + if x != nil { + return x.Reason + } + return RowRefusal_ROW_REFUSAL_UNSPECIFIED +} + type CapabilitiesRequest struct { state protoimpl.MessageState `protogen:"open.v1"` unknownFields protoimpl.UnknownFields @@ -2801,10 +2879,11 @@ const file_codefly_warehouse_v0_warehouse_proto_rawDesc = "" + "\x10arrow_ipc_schema\x18\x02 \x01(\fR\x0earrowIpcSchema\"o\n" + "\x10InsertRowsResult\x12#\n" + "\rrows_inserted\x18\x01 \x01(\x03R\frowsInserted\x126\n" + - "\x06errors\x18\x02 \x03(\v2\x1e.codefly.warehouse.v0.RowErrorR\x06errors\"=\n" + + "\x06errors\x18\x02 \x03(\v2\x1e.codefly.warehouse.v0.RowErrorR\x06errors\"w\n" + "\bRowError\x12\x1b\n" + "\trow_index\x18\x01 \x01(\x03R\browIndex\x12\x14\n" + - "\x05error\x18\x02 \x01(\tR\x05error\"\x15\n" + + "\x05error\x18\x02 \x01(\tR\x05error\x128\n" + + "\x06reason\x18\x03 \x01(\x0e2 .codefly.warehouse.v0.RowRefusalR\x06reason\"\x15\n" + "\x13CapabilitiesRequest\"\xa5\x03\n" + "\x13BackendCapabilities\x12\x18\n" + "\abackend\x18\x01 \x01(\tR\abackend\x12\x1d\n" + @@ -2879,7 +2958,13 @@ const file_codefly_warehouse_v0_warehouse_proto_rawDesc = "" + "\x1dWRITE_DISPOSITION_UNSPECIFIED\x10\x00\x12\x1c\n" + "\x18WRITE_DISPOSITION_APPEND\x10\x01\x12\x1e\n" + "\x1aWRITE_DISPOSITION_TRUNCATE\x10\x02\x12\x1b\n" + - "\x17WRITE_DISPOSITION_EMPTY\x10\x032\xb3\n" + + "\x17WRITE_DISPOSITION_EMPTY\x10\x03*\x88\x01\n" + + "\n" + + "RowRefusal\x12\x1b\n" + + "\x17ROW_REFUSAL_UNSPECIFIED\x10\x00\x12\x1d\n" + + "\x19ROW_REFUSAL_INVALID_VALUE\x10\x01\x12\x1d\n" + + "\x19ROW_REFUSAL_ROW_TOO_LARGE\x10\x02\x12\x1f\n" + + "\x1bROW_REFUSAL_SCHEMA_MISMATCH\x10\x032\xb3\n" + "\n" + "\tWarehouse\x12R\n" + "\x05Query\x12\".codefly.warehouse.v0.QueryRequest\x1a#.codefly.warehouse.v0.QueryResponse0\x01\x12H\n" + @@ -2912,7 +2997,7 @@ func file_codefly_warehouse_v0_warehouse_proto_rawDescGZIP() []byte { return file_codefly_warehouse_v0_warehouse_proto_rawDescData } -var file_codefly_warehouse_v0_warehouse_proto_enumTypes = make([]protoimpl.EnumInfo, 5) +var file_codefly_warehouse_v0_warehouse_proto_enumTypes = make([]protoimpl.EnumInfo, 6) var file_codefly_warehouse_v0_warehouse_proto_msgTypes = make([]protoimpl.MessageInfo, 40) var file_codefly_warehouse_v0_warehouse_proto_goTypes = []any{ (ColumnType)(0), // 0: codefly.warehouse.v0.ColumnType @@ -2920,118 +3005,120 @@ var file_codefly_warehouse_v0_warehouse_proto_goTypes = []any{ (TableKind)(0), // 2: codefly.warehouse.v0.TableKind (LoadFormat)(0), // 3: codefly.warehouse.v0.LoadFormat (WriteDisposition)(0), // 4: codefly.warehouse.v0.WriteDisposition - (*Column)(nil), // 5: codefly.warehouse.v0.Column - (*TableRef)(nil), // 6: codefly.warehouse.v0.TableRef - (*JobStats)(nil), // 7: codefly.warehouse.v0.JobStats - (*Job)(nil), // 8: codefly.warehouse.v0.Job - (*QueryParam)(nil), // 9: codefly.warehouse.v0.QueryParam - (*QueryRequest)(nil), // 10: codefly.warehouse.v0.QueryRequest - (*QueryResponse)(nil), // 11: codefly.warehouse.v0.QueryResponse - (*QueryHeader)(nil), // 12: codefly.warehouse.v0.QueryHeader - (*GetJobRequest)(nil), // 13: codefly.warehouse.v0.GetJobRequest - (*CancelJobRequest)(nil), // 14: codefly.warehouse.v0.CancelJobRequest - (*CancelJobResult)(nil), // 15: codefly.warehouse.v0.CancelJobResult - (*ListDatasetsRequest)(nil), // 16: codefly.warehouse.v0.ListDatasetsRequest - (*Dataset)(nil), // 17: codefly.warehouse.v0.Dataset - (*ListDatasetsResult)(nil), // 18: codefly.warehouse.v0.ListDatasetsResult - (*ListTablesRequest)(nil), // 19: codefly.warehouse.v0.ListTablesRequest - (*TableInfo)(nil), // 20: codefly.warehouse.v0.TableInfo - (*ListTablesResult)(nil), // 21: codefly.warehouse.v0.ListTablesResult - (*GetTableRequest)(nil), // 22: codefly.warehouse.v0.GetTableRequest - (*TableSchema)(nil), // 23: codefly.warehouse.v0.TableSchema - (*CreateDatasetRequest)(nil), // 24: codefly.warehouse.v0.CreateDatasetRequest - (*DropDatasetRequest)(nil), // 25: codefly.warehouse.v0.DropDatasetRequest - (*DropResult)(nil), // 26: codefly.warehouse.v0.DropResult - (*CreateTableRequest)(nil), // 27: codefly.warehouse.v0.CreateTableRequest - (*DropTableRequest)(nil), // 28: codefly.warehouse.v0.DropTableRequest - (*LoadRequest)(nil), // 29: codefly.warehouse.v0.LoadRequest - (*UnloadRequest)(nil), // 30: codefly.warehouse.v0.UnloadRequest - (*InsertRowsRequest)(nil), // 31: codefly.warehouse.v0.InsertRowsRequest - (*InsertHeader)(nil), // 32: codefly.warehouse.v0.InsertHeader - (*InsertRowsResult)(nil), // 33: codefly.warehouse.v0.InsertRowsResult - (*RowError)(nil), // 34: codefly.warehouse.v0.RowError - (*CapabilitiesRequest)(nil), // 35: codefly.warehouse.v0.CapabilitiesRequest - (*BackendCapabilities)(nil), // 36: codefly.warehouse.v0.BackendCapabilities - (*NativeRequest)(nil), // 37: codefly.warehouse.v0.NativeRequest - (*NativeResult)(nil), // 38: codefly.warehouse.v0.NativeResult - nil, // 39: codefly.warehouse.v0.Dataset.LabelsEntry - nil, // 40: codefly.warehouse.v0.CreateDatasetRequest.LabelsEntry - nil, // 41: codefly.warehouse.v0.LoadRequest.OptionsEntry - nil, // 42: codefly.warehouse.v0.UnloadRequest.OptionsEntry - nil, // 43: codefly.warehouse.v0.NativeRequest.ParamsEntry - nil, // 44: codefly.warehouse.v0.NativeResult.ValuesEntry + (RowRefusal)(0), // 5: codefly.warehouse.v0.RowRefusal + (*Column)(nil), // 6: codefly.warehouse.v0.Column + (*TableRef)(nil), // 7: codefly.warehouse.v0.TableRef + (*JobStats)(nil), // 8: codefly.warehouse.v0.JobStats + (*Job)(nil), // 9: codefly.warehouse.v0.Job + (*QueryParam)(nil), // 10: codefly.warehouse.v0.QueryParam + (*QueryRequest)(nil), // 11: codefly.warehouse.v0.QueryRequest + (*QueryResponse)(nil), // 12: codefly.warehouse.v0.QueryResponse + (*QueryHeader)(nil), // 13: codefly.warehouse.v0.QueryHeader + (*GetJobRequest)(nil), // 14: codefly.warehouse.v0.GetJobRequest + (*CancelJobRequest)(nil), // 15: codefly.warehouse.v0.CancelJobRequest + (*CancelJobResult)(nil), // 16: codefly.warehouse.v0.CancelJobResult + (*ListDatasetsRequest)(nil), // 17: codefly.warehouse.v0.ListDatasetsRequest + (*Dataset)(nil), // 18: codefly.warehouse.v0.Dataset + (*ListDatasetsResult)(nil), // 19: codefly.warehouse.v0.ListDatasetsResult + (*ListTablesRequest)(nil), // 20: codefly.warehouse.v0.ListTablesRequest + (*TableInfo)(nil), // 21: codefly.warehouse.v0.TableInfo + (*ListTablesResult)(nil), // 22: codefly.warehouse.v0.ListTablesResult + (*GetTableRequest)(nil), // 23: codefly.warehouse.v0.GetTableRequest + (*TableSchema)(nil), // 24: codefly.warehouse.v0.TableSchema + (*CreateDatasetRequest)(nil), // 25: codefly.warehouse.v0.CreateDatasetRequest + (*DropDatasetRequest)(nil), // 26: codefly.warehouse.v0.DropDatasetRequest + (*DropResult)(nil), // 27: codefly.warehouse.v0.DropResult + (*CreateTableRequest)(nil), // 28: codefly.warehouse.v0.CreateTableRequest + (*DropTableRequest)(nil), // 29: codefly.warehouse.v0.DropTableRequest + (*LoadRequest)(nil), // 30: codefly.warehouse.v0.LoadRequest + (*UnloadRequest)(nil), // 31: codefly.warehouse.v0.UnloadRequest + (*InsertRowsRequest)(nil), // 32: codefly.warehouse.v0.InsertRowsRequest + (*InsertHeader)(nil), // 33: codefly.warehouse.v0.InsertHeader + (*InsertRowsResult)(nil), // 34: codefly.warehouse.v0.InsertRowsResult + (*RowError)(nil), // 35: codefly.warehouse.v0.RowError + (*CapabilitiesRequest)(nil), // 36: codefly.warehouse.v0.CapabilitiesRequest + (*BackendCapabilities)(nil), // 37: codefly.warehouse.v0.BackendCapabilities + (*NativeRequest)(nil), // 38: codefly.warehouse.v0.NativeRequest + (*NativeResult)(nil), // 39: codefly.warehouse.v0.NativeResult + nil, // 40: codefly.warehouse.v0.Dataset.LabelsEntry + nil, // 41: codefly.warehouse.v0.CreateDatasetRequest.LabelsEntry + nil, // 42: codefly.warehouse.v0.LoadRequest.OptionsEntry + nil, // 43: codefly.warehouse.v0.UnloadRequest.OptionsEntry + nil, // 44: codefly.warehouse.v0.NativeRequest.ParamsEntry + nil, // 45: codefly.warehouse.v0.NativeResult.ValuesEntry } var file_codefly_warehouse_v0_warehouse_proto_depIdxs = []int32{ 0, // 0: codefly.warehouse.v0.Column.type:type_name -> codefly.warehouse.v0.ColumnType - 5, // 1: codefly.warehouse.v0.Column.fields:type_name -> codefly.warehouse.v0.Column + 6, // 1: codefly.warehouse.v0.Column.fields:type_name -> codefly.warehouse.v0.Column 1, // 2: codefly.warehouse.v0.Job.state:type_name -> codefly.warehouse.v0.JobState - 7, // 3: codefly.warehouse.v0.Job.stats:type_name -> codefly.warehouse.v0.JobStats + 8, // 3: codefly.warehouse.v0.Job.stats:type_name -> codefly.warehouse.v0.JobStats 0, // 4: codefly.warehouse.v0.QueryParam.type:type_name -> codefly.warehouse.v0.ColumnType - 9, // 5: codefly.warehouse.v0.QueryRequest.params:type_name -> codefly.warehouse.v0.QueryParam - 12, // 6: codefly.warehouse.v0.QueryResponse.header:type_name -> codefly.warehouse.v0.QueryHeader - 8, // 7: codefly.warehouse.v0.QueryHeader.job:type_name -> codefly.warehouse.v0.Job - 5, // 8: codefly.warehouse.v0.QueryHeader.schema:type_name -> codefly.warehouse.v0.Column - 39, // 9: codefly.warehouse.v0.Dataset.labels:type_name -> codefly.warehouse.v0.Dataset.LabelsEntry - 17, // 10: codefly.warehouse.v0.ListDatasetsResult.datasets:type_name -> codefly.warehouse.v0.Dataset - 6, // 11: codefly.warehouse.v0.TableInfo.ref:type_name -> codefly.warehouse.v0.TableRef + 10, // 5: codefly.warehouse.v0.QueryRequest.params:type_name -> codefly.warehouse.v0.QueryParam + 13, // 6: codefly.warehouse.v0.QueryResponse.header:type_name -> codefly.warehouse.v0.QueryHeader + 9, // 7: codefly.warehouse.v0.QueryHeader.job:type_name -> codefly.warehouse.v0.Job + 6, // 8: codefly.warehouse.v0.QueryHeader.schema:type_name -> codefly.warehouse.v0.Column + 40, // 9: codefly.warehouse.v0.Dataset.labels:type_name -> codefly.warehouse.v0.Dataset.LabelsEntry + 18, // 10: codefly.warehouse.v0.ListDatasetsResult.datasets:type_name -> codefly.warehouse.v0.Dataset + 7, // 11: codefly.warehouse.v0.TableInfo.ref:type_name -> codefly.warehouse.v0.TableRef 2, // 12: codefly.warehouse.v0.TableInfo.kind:type_name -> codefly.warehouse.v0.TableKind - 20, // 13: codefly.warehouse.v0.ListTablesResult.tables:type_name -> codefly.warehouse.v0.TableInfo - 6, // 14: codefly.warehouse.v0.GetTableRequest.ref:type_name -> codefly.warehouse.v0.TableRef - 20, // 15: codefly.warehouse.v0.TableSchema.info:type_name -> codefly.warehouse.v0.TableInfo - 5, // 16: codefly.warehouse.v0.TableSchema.columns:type_name -> codefly.warehouse.v0.Column - 40, // 17: codefly.warehouse.v0.CreateDatasetRequest.labels:type_name -> codefly.warehouse.v0.CreateDatasetRequest.LabelsEntry - 6, // 18: codefly.warehouse.v0.CreateTableRequest.ref:type_name -> codefly.warehouse.v0.TableRef - 5, // 19: codefly.warehouse.v0.CreateTableRequest.columns:type_name -> codefly.warehouse.v0.Column - 6, // 20: codefly.warehouse.v0.DropTableRequest.ref:type_name -> codefly.warehouse.v0.TableRef - 6, // 21: codefly.warehouse.v0.LoadRequest.dest:type_name -> codefly.warehouse.v0.TableRef + 21, // 13: codefly.warehouse.v0.ListTablesResult.tables:type_name -> codefly.warehouse.v0.TableInfo + 7, // 14: codefly.warehouse.v0.GetTableRequest.ref:type_name -> codefly.warehouse.v0.TableRef + 21, // 15: codefly.warehouse.v0.TableSchema.info:type_name -> codefly.warehouse.v0.TableInfo + 6, // 16: codefly.warehouse.v0.TableSchema.columns:type_name -> codefly.warehouse.v0.Column + 41, // 17: codefly.warehouse.v0.CreateDatasetRequest.labels:type_name -> codefly.warehouse.v0.CreateDatasetRequest.LabelsEntry + 7, // 18: codefly.warehouse.v0.CreateTableRequest.ref:type_name -> codefly.warehouse.v0.TableRef + 6, // 19: codefly.warehouse.v0.CreateTableRequest.columns:type_name -> codefly.warehouse.v0.Column + 7, // 20: codefly.warehouse.v0.DropTableRequest.ref:type_name -> codefly.warehouse.v0.TableRef + 7, // 21: codefly.warehouse.v0.LoadRequest.dest:type_name -> codefly.warehouse.v0.TableRef 3, // 22: codefly.warehouse.v0.LoadRequest.format:type_name -> codefly.warehouse.v0.LoadFormat 4, // 23: codefly.warehouse.v0.LoadRequest.write_disposition:type_name -> codefly.warehouse.v0.WriteDisposition - 5, // 24: codefly.warehouse.v0.LoadRequest.schema:type_name -> codefly.warehouse.v0.Column - 41, // 25: codefly.warehouse.v0.LoadRequest.options:type_name -> codefly.warehouse.v0.LoadRequest.OptionsEntry - 6, // 26: codefly.warehouse.v0.UnloadRequest.table:type_name -> codefly.warehouse.v0.TableRef + 6, // 24: codefly.warehouse.v0.LoadRequest.schema:type_name -> codefly.warehouse.v0.Column + 42, // 25: codefly.warehouse.v0.LoadRequest.options:type_name -> codefly.warehouse.v0.LoadRequest.OptionsEntry + 7, // 26: codefly.warehouse.v0.UnloadRequest.table:type_name -> codefly.warehouse.v0.TableRef 3, // 27: codefly.warehouse.v0.UnloadRequest.format:type_name -> codefly.warehouse.v0.LoadFormat - 42, // 28: codefly.warehouse.v0.UnloadRequest.options:type_name -> codefly.warehouse.v0.UnloadRequest.OptionsEntry - 32, // 29: codefly.warehouse.v0.InsertRowsRequest.header:type_name -> codefly.warehouse.v0.InsertHeader - 6, // 30: codefly.warehouse.v0.InsertHeader.table:type_name -> codefly.warehouse.v0.TableRef - 34, // 31: codefly.warehouse.v0.InsertRowsResult.errors:type_name -> codefly.warehouse.v0.RowError - 3, // 32: codefly.warehouse.v0.BackendCapabilities.load_formats:type_name -> codefly.warehouse.v0.LoadFormat - 43, // 33: codefly.warehouse.v0.NativeRequest.params:type_name -> codefly.warehouse.v0.NativeRequest.ParamsEntry - 44, // 34: codefly.warehouse.v0.NativeResult.values:type_name -> codefly.warehouse.v0.NativeResult.ValuesEntry - 10, // 35: codefly.warehouse.v0.Warehouse.Query:input_type -> codefly.warehouse.v0.QueryRequest - 13, // 36: codefly.warehouse.v0.Warehouse.GetJob:input_type -> codefly.warehouse.v0.GetJobRequest - 14, // 37: codefly.warehouse.v0.Warehouse.CancelJob:input_type -> codefly.warehouse.v0.CancelJobRequest - 16, // 38: codefly.warehouse.v0.Warehouse.ListDatasets:input_type -> codefly.warehouse.v0.ListDatasetsRequest - 19, // 39: codefly.warehouse.v0.Warehouse.ListTables:input_type -> codefly.warehouse.v0.ListTablesRequest - 22, // 40: codefly.warehouse.v0.Warehouse.GetTable:input_type -> codefly.warehouse.v0.GetTableRequest - 24, // 41: codefly.warehouse.v0.Warehouse.CreateDataset:input_type -> codefly.warehouse.v0.CreateDatasetRequest - 25, // 42: codefly.warehouse.v0.Warehouse.DropDataset:input_type -> codefly.warehouse.v0.DropDatasetRequest - 27, // 43: codefly.warehouse.v0.Warehouse.CreateTable:input_type -> codefly.warehouse.v0.CreateTableRequest - 28, // 44: codefly.warehouse.v0.Warehouse.DropTable:input_type -> codefly.warehouse.v0.DropTableRequest - 29, // 45: codefly.warehouse.v0.Warehouse.Load:input_type -> codefly.warehouse.v0.LoadRequest - 30, // 46: codefly.warehouse.v0.Warehouse.Unload:input_type -> codefly.warehouse.v0.UnloadRequest - 31, // 47: codefly.warehouse.v0.Warehouse.InsertRows:input_type -> codefly.warehouse.v0.InsertRowsRequest - 35, // 48: codefly.warehouse.v0.Warehouse.Capabilities:input_type -> codefly.warehouse.v0.CapabilitiesRequest - 37, // 49: codefly.warehouse.v0.Warehouse.Native:input_type -> codefly.warehouse.v0.NativeRequest - 11, // 50: codefly.warehouse.v0.Warehouse.Query:output_type -> codefly.warehouse.v0.QueryResponse - 8, // 51: codefly.warehouse.v0.Warehouse.GetJob:output_type -> codefly.warehouse.v0.Job - 15, // 52: codefly.warehouse.v0.Warehouse.CancelJob:output_type -> codefly.warehouse.v0.CancelJobResult - 18, // 53: codefly.warehouse.v0.Warehouse.ListDatasets:output_type -> codefly.warehouse.v0.ListDatasetsResult - 21, // 54: codefly.warehouse.v0.Warehouse.ListTables:output_type -> codefly.warehouse.v0.ListTablesResult - 23, // 55: codefly.warehouse.v0.Warehouse.GetTable:output_type -> codefly.warehouse.v0.TableSchema - 17, // 56: codefly.warehouse.v0.Warehouse.CreateDataset:output_type -> codefly.warehouse.v0.Dataset - 26, // 57: codefly.warehouse.v0.Warehouse.DropDataset:output_type -> codefly.warehouse.v0.DropResult - 23, // 58: codefly.warehouse.v0.Warehouse.CreateTable:output_type -> codefly.warehouse.v0.TableSchema - 26, // 59: codefly.warehouse.v0.Warehouse.DropTable:output_type -> codefly.warehouse.v0.DropResult - 8, // 60: codefly.warehouse.v0.Warehouse.Load:output_type -> codefly.warehouse.v0.Job - 8, // 61: codefly.warehouse.v0.Warehouse.Unload:output_type -> codefly.warehouse.v0.Job - 33, // 62: codefly.warehouse.v0.Warehouse.InsertRows:output_type -> codefly.warehouse.v0.InsertRowsResult - 36, // 63: codefly.warehouse.v0.Warehouse.Capabilities:output_type -> codefly.warehouse.v0.BackendCapabilities - 38, // 64: codefly.warehouse.v0.Warehouse.Native:output_type -> codefly.warehouse.v0.NativeResult - 50, // [50:65] is the sub-list for method output_type - 35, // [35:50] is the sub-list for method input_type - 35, // [35:35] is the sub-list for extension type_name - 35, // [35:35] is the sub-list for extension extendee - 0, // [0:35] is the sub-list for field type_name + 43, // 28: codefly.warehouse.v0.UnloadRequest.options:type_name -> codefly.warehouse.v0.UnloadRequest.OptionsEntry + 33, // 29: codefly.warehouse.v0.InsertRowsRequest.header:type_name -> codefly.warehouse.v0.InsertHeader + 7, // 30: codefly.warehouse.v0.InsertHeader.table:type_name -> codefly.warehouse.v0.TableRef + 35, // 31: codefly.warehouse.v0.InsertRowsResult.errors:type_name -> codefly.warehouse.v0.RowError + 5, // 32: codefly.warehouse.v0.RowError.reason:type_name -> codefly.warehouse.v0.RowRefusal + 3, // 33: codefly.warehouse.v0.BackendCapabilities.load_formats:type_name -> codefly.warehouse.v0.LoadFormat + 44, // 34: codefly.warehouse.v0.NativeRequest.params:type_name -> codefly.warehouse.v0.NativeRequest.ParamsEntry + 45, // 35: codefly.warehouse.v0.NativeResult.values:type_name -> codefly.warehouse.v0.NativeResult.ValuesEntry + 11, // 36: codefly.warehouse.v0.Warehouse.Query:input_type -> codefly.warehouse.v0.QueryRequest + 14, // 37: codefly.warehouse.v0.Warehouse.GetJob:input_type -> codefly.warehouse.v0.GetJobRequest + 15, // 38: codefly.warehouse.v0.Warehouse.CancelJob:input_type -> codefly.warehouse.v0.CancelJobRequest + 17, // 39: codefly.warehouse.v0.Warehouse.ListDatasets:input_type -> codefly.warehouse.v0.ListDatasetsRequest + 20, // 40: codefly.warehouse.v0.Warehouse.ListTables:input_type -> codefly.warehouse.v0.ListTablesRequest + 23, // 41: codefly.warehouse.v0.Warehouse.GetTable:input_type -> codefly.warehouse.v0.GetTableRequest + 25, // 42: codefly.warehouse.v0.Warehouse.CreateDataset:input_type -> codefly.warehouse.v0.CreateDatasetRequest + 26, // 43: codefly.warehouse.v0.Warehouse.DropDataset:input_type -> codefly.warehouse.v0.DropDatasetRequest + 28, // 44: codefly.warehouse.v0.Warehouse.CreateTable:input_type -> codefly.warehouse.v0.CreateTableRequest + 29, // 45: codefly.warehouse.v0.Warehouse.DropTable:input_type -> codefly.warehouse.v0.DropTableRequest + 30, // 46: codefly.warehouse.v0.Warehouse.Load:input_type -> codefly.warehouse.v0.LoadRequest + 31, // 47: codefly.warehouse.v0.Warehouse.Unload:input_type -> codefly.warehouse.v0.UnloadRequest + 32, // 48: codefly.warehouse.v0.Warehouse.InsertRows:input_type -> codefly.warehouse.v0.InsertRowsRequest + 36, // 49: codefly.warehouse.v0.Warehouse.Capabilities:input_type -> codefly.warehouse.v0.CapabilitiesRequest + 38, // 50: codefly.warehouse.v0.Warehouse.Native:input_type -> codefly.warehouse.v0.NativeRequest + 12, // 51: codefly.warehouse.v0.Warehouse.Query:output_type -> codefly.warehouse.v0.QueryResponse + 9, // 52: codefly.warehouse.v0.Warehouse.GetJob:output_type -> codefly.warehouse.v0.Job + 16, // 53: codefly.warehouse.v0.Warehouse.CancelJob:output_type -> codefly.warehouse.v0.CancelJobResult + 19, // 54: codefly.warehouse.v0.Warehouse.ListDatasets:output_type -> codefly.warehouse.v0.ListDatasetsResult + 22, // 55: codefly.warehouse.v0.Warehouse.ListTables:output_type -> codefly.warehouse.v0.ListTablesResult + 24, // 56: codefly.warehouse.v0.Warehouse.GetTable:output_type -> codefly.warehouse.v0.TableSchema + 18, // 57: codefly.warehouse.v0.Warehouse.CreateDataset:output_type -> codefly.warehouse.v0.Dataset + 27, // 58: codefly.warehouse.v0.Warehouse.DropDataset:output_type -> codefly.warehouse.v0.DropResult + 24, // 59: codefly.warehouse.v0.Warehouse.CreateTable:output_type -> codefly.warehouse.v0.TableSchema + 27, // 60: codefly.warehouse.v0.Warehouse.DropTable:output_type -> codefly.warehouse.v0.DropResult + 9, // 61: codefly.warehouse.v0.Warehouse.Load:output_type -> codefly.warehouse.v0.Job + 9, // 62: codefly.warehouse.v0.Warehouse.Unload:output_type -> codefly.warehouse.v0.Job + 34, // 63: codefly.warehouse.v0.Warehouse.InsertRows:output_type -> codefly.warehouse.v0.InsertRowsResult + 37, // 64: codefly.warehouse.v0.Warehouse.Capabilities:output_type -> codefly.warehouse.v0.BackendCapabilities + 39, // 65: codefly.warehouse.v0.Warehouse.Native:output_type -> codefly.warehouse.v0.NativeResult + 51, // [51:66] is the sub-list for method output_type + 36, // [36:51] is the sub-list for method input_type + 36, // [36:36] is the sub-list for extension type_name + 36, // [36:36] is the sub-list for extension extendee + 0, // [0:36] is the sub-list for field type_name } func init() { file_codefly_warehouse_v0_warehouse_proto_init() } @@ -3056,7 +3143,7 @@ func file_codefly_warehouse_v0_warehouse_proto_init() { File: protoimpl.DescBuilder{ GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_codefly_warehouse_v0_warehouse_proto_rawDesc), len(file_codefly_warehouse_v0_warehouse_proto_rawDesc)), - NumEnums: 5, + NumEnums: 6, NumMessages: 40, NumExtensions: 0, NumServices: 1, diff --git a/internal/backend/backend.go b/internal/backend/backend.go index 13e9d34..2734226 100644 --- a/internal/backend/backend.go +++ b/internal/backend/backend.go @@ -292,10 +292,37 @@ type UnloadOptions struct { Options map[string]string } -// RowError is a per-row rejection from a streaming insert. +// RowRefusal is the portable reason a backend permanently refused one row of a +// streaming insert. It tells a caller what to do next: retrying the same row +// against the same table gets the same answer, unless the reason says the +// table, not the row, is what disagrees. A failure that says nothing against the +// row (a throttle, a quota, a timeout, a dead backend) is never a row refusal; +// the insert fails with a serr code instead. +type RowRefusal int + +const ( + // RefusalUnspecified is a refusal the backend does not explain, or one no + // other reason fits. It is still a refusal of the row. + RefusalUnspecified RowRefusal = iota + // RefusalInvalidValue means a value cannot be stored in its column: the + // wrong type, out of range, malformed, or null in a required column. + RefusalInvalidValue + // RefusalRowTooLarge means the row alone exceeds what the backend accepts. + RefusalRowTooLarge + // RefusalSchemaMismatch means the row's columns disagree with the table's: + // a column the table lacks, or a required column the row does not carry. The + // table would have to change for the same row to be accepted. + RefusalSchemaMismatch +) + +// RowError is a per-row rejection from a streaming insert. RowIndex counts rows +// from 0 across every batch of the call, in the order they were read. Error is +// a description the backend wrote itself, never a vendor's own text; callers +// branch on Reason. type RowError struct { RowIndex int64 Error string + Reason RowRefusal } // InsertResult reports a streaming insert outcome. diff --git a/internal/backend/mem/mem_test.go b/internal/backend/mem/mem_test.go index 0877259..42fc6b9 100644 --- a/internal/backend/mem/mem_test.go +++ b/internal/backend/mem/mem_test.go @@ -62,3 +62,13 @@ func TestCapabilities(t *testing.T) { require.True(t, caps.DDL) require.False(t, caps.ArrowResults) } + +// mem has no engine to ingest rows, so it says so instead of accepting rows it +// would drop, and it never reports a per-row refusal it did not make. +func TestInsertRowsUnsupported(t *testing.T) { + b := New(backend.Config{}) + res, err := b.InsertRows(context.Background(), backend.TableRef{Dataset: "d", Table: "t"}, nil, nil) + require.Nil(t, res) + require.True(t, serr.Is(err, serr.Unsupported)) + require.False(t, b.Capabilities().StreamingInsert) +} diff --git a/internal/server/convert.go b/internal/server/convert.go index a67ac03..6971377 100644 --- a/internal/server/convert.go +++ b/internal/server/convert.go @@ -245,6 +245,26 @@ var protoToWriteDisp = map[whv0.WriteDisposition]backend.WriteDisposition{ whv0.WriteDisposition_WRITE_DISPOSITION_EMPTY: backend.WriteEmpty, } +// --- RowError --- + +var rowRefusalToProto = map[backend.RowRefusal]whv0.RowRefusal{ + backend.RefusalUnspecified: whv0.RowRefusal_ROW_REFUSAL_UNSPECIFIED, + backend.RefusalInvalidValue: whv0.RowRefusal_ROW_REFUSAL_INVALID_VALUE, + backend.RefusalRowTooLarge: whv0.RowRefusal_ROW_REFUSAL_ROW_TOO_LARGE, + backend.RefusalSchemaMismatch: whv0.RowRefusal_ROW_REFUSAL_SCHEMA_MISMATCH, +} + +// rowErrorOut maps a refusal a backend reported to the wire. A reason this +// server does not know leaves the refusal unspecified: the row is still +// refused, and nothing false is said about why. +func rowErrorOut(e backend.RowError) *whv0.RowError { + return &whv0.RowError{ + RowIndex: e.RowIndex, + Error: e.Error, + Reason: rowRefusalToProto[e.Reason], + } +} + // --- Capabilities --- func capabilitiesOut(c backend.Capabilities) *whv0.BackendCapabilities { diff --git a/internal/server/server.go b/internal/server/server.go index 2611117..3909a99 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -212,7 +212,7 @@ func (s *Server) InsertRows(stream whv0.Warehouse_InsertRowsServer) error { } out := &whv0.InsertRowsResult{RowsInserted: res.RowsInserted} for _, e := range res.Errors { - out.Errors = append(out.Errors, &whv0.RowError{RowIndex: e.RowIndex, Error: e.Error}) + out.Errors = append(out.Errors, rowErrorOut(e)) } return stream.SendAndClose(out) } diff --git a/internal/server/server_test.go b/internal/server/server_test.go index 5f231cb..06e2d68 100644 --- a/internal/server/server_test.go +++ b/internal/server/server_test.go @@ -19,10 +19,16 @@ import ( // dial spins up the Warehouse service over an in-memory listener backed by the // mem backend and returns a connected client. func dial(t *testing.T) whv0.WarehouseClient { + t.Helper() + return dialBackend(t, mem.New(backend.Config{})) +} + +// dialBackend is dial over any backend. +func dialBackend(t *testing.T, be backend.Backend) whv0.WarehouseClient { t.Helper() lis := bufconn.Listen(1 << 20) srv := grpc.NewServer() - whv0.RegisterWarehouseServer(srv, New(mem.New(backend.Config{}))) + whv0.RegisterWarehouseServer(srv, New(be)) go func() { _ = srv.Serve(lis) }() t.Cleanup(srv.Stop) @@ -73,3 +79,68 @@ func TestServerCapabilities(t *testing.T) { require.Equal(t, "mem", caps.GetBackend()) require.True(t, caps.GetDdl()) } + +// refusingBackend is mem with an InsertRows that refuses one row for each +// portable reason, so the wire mapping of RowError is exercised without a +// warehouse. +type refusingBackend struct{ *mem.Backend } + +func (refusingBackend) InsertRows(_ context.Context, _ backend.TableRef, _ []byte, batches backend.BatchReader) (*backend.InsertResult, error) { + for { + if _, err := batches.Next(); err != nil { + break + } + } + return &backend.InsertResult{ + RowsInserted: 1, + Errors: []backend.RowError{ + {RowIndex: 0, Error: "no reason given", Reason: backend.RefusalUnspecified}, + {RowIndex: 1, Error: "bad value", Reason: backend.RefusalInvalidValue}, + {RowIndex: 2, Error: "too big", Reason: backend.RefusalRowTooLarge}, + {RowIndex: 3, Error: "table differs", Reason: backend.RefusalSchemaMismatch}, + {RowIndex: 4, Error: "a reason this server was never taught", Reason: backend.RowRefusal(99)}, + }, + }, nil +} + +func TestServerInsertRowsCarriesPerRowRefusalReasons(t *testing.T) { + ctx := context.Background() + c := dialBackend(t, refusingBackend{mem.New(backend.Config{})}) + + stream, err := c.InsertRows(ctx) + require.NoError(t, err) + require.NoError(t, stream.Send(&whv0.InsertRowsRequest{Kind: &whv0.InsertRowsRequest_Header{ + Header: &whv0.InsertHeader{Table: &whv0.TableRef{Dataset: "d", Table: "t"}}, + }})) + require.NoError(t, stream.Send(&whv0.InsertRowsRequest{Kind: &whv0.InsertRowsRequest_ArrowBatch{ArrowBatch: []byte("batch")}})) + res, err := stream.CloseAndRecv() + require.NoError(t, err) + + require.EqualValues(t, 1, res.GetRowsInserted()) + got := make([]whv0.RowRefusal, 0, len(res.GetErrors())) + for i, e := range res.GetErrors() { + require.EqualValues(t, i, e.GetRowIndex()) + got = append(got, e.GetReason()) + } + require.Equal(t, []whv0.RowRefusal{ + whv0.RowRefusal_ROW_REFUSAL_UNSPECIFIED, + whv0.RowRefusal_ROW_REFUSAL_INVALID_VALUE, + whv0.RowRefusal_ROW_REFUSAL_ROW_TOO_LARGE, + whv0.RowRefusal_ROW_REFUSAL_SCHEMA_MISMATCH, + // A reason the server does not know is left unspecified, never guessed. + whv0.RowRefusal_ROW_REFUSAL_UNSPECIFIED, + }, got) + require.Equal(t, "bad value", res.GetErrors()[1].GetError()) +} + +func TestServerInsertRowsUnimplementedOnMem(t *testing.T) { + ctx := context.Background() + c := dial(t) + stream, err := c.InsertRows(ctx) + require.NoError(t, err) + require.NoError(t, stream.Send(&whv0.InsertRowsRequest{Kind: &whv0.InsertRowsRequest_Header{ + Header: &whv0.InsertHeader{Table: &whv0.TableRef{Dataset: "d", Table: "t"}}, + }})) + _, err = stream.CloseAndRecv() + require.Equal(t, codes.Unimplemented, status.Code(err)) +} diff --git a/proto/codefly/warehouse/v0/warehouse.proto b/proto/codefly/warehouse/v0/warehouse.proto index c3e9dfe..223fd68 100644 --- a/proto/codefly/warehouse/v0/warehouse.proto +++ b/proto/codefly/warehouse/v0/warehouse.proto @@ -311,11 +311,37 @@ message InsertRowsResult { repeated RowError errors = 2; } +// RowRefusal is the portable reason a backend permanently refused one row of an +// InsertRows. It answers what a caller may do next: retrying the same row against +// the same table gets the same answer, unless the reason says the table, not the +// row, is what disagrees. A failure that says nothing against the row (a +// throttle, a quota, a timeout, a dead backend) is never a row refusal; it fails +// the call with a status code instead. +enum RowRefusal { + // The backend refused the row without saying which of the reasons below + // applies, or none of them does. Treat it as a refusal of the row. + ROW_REFUSAL_UNSPECIFIED = 0; + // A value cannot be stored in its column: the wrong type, out of range, + // malformed, not valid text, or null in a column that is required. + ROW_REFUSAL_INVALID_VALUE = 1; + // The row alone is bigger than the backend accepts for one row. + ROW_REFUSAL_ROW_TOO_LARGE = 2; + // The row's columns disagree with the table's: a column the table does not + // have, or a required column the row does not carry. The table would have to + // change for the same row to be accepted, so a caller must not conclude the + // row itself is bad. + ROW_REFUSAL_SCHEMA_MISMATCH = 3; +} + // RowError reports a per-row rejection — streaming insert is not atomic, so -// partial success is normal and surfaced per offset. +// partial success is normal and surfaced per offset. row_index counts rows from +// 0 across every batch of the call, in the order they were sent. error is a +// description written by this server, never a backend's own text; branch on +// reason. message RowError { int64 row_index = 1; string error = 2; + RowRefusal reason = 3; } message CapabilitiesRequest {} From 37e5fcca1a95997308ed65aac2e02ec6045bb624 Mon Sep 17 00:00:00 2001 From: crisfiorini-oss <301773486+crisfiorini-oss@users.noreply.github.com> Date: Thu, 8 Oct 2026 16:04:59 -0400 Subject: [PATCH 02/14] feat(serr): report an exceeded time limit as DEADLINE_EXCEEDED The normalized error model had no code for "the query ran past its time limit", so a backend that enforces timeout_ms or SWH_QUERY_TIMEOUT could only report it as Internal, which tells a client the gateway is broken when the query was merely too slow. Add serr.DeadlineExceeded and map it to the gRPC code of the same name, so a client can branch on it and retry with a longer limit. The service comment in the proto lists the status codes a client may see; it now names DEADLINE_EXCEEDED too, and gen/ is regenerated through `codefly generate proto` (comment-only change, generator banners unchanged). A test pins every normalized code to its gRPC code. Co-Authored-By: Claude Sonnet 5.5 --- gen/codefly/warehouse/v0/warehouse_grpc.pb.go | 6 ++-- internal/serr/serr.go | 5 +++ internal/server/errors.go | 2 ++ internal/server/errors_test.go | 32 +++++++++++++++++++ proto/codefly/warehouse/v0/warehouse.proto | 3 +- 5 files changed, 45 insertions(+), 3 deletions(-) create mode 100644 internal/server/errors_test.go diff --git a/gen/codefly/warehouse/v0/warehouse_grpc.pb.go b/gen/codefly/warehouse/v0/warehouse_grpc.pb.go index 9fee995..4462838 100644 --- a/gen/codefly/warehouse/v0/warehouse_grpc.pb.go +++ b/gen/codefly/warehouse/v0/warehouse_grpc.pb.go @@ -51,7 +51,8 @@ const ( // Warehouse is the generic data plane. Errors are returned as gRPC status codes // normalized across backends: NOT_FOUND, ALREADY_EXISTS, FAILED_PRECONDITION // (conditional DDL), INVALID_ARGUMENT (bad SQL), UNIMPLEMENTED (unsupported), -// PERMISSION_DENIED, RESOURCE_EXHAUSTED (quota/bytes-billed), INTERNAL. +// PERMISSION_DENIED, RESOURCE_EXHAUSTED (quota/bytes-billed), DEADLINE_EXCEEDED +// (the query or server time limit), INTERNAL. type WarehouseClient interface { // Query plane. Query(ctx context.Context, in *QueryRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[QueryResponse], error) @@ -251,7 +252,8 @@ func (c *warehouseClient) Native(ctx context.Context, in *NativeRequest, opts .. // Warehouse is the generic data plane. Errors are returned as gRPC status codes // normalized across backends: NOT_FOUND, ALREADY_EXISTS, FAILED_PRECONDITION // (conditional DDL), INVALID_ARGUMENT (bad SQL), UNIMPLEMENTED (unsupported), -// PERMISSION_DENIED, RESOURCE_EXHAUSTED (quota/bytes-billed), INTERNAL. +// PERMISSION_DENIED, RESOURCE_EXHAUSTED (quota/bytes-billed), DEADLINE_EXCEEDED +// (the query or server time limit), INTERNAL. type WarehouseServer interface { // Query plane. Query(*QueryRequest, grpc.ServerStreamingServer[QueryResponse]) error diff --git a/internal/serr/serr.go b/internal/serr/serr.go index c23a923..71d4eca 100644 --- a/internal/serr/serr.go +++ b/internal/serr/serr.go @@ -29,6 +29,9 @@ const ( Throttled // InvalidArgument means the request (or SQL) was malformed. InvalidArgument + // DeadlineExceeded means the operation did not finish within its time + // limit: the caller's or the server's configured one. + DeadlineExceeded ) func (c Code) String() string { @@ -47,6 +50,8 @@ func (c Code) String() string { return "Throttled" case InvalidArgument: return "InvalidArgument" + case DeadlineExceeded: + return "DeadlineExceeded" default: return "Internal" } diff --git a/internal/server/errors.go b/internal/server/errors.go index 2399181..d7022fe 100644 --- a/internal/server/errors.go +++ b/internal/server/errors.go @@ -28,6 +28,8 @@ func toStatus(err error) error { c = codes.ResourceExhausted case serr.InvalidArgument: c = codes.InvalidArgument + case serr.DeadlineExceeded: + c = codes.DeadlineExceeded default: c = codes.Internal } diff --git a/internal/server/errors_test.go b/internal/server/errors_test.go new file mode 100644 index 0000000..327a527 --- /dev/null +++ b/internal/server/errors_test.go @@ -0,0 +1,32 @@ +package server + +import ( + "errors" + "testing" + + "github.com/codefly-dev/service-warehouse/internal/serr" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +// Every normalized code reaches the client as one gRPC code, and only that code +// is what a client branches on. +func TestToStatusMapsEveryNormalizedCode(t *testing.T) { + for code, want := range map[serr.Code]codes.Code{ + serr.Internal: codes.Internal, + serr.NotFound: codes.NotFound, + serr.AlreadyExists: codes.AlreadyExists, + serr.PreconditionFailed: codes.FailedPrecondition, + serr.Unsupported: codes.Unimplemented, + serr.PermissionDenied: codes.PermissionDenied, + serr.Throttled: codes.ResourceExhausted, + serr.InvalidArgument: codes.InvalidArgument, + serr.DeadlineExceeded: codes.DeadlineExceeded, + } { + require.Equal(t, want, status.Code(toStatus(serr.New(code, "Op", "msg"))), code.String()) + } + // An error that carries no code is an internal failure, never a guess. + require.Equal(t, codes.Internal, status.Code(toStatus(errors.New("plain")))) + require.NoError(t, toStatus(nil)) +} diff --git a/proto/codefly/warehouse/v0/warehouse.proto b/proto/codefly/warehouse/v0/warehouse.proto index 223fd68..696dbec 100644 --- a/proto/codefly/warehouse/v0/warehouse.proto +++ b/proto/codefly/warehouse/v0/warehouse.proto @@ -378,7 +378,8 @@ message NativeResult { // Warehouse is the generic data plane. Errors are returned as gRPC status codes // normalized across backends: NOT_FOUND, ALREADY_EXISTS, FAILED_PRECONDITION // (conditional DDL), INVALID_ARGUMENT (bad SQL), UNIMPLEMENTED (unsupported), -// PERMISSION_DENIED, RESOURCE_EXHAUSTED (quota/bytes-billed), INTERNAL. +// PERMISSION_DENIED, RESOURCE_EXHAUSTED (quota/bytes-billed), DEADLINE_EXCEEDED +// (the query or server time limit), INTERNAL. service Warehouse { // Query plane. rpc Query(QueryRequest) returns (stream QueryResponse); From 2d563095c87f1444bc26e0d0fe8e826d7967f980 Mon Sep 17 00:00:00 2001 From: crisfiorini-oss <301773486+crisfiorini-oss@users.noreply.github.com> Date: Thu, 8 Oct 2026 16:28:25 -0400 Subject: [PATCH 03/14] feat(config): ask for the database alone from a backend that reads no DSN A backend bound to a database was refused with "SWH_DATABASE (or SWH_DSN) is required" whatever it was. For a backend with no connection string, such as one bound to a cloud project, the message sends the operator to set a variable the backend then refuses, and a DSN given alone passed the check only to fail later. Registration gains DSN, which says whether the backend reads SWH_DSN. Only such a backend may stand a DSN in for the database, and only it is told it may. The two kinds are covered by test backends, and the unregistered-kind test now uses snowflake so it does not name a backend this repository implements. Co-Authored-By: Claude Sonnet 5.5 --- internal/backend/registry.go | 5 +++++ internal/config/config.go | 8 ++++++-- internal/config/config_test.go | 36 ++++++++++++++++++++++++++++------ 3 files changed, 41 insertions(+), 8 deletions(-) diff --git a/internal/backend/registry.go b/internal/backend/registry.go index 732fcc7..c0ddfb7 100644 --- a/internal/backend/registry.go +++ b/internal/backend/registry.go @@ -14,6 +14,11 @@ type Constructor func(ctx context.Context, cfg Config) (Backend, error) // own configuration requirements; nothing outside this package should restate // them. type Registration struct { + // DSN marks a backend that reads SWH_DSN, a backend-native connection string + // that stands in for the database. Without it the configuration error asks + // for SWH_DATABASE alone, instead of offering a setting the backend would + // refuse. + DSN bool // SelfContained marks a backend that is not bound to a database — mem holds // its catalog in process, and DuckDB's database is an optional file path. // The zero value means bound, so a backend that says nothing is treated as diff --git a/internal/config/config.go b/internal/config/config.go index 6597b6d..47983e4 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -79,8 +79,12 @@ func FromEnv() (Config, error) { return Config{}, fmt.Errorf("unknown backend kind %q (registered: %v)", cfg.Backend.Kind, backend.Registered()) } - if !reg.SelfContained && cfg.Backend.Database == "" && cfg.Backend.DSN == "" { - return Config{}, fmt.Errorf("SWH_DATABASE (or SWH_DSN) is required for backend %q", cfg.Backend.Kind) + if !reg.SelfContained && cfg.Backend.Database == "" && !(reg.DSN && cfg.Backend.DSN != "") { + need := "SWH_DATABASE" + if reg.DSN { + need += " (or SWH_DSN)" + } + return Config{}, fmt.Errorf("%s is required for backend %q", need, cfg.Backend.Kind) } // Last, so a configuration broken in more than one way reports the backend // first: the token cannot be what stops a server that names no real backend. diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 552e952..092be3c 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -16,13 +16,22 @@ import ( _ "github.com/codefly-dev/service-warehouse/internal/backend/mem" ) -// boundKind registers a backend that is bound to a database, standing in for -// the cloud backends until one is implemented. Registering a real kind's name -// would assert a support claim this repository does not make. -const boundKind = "test-bound" +// boundKind registers a backend that is bound to a database and can be given a +// DSN instead, standing in for the cloud backends that are not implemented. +// Registering a real kind's name would assert a support claim this repository +// does not make. noDSNKind is the same without a DSN, which is what a project- +// bound backend such as BigQuery looks like to the configuration. +const ( + boundKind = "test-bound" + noDSNKind = "test-bound-no-dsn" +) func init() { backend.Register(boundKind, backend.Registration{ + DSN: true, + Open: func(context.Context, backend.Config) (backend.Backend, error) { return nil, nil }, + }) + backend.Register(noDSNKind, backend.Registration{ Open: func(context.Context, backend.Config) (backend.Backend, error) { return nil, nil }, }) } @@ -68,6 +77,21 @@ func TestBoundBackendRequiresDatabaseOrDSN(t *testing.T) { require.ErrorContains(t, err, boundKind) } +// A backend that reads no DSN must not be told to supply one: that sends the +// operator to set a variable the backend then refuses. +func TestBackendWithoutDSNAsksForTheDatabaseOnly(t *testing.T) { + clearEnv(t) + t.Setenv("SWH_BACKEND", noDSNKind) + + _, err := FromEnv() + require.ErrorContains(t, err, "SWH_DATABASE") + require.NotContains(t, err.Error(), "SWH_DSN") + + t.Setenv("SWH_DSN", "user:pass@account/db") + _, err = FromEnv() + require.ErrorContains(t, err, "SWH_DATABASE", "a DSN does not stand in for the database of a backend that has none") +} + func TestBoundBackendAcceptsDSNAlone(t *testing.T) { clearEnv(t) t.Setenv("SWH_BACKEND", boundKind) @@ -83,10 +107,10 @@ func TestBoundBackendAcceptsDSNAlone(t *testing.T) { // backend that does not exist. func TestUnknownBackendIsRejectedBeforeItsDatabase(t *testing.T) { clearEnv(t) - t.Setenv("SWH_BACKEND", "bigquery") + t.Setenv("SWH_BACKEND", "snowflake") _, err := FromEnv() - require.ErrorContains(t, err, `unknown backend kind "bigquery"`) + require.ErrorContains(t, err, `unknown backend kind "snowflake"`) require.ErrorContains(t, err, "mem") require.NotContains(t, err.Error(), "SWH_DATABASE") } From 7d58781269f77a88a81cf5e2a7b82a2c1a7d8401 Mon Sep 17 00:00:00 2001 From: crisfiorini-oss <301773486+crisfiorini-oss@users.noreply.github.com> Date: Thu, 8 Oct 2026 16:29:01 -0400 Subject: [PATCH 04/14] feat(backend): add the bigquery backend A BigQuery backend behind the uniform contract, so a client sends SQL and reads Arrow without linking the BigQuery SDK. It is registered as `bigquery` and compiled into the one binary. Implemented: the catalog and DDL plane (list/get/create/drop of datasets and tables, with partitioning, clustering and CTAS), Query with named parameters, a dry run, a byte cap and a timeout and Arrow-IPC batches, GetJob and CancelJob, and InsertRows from Arrow-IPC batches with streaming inserts. Load, Unload and Native return serr.Unsupported and Capabilities says so. - Every BigQuery failure is mapped to a serr.Code in one place (errors.go). A client is told a fixed sentence per code; the one vendor text it can read is the diagnosis of its own invalid SQL, and the vendor's own error goes to the operator's log and is never wrapped. - The type mapping is explicit in both directions and refuses what it cannot carry (a RANGE result, an Arrow map, a nanosecond timestamp) with Unsupported or an invalid argument, never a lookalike type or a rounded value. - InsertRows is not atomic: BigQuery is asked to skip invalid rows and each refused row is reported with its index across the whole stream and a portable reason. A row whose content BigQuery refuses is INVALID_VALUE, a row over the request limit is ROW_TOO_LARGE, and a batch whose columns disagree with the table's, found by comparing it with the table's metadata and not by reading vendor text, refuses every row as SCHEMA_MISMATCH without sending any. A failure that says nothing against a row fails the call. - Names that become request paths are checked first: the client library expands them without escaping, so a dataset named "d/tables/t" would address a table. - SWH_DATABASE, SWH_DATASET, SWH_LOCATION, SWH_CREDENTIALS_FILE, SWH_MAX_QUERY_BYTES and SWH_QUERY_TIMEOUT already existed and cover what BigQuery needs, so no variable is added; settings it never reads (a DSN, a host, a login) are refused at startup instead of ignored. Credentials are Application Default Credentials, or the service-account file. Adds cloud.google.com/go/bigquery and github.com/apache/arrow-go/v18, the maintained Arrow module (v18.8.0 is its latest; apache/arrow/go/v15 is the old one). grpc moves to v1.83.2, which the BigQuery client requires. Unit tests run the real client against a fake of BigQuery's REST API served from the test process: they check what is sent and how answers are read, not what BigQuery does. One integration test runs against a real project when SWH_TEST_BIGQUERY_PROJECT is set and is skipped with a message otherwise; it has not been run. Co-Authored-By: Claude Sonnet 5.5 --- cmd/service-warehouse/main.go | 1 + go.mod | 28 +- go.sum | 86 +++- internal/backend/bigquery/arrow.go | 449 ++++++++++++++++++ internal/backend/bigquery/arrow_test.go | 339 +++++++++++++ internal/backend/bigquery/bigquery.go | 184 +++++++ internal/backend/bigquery/bigquery_test.go | 121 +++++ internal/backend/bigquery/catalog.go | 380 +++++++++++++++ internal/backend/bigquery/catalog_test.go | 374 +++++++++++++++ internal/backend/bigquery/errors.go | 225 +++++++++ internal/backend/bigquery/errors_test.go | 153 ++++++ internal/backend/bigquery/fake_test.go | 128 +++++ internal/backend/bigquery/ident.go | 52 ++ internal/backend/bigquery/insert.go | 348 ++++++++++++++ internal/backend/bigquery/insert_test.go | 446 +++++++++++++++++ internal/backend/bigquery/integration_test.go | 275 +++++++++++ internal/backend/bigquery/params.go | 238 ++++++++++ internal/backend/bigquery/params_test.go | 122 +++++ internal/backend/bigquery/query.go | 297 ++++++++++++ internal/backend/bigquery/query_test.go | 447 +++++++++++++++++ internal/backend/bigquery/types.go | 362 ++++++++++++++ internal/backend/bigquery/types_test.go | 160 +++++++ internal/backend/bigquery/values.go | 254 ++++++++++ 23 files changed, 5467 insertions(+), 2 deletions(-) create mode 100644 internal/backend/bigquery/arrow.go create mode 100644 internal/backend/bigquery/arrow_test.go create mode 100644 internal/backend/bigquery/bigquery.go create mode 100644 internal/backend/bigquery/bigquery_test.go create mode 100644 internal/backend/bigquery/catalog.go create mode 100644 internal/backend/bigquery/catalog_test.go create mode 100644 internal/backend/bigquery/errors.go create mode 100644 internal/backend/bigquery/errors_test.go create mode 100644 internal/backend/bigquery/fake_test.go create mode 100644 internal/backend/bigquery/ident.go create mode 100644 internal/backend/bigquery/insert.go create mode 100644 internal/backend/bigquery/insert_test.go create mode 100644 internal/backend/bigquery/integration_test.go create mode 100644 internal/backend/bigquery/params.go create mode 100644 internal/backend/bigquery/params_test.go create mode 100644 internal/backend/bigquery/query.go create mode 100644 internal/backend/bigquery/query_test.go create mode 100644 internal/backend/bigquery/types.go create mode 100644 internal/backend/bigquery/types_test.go create mode 100644 internal/backend/bigquery/values.go diff --git a/cmd/service-warehouse/main.go b/cmd/service-warehouse/main.go index 76abd0f..7793067 100644 --- a/cmd/service-warehouse/main.go +++ b/cmd/service-warehouse/main.go @@ -26,6 +26,7 @@ import ( // Backends register themselves via init(); importing them compiles each into // the single binary and makes its kind selectable by config. + _ "github.com/codefly-dev/service-warehouse/internal/backend/bigquery" _ "github.com/codefly-dev/service-warehouse/internal/backend/duckdb" _ "github.com/codefly-dev/service-warehouse/internal/backend/mem" ) diff --git a/go.mod b/go.mod index 23f5033..7ec7d4c 100644 --- a/go.mod +++ b/go.mod @@ -3,8 +3,14 @@ module github.com/codefly-dev/service-warehouse go 1.27.3 require ( + cloud.google.com/go v0.123.0 + cloud.google.com/go/bigquery v1.85.0 + github.com/apache/arrow-go/v18 v18.8.0 github.com/codefly-dev/core v0.19.4 github.com/stretchr/testify v1.12.1 + golang.org/x/oauth2 v0.36.0 + golang.org/x/sync v0.23.0 + google.golang.org/api v0.287.1 google.golang.org/grpc v1.84.0 google.golang.org/protobuf v1.36.12 gopkg.in/yaml.v3 v3.0.1 @@ -15,12 +21,17 @@ require ( buf.build/go/protovalidate v1.4.0 // indirect cel.dev/cel-go v0.32.0 // indirect cel.dev/expr v0.25.3 // indirect + cloud.google.com/go/auth v0.20.0 // indirect + cloud.google.com/go/auth/oauth2adapt v0.2.8 // indirect + cloud.google.com/go/compute/metadata v0.9.0 // indirect + cloud.google.com/go/iam v1.11.0 // indirect github.com/Azure/go-ansiterm v0.0.0-20261006220739-8c912ac31dd4 // indirect github.com/Masterminds/semver v1.5.0 // indirect github.com/Masterminds/semver/v3 v3.5.0 // indirect github.com/Microsoft/go-winio v0.6.2 // indirect github.com/VividCortex/ewma v1.2.0 // indirect github.com/antlr4-go/antlr/v4 v4.13.1 // indirect + github.com/apache/arrow/go/v15 v15.0.2 // indirect github.com/blang/semver v3.5.1+incompatible // indirect github.com/blang/semver/v4 v4.0.0 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect @@ -61,14 +72,20 @@ require ( github.com/go-openapi/swag/typeutils v0.29.2 // indirect github.com/go-openapi/swag/yamlutils v0.29.2 // indirect github.com/go-viper/mapstructure/v2 v2.5.0 // indirect + github.com/goccy/go-json v0.10.6 // indirect github.com/gofrs/flock v0.13.1 // indirect + github.com/google/flatbuffers v25.12.19+incompatible // indirect github.com/google/gnostic-models v0.7.1 // indirect github.com/google/go-github/v89 v89.0.0 // indirect github.com/google/go-querystring v1.2.0 // indirect + github.com/google/s2a-go v0.1.9 // indirect github.com/google/uuid v1.6.0 // indirect + github.com/googleapis/enterprise-certificate-proxy v0.3.17 // indirect + github.com/googleapis/gax-go/v2 v2.23.0 // indirect github.com/hashicorp/errwrap v1.1.0 // indirect github.com/hashicorp/go-multierror v1.1.1 // indirect - github.com/kr/text v0.2.0 // indirect + github.com/klauspost/compress v1.19.2 // indirect + github.com/klauspost/cpuid/v2 v2.4.0 // indirect github.com/lufia/plan9stats v0.0.0-20260330125221-c963978e514e // indirect github.com/mattn/go-colorable v0.1.15 // indirect github.com/mattn/go-isatty v0.0.24 // indirect @@ -80,6 +97,7 @@ require ( github.com/oklog/ulid/v2 v2.1.2 // indirect github.com/opencontainers/go-digest v1.0.0 // indirect github.com/opencontainers/image-spec v1.1.1 // indirect + github.com/pierrec/lz4/v4 v4.1.29 // indirect github.com/pkg/errors v0.9.1 // indirect github.com/power-devops/perfstat v0.0.0-20240221224432-82ca36839d55 // indirect github.com/santhosh-tekuri/jsonschema/v6 v6.0.3 // indirect @@ -90,6 +108,7 @@ require ( github.com/xlab/treeprint v1.2.0 // indirect github.com/yoheimuta/go-protoparser/v4 v4.14.2 // indirect github.com/yusufpapurcu/wmi v1.2.4 // indirect + github.com/zeebo/xxh3 v1.1.0 // indirect go.opentelemetry.io/auto/sdk v1.2.1 // indirect go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.72.0 // indirect go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.71.0 // indirect @@ -99,11 +118,18 @@ require ( go.opentelemetry.io/otel/trace v1.47.0 // indirect go.yaml.in/yaml/v2 v2.4.4 // indirect go.yaml.in/yaml/v3 v3.0.5 // indirect + golang.org/x/crypto v0.57.0 // indirect golang.org/x/exp v0.0.0-20260908205506-85c1c2202aba // indirect + golang.org/x/mod v0.41.0 // indirect golang.org/x/net v0.60.0 // indirect golang.org/x/sys v0.49.0 // indirect + golang.org/x/telemetry v0.0.0-20260908163034-4bcc4b2ee518 // indirect golang.org/x/term v0.46.0 // indirect golang.org/x/text v0.42.0 // indirect + golang.org/x/time v0.15.0 // indirect + golang.org/x/tools v0.50.0 // indirect + golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da // indirect + google.golang.org/genproto v0.0.0-20260319201613-d00831a3d3e7 // indirect google.golang.org/genproto/googleapis/api v0.0.0-20260928230214-8a89bd6388cc // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260928230214-8a89bd6388cc // indirect gopkg.in/evanphx/json-patch.v4 v4.13.0 // indirect diff --git a/go.sum b/go.sum index f5aeeda..92435b3 100644 --- a/go.sum +++ b/go.sum @@ -6,10 +6,36 @@ cel.dev/cel-go v0.32.0 h1:irvpFKr5EuGPyxeME03ERh0rii1TX+BDAnB9eL3IvNk= cel.dev/cel-go v0.32.0/go.mod h1:DnVip7tpJSsgZymwfT+m1tnEVy3ivAjSMXPx12YrMkU= cel.dev/expr v0.25.3 h1:A2jO8jwOugrrovveCWfj0KEZOfqiLgAcwjpHPhzIGw0= cel.dev/expr v0.25.3/go.mod h1:hrXvqGP6G6gyx8UAHSHJ5RGk//1Oj5nXQ2NI02Nrsg4= +cloud.google.com/go v0.123.0 h1:2NAUJwPR47q+E35uaJeYoNhuNEM9kM8SjgRgdeOJUSE= +cloud.google.com/go v0.123.0/go.mod h1:xBoMV08QcqUGuPW65Qfm1o9Y4zKZBpGS+7bImXLTAZU= +cloud.google.com/go/auth v0.20.0 h1:kXTssoVb4azsVDoUiF8KvxAqrsQcQtB53DcSgta74CA= +cloud.google.com/go/auth v0.20.0/go.mod h1:942/yi/itH1SsmpyrbnTMDgGfdy2BUqIKyd0cyYLc5Q= +cloud.google.com/go/auth/oauth2adapt v0.2.8 h1:keo8NaayQZ6wimpNSmW5OPc283g65QNIiLpZnkHRbnc= +cloud.google.com/go/auth/oauth2adapt v0.2.8/go.mod h1:XQ9y31RkqZCcwJWNSx2Xvric3RrU88hAYYbjDWYDL+c= +cloud.google.com/go/bigquery v1.85.0 h1:zsFsa8jOVkU4c7CWE1cbrfsemtNbM3YRUmtFRYXYN58= +cloud.google.com/go/bigquery v1.85.0/go.mod h1:oBma1P5/b1Jtd8xRLKoyTeNIMlACGHbSMLudzxHGHgc= +cloud.google.com/go/compute/metadata v0.9.0 h1:pDUj4QMoPejqq20dK0Pg2N4yG9zIkYGdBtwLoEkH9Zs= +cloud.google.com/go/compute/metadata v0.9.0/go.mod h1:E0bWwX5wTnLPedCKqk3pJmVgCBSM6qQI1yTBdEb3C10= +cloud.google.com/go/datacatalog v1.32.0 h1:fyYn8ODkGil5y3zTIqgIhOfzTu1ACaU2o+C750CO6Ac= +cloud.google.com/go/datacatalog v1.32.0/go.mod h1:DE272tynQUwheJeQAyVfV+nO8yrdkuDyOgH2LtOrkWM= +cloud.google.com/go/iam v1.11.0 h1:KieQ9Pb+LLPak1O3Rv3GgCxhnmkYf7Xyh0P5HfF1jFM= +cloud.google.com/go/iam v1.11.0/go.mod h1:KP+nKGugNJW4LcLx1uEZcq1ok5sQHFaQehQNl4QDgV4= +cloud.google.com/go/longrunning v1.2.0 h1:WjYH3YHBGCxGJP9M4dWGHBfXr/cFIjMkNgWcJj7/iMM= +cloud.google.com/go/longrunning v1.2.0/go.mod h1:5KMQALFGOCtFoi2xSOA1u3H7WKlhmckgiyFw7+LGQp0= +cloud.google.com/go/monitoring v1.24.3 h1:dde+gMNc0UhPZD1Azu6at2e79bfdztVDS5lvhOdsgaE= +cloud.google.com/go/monitoring v1.24.3/go.mod h1:nYP6W0tm3N9H/bOw8am7t62YTzZY+zUeQ+Bi6+2eonI= +cloud.google.com/go/storage v1.62.3 h1:SZq1t23NCI+e96dH77Dg3PEfsNNEjqO8zE5AnD8gVD0= +cloud.google.com/go/storage v1.62.3/go.mod h1:cpYz/kRVZ+UQAF1uHeea10/9ewcRbxGoGNKsS9daSXA= dario.cat/mergo v1.0.2 h1:85+piFYR1tMbRrLcDwR18y4UKJ3aH1Tbzi24VRW1TK8= dario.cat/mergo v1.0.2/go.mod h1:E/hbnu0NxMFBjpMIE34DRGLWqDy0g5FuKDhCb31ngxA= github.com/Azure/go-ansiterm v0.0.0-20261006220739-8c912ac31dd4 h1:m2xcjRABtG288vnlV3U/PKh/HnPURxZp50RKNUwBxM8= github.com/Azure/go-ansiterm v0.0.0-20261006220739-8c912ac31dd4/go.mod h1:MTUE7ZuZO89FAtAwozGg6AWD83pcdefnEq5gF82MB8s= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.34.0 h1:yzIYdwuro811Z27D3T80Wkd3rqZzb0K43nner7Eh1yE= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.34.0/go.mod h1:pJTkW8hEUIIi3Pf65lPZOnn4Y81yCllX6IWk2jNXdkM= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.55.0 h1:UnDZ/zFfG1JhH/DqxIZYU/1CUAlTUScoXD/LcM2Ykk8= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.55.0/go.mod h1:IA1C1U7jO/ENqm/vhi7V9YYpBsp+IMyqNrEN94N7tVc= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.55.0 h1:0s6TxfCu2KHkkZPnBfsQ2y5qia0jl3MMrmBhu3nCOYk= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.55.0/go.mod h1:Mf6O40IAyB9zR/1J8nGDDPirZQQPbYJni8Yisy7NTMc= github.com/Masterminds/semver v1.5.0 h1:H65muMkzWKEuNDnfl9d70GUjFniHKHRbFPGBuZ3QEww= github.com/Masterminds/semver v1.5.0/go.mod h1:MB6lktGJrhw8PrUyiEoblNEGEQ+RzHPF078ddwwvV3Y= github.com/Masterminds/semver/v3 v3.5.0 h1:kQceYJfbupGfZOKZQg0kou0DgAKhzDg2NZPAwZ/2OOE= @@ -20,8 +46,16 @@ github.com/ProtonMail/go-crypto v1.4.1 h1:9RfcZHqEQUvP8RzecWEUafnZVtEvrBVL9BiF67 github.com/ProtonMail/go-crypto v1.4.1/go.mod h1:e1OaTyu5SYVrO9gKOEhTc+5UcXtTUa+P3uLudwcgPqo= github.com/VividCortex/ewma v1.2.0 h1:f58SaIzcDXrSy3kWaHNvuJgJ3Nmz59Zji6XoJR/q1ow= github.com/VividCortex/ewma v1.2.0/go.mod h1:nz4BbCtbLyFDeC9SUHbtcT5644juEuWfUAUnGx7j5l4= +github.com/andybalholm/brotli v1.2.3 h1:8H1qwOkl2LPfjf3YezB90JnCliZb6SInJ/OJkEbA5NQ= +github.com/andybalholm/brotli v1.2.3/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY= github.com/antlr4-go/antlr/v4 v4.13.1 h1:SqQKkuVZ+zWkMMNkjy5FZe5mr5WURWnlpmOuzYWrPrQ= github.com/antlr4-go/antlr/v4 v4.13.1/go.mod h1:GKmUxMtwp6ZgGwZSva4eWPC5mS6vUAmOABFgjdkM7Nw= +github.com/apache/arrow-go/v18 v18.8.0 h1:BLOzbPv7bxMPgXPacAg6HQjnxupYsZzC4tf+FkqPU/M= +github.com/apache/arrow-go/v18 v18.8.0/go.mod h1:uJCFfCwq0KsxCmsCfQg4ft+LsW+iHYzAXiSDh5ug/8U= +github.com/apache/arrow/go/v15 v15.0.2 h1:60IliRbiyTWCWjERBCkO1W4Qun9svcYoZrSLcyOsMLE= +github.com/apache/arrow/go/v15 v15.0.2/go.mod h1:DGXsR3ajT524njufqf95822i+KTh+yea1jass9YXgjA= +github.com/apache/thrift v0.24.0 h1:zy31L1a49QTNB2bG1BBfMXol3yJrTH975G3pPubQVLQ= +github.com/apache/thrift v0.24.0/go.mod h1:zPt6WxgvTOM6hF92y8C+MkEM5LMxZuk4JcQOiU4Esvs= github.com/blang/semver v3.5.1+incompatible h1:cQNTCjp13qL8KC3Nbxr/y2Bqb63oX6wdnnjpJbkM4JQ= github.com/blang/semver v3.5.1+incompatible/go.mod h1:kRBLl5iJ+tD4TcOOxsy/0fnwebNt5EWlYSAyrTnjyyk= github.com/blang/semver/v4 v4.0.0 h1:1PFHFE6yCCTv8C1TeyNNarDzntLi7wMI5i/pzqYIsAM= @@ -38,6 +72,8 @@ github.com/clipperhouse/uax29/v2 v2.7.0 h1:+gs4oBZ2gPfVrKPthwbMzWZDaAFPGYK72F0NJ github.com/clipperhouse/uax29/v2 v2.7.0/go.mod h1:EFJ2TJMRUaplDxHKj1qAEhCtQPW2tJSwu5BF98AuoVM= github.com/cloudflare/circl v1.6.3 h1:9GPOhQGF9MCYUeXyMYlqTR6a5gTrgR/fBLXvUgtVcg8= github.com/cloudflare/circl v1.6.3/go.mod h1:2eXP6Qfat4O/Yhh8BznvKnJ+uzEoTQ6jVKJRn81BiS4= +github.com/cncf/xds/go v0.0.0-20260202195803-dba9d589def2 h1:aBangftG7EVZoUb69Os8IaYg++6uMOdKK83QtkkvJik= +github.com/cncf/xds/go v0.0.0-20260202195803-dba9d589def2/go.mod h1:qwXFYgsP6T7XnJtbKlf1HP8AjxZZyzxMmc+Lq5GjlU4= github.com/codefly-dev/core v0.19.4 h1:vUvaFrWkB8B3ksQY/qlgBTVj3jMNhdEpJca8aofaals= github.com/codefly-dev/core v0.19.4/go.mod h1:0uX/7lTBVaGV0RI2Mb3mpEr/PPdx/jZqXHym/U8FVDI= github.com/codefly-dev/gortk v0.2.0 h1:7bOlS5valYz2zil+fZctQNcPCYBcPj86abcw9N8h1hQ= @@ -48,7 +84,6 @@ github.com/containerd/errdefs/pkg v0.3.0 h1:9IKJ06FvyNlexW690DXuQNx2KA2cUJXx151X github.com/containerd/errdefs/pkg v0.3.0/go.mod h1:NJw6s9HwNuRhnjJhM7pylWwMyAkmCQvQ4GpJHEqRLVk= github.com/containerd/log v0.1.0 h1:TCJt7ioM2cr/tfR8GPbGf9/VRAX8D2B4PjzCpfX540I= github.com/containerd/log v0.1.0/go.mod h1:VRRf09a7mHDIRezVKTRCrOq78v577GXq3bSa3EhrzVo= -github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= github.com/cyphar/filepath-securejoin v0.7.0 h1:s0Y3ITPy6sQn5xt54DuYvTF8hu134ooYLUb58DX/HjE= github.com/cyphar/filepath-securejoin v0.7.0/go.mod h1:ymLGms/u3BYaviIiuKFnUx8EkQEZeK6cInNoAPJA3o4= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= @@ -66,6 +101,11 @@ github.com/docker/go-units v0.5.0 h1:69rxXcBk27SvSaaxTtLh/8llcHD8vYHT7WSdRZ/jvr4 github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk= github.com/emirpasic/gods v1.18.1 h1:FXtiHYKDGKCW2KzwZKx0iC0PQmdlorYgdFG9jPXJ1Bc= github.com/emirpasic/gods v1.18.1/go.mod h1:8tpGGwCnJ5H4r6BWwaV6OrWmMoPhUl5jm/FMNAnJvWQ= +github.com/envoyproxy/go-control-plane v0.14.0 h1:hbG2kr4RuFj222B6+7T83thSPqLjwBIfQawTkC++2HA= +github.com/envoyproxy/go-control-plane/envoy v1.37.0 h1:u3riX6BoYRfF4Dr7dwSOroNfdSbEPe9Yyl09/B6wBrQ= +github.com/envoyproxy/go-control-plane/envoy v1.37.0/go.mod h1:DReE9MMrmecPy+YvQOAOHNYMALuowAnbjjEMkkWOi6A= +github.com/envoyproxy/protoc-gen-validate v1.3.3 h1:MVQghNeW+LZcmXe7SY1V36Z+WFMDjpqGAGacLe2T0ds= +github.com/envoyproxy/protoc-gen-validate v1.3.3/go.mod h1:TsndJ/ngyIdQRhMcVVGDDHINPLWB7C82oDArY51KfB0= github.com/fatih/color v1.19.0 h1:Zp3PiM21/9Ld6FzSKyL5c/BULoe/ONr9KlbYVOfG8+w= github.com/fatih/color v1.19.0/go.mod h1:zNk67I0ZUT1bEGsSGyCZYZNrHuTkJJB+r6Q9VuMi0LE= github.com/felixge/httpsnoop v1.1.0 h1:3YtUj32ZZkqZtt3sZZsClsymw/QDuVfpNhoA31zeORc= @@ -80,6 +120,8 @@ github.com/go-git/go-billy/v5 v5.9.2 h1:OXFSRyz4g20upsGDJgQG9Bak1l/ZEv8GHVYB52O7 github.com/go-git/go-billy/v5 v5.9.2/go.mod h1:ExsU+jcGwXTBOnyilvAnEM1wug1IxHr4yP2ZXsNRtV0= github.com/go-git/go-git/v5 v5.19.3 h1:qHttbZ+Am7wEjdTpR1XMQ4rl42FK9SIXzZ4o9x7s6M4= github.com/go-git/go-git/v5 v5.19.3/go.mod h1:Ye4C8dVigkqrv9UsB4NMdSLV4yAIkLiIhW61gcOyhl4= +github.com/go-jose/go-jose/v4 v4.1.4 h1:moDMcTHmvE6Groj34emNPLs/qtYXRVcd6S7NHbHz3kA= +github.com/go-jose/go-jose/v4 v4.1.4/go.mod h1:x4oUasVrzR7071A4TnHLGSPpNOm2a21K9Kf04k1rs08= github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= github.com/go-logr/logr v1.4.4 h1:tG4xh9yMsRCAiodLVTxyrkzSZ9+o0L1Kg/+cPVcbP/8= github.com/go-logr/logr v1.4.4/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= @@ -134,12 +176,16 @@ github.com/go-openapi/testify/v2 v2.8.0 h1:19QDx5b57p8KjO2E4sgs/woL4Akp+FGLRlVm2 github.com/go-openapi/testify/v2 v2.8.0/go.mod h1:4gUN8jC+RqOE2qjW/vDmfyTvl5BCht4qL7oca0odrak= github.com/go-viper/mapstructure/v2 v2.5.0 h1:vM5IJoUAy3d7zRSVtIwQgBj7BiWtMPfmPEgAXnvj1Ro= github.com/go-viper/mapstructure/v2 v2.5.0/go.mod h1:oJDH3BJKyqBA2TXFhDsKDGDTlndYOZ6rGS0BRZIxGhM= +github.com/goccy/go-json v0.10.6 h1:p8HrPJzOakx/mn/bQtjgNjdTcN+/S6FcG2CTtQOrHVU= +github.com/goccy/go-json v0.10.6/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M= github.com/gofrs/flock v0.13.1 h1:jjREztyBeSKBZYAC+mgc1laB+xsgy4kYMf3FbKF2UBo= github.com/gofrs/flock v0.13.1/go.mod h1:sf4BFiHwnvgxa25DlQoDqXQnwRMEOwqxRq37P6MzzmE= github.com/golang/groupcache v0.0.0-20241129210726-2c02b8208cf8 h1:f+oWsMOmNPc8JmEHVZIycC7hBoQxHH9pNKQORJNozsQ= github.com/golang/groupcache v0.0.0-20241129210726-2c02b8208cf8/go.mod h1:wcDNUvekVysuuOpQKo3191zZyTpiI6se1N1ULghS0sw= github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= +github.com/google/flatbuffers v25.12.19+incompatible h1:haMV2JRRJCe1998HeW/p0X9UaMTK6SDo0ffLn2+DbLs= +github.com/google/flatbuffers v25.12.19+incompatible/go.mod h1:1AeVuKshWv4vARoZatz6mlQ0JxURH0Kv5+zNeJKJCa8= github.com/google/gnostic-models v0.7.1 h1:SisTfuFKJSKM5CPZkffwi6coztzzeYUhc3v4yxLWH8c= github.com/google/gnostic-models v0.7.1/go.mod h1:whL5G0m6dmc5cPxKc5bdKdEN3UjI7OUGxBlw57miDrQ= github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= @@ -149,8 +195,16 @@ github.com/google/go-github/v89 v89.0.0 h1:35bEK5XoEcF3PZrlVbl9XN63f5BcJRA/UGkxe github.com/google/go-github/v89 v89.0.0/go.mod h1:QLcbU0ipeAqQuR5KSg8c2lql4Qk1EwJ2dWz/0rP4Nho= github.com/google/go-querystring v1.2.0 h1:yhqkPbu2/OH+V9BfpCVPZkNmUXhb2gBxJArfhIxNtP0= github.com/google/go-querystring v1.2.0/go.mod h1:8IFJqpSRITyJ8QhQ13bmbeMBDfmeEJZD5A0egEOmkqU= +github.com/google/martian/v3 v3.3.3 h1:DIhPTQrbPkgs2yJYdXU/eNACCG5DVQjySNRNlflZ9Fc= +github.com/google/martian/v3 v3.3.3/go.mod h1:iEPrYcgCF7jA9OtScMFQyAlZZ4YXTKEtJ1E6RWzmBA0= +github.com/google/s2a-go v0.1.9 h1:LGD7gtMgezd8a/Xak7mEWL0PjoTQFvpRudN895yqKW0= +github.com/google/s2a-go v0.1.9/go.mod h1:YA0Ei2ZQL3acow2O62kdp9UlnvMmU7kA6Eutn0dXayM= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/googleapis/enterprise-certificate-proxy v0.3.17 h1:73NfMHdiqo9JFU9+7a5ExpVa10/R29pXfZIaW559nrg= +github.com/googleapis/enterprise-certificate-proxy v0.3.17/go.mod h1:rSEsBUemEBZEexP2y6jPp16LUmUbjmSbcPMQizR0o4k= +github.com/googleapis/gax-go/v2 v2.23.0 h1:Tchl7qkvE7Ip3y+ztvNufYFvkfqTe7NfLTYGIdJRLuE= +github.com/googleapis/gax-go/v2 v2.23.0/go.mod h1:rBQKOVJCdb8IFEzg+FCwlt1LP/xMDGuqUXhUG+XMXEg= github.com/grpc-ecosystem/grpc-gateway/v2 v2.31.0 h1:Bd7KaOxzULLxtZ/K5s1aLbWhR0+5RToO65TXHsf3bqQ= github.com/grpc-ecosystem/grpc-gateway/v2 v2.31.0/go.mod h1:nN7ts3dFXKtCZWc//yfkpcQNKJABg16/uDVAZpLDalo= github.com/hashicorp/errwrap v1.0.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4= @@ -162,6 +216,10 @@ github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99 h1:BQSFePA1RWJOl github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99/go.mod h1:1lJo3i6rXxKeerYnT8Nvf0QmHCRC1n8sfWVwXF2Frvo= github.com/kevinburke/ssh_config v1.6.0 h1:J1FBfmuVosPHf5GRdltRLhPJtJpTlMdKTBjRgTaQBFY= github.com/kevinburke/ssh_config v1.6.0/go.mod h1:q2RIzfka+BXARoNexmF9gkxEX7DmvbW9P4hIVx2Kg4M= +github.com/klauspost/compress v1.19.2 h1:hMRETovs/pu/dVWN7zIT1PGG8t509MwT6bO7XSi26R8= +github.com/klauspost/compress v1.19.2/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= +github.com/klauspost/cpuid/v2 v2.4.0 h1:S6Hrbc7+ywsr0r+RLapfGBHfyefhCTwEh3A0tV913Dw= +github.com/klauspost/cpuid/v2 v2.4.0/go.mod h1:19jmZ9mjzoF//ddRSUsv0zfBTJWh3QJh9FNxZTMrGxU= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= @@ -197,10 +255,14 @@ github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3I github.com/opencontainers/image-spec v1.1.1 h1:y0fUlFfIZhPF1W537XOLg0/fcx6zcHCJwooC2xJA040= github.com/opencontainers/image-spec v1.1.1/go.mod h1:qpqAh3Dmcf36wStyyWU+kCeDgrGnAve2nCC8+7h8Q0M= github.com/pborman/getopt v0.0.0-20170112200414-7148bc3a4c30/go.mod h1:85jBQOZwpVEaDAr341tbn15RS4fCAsIst0qp7i8ex1o= +github.com/pierrec/lz4/v4 v4.1.29 h1:CDQY6qZOLI4DW0Nx6R1vRrifrCeQHnNXkMb0hZWXFjg= +github.com/pierrec/lz4/v4 v4.1.29/go.mod h1:EoQMVJgeeEOMsCqCzqFm2O0cJvljX2nGZjcRIPL34O4= github.com/pjbgf/sha1cd v0.7.0 h1:ZRNPKHj+gfkLBf0KJv/p1Hmz+7bqCT5o311YX0nq+DA= github.com/pjbgf/sha1cd v0.7.0/go.mod h1:pKR5Li+qTCo+ebqITZfwPlJGoCFCvSO00qxjbNtTUFQ= github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 h1:GFCKgmp0tecUJ0sJuv4pzYCqS9+RGSn52M3FUwPs+uo= +github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10/go.mod h1:t/avpk3KcrXxUnYOhZhMXJlSEyie6gQbtLq5NM3loB8= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/power-devops/perfstat v0.0.0-20240221224432-82ca36839d55 h1:o4JXh1EVt9k/+g42oCprj/FisM4qX9L3sZB3upGN2ZU= github.com/power-devops/perfstat v0.0.0-20240221224432-82ca36839d55/go.mod h1:OmDBASR4679mdNQnz2pUhc2G8CO2JrUAVFDRBDP/hJE= @@ -222,6 +284,8 @@ github.com/sirupsen/logrus v1.9.4 h1:TsZE7l11zFCLZnZ+teH4Umoq5BhEIfIzfRDZ1Uzql2w github.com/sirupsen/logrus v1.9.4/go.mod h1:ftWc9WdOfJ0a92nsE2jF5u5ZwH8Bv2zdeOC42RjbV2g= github.com/skeema/knownhosts v1.3.2 h1:EDL9mgf4NzwMXCTfaxSD/o/a5fxDw/xL9nkU28JjdBg= github.com/skeema/knownhosts v1.3.2/go.mod h1:bEg3iQAuw+jyiw+484wwFJoKSLwcfd7fqRy+N0QTiow= +github.com/spiffe/go-spiffe/v2 v2.8.1 h1:eXZMLsu+3MLEPJyGJkolqtVrteZfQdUpOWj6LTiDl/E= +github.com/spiffe/go-spiffe/v2 v2.8.1/go.mod h1:47Q0Q9/AqGha8QLHp+kxpH4Wca7X7EnOtlIJy3mxZ3U= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.5.3 h1:jmXUvGomnU1o3W/V5h2VEradbpJDwGrzugQQvL0POH4= github.com/stretchr/objx v0.5.3/go.mod h1:rDQraq+vQZU7Fde9LOZLr8Tax6zZvy4kuNKF+QYS+U0= @@ -240,8 +304,14 @@ github.com/yoheimuta/go-protoparser/v4 v4.14.2 h1:/P/LlX1CF9NaTWEltGcIZVvNlPbhAB github.com/yoheimuta/go-protoparser/v4 v4.14.2/go.mod h1:AHNNnSWnb0UoL4QgHPiOAg2BniQceFscPI5X/BZNHl8= github.com/yusufpapurcu/wmi v1.2.4 h1:zFUKzehAFReQwLys1b/iSMl+JQGSCSjtVqQn9bBrPo0= github.com/yusufpapurcu/wmi v1.2.4/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0= +github.com/zeebo/assert v1.3.0 h1:g7C04CbJuIDKNPFHmsk4hwZDO5O+kntRxzaUoNXj+IQ= +github.com/zeebo/assert v1.3.0/go.mod h1:Pq9JiuJQpG8JLJdtkwrJESF0Foym2/D9XMU5ciN/wJ0= +github.com/zeebo/xxh3 v1.1.0 h1:s7DLGDK45Dyfg7++yxI0khrfwq9661w9EN78eP/UZVs= +github.com/zeebo/xxh3 v1.1.0/go.mod h1:IisAie1LELR4xhVinxWS5+zf1lA4p0MW4T+w+W07F5s= go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= +go.opentelemetry.io/contrib/detectors/gcp v1.44.0 h1:NmLfL734pJhM0JKaYd2Y28+nY9dPRWYAAbxhRCrKXPw= +go.opentelemetry.io/contrib/detectors/gcp v1.44.0/go.mod h1:tNAsgd8avTGke1+MndXlU5Cru4PQ9Ai/cCNWQv/ZJ/s= go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.72.0 h1:Tq+E65HaNbZTeqcLw6QUhsV+/EIWdkfUH2CQzcrlE5M= go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.72.0/go.mod h1:A+EVhDakAj1waJBPhmTszZaSr08ww8MRRKy1nB1rqck= go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.71.0 h1:3g7B90UzBltIDKq1/5mrTGxTnOFDV0ICOhLoxiZ8jlg= @@ -278,19 +348,33 @@ golang.org/x/mod v0.41.0 h1:qJmnOUb4YB+FsEuM3HcWucdZASCPGhsX6uljO6pog0c= golang.org/x/mod v0.41.0/go.mod h1:Ek9pY8RKWXwsWvd3rQiHYtMqkjSUV+s1Rj7j4H5Ur6o= golang.org/x/net v0.60.0 h1:79p50tfZlm0J9YfoDsSi639qSXNGVwEzOPLCxM2FsYU= golang.org/x/net v0.60.0/go.mod h1:2DA/G1UfVbCpQPeWTmMPGY7Cs2PkBkwu743bVX5PIVg= +golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs= +golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q= +golang.org/x/sync v0.23.0 h1:KameEIfc1IkluZyXWLn39Wd4tURc6GbCiISGiZm2bQk= +golang.org/x/sync v0.23.0/go.mod h1:sUUOizhqBxiL6pEWpqNLUiaJn1ShEbZ6BBqskPbjZm0= golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20201204225414-ed752295db88/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.49.0 h1:XbzkgYJHdqh/8m2Uu0W/dQv8nktxx4BFHp1M0gROTXA= golang.org/x/sys v0.49.0/go.mod h1:hNLxWAXmnKAxqDtdwIYC4bM9oQPEecfsnNMuSxOs3og= +golang.org/x/telemetry v0.0.0-20260908163034-4bcc4b2ee518 h1:F5BWKvW126NXR74uxkxuc1jQHhm/rwm/J3rSiFyuRs4= +golang.org/x/telemetry v0.0.0-20260908163034-4bcc4b2ee518/go.mod h1:i+ivNqjDnTF3WTElsdk5g9V5DTSBYgdNo7xTU9SDwYA= golang.org/x/term v0.46.0 h1:3+OXuTbaKDgwk8jTi3aSLHRlmWqHEUDUtxnbFigO4YE= golang.org/x/term v0.46.0/go.mod h1:+K02xbkittuwc0Am4abfA3Fc+XRGXkvBXNO88NCXPoc= golang.org/x/text v0.42.0 h1:JbOZXgfeCPU9gacVtYliJqOhD+zhrEqK4LfdpmlUZqI= golang.org/x/text v0.42.0/go.mod h1:ojzP1Z+2QtioaF8DTtO8K5q7JWVVYwZKenzujK0Zd0E= golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U= golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno= +golang.org/x/tools v0.50.0 h1:c2ifzfcuY7L90lZ2aKd8S4K2NpASF08SZx9ZuJkHmSU= +golang.org/x/tools v0.50.0/go.mod h1:7ulVMw3831Mwi5EZD6RomGyffr4VFjuNYXf2BbCEAV0= +golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da h1:noIWHXmPHxILtqtCOPIhSt0ABwskkZKjD3bXGnZGpNY= +golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da/go.mod h1:NDW/Ps6MPRej6fsCIbMTohpP40sJ/P/vI1MoTEGwX90= gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E= +google.golang.org/api v0.287.1 h1:LiyJx32VU3cwQfLchn/513qKhc25hq0pEANYJoWNnnI= +google.golang.org/api v0.287.1/go.mod h1:lM2kYRzYUCBY91P9h6VF1PYmvhxii3O5hji37qRvIcY= +google.golang.org/genproto v0.0.0-20260319201613-d00831a3d3e7 h1:XzmzkmB14QhVhgnawEVsOn6OFsnpyxNPRY9QV01dNB0= +google.golang.org/genproto v0.0.0-20260319201613-d00831a3d3e7/go.mod h1:L43LFes82YgSonw6iTXTxXUX1OlULt4AQtkik4ULL/I= google.golang.org/genproto/googleapis/api v0.0.0-20260928230214-8a89bd6388cc h1:rZHRz0ogq4Pid1IKtA/ivHi8XONVF5tqUQXA2TtgpkE= google.golang.org/genproto/googleapis/api v0.0.0-20260928230214-8a89bd6388cc/go.mod h1:K/n0XLazCJHjD73YuMyNJsN38BPPy1kfFnKxUXaThFU= google.golang.org/genproto/googleapis/rpc v0.0.0-20260928230214-8a89bd6388cc h1:4bNTbnb44EqGVy9HaQxSY2jnafSUdgV3qHtedvNpDKg= diff --git a/internal/backend/bigquery/arrow.go b/internal/backend/bigquery/arrow.go new file mode 100644 index 0000000..78b84e3 --- /dev/null +++ b/internal/backend/bigquery/arrow.go @@ -0,0 +1,449 @@ +package bigquery + +import ( + "bytes" + "errors" + "fmt" + "io" + "math/big" + "time" + + bq "cloud.google.com/go/bigquery" + "cloud.google.com/go/civil" + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/decimal128" + "github.com/apache/arrow-go/v18/arrow/decimal256" + "github.com/apache/arrow-go/v18/arrow/ipc" + "github.com/apache/arrow-go/v18/arrow/memory" + + "github.com/codefly-dev/service-warehouse/internal/backend" + "github.com/codefly-dev/service-warehouse/internal/serr" +) + +// This file is the Arrow side of the backend, in both directions: BigQuery +// result values into Arrow-IPC batches for Query, and Arrow-IPC batches into +// the values a streaming insert sends for InsertRows. +// +// Framing, which is the contract a client reads against: the header carries one +// encapsulated IPC Schema message, and every batch is one encapsulated IPC +// RecordBatch message without the schema. The schema followed by the batches is +// a valid Arrow IPC stream (an end-of-stream marker is optional), which is how +// BigQuery's own Storage Read API frames Arrow too. + +const ( + // batchMaxRows and batchTargetBytes bound one result batch, so a wide table + // does not become one message a client cannot receive; the byte bound + // counts the text and binary a batch holds, which is what dominates. + batchMaxRows = 8192 + batchTargetBytes = 1 << 20 +) + +// eos is the 8-byte end-of-stream marker an IPC stream writer appends: the +// continuation indicator followed by a zero metadata length. +var eos = []byte{0xff, 0xff, 0xff, 0xff, 0, 0, 0, 0} + +// framer writes the schema message once and each record batch as a message of +// its own, by running one IPC stream writer and cutting what it emits. +type framer struct { + schema *arrow.Schema + schemaMsg []byte + buf bytes.Buffer + writer *ipc.Writer + primed bool +} + +func newFramer(schema *arrow.Schema, mem memory.Allocator) (*framer, error) { + // A stream with no batches is the schema message and the end marker. + var only bytes.Buffer + empty := ipc.NewWriter(&only, ipc.WithSchema(schema), ipc.WithAllocator(mem)) + if err := empty.Close(); err != nil { + return nil, serr.Wrap(serr.Internal, "arrow", fmt.Errorf("encode the schema: %w", err)) + } + encoded := only.Bytes() + if !bytes.HasSuffix(encoded, eos) { + return nil, serr.New(serr.Internal, "arrow", "the IPC writer's stream does not end with the end-of-stream marker") + } + f := &framer{schema: schema, schemaMsg: append([]byte(nil), encoded[:len(encoded)-len(eos)]...)} + f.writer = ipc.NewWriter(&f.buf, ipc.WithSchema(schema), ipc.WithAllocator(mem)) + return f, nil +} + +// batch returns rec as one RecordBatch message. +func (f *framer) batch(rec arrow.RecordBatch) ([]byte, error) { + f.buf.Reset() + if err := f.writer.Write(rec); err != nil { + return nil, serr.Wrap(serr.Internal, "arrow", fmt.Errorf("encode a batch: %w", err)) + } + out := f.buf.Bytes() + if !f.primed { + // The first write also carries the schema message, which the header + // already sent. + if !bytes.HasPrefix(out, f.schemaMsg) { + return nil, serr.New(serr.Internal, "arrow", "the IPC writer's first batch does not start with the schema message") + } + out = out[len(f.schemaMsg):] + f.primed = true + } + return append([]byte(nil), out...), nil +} + +// resultEncoder turns BigQuery rows into Arrow batches. +type resultEncoder struct { + cols []backend.Column + mem memory.Allocator + builder *array.RecordBuilder + framer *framer + rows int + size int +} + +func newResultEncoder(cols []backend.Column) (*resultEncoder, error) { + schema, err := arrowSchema(cols) + if err != nil { + return nil, err + } + mem := memory.NewGoAllocator() + fr, err := newFramer(schema, mem) + if err != nil { + return nil, err + } + return &resultEncoder{ + cols: cols, + mem: mem, + builder: array.NewRecordBuilder(mem, schema), + framer: fr, + }, nil +} + +// schemaMessage is the IPC Schema message for the result header. +func (e *resultEncoder) schemaMessage() []byte { return e.framer.schemaMsg } + +// add appends one row, which must carry one value per column. +func (e *resultEncoder) add(row []bq.Value) error { + if len(row) != len(e.cols) { + return serr.New(serr.Internal, "Query", fmt.Sprintf("a result row has %d values for %d columns", len(row), len(e.cols))) + } + for i, c := range e.cols { + if err := appendValue(e.builder.Field(i), c, row[i]); err != nil { + return err + } + e.size += valueSize(row[i]) + } + e.rows++ + return nil +} + +func (e *resultEncoder) pending() bool { return e.rows > 0 } + +func (e *resultEncoder) full() bool { return e.rows >= batchMaxRows || e.size >= batchTargetBytes } + +// flush returns the rows added since the last flush as one batch. +func (e *resultEncoder) flush() ([]byte, error) { + rec := e.builder.NewRecordBatch() + defer rec.Release() + e.rows, e.size = 0, 0 + return e.framer.batch(rec) +} + +func (e *resultEncoder) release() { e.builder.Release() } + +// valueSize is a cheap estimate of a value's weight, enough to bound a batch. +func valueSize(v bq.Value) int { + switch x := v.(type) { + case string: + return len(x) + case []byte: + return len(x) + case []bq.Value: + n := 8 + for _, e := range x { + n += valueSize(e) + } + return n + default: + return 8 + } +} + +// mismatch is a BigQuery value of a Go type the column's type does not produce. +func mismatch(c backend.Column, v bq.Value) error { + return serr.New(serr.Internal, "Query", fmt.Sprintf("column %q (%s) produced a value of type %T", c.Name, c.NativeType, v)) +} + +// appendValue adds one BigQuery value to the Arrow builder of its column. +func appendValue(b array.Builder, c backend.Column, v bq.Value) error { + if v == nil { + if c.Type == backend.TypeArray && !c.Nullable { + // BigQuery has no NULL array; an absent one is an empty array. + b.(*array.ListBuilder).Append(true) + return nil + } + b.AppendNull() + return nil + } + switch c.Type { + case backend.TypeBool: + x, ok := v.(bool) + if !ok { + return mismatch(c, v) + } + b.(*array.BooleanBuilder).Append(x) + case backend.TypeInt64: + x, ok := v.(int64) + if !ok { + return mismatch(c, v) + } + b.(*array.Int64Builder).Append(x) + case backend.TypeFloat64: + x, ok := v.(float64) + if !ok { + return mismatch(c, v) + } + b.(*array.Float64Builder).Append(x) + case backend.TypeNumeric: + x, ok := v.(*big.Rat) + if !ok { + return mismatch(c, v) + } + return appendDecimal(b, c, x) + case backend.TypeString, backend.TypeJSON, backend.TypeGeography: + x, ok := v.(string) + if !ok { + return mismatch(c, v) + } + b.(*array.StringBuilder).Append(x) + case backend.TypeBytes: + x, ok := v.([]byte) + if !ok { + return mismatch(c, v) + } + b.(*array.BinaryBuilder).Append(x) + case backend.TypeDate: + x, ok := v.(civil.Date) + if !ok { + return mismatch(c, v) + } + b.(*array.Date32Builder).Append(arrow.Date32(x.DaysSince(civil.Date{Year: 1970, Month: time.January, Day: 1}))) + case backend.TypeTime: + x, ok := v.(civil.Time) + if !ok { + return mismatch(c, v) + } + micros := ((int64(x.Hour)*60+int64(x.Minute))*60+int64(x.Second))*1_000_000 + int64(x.Nanosecond)/1000 + b.(*array.Time64Builder).Append(arrow.Time64(micros)) + case backend.TypeTimestamp: + x, ok := v.(civil.DateTime) + if !ok { + return mismatch(c, v) + } + b.(*array.TimestampBuilder).Append(arrow.Timestamp(x.In(time.UTC).UnixMicro())) + case backend.TypeTimestampTZ: + x, ok := v.(time.Time) + if !ok { + return mismatch(c, v) + } + b.(*array.TimestampBuilder).Append(arrow.Timestamp(x.UnixMicro())) + case backend.TypeInterval: + x, ok := v.(*bq.IntervalValue) + if !ok { + return mismatch(c, v) + } + b.(*array.MonthDayNanoIntervalBuilder).Append(arrow.MonthDayNanoInterval{ + Months: x.Years*12 + x.Months, + Days: x.Days, + Nanoseconds: ((int64(x.Hours)*60+int64(x.Minutes))*60+int64(x.Seconds))*1_000_000_000 + int64(x.SubSecondNanos), + }) + case backend.TypeArray: + x, ok := v.([]bq.Value) + if !ok { + return mismatch(c, v) + } + lb := b.(*array.ListBuilder) + lb.Append(true) + for _, item := range x { + if err := appendValue(lb.ValueBuilder(), c.Fields[0], item); err != nil { + return err + } + } + case backend.TypeStruct: + x, ok := v.([]bq.Value) + if !ok || len(x) != len(c.Fields) { + return mismatch(c, v) + } + sb := b.(*array.StructBuilder) + sb.Append(true) + for i, sub := range c.Fields { + if err := appendValue(sb.FieldBuilder(i), sub, x[i]); err != nil { + return err + } + } + default: + return serr.New(serr.Unsupported, "Query", fmt.Sprintf("column %q has no portable type", c.Name)) + } + return nil +} + +// unscaled is r as an integer count of 10^-scale, which BigQuery guarantees for a +// value of its own column. A value that is not one means the column's declared +// scale and its data disagree, so it is an error and never a rounded number. +func unscaled(r *big.Rat, scale int32) (*big.Int, error) { + if scale < 0 { + return nil, fmt.Errorf("negative decimal scale %d", scale) + } + factor := new(big.Int).Exp(big.NewInt(10), big.NewInt(int64(scale)), nil) + scaled := new(big.Int).Mul(r.Num(), factor) + quotient, remainder := new(big.Int).QuoRem(scaled, r.Denom(), new(big.Int)) + if remainder.Sign() != 0 { + return nil, fmt.Errorf("%s has more than %d fractional digits", r.FloatString(40), scale) + } + return quotient, nil +} + +func appendDecimal(b array.Builder, c backend.Column, r *big.Rat) error { + switch t := b.Type().(type) { + case *arrow.Decimal128Type: + n, err := unscaled(r, t.Scale) + if err != nil { + return serr.Wrap(serr.Internal, "Query", fmt.Errorf("column %q: %w", c.Name, err)) + } + if n.BitLen() > 126 { + return serr.New(serr.Internal, "Query", fmt.Sprintf("column %q: a value does not fit %s", c.Name, t)) + } + num := decimal128.FromBigInt(n) + if !num.FitsInPrecision(t.Precision) { + return serr.New(serr.Internal, "Query", fmt.Sprintf("column %q: a value does not fit %s", c.Name, t)) + } + b.(*array.Decimal128Builder).Append(num) + case *arrow.Decimal256Type: + n, err := unscaled(r, t.Scale) + if err != nil { + return serr.Wrap(serr.Internal, "Query", fmt.Errorf("column %q: %w", c.Name, err)) + } + if n.BitLen() > 254 { + return serr.New(serr.Internal, "Query", fmt.Sprintf("column %q: a value does not fit %s", c.Name, t)) + } + num := decimal256.FromBigInt(n) + if !num.FitsInPrecision(t.Precision) { + return serr.New(serr.Internal, "Query", fmt.Sprintf("column %q: a value does not fit %s", c.Name, t)) + } + b.(*array.Decimal256Builder).Append(num) + default: + return mismatch(c, r) + } + return nil +} + +// batchStream reads the Arrow batches of an InsertRows call. The header's schema +// message, when there is one, completes each batch; a batch that carries its own +// schema is read as the stream it is. +type batchStream struct { + mem memory.Allocator + header []byte // the header's Schema message with any end marker removed; nil when absent + src backend.BatchReader + + schema *arrow.Schema // the schema every record must have, known after the first + cur *ipc.Reader +} + +func newBatchStream(header []byte, src backend.BatchReader) (*batchStream, error) { + s := &batchStream{mem: memory.NewGoAllocator(), src: src} + if len(header) == 0 { + return s, nil + } + s.header = bytes.TrimSuffix(header, eos) + schema, err := s.readSchema(bytes.NewReader(s.header)) + if err != nil { + return nil, err + } + s.schema = schema + return s, nil +} + +func invalidBatch(err error) error { + return serr.Wrap(serr.InvalidArgument, "InsertRows", fmt.Errorf("the Arrow data is malformed: %w", err)) +} + +// guard turns a panic inside the Arrow reader, which is fed bytes from the +// network, into an invalid-argument error. +func guard(err *error) { + if r := recover(); r != nil { + *err = invalidBatch(fmt.Errorf("%v", r)) + } +} + +func (s *batchStream) readSchema(r io.Reader) (schema *arrow.Schema, err error) { + defer guard(&err) + reader, err := ipc.NewReader(r, ipc.WithAllocator(s.mem)) + if err != nil { + return nil, invalidBatch(err) + } + defer reader.Release() + return reader.Schema(), nil +} + +// carriesSchema reports whether an IPC payload begins with a Schema message. +func carriesSchema(payload []byte) bool { + messages := ipc.NewMessageReader(bytes.NewReader(payload)) + defer messages.Release() + message, err := messages.Message() + if err != nil { + return false + } + defer message.Release() + return message.Type() == ipc.MessageSchema +} + +// next returns the next record batch, or io.EOF when the stream is exhausted. The +// caller releases the record. +func (s *batchStream) next() (rec arrow.RecordBatch, err error) { + defer guard(&err) + for { + if s.cur != nil { + if s.cur.Next() { + rec := s.cur.RecordBatch() + if s.schema == nil { + s.schema = rec.Schema() + } else if !s.schema.Equal(rec.Schema()) { + return nil, serr.New(serr.InvalidArgument, "InsertRows", "a batch's schema differs from the schema of the batches before it") + } + rec.Retain() + return rec, nil + } + err := s.cur.Err() + s.cur.Release() + s.cur = nil + if err != nil { + return nil, invalidBatch(err) + } + } + payload, err := s.src.Next() + if errors.Is(err, io.EOF) { + return nil, io.EOF + } + if err != nil { + return nil, err + } + var stream io.Reader + switch { + case carriesSchema(payload): + stream = bytes.NewReader(payload) + case s.header != nil: + stream = io.MultiReader(bytes.NewReader(s.header), bytes.NewReader(payload)) + default: + return nil, serr.New(serr.InvalidArgument, "InsertRows", "a batch carries no schema and the header gave none") + } + reader, err := ipc.NewReader(stream, ipc.WithAllocator(s.mem)) + if err != nil { + return nil, invalidBatch(err) + } + s.cur = reader + } +} + +func (s *batchStream) close() { + if s.cur != nil { + s.cur.Release() + s.cur = nil + } +} diff --git a/internal/backend/bigquery/arrow_test.go b/internal/backend/bigquery/arrow_test.go new file mode 100644 index 0000000..143f0c6 --- /dev/null +++ b/internal/backend/bigquery/arrow_test.go @@ -0,0 +1,339 @@ +package bigquery + +import ( + "bytes" + "errors" + "io" + "math" + "math/big" + "testing" + "time" + + bq "cloud.google.com/go/bigquery" + "cloud.google.com/go/civil" + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/decimal128" + "github.com/apache/arrow-go/v18/arrow/ipc" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/stretchr/testify/require" + + "github.com/codefly-dev/service-warehouse/internal/serr" +) + +func rat(s string) *big.Rat { + r, ok := new(big.Rat).SetString(s) + if !ok { + panic("bad rat " + s) + } + return r +} + +// everyType is a BigQuery schema with one column of each type this backend maps, +// and the values a result of it holds: a full row, and a row of NULLs. +func everyType() (bq.Schema, [][]bq.Value, [][]any) { + schema := bq.Schema{ + {Name: "b", Type: bq.BooleanFieldType}, + {Name: "i", Type: bq.IntegerFieldType}, + {Name: "f", Type: bq.FloatFieldType}, + {Name: "n", Type: bq.NumericFieldType}, + {Name: "bn", Type: bq.BigNumericFieldType}, + {Name: "s", Type: bq.StringFieldType}, + {Name: "y", Type: bq.BytesFieldType}, + {Name: "d", Type: bq.DateFieldType}, + {Name: "tm", Type: bq.TimeFieldType}, + {Name: "dt", Type: bq.DateTimeFieldType}, + {Name: "ts", Type: bq.TimestampFieldType}, + {Name: "iv", Type: bq.IntervalFieldType}, + {Name: "j", Type: bq.JSONFieldType}, + {Name: "g", Type: bq.GeographyFieldType}, + {Name: "arr", Type: bq.IntegerFieldType, Repeated: true}, + {Name: "st", Type: bq.RecordFieldType, Schema: bq.Schema{ + {Name: "a", Type: bq.IntegerFieldType}, {Name: "z", Type: bq.StringFieldType}, + }}, + {Name: "sa", Type: bq.RecordFieldType, Repeated: true, Schema: bq.Schema{{Name: "x", Type: bq.IntegerFieldType}}}, + } + values := [][]bq.Value{ + { + true, int64(-42), 1.5, rat("12.34"), rat("-98765432109876543210.12345678901234567890"), + "héllo", []byte{0, 1, 2, 255}, + civil.Date{Year: 2024, Month: time.February, Day: 29}, + civil.Time{Hour: 13, Minute: 4, Second: 5, Nanosecond: 123456000}, + civil.DateTime{Date: civil.Date{Year: 2024, Month: time.February, Day: 29}, Time: civil.Time{Hour: 13, Minute: 4, Second: 5, Nanosecond: 123456000}}, + time.Date(2024, time.February, 29, 13, 4, 5, 123456000, time.UTC), + &bq.IntervalValue{Years: 1, Months: 2, Days: 3, Hours: 4, Minutes: 5, Seconds: 6, SubSecondNanos: 789000000}, + `{"k":[1,2]}`, "POINT(1 2)", + []bq.Value{int64(1), int64(2), int64(3)}, + []bq.Value{int64(7), "seven"}, + []bq.Value{[]bq.Value{int64(1)}, []bq.Value{int64(2)}}, + }, + {nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, []bq.Value(nil), nil, []bq.Value(nil)}, + } + // What a client reading the batches back sees, as the JSON a streaming insert + // would send for the same cell: the reverse direction of the same mapping. + wantJSON := [][]any{ + { + true, int64(-42), 1.5, "12.340000000", "-98765432109876543210.12345678901234567890000000000000000000", + "héllo", "AAEC/w==", "2024-02-29", "13:04:05.123456", + "2024-02-29T13:04:05.123456", "2024-02-29T13:04:05.123456Z", + "1-2 3 4:5:6.789", `{"k":[1,2]}`, "POINT(1 2)", + []any{int64(1), int64(2), int64(3)}, + map[string]any{"a": int64(7), "z": "seven"}, + []any{map[string]any{"x": int64(1)}, map[string]any{"x": int64(2)}}, + }, + {nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, []any{}, nil, []any{}}, + } + return schema, values, wantJSON +} + +func readBack(t *testing.T, schemaMsg []byte, batch []byte) arrow.RecordBatch { + t.Helper() + reader, err := ipc.NewReader(io.MultiReader(bytes.NewReader(schemaMsg), bytes.NewReader(batch))) + require.NoError(t, err) + t.Cleanup(reader.Release) + require.True(t, reader.Next(), "the batch holds a record: %v", reader.Err()) + rec := reader.RecordBatch() + rec.Retain() + t.Cleanup(rec.Release) + return rec +} + +func TestResultEncoderRoundTripsEveryType(t *testing.T) { + schema, rows, want := everyType() + cols := columnsFromSchema(schema) + enc, err := newResultEncoder(cols) + require.NoError(t, err) + defer enc.release() + for _, row := range rows { + require.NoError(t, enc.add(row)) + } + batch, err := enc.flush() + require.NoError(t, err) + rec := readBack(t, enc.schemaMessage(), batch) + + require.EqualValues(t, len(rows), rec.NumRows()) + for c, col := range cols { + for r := range rows { + got, err := jsonValue(col.Name, rec.Column(c), r) + require.NoError(t, err, "column %s row %d", col.Name, r) + require.Equal(t, want[r][c], got, "column %s row %d", col.Name, r) + } + } +} + +func TestResultEncoderArrowTypes(t *testing.T) { + schema, _, _ := everyType() + cols := columnsFromSchema(schema) + got, err := arrowSchema(cols) + require.NoError(t, err) + want := map[string]arrow.DataType{ + "b": arrow.FixedWidthTypes.Boolean, "i": arrow.PrimitiveTypes.Int64, "f": arrow.PrimitiveTypes.Float64, + "n": &arrow.Decimal128Type{Precision: 38, Scale: 9}, + "bn": &arrow.Decimal256Type{Precision: 76, Scale: 38}, + "s": arrow.BinaryTypes.String, "y": arrow.BinaryTypes.Binary, + "d": arrow.FixedWidthTypes.Date32, "tm": arrow.FixedWidthTypes.Time64us, + "dt": &arrow.TimestampType{Unit: arrow.Microsecond}, + "ts": &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: "UTC"}, + "iv": arrow.FixedWidthTypes.MonthDayNanoInterval, + "j": arrow.BinaryTypes.String, "g": arrow.BinaryTypes.String, + } + for name, dt := range want { + idx := got.FieldIndices(name) + require.Len(t, idx, 1, name) + require.True(t, arrow.TypeEqual(dt, got.Field(idx[0]).Type), "%s: got %s want %s", name, got.Field(idx[0]).Type, dt) + } + arr := got.Field(got.FieldIndices("arr")[0]) + require.Equal(t, "list", arr.Type.String(), "an ARRAY is a list of elements that are never NULL") + require.False(t, arr.Nullable, "BigQuery has no NULL array") + st := got.Field(got.FieldIndices("st")[0]).Type.(*arrow.StructType) + require.Equal(t, []string{"a", "z"}, []string{st.Field(0).Name, st.Field(1).Name}) +} + +func TestNumericIsCarriedWithItsDeclaredPrecisionAndScale(t *testing.T) { + cols := columnsFromSchema(bq.Schema{{Name: "p", Type: bq.NumericFieldType, Precision: 10, Scale: 2}}) + schema, err := arrowSchema(cols) + require.NoError(t, err) + require.True(t, arrow.TypeEqual(&arrow.Decimal128Type{Precision: 10, Scale: 2}, schema.Field(0).Type)) + + enc, err := newResultEncoder(cols) + require.NoError(t, err) + defer enc.release() + require.NoError(t, enc.add([]bq.Value{rat("-123.45")})) + batch, err := enc.flush() + require.NoError(t, err) + rec := readBack(t, enc.schemaMessage(), batch) + require.Equal(t, decimal128.FromI64(-12345), rec.Column(0).(*array.Decimal128).Value(0)) +} + +func TestADecimalThatDoesNotFitItsColumnIsAnErrorNotARoundedNumber(t *testing.T) { + cols := columnsFromSchema(bq.Schema{{Name: "p", Type: bq.NumericFieldType, Precision: 10, Scale: 2}}) + enc, err := newResultEncoder(cols) + require.NoError(t, err) + defer enc.release() + require.Error(t, enc.add([]bq.Value{rat("1.005")}), "a third fractional digit would be rounded away") + require.Error(t, enc.add([]bq.Value{rat("123456789.01")}), "eleven digits do not fit precision 10") +} + +func TestResultEncoderRefusesAValueOfTheWrongType(t *testing.T) { + enc, err := newResultEncoder(columnsFromSchema(bq.Schema{{Name: "i", Type: bq.IntegerFieldType}})) + require.NoError(t, err) + defer enc.release() + err = enc.add([]bq.Value{"not an int"}) + require.True(t, serr.Is(err, serr.Internal)) + require.Error(t, enc.add([]bq.Value{}), "a row with too few values") +} + +func TestFramingIsTheSchemaMessageThenEachBatchMessage(t *testing.T) { + schema := arrow.NewSchema([]arrow.Field{{Name: "id", Type: arrow.PrimitiveTypes.Int64}}, nil) + mem := memory.NewGoAllocator() + fr, err := newFramer(schema, mem) + require.NoError(t, err) + require.True(t, carriesSchema(fr.schemaMsg)) + require.False(t, bytes.HasSuffix(fr.schemaMsg, eos), "the schema message is not terminated") + + build := func(vals ...int64) []byte { + b := array.NewRecordBuilder(mem, schema) + defer b.Release() + for _, v := range vals { + b.Field(0).(*array.Int64Builder).Append(v) + } + rec := b.NewRecordBatch() + defer rec.Release() + out, err := fr.batch(rec) + require.NoError(t, err) + return out + } + first, second := build(1, 2), build(3) + require.False(t, carriesSchema(first), "a batch does not repeat the schema") + require.False(t, carriesSchema(second)) + require.False(t, bytes.HasPrefix(second, fr.schemaMsg)) + + // The schema and the batches together are one valid Arrow IPC stream, which + // is what lets a client hand them to any Arrow reader. + stream := io.MultiReader(bytes.NewReader(fr.schemaMsg), bytes.NewReader(first), bytes.NewReader(second), bytes.NewReader(eos)) + reader, err := ipc.NewReader(stream) + require.NoError(t, err) + defer reader.Release() + var seen []int64 + for reader.Next() { + col := reader.RecordBatch().Column(0).(*array.Int64) + for i := 0; i < col.Len(); i++ { + seen = append(seen, col.Value(i)) + } + } + require.NoError(t, reader.Err()) + require.Equal(t, []int64{1, 2, 3}, seen) +} + +func TestBatchStreamReadsHeaderSchemaWithBatchesAndSelfDescribingStreams(t *testing.T) { + header, batch := eventBatch(t, []int64{1, 2}, []string{"a", "b"}) + + read := func(header []byte, payloads ...[]byte) (rows int64, err error) { + s, err := newBatchStream(header, &sliceReader{payloads: payloads}) + if err != nil { + return 0, err + } + defer s.close() + for { + rec, err := s.next() + if errors.Is(err, io.EOF) { + return rows, nil + } + if err != nil { + return rows, err + } + rows += rec.NumRows() + rec.Release() + } + } + + n, err := read(header, batch, batch) + require.NoError(t, err) + require.EqualValues(t, 4, n) + + n, err = read(nil, append(append([]byte{}, header...), batch...), batch[:0:0]) + require.Error(t, err, "an empty payload is not a batch") + _ = n + + n, err = read(nil, append(append(append([]byte{}, header...), batch...), eos...)) + require.NoError(t, err) + require.EqualValues(t, 2, n) + + n, err = read(header) + require.NoError(t, err) + require.Zero(t, n, "no batches is no rows") +} + +func TestJSONValuesOfTheCellsAStreamingInsertSends(t *testing.T) { + mem := memory.NewGoAllocator() + build := func(dt arrow.DataType, fill func(b array.Builder)) arrow.Array { + b := array.NewBuilder(mem, dt) + defer b.Release() + fill(b) + return b.NewArray() + } + cell := func(arr arrow.Array) (any, error) { + defer arr.Release() + return jsonValue("c", arr, 0) + } + + for name, tc := range map[string]struct { + arr arrow.Array + want any + }{ + "NaN is the string BigQuery reads as NaN": {build(arrow.PrimitiveTypes.Float64, func(b array.Builder) { b.(*array.Float64Builder).Append(math.NaN()) }), "NaN"}, + "+Inf": {build(arrow.PrimitiveTypes.Float64, func(b array.Builder) { b.(*array.Float64Builder).Append(math.Inf(1)) }), "Infinity"}, + "-Inf": {build(arrow.PrimitiveTypes.Float64, func(b array.Builder) { b.(*array.Float64Builder).Append(math.Inf(-1)) }), "-Infinity"}, + "float32": {build(arrow.PrimitiveTypes.Float32, func(b array.Builder) { b.(*array.Float32Builder).Append(0.5) }), 0.5}, + "small uint": {build(arrow.PrimitiveTypes.Uint32, func(b array.Builder) { b.(*array.Uint32Builder).Append(4000000000) }), int64(4000000000)}, + "uint64 in range": {build(arrow.PrimitiveTypes.Uint64, func(b array.Builder) { b.(*array.Uint64Builder).Append(math.MaxInt64) }), int64(math.MaxInt64)}, + "time in seconds": {build(&arrow.Time32Type{Unit: arrow.Second}, func(b array.Builder) { b.(*array.Time32Builder).Append(arrow.Time32(3661)) }), "01:01:01.000000"}, + "naive timestamp": {build(&arrow.TimestampType{Unit: arrow.Millisecond}, func(b array.Builder) { b.(*array.TimestampBuilder).Append(arrow.Timestamp(1700000000123)) }), "2023-11-14T22:13:20.123000"}, + "timestamp in ns": {build(&arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}, func(b array.Builder) { b.(*array.TimestampBuilder).Append(arrow.Timestamp(1700000000123456000)) }), "2023-11-14T22:13:20.123456Z"}, + "date64": {build(arrow.FixedWidthTypes.Date64, func(b array.Builder) { b.(*array.Date64Builder).Append(arrow.Date64(86400000 * 19000)) }), "2022-01-08"}, + "fixed size bytes": {build(&arrow.FixedSizeBinaryType{ByteWidth: 2}, func(b array.Builder) { b.(*array.FixedSizeBinaryBuilder).Append([]byte{1, 2}) }), "AQI="}, + "negative decimal": {build(&arrow.Decimal128Type{Precision: 10, Scale: 2}, func(b array.Builder) { b.(*array.Decimal128Builder).Append(decimal128.FromI64(-5)) }), "-0.05"}, + "interval": {build(arrow.FixedWidthTypes.MonthDayNanoInterval, func(b array.Builder) { + b.(*array.MonthDayNanoIntervalBuilder).Append(arrow.MonthDayNanoInterval{Months: 14, Days: 1, Nanoseconds: 3_661_500_000_000}) + }), "1-2 1 1:1:1.5"}, + } { + t.Run(name, func(t *testing.T) { + got, err := cell(tc.arr) + require.NoError(t, err) + require.Equal(t, tc.want, got) + }) + } + + for name, arr := range map[string]arrow.Array{ + "uint64 beyond int64": build(arrow.PrimitiveTypes.Uint64, func(b array.Builder) { b.(*array.Uint64Builder).Append(math.MaxUint64) }), + "text that is not UTF-8": build(arrow.BinaryTypes.String, func(b array.Builder) { b.(*array.StringBuilder).Append("a\xffb") }), + "sub-microsecond": build(&arrow.TimestampType{Unit: arrow.Nanosecond}, func(b array.Builder) { b.(*array.TimestampBuilder).Append(arrow.Timestamp(1_700_000_000_000_000_001)) }), + "time past a day": build(arrow.FixedWidthTypes.Time64us, func(b array.Builder) { b.(*array.Time64Builder).Append(arrow.Time64(25 * 3600 * 1_000_000)) }), + "negative time": build(arrow.FixedWidthTypes.Time64us, func(b array.Builder) { b.(*array.Time64Builder).Append(arrow.Time64(-1)) }), + } { + t.Run(name, func(t *testing.T) { + _, err := cell(arr) + ve, ok := asValueError(err) + require.True(t, ok, "a cell BigQuery cannot be given is a value error, which refuses its row: %v", err) + require.Equal(t, "c", ve.column) + }) + } +} + +func TestAnArrowTypeWithNoBigQueryEquivalentFailsTheWholeStream(t *testing.T) { + for name, dt := range map[string]arrow.DataType{ + "map": arrow.MapOf(arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int64), + "duration": &arrow.DurationType{Unit: arrow.Second}, + "dictionary": &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int8, ValueType: arrow.BinaryTypes.String}, + "nested": arrow.StructOf(arrow.Field{Name: "d", Type: &arrow.DurationType{Unit: arrow.Second}}), + "in a list": arrow.ListOf(&arrow.DurationType{Unit: arrow.Second}), + } { + t.Run(name, func(t *testing.T) { + err := checkArrowSchema(arrow.NewSchema([]arrow.Field{{Name: "c", Type: dt}}, nil)) + var typeErr *arrowTypeError + require.ErrorAs(t, err, &typeErr) + require.Equal(t, "c", typeErr.column[:1]) + }) + } + require.NoError(t, checkArrowSchema(eventSchema)) +} diff --git a/internal/backend/bigquery/bigquery.go b/internal/backend/bigquery/bigquery.go new file mode 100644 index 0000000..feae773 --- /dev/null +++ b/internal/backend/bigquery/bigquery.go @@ -0,0 +1,184 @@ +// Package bigquery is the BigQuery backend: the whole catalog and DDL plane, +// SQL queries with bound parameters and a dry run, job inspection and +// cancellation, and streaming inserts, all in Arrow. Bulk load and unload and +// the Native verbs are not offered; they return serr.Unsupported and +// Capabilities says so, so a client is never told something false. +// +// The server is bound to one project (SWH_DATABASE). Credentials are +// Application Default Credentials, or the service-account file named by +// SWH_CREDENTIALS_FILE; nothing else is read from the environment and no +// credential is ever written into a request this server makes on its own. +// +// Every BigQuery failure is mapped to a serr.Code before it leaves this package +// (errors.go); a client branches on the code and never sees a reason token, an +// HTTP status or a vendor message, except that an invalid request keeps +// BigQuery's diagnosis of the caller's own SQL. +package bigquery + +import ( + "context" + "fmt" + "regexp" + "strings" + + bq "cloud.google.com/go/bigquery" + "google.golang.org/api/option" + + "github.com/codefly-dev/service-warehouse/internal/backend" + "github.com/codefly-dev/service-warehouse/internal/serr" +) + +// Kind is the name this backend registers under, and the value of SWH_BACKEND +// that selects it. +const Kind = "bigquery" + +func init() { + backend.Register(Kind, backend.Registration{ + // A BigQuery server is bound to a project, which SWH_DATABASE names; it has + // no connection string. + SelfContained: false, + DSN: false, + Open: func(ctx context.Context, cfg backend.Config) (backend.Backend, error) { + b, err := Open(ctx, cfg) + if err != nil { + // Not `return b, err`: a nil *Backend in a backend.Backend is not nil. + return nil, err + } + return b, nil + }, + }) +} + +// Backend is the BigQuery warehouse backend. It is safe for concurrent use. +type Backend struct { + cfg backend.Config + project string + client *bq.Client +} + +var _ backend.Backend = (*Backend)(nil) + +// projectID is the shape of a project id, including the legacy +// domain-scoped form (example.com:project). It is checked because the id is put +// into request paths. +var projectID = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._:-]*$`) + +// Open builds the backend from cfg. Extra client options exist for tests, which +// point the client at a fake endpoint; the server passes none. +func Open(ctx context.Context, cfg backend.Config, opts ...option.ClientOption) (*Backend, error) { + const op = "bigquery.Open" + if cfg.Database == "" { + return nil, serr.New(serr.InvalidArgument, op, "SWH_DATABASE must name the BigQuery project this server is bound to") + } + if !projectID.MatchString(cfg.Database) { + return nil, serr.New(serr.InvalidArgument, op, fmt.Sprintf("SWH_DATABASE %q is not a valid BigQuery project id", cfg.Database)) + } + // A setting this backend never reads must not be accepted: silently ignoring + // SWH_USER would let an operator believe a login was configured. + if unused := unusedSettings(cfg); len(unused) > 0 { + return nil, serr.New(serr.InvalidArgument, op, + fmt.Sprintf("%s is not used by the bigquery backend; remove it (credentials are Application Default Credentials or SWH_CREDENTIALS_FILE)", + strings.Join(unused, ", "))) + } + if cfg.DefaultDataset != "" { + if err := checkDatasetID(op, cfg.DefaultDataset); err != nil { + return nil, err + } + } + if cfg.Location != "" && !locationID.MatchString(cfg.Location) { + return nil, serr.New(serr.InvalidArgument, op, fmt.Sprintf("SWH_LOCATION %q is not a BigQuery location", cfg.Location)) + } + if cfg.CredentialsFile != "" { + // Typed as a service account, so the file cannot silently be some other + // kind of credential configuration. + opts = append(opts, option.WithAuthCredentialsFile(option.ServiceAccount, cfg.CredentialsFile)) + } + client, err := bq.NewClient(ctx, cfg.Database, opts...) + if err != nil { + // Startup failures are the operator's to read, so the cause is kept. + return nil, serr.Wrap(serr.Internal, op, fmt.Errorf("create the BigQuery client: %w", err)) + } + client.Location = cfg.Location + return &Backend{cfg: cfg, project: cfg.Database, client: client}, nil +} + +// unusedSettings names the SWH_* settings that are set and mean nothing to +// BigQuery. +func unusedSettings(cfg backend.Config) []string { + var unused []string + for _, s := range []struct { + name string + set bool + }{ + {"SWH_DSN", cfg.DSN != ""}, + {"SWH_HOST", cfg.Host != ""}, + {"SWH_PORT", cfg.Port != 0}, + {"SWH_USER", cfg.User != ""}, + {"SWH_PASSWORD", cfg.Password != ""}, + {"SWH_ACCOUNT", cfg.Account != ""}, + } { + if s.set { + unused = append(unused, s.name) + } + } + return unused +} + +// Name is the backend kind. +func (b *Backend) Name() string { return Kind } + +// Capabilities is what this backend honors. Load, unload and the native verbs +// are absent, and say so. +func (b *Backend) Capabilities() backend.Capabilities { + return backend.Capabilities{ + Backend: Kind, + AsyncJobs: true, + DryRun: true, + Parameters: true, + DDL: true, + Load: false, + Unload: false, + StreamingInsert: true, + ArrowResults: true, + MaxQueryBytes: b.cfg.MaxQueryBytes, + LoadFormats: nil, + NativeVerbs: nil, + } +} + +// Load is not offered: bulk loads from object storage are a gap, not an +// emulation. +func (b *Backend) Load(context.Context, backend.LoadOptions) (*backend.Job, error) { + return nil, serr.New(serr.Unsupported, "Load", "the bigquery backend does not offer bulk load") +} + +// Unload is not offered. +func (b *Backend) Unload(context.Context, backend.UnloadOptions) (*backend.Job, error) { + return nil, serr.New(serr.Unsupported, "Unload", "the bigquery backend does not offer unload") +} + +// Native has no verbs on this backend. +func (b *Backend) Native(_ context.Context, verb string, _ map[string]string) (map[string]string, error) { + return nil, serr.New(serr.Unsupported, "Native", "the bigquery backend has no native verbs: "+verb) +} + +// Close releases the client. +func (b *Backend) Close() error { return b.client.Close() } + +// resolve fills the default dataset into a table reference and validates the +// names, which are put into request paths. +func (b *Backend) resolve(op string, ref backend.TableRef) (backend.TableRef, error) { + if ref.Dataset == "" { + ref.Dataset = b.cfg.DefaultDataset + } + if ref.Dataset == "" { + return ref, serr.New(serr.InvalidArgument, op, "the table names no dataset and SWH_DATASET is not set") + } + if err := checkDatasetID(op, ref.Dataset); err != nil { + return ref, err + } + if err := checkTableID(op, ref.Table); err != nil { + return ref, err + } + return ref, nil +} diff --git a/internal/backend/bigquery/bigquery_test.go b/internal/backend/bigquery/bigquery_test.go new file mode 100644 index 0000000..b92ff95 --- /dev/null +++ b/internal/backend/bigquery/bigquery_test.go @@ -0,0 +1,121 @@ +package bigquery + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/codefly-dev/service-warehouse/internal/backend" + "github.com/codefly-dev/service-warehouse/internal/serr" +) + +func TestRegisteredAsAProjectBoundBackendWithoutADSN(t *testing.T) { + reg, ok := backend.Lookup("bigquery") + require.True(t, ok) + require.False(t, reg.SelfContained, "a BigQuery server is bound to a project") + require.False(t, reg.DSN, "BigQuery has no connection string") + require.Contains(t, backend.Registered(), "bigquery") +} + +// A failed open must hand back a nil interface: a nil *Backend inside a +// backend.Backend compares unequal to nil, and a caller checking the backend +// instead of the error would use it. +func TestRegisteredOpenReturnsANilBackendOnFailure(t *testing.T) { + reg, ok := backend.Lookup("bigquery") + require.True(t, ok) + be, err := reg.Open(context.Background(), backend.Config{}) + require.Error(t, err) + require.True(t, be == nil, "the backend is a true nil, not a typed one") +} + +func TestOpenRefusesAConfigurationItCannotHonor(t *testing.T) { + for name, tc := range map[string]struct { + cfg backend.Config + want string + }{ + "no project": {backend.Config{}, "SWH_DATABASE"}, + "a DSN instead": {backend.Config{DSN: "bigquery://p"}, "SWH_DATABASE"}, + "a project with a slash": {backend.Config{Database: "p/jobs"}, "project id"}, + "a project with a space": {backend.Config{Database: "my project"}, "project id"}, + "a DSN beside a project": {backend.Config{Database: "p-1", DSN: "x"}, "SWH_DSN"}, + "a host": {backend.Config{Database: "p-1", Host: "h"}, "SWH_HOST"}, + "a port": {backend.Config{Database: "p-1", Port: 5432}, "SWH_PORT"}, + "a user": {backend.Config{Database: "p-1", User: "u"}, "SWH_USER"}, + "a password": {backend.Config{Database: "p-1", Password: "pw"}, "SWH_PASSWORD"}, + "an account": {backend.Config{Database: "p-1", Account: "a"}, "SWH_ACCOUNT"}, + "a bad default dataset": {backend.Config{Database: "p-1", DefaultDataset: "a/b"}, "dataset"}, + "a bad location": {backend.Config{Database: "p-1", Location: "eu west"}, "SWH_LOCATION"}, + } { + t.Run(name, func(t *testing.T) { + b, err := Open(context.Background(), tc.cfg) + require.Nil(t, b) + require.True(t, serr.Is(err, serr.InvalidArgument), "got %v", err) + require.Contains(t, err.Error(), tc.want) + }) + } +} + +func TestOpenNamesEverySettingItRejects(t *testing.T) { + _, err := Open(context.Background(), backend.Config{Database: "p-1", User: "u", Password: "pw", Host: "h"}) + require.ErrorContains(t, err, "SWH_HOST, SWH_USER, SWH_PASSWORD") + require.NotContains(t, err.Error(), "pw", "a secret is never repeated back") +} + +func TestOpenFailsOnAnUnreadableCredentialsFileAndSaysWhy(t *testing.T) { + _, err := Open(context.Background(), backend.Config{Database: "p-1", CredentialsFile: t.TempDir() + "/missing.json"}) + require.Error(t, err) + require.Contains(t, err.Error(), "missing.json", "the operator is told which file") +} + +func TestCapabilitiesSayExactlyWhatWorks(t *testing.T) { + b := newFake(t).open(t, backend.Config{MaxQueryBytes: 1 << 30}) + require.Equal(t, "bigquery", b.Name()) + require.Equal(t, backend.Capabilities{ + Backend: "bigquery", + AsyncJobs: true, + DryRun: true, + Parameters: true, + DDL: true, + Load: false, + Unload: false, + StreamingInsert: true, + ArrowResults: true, + MaxQueryBytes: 1 << 30, + }, b.Capabilities()) +} + +func TestWhatIsNotOfferedIsUnsupportedAndNeverEmpty(t *testing.T) { + f := newFake(t) + b := f.open(t, backend.Config{}) + ctx := context.Background() + + job, err := b.Load(ctx, backend.LoadOptions{Dest: backend.TableRef{Dataset: "d", Table: "t"}, SourceURIs: []string{"gs://b/o"}}) + require.Nil(t, job) + require.True(t, serr.Is(err, serr.Unsupported)) + + job, err = b.Unload(ctx, backend.UnloadOptions{SQL: "SELECT 1", DestURI: "gs://b/o"}) + require.Nil(t, job) + require.True(t, serr.Is(err, serr.Unsupported)) + + values, err := b.Native(ctx, "time-travel", map[string]string{"at": "now"}) + require.Nil(t, values) + require.True(t, serr.Is(err, serr.Unsupported)) + + require.Empty(t, f.requests) + caps := b.Capabilities() + require.False(t, caps.Load) + require.False(t, caps.Unload) + require.Empty(t, caps.LoadFormats) + require.Empty(t, caps.NativeVerbs) +} + +func TestLimitIsTheTighterOfTheCallersAndTheServers(t *testing.T) { + require.Equal(t, int64(10), limit(int64(10), int64(100))) + require.Equal(t, int64(100), limit(int64(1000), int64(100))) + require.Equal(t, int64(100), limit(int64(0), int64(100))) + require.Equal(t, int64(10), limit(int64(10), int64(0))) + require.Equal(t, int64(0), limit(int64(0), int64(0))) + require.Equal(t, time.Second, limit(time.Second, time.Minute)) +} diff --git a/internal/backend/bigquery/catalog.go b/internal/backend/bigquery/catalog.go new file mode 100644 index 0000000..901e5b7 --- /dev/null +++ b/internal/backend/bigquery/catalog.go @@ -0,0 +1,380 @@ +package bigquery + +import ( + "context" + "fmt" + + bq "cloud.google.com/go/bigquery" + "golang.org/x/sync/errgroup" + "google.golang.org/api/iterator" + + "github.com/codefly-dev/service-warehouse/internal/backend" + "github.com/codefly-dev/service-warehouse/internal/serr" +) + +const ( + // defaultPageSize is the page of a catalog listing that names no limit, and + // maxPageSize the most a page holds: a larger request gets a full page and a + // token to continue, which is how a client already reads every page. + defaultPageSize = 50 + maxPageSize = 500 + // metadataParallelism bounds the metadata reads one listing makes at once. + // BigQuery lists only names, so each entry's detail is a request of its own, + // which is also why a page is kept modest. + metadataParallelism = 8 +) + +func pageSize(op string, requested int32) (int, error) { + switch { + case requested < 0: + return 0, serr.New(serr.InvalidArgument, op, "the page limit must not be negative") + case requested == 0: + return defaultPageSize, nil + case requested > maxPageSize: + return maxPageSize, nil + default: + return int(requested), nil + } +} + +// describeEach reads the detail of every item of a page, a few at a time. An item +// that disappeared between the listing and the read (NotFound) is left out of the +// page; any other failure fails it. +func describeEach[In, Out any](ctx context.Context, op string, items []In, read func(context.Context, In) (Out, error)) ([]Out, error) { + results := make([]Out, len(items)) + kept := make([]bool, len(items)) + group, ctx := errgroup.WithContext(ctx) + group.SetLimit(metadataParallelism) + for i, item := range items { + group.Go(func() error { + out, err := read(ctx, item) + if err != nil { + if is(op, err, serr.NotFound) { + return nil + } + return classify(op, err) + } + results[i], kept[i] = out, true + return nil + }) + } + if err := group.Wait(); err != nil { + return nil, err + } + out := make([]Out, 0, len(items)) + for i := range items { + if kept[i] { + out = append(out, results[i]) + } + } + return out, nil +} + +// ListDatasets lists the datasets of the project, a page at a time. +func (b *Backend) ListDatasets(ctx context.Context, opts backend.ListDatasetsOptions) (*backend.ListDatasetsResult, error) { + const op = "ListDatasets" + size, err := pageSize(op, opts.Limit) + if err != nil { + return nil, err + } + var page []*bq.Dataset + next, err := iterator.NewPager(b.client.Datasets(ctx), size, opts.PageToken).NextPage(&page) + if err != nil { + return nil, classify(op, err) + } + datasets, err := describeEach(ctx, op, page, func(ctx context.Context, d *bq.Dataset) (backend.Dataset, error) { + md, err := d.Metadata(ctx) + if err != nil { + return backend.Dataset{}, err + } + return datasetOut(d.DatasetID, md), nil + }) + if err != nil { + return nil, err + } + return &backend.ListDatasetsResult{Datasets: datasets, NextPageToken: next}, nil +} + +func datasetOut(name string, md *bq.DatasetMetadata) backend.Dataset { + return backend.Dataset{Name: name, Location: md.Location, Description: md.Description, Labels: md.Labels} +} + +// ListTables lists the tables of a dataset, a page at a time. +func (b *Backend) ListTables(ctx context.Context, opts backend.ListTablesOptions) (*backend.ListTablesResult, error) { + const op = "ListTables" + dataset := opts.Dataset + if dataset == "" { + dataset = b.cfg.DefaultDataset + } + if dataset == "" { + return nil, serr.New(serr.InvalidArgument, op, "no dataset is named and SWH_DATASET is not set") + } + if err := checkDatasetID(op, dataset); err != nil { + return nil, err + } + size, err := pageSize(op, opts.Limit) + if err != nil { + return nil, err + } + var page []*bq.Table + next, err := iterator.NewPager(b.client.Dataset(dataset).Tables(ctx), size, opts.PageToken).NextPage(&page) + if err != nil { + return nil, classify(op, err) + } + tables, err := describeEach(ctx, op, page, func(ctx context.Context, t *bq.Table) (backend.TableInfo, error) { + md, err := t.Metadata(ctx) + if err != nil { + return backend.TableInfo{}, err + } + return tableInfoOut(backend.TableRef{Dataset: dataset, Table: t.TableID}, md), nil + }) + if err != nil { + return nil, err + } + return &backend.ListTablesResult{Tables: tables, NextPageToken: next}, nil +} + +func tableKind(t bq.TableType) backend.TableKind { + switch t { + case bq.RegularTable: + return backend.KindTable + case bq.ViewTable: + return backend.KindView + case bq.MaterializedView: + return backend.KindMaterializedView + case bq.ExternalTable: + return backend.KindExternal + default: + // A snapshot, or a kind added since: reported as what it is not known to + // be, never as a table. + return backend.KindUnspecified + } +} + +func tableInfoOut(ref backend.TableRef, md *bq.TableMetadata) backend.TableInfo { + return backend.TableInfo{ + Ref: ref, + Kind: tableKind(md.Type), + NumRows: int64(md.NumRows), + NumBytes: md.NumBytes, + Created: md.CreationTime, + Modified: md.LastModifiedTime, + } +} + +// partitionColumns names the partitioning column of a table. A table partitioned +// by ingestion time has none, and reports BigQuery's pseudo-column. +func partitionColumns(md *bq.TableMetadata) []string { + switch { + case md.TimePartitioning != nil && md.TimePartitioning.Field != "": + return []string{md.TimePartitioning.Field} + case md.TimePartitioning != nil: + return []string{"_PARTITIONTIME"} + case md.RangePartitioning != nil: + return []string{md.RangePartitioning.Field} + default: + return nil + } +} + +func tableSchemaOut(ref backend.TableRef, md *bq.TableMetadata) backend.TableSchema { + out := backend.TableSchema{ + Info: tableInfoOut(ref, md), + Columns: columnsFromSchema(md.Schema), + PartitionBy: partitionColumns(md), + } + if md.Clustering != nil { + out.ClusterBy = md.Clustering.Fields + } + return out +} + +// GetTable returns a table's columns, partitioning and clustering. +func (b *Backend) GetTable(ctx context.Context, ref backend.TableRef) (*backend.TableSchema, error) { + const op = "GetTable" + ref, err := b.resolve(op, ref) + if err != nil { + return nil, err + } + md, err := b.client.Dataset(ref.Dataset).Table(ref.Table).Metadata(ctx) + if err != nil { + return nil, classify(op, err) + } + out := tableSchemaOut(ref, md) + return &out, nil +} + +// CreateDataset creates a dataset, in the configured location when the request +// names none. +func (b *Backend) CreateDataset(ctx context.Context, name string, opts backend.CreateDatasetOptions) (*backend.Dataset, error) { + const op = "CreateDataset" + if err := checkDatasetID(op, name); err != nil { + return nil, err + } + location := opts.Location + if location == "" { + location = b.cfg.Location + } + if location != "" && !locationID.MatchString(location) { + return nil, serr.New(serr.InvalidArgument, op, fmt.Sprintf("%q is not a BigQuery location", location)) + } + ds := b.client.Dataset(name) + md := &bq.DatasetMetadata{Location: location, Description: opts.Description, Labels: opts.Labels} + if err := ds.Create(ctx, md); err != nil { + if !opts.IfNotExists || !is(op, err, serr.AlreadyExists) { + return nil, classify(op, err) + } + existing, err := ds.Metadata(ctx) + if err != nil { + return nil, classify(op, err) + } + out := datasetOut(name, existing) + return &out, nil + } + created, err := ds.Metadata(ctx) + if err != nil { + return nil, classify(op, err) + } + out := datasetOut(name, created) + return &out, nil +} + +// DropDataset deletes a dataset. Without cascade a dataset that still holds +// tables is not deleted. +func (b *Backend) DropDataset(ctx context.Context, name string, opts backend.DropDatasetOptions) error { + const op = "DropDataset" + if err := checkDatasetID(op, name); err != nil { + return err + } + ds := b.client.Dataset(name) + var err error + if opts.Cascade { + err = ds.DeleteWithContents(ctx) + } else { + err = ds.Delete(ctx) + } + if err != nil && opts.IfExists && is(op, err, serr.NotFound) { + return nil + } + return classify(op, err) +} + +// partitioning is the time partitioning a portable PartitionBy asks for. BigQuery +// partitions by one column, at day granularity here; a second column is not +// something it can do, so it is refused and not dropped. +func partitioning(op string, columns []string) (*bq.TimePartitioning, error) { + switch len(columns) { + case 0: + return nil, nil + case 1: + field := columns[0] + if field == "_PARTITIONTIME" || field == "_PARTITIONDATE" { + return &bq.TimePartitioning{Type: bq.DayPartitioningType}, nil + } + return &bq.TimePartitioning{Type: bq.DayPartitioningType, Field: field}, nil + default: + return nil, serr.New(serr.Unsupported, op, "BigQuery partitions a table by one column") + } +} + +func clustering(columns []string) *bq.Clustering { + if len(columns) == 0 { + return nil + } + return &bq.Clustering{Fields: columns} +} + +// CreateTable creates a table from columns, or from a query (CTAS). +func (b *Backend) CreateTable(ctx context.Context, ref backend.TableRef, opts backend.CreateTableOptions) (*backend.TableSchema, error) { + const op = "CreateTable" + ref, err := b.resolve(op, ref) + if err != nil { + return nil, err + } + if opts.AsSelect != "" && len(opts.Columns) > 0 { + return nil, serr.New(serr.InvalidArgument, op, "a table is created from columns or from a query, not both") + } + if opts.AsSelect == "" && len(opts.Columns) == 0 { + return nil, serr.New(serr.InvalidArgument, op, "a table needs columns or a query to be created from") + } + part, err := partitioning(op, opts.PartitionBy) + if err != nil { + return nil, err + } + table := b.client.Dataset(ref.Dataset).Table(ref.Table) + + if opts.IfNotExists { + // Idempotent: an existing table is the answer, whatever it holds. + if md, err := table.Metadata(ctx); err == nil { + out := tableSchemaOut(ref, md) + return &out, nil + } else if !is(op, err, serr.NotFound) { + return nil, classify(op, err) + } + } + + if opts.AsSelect != "" { + if err := b.createAsSelect(ctx, table, opts, part); err != nil { + return nil, err + } + } else { + schema, err := schemaFromColumns(opts.Columns) + if err != nil { + return nil, err + } + meta := &bq.TableMetadata{Schema: schema, TimePartitioning: part, Clustering: clustering(opts.ClusterBy)} + if err := table.Create(ctx, meta); err != nil { + return nil, classify(op, err) + } + } + return b.GetTable(ctx, ref) +} + +// createAsSelect runs the query into the new table, which fails if the table +// already holds rows. +func (b *Backend) createAsSelect(ctx context.Context, table *bq.Table, opts backend.CreateTableOptions, part *bq.TimePartitioning) error { + const op = "CreateTable" + timeout := b.cfg.QueryTimeout + if timeout > 0 { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, timeout) + defer cancel() + } + q := b.client.Query(opts.AsSelect) + q.Dst = table + q.CreateDisposition = bq.CreateIfNeeded + q.WriteDisposition = bq.WriteEmpty + q.TimePartitioning = part + q.Clustering = clustering(opts.ClusterBy) + q.MaxBytesBilled = b.cfg.MaxQueryBytes + if b.cfg.DefaultDataset != "" { + q.DefaultProjectID, q.DefaultDatasetID = b.project, b.cfg.DefaultDataset + } + if timeout > 0 { + q.JobTimeout = timeout + } + job, err := q.Run(ctx) + if err != nil { + return classify(op, err) + } + status, err := job.Wait(ctx) + if err != nil { + cancelAbandoned(ctx, job) + return classify(op, err) + } + return classify(op, status.Err()) +} + +// DropTable deletes a table. +func (b *Backend) DropTable(ctx context.Context, ref backend.TableRef, opts backend.DropTableOptions) error { + const op = "DropTable" + ref, err := b.resolve(op, ref) + if err != nil { + return err + } + err = b.client.Dataset(ref.Dataset).Table(ref.Table).Delete(ctx) + if err != nil && opts.IfExists && is(op, err, serr.NotFound) { + return nil + } + return classify(op, err) +} diff --git a/internal/backend/bigquery/catalog_test.go b/internal/backend/bigquery/catalog_test.go new file mode 100644 index 0000000..a24f362 --- /dev/null +++ b/internal/backend/bigquery/catalog_test.go @@ -0,0 +1,374 @@ +package bigquery + +import ( + "context" + "net/http" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/codefly-dev/service-warehouse/internal/backend" + "github.com/codefly-dev/service-warehouse/internal/serr" +) + +const ( + datasetsPath = "/projects/test-project/datasets" + tablesPath = "/projects/test-project/datasets/d/tables" + tablePath = "/projects/test-project/datasets/d/tables/t" +) + +func datasetRef(id string) map[string]any { + return map[string]any{"projectId": "test-project", "datasetId": id} +} + +func tableRef(dataset, id string) map[string]any { + return map[string]any{"projectId": "test-project", "datasetId": dataset, "tableId": id} +} + +func TestListDatasetsReadsEachDatasetsDetailAndPages(t *testing.T) { + f := newFake(t) + f.handle("GET "+datasetsPath, 200, map[string]any{ + "datasets": []map[string]any{{"datasetReference": datasetRef("a")}, {"datasetReference": datasetRef("gone")}, {"datasetReference": datasetRef("b")}}, + "nextPageToken": "next-page", + }) + f.handle("GET "+datasetsPath+"/a", 200, map[string]any{ + "datasetReference": datasetRef("a"), "location": "EU", "description": "analytics", "labels": map[string]string{"team": "x"}, + }) + f.handle("GET "+datasetsPath+"/b", 200, map[string]any{"datasetReference": datasetRef("b"), "location": "US"}) + f.handle("GET "+datasetsPath+"/gone", 404, apiError(404, "notFound", "Not found")) + b := f.open(t, backend.Config{}) + + res, err := b.ListDatasets(context.Background(), backend.ListDatasetsOptions{PageToken: "this-page", Limit: 3}) + require.NoError(t, err) + require.Equal(t, []backend.Dataset{ + {Name: "a", Location: "EU", Description: "analytics", Labels: map[string]string{"team": "x"}}, + {Name: "b", Location: "US"}, + }, res.Datasets, "a dataset deleted between the listing and the read is left out") + require.Equal(t, "next-page", res.NextPageToken) + + listed := f.seen("GET", datasetsPath)[0] + require.Equal(t, "this-page", listed.Query["pageToken"][0]) + require.Equal(t, "3", listed.Query["maxResults"][0]) +} + +func TestAPageIsNeverLargerThanTheMaximumAndNeverNegative(t *testing.T) { + f := newFake(t) + f.handle("GET "+datasetsPath, 200, map[string]any{}) + b := f.open(t, backend.Config{}) + + _, err := b.ListDatasets(context.Background(), backend.ListDatasetsOptions{Limit: 1 << 20}) + require.NoError(t, err) + require.Equal(t, "500", f.seen("GET", datasetsPath)[0].Query["maxResults"][0]) + + _, err = b.ListDatasets(context.Background(), backend.ListDatasetsOptions{}) + require.NoError(t, err) + require.Equal(t, "50", f.seen("GET", datasetsPath)[1].Query["maxResults"][0], "no limit asks for the default page") + + _, err = b.ListDatasets(context.Background(), backend.ListDatasetsOptions{Limit: -1}) + require.True(t, serr.Is(err, serr.InvalidArgument)) +} + +func TestListTablesReportsKindAndSize(t *testing.T) { + f := newFake(t) + f.handle("GET "+tablesPath, 200, map[string]any{"tables": []map[string]any{ + {"tableReference": tableRef("d", "events"), "type": "TABLE"}, + {"tableReference": tableRef("d", "v"), "type": "VIEW"}, + {"tableReference": tableRef("d", "snap"), "type": "SNAPSHOT"}, + }}) + f.handle("GET "+tablesPath+"/events", 200, map[string]any{ + "tableReference": tableRef("d", "events"), "type": "TABLE", "numRows": "42", "numBytes": "4200", + "creationTime": "1700000000000", "lastModifiedTime": "1700000005000", + }) + f.handle("GET "+tablesPath+"/v", 200, map[string]any{"tableReference": tableRef("d", "v"), "type": "VIEW"}) + f.handle("GET "+tablesPath+"/snap", 200, map[string]any{"tableReference": tableRef("d", "snap"), "type": "SNAPSHOT"}) + b := f.open(t, backend.Config{}) + + res, err := b.ListTables(context.Background(), backend.ListTablesOptions{Dataset: "d"}) + require.NoError(t, err) + require.Len(t, res.Tables, 3) + require.Equal(t, backend.TableRef{Dataset: "d", Table: "events"}, res.Tables[0].Ref) + require.Equal(t, backend.KindTable, res.Tables[0].Kind) + require.EqualValues(t, 42, res.Tables[0].NumRows) + require.EqualValues(t, 4200, res.Tables[0].NumBytes) + require.EqualValues(t, 1700000000000, res.Tables[0].Created.UnixMilli()) + require.Equal(t, backend.KindView, res.Tables[1].Kind) + require.Equal(t, backend.KindUnspecified, res.Tables[2].Kind, "a snapshot is not reported as a table") + require.Empty(t, res.NextPageToken) + + _, err = b.ListTables(context.Background(), backend.ListTablesOptions{}) + require.True(t, serr.Is(err, serr.InvalidArgument), "no dataset is named and none is configured") +} + +func TestGetTableDescribesColumnsPartitioningAndClustering(t *testing.T) { + f := newFake(t) + f.handle("GET "+tablePath, 200, map[string]any{ + "tableReference": tableRef("d", "t"), "type": "TABLE", "numRows": "7", + "schema": map[string]any{"fields": []map[string]any{ + {"name": "id", "type": "INTEGER", "mode": "REQUIRED", "description": "the key"}, + {"name": "amount", "type": "NUMERIC", "mode": "NULLABLE", "precision": "10", "scale": "2"}, + {"name": "tags", "type": "STRING", "mode": "REPEATED"}, + {"name": "at", "type": "TIMESTAMP", "mode": "NULLABLE"}, + {"name": "items", "type": "RECORD", "mode": "REPEATED", "fields": []map[string]any{ + {"name": "sku", "type": "STRING", "mode": "NULLABLE", "maxLength": "12"}, + {"name": "qty", "type": "INTEGER", "mode": "NULLABLE"}, + }}, + {"name": "span", "type": "RANGE", "mode": "NULLABLE", "rangeElementType": map[string]any{"type": "DATE"}}, + }}, + "timePartitioning": map[string]any{"type": "DAY", "field": "at"}, + "clustering": map[string]any{"fields": []string{"id"}}, + }) + b := f.open(t, backend.Config{}) + + got, err := b.GetTable(context.Background(), backend.TableRef{Dataset: "d", Table: "t"}) + require.NoError(t, err) + require.Equal(t, []string{"at"}, got.PartitionBy) + require.Equal(t, []string{"id"}, got.ClusterBy) + require.EqualValues(t, 7, got.Info.NumRows) + require.Len(t, got.Columns, 6) + + require.Equal(t, backend.Column{Name: "id", Type: backend.TypeInt64, NativeType: "INT64", Description: "the key"}, got.Columns[0]) + require.Equal(t, backend.Column{Name: "amount", Type: backend.TypeNumeric, Nullable: true, NativeType: "NUMERIC(10, 2)", Precision: 10, Scale: 2}, got.Columns[1]) + require.Equal(t, backend.Column{ + Name: "tags", Type: backend.TypeArray, NativeType: "ARRAY", + Fields: []backend.Column{{Name: "item", Type: backend.TypeString, NativeType: "STRING"}}, + }, got.Columns[2]) + require.Equal(t, backend.TypeTimestampTZ, got.Columns[3].Type) + items := got.Columns[4] + require.Equal(t, backend.TypeArray, items.Type) + require.Equal(t, "ARRAY>", items.NativeType) + require.Equal(t, backend.TypeStruct, items.Fields[0].Type) + require.Equal(t, []string{"sku", "qty"}, []string{items.Fields[0].Fields[0].Name, items.Fields[0].Fields[1].Name}) + require.Equal(t, backend.Column{Name: "span", Type: backend.TypeUnknown, Nullable: true, NativeType: "RANGE"}, got.Columns[5], + "a type with no portable bucket is UNKNOWN with BigQuery's spelling kept") +} + +func TestGetTableReportsIngestionTimePartitioningByItsPseudoColumn(t *testing.T) { + f := newFake(t) + f.handle("GET "+tablePath, 200, map[string]any{"tableReference": tableRef("d", "t"), "type": "TABLE", "timePartitioning": map[string]any{"type": "DAY"}}) + b := f.open(t, backend.Config{}) + got, err := b.GetTable(context.Background(), backend.TableRef{Dataset: "d", Table: "t"}) + require.NoError(t, err) + require.Equal(t, []string{"_PARTITIONTIME"}, got.PartitionBy) +} + +func TestCreateDatasetUsesTheConfiguredLocationAndIsIdempotentOnRequest(t *testing.T) { + quiet(t) + f := newFake(t) + created := map[string]any{"datasetReference": datasetRef("fresh"), "location": "EU", "description": "x"} + f.handle("POST "+datasetsPath, 200, created) + f.handle("GET "+datasetsPath+"/fresh", 200, created) + b := f.open(t, backend.Config{Location: "EU"}) + + got, err := b.CreateDataset(context.Background(), "fresh", backend.CreateDatasetOptions{Description: "x", Labels: map[string]string{"k": "v"}}) + require.NoError(t, err) + require.Equal(t, backend.Dataset{Name: "fresh", Location: "EU", Description: "x"}, *got) + body := f.seen("POST", datasetsPath)[0].Body + require.Equal(t, "EU", body["location"]) + require.Equal(t, map[string]any{"k": "v"}, body["labels"]) + require.Equal(t, "fresh", body["datasetReference"].(map[string]any)["datasetId"]) + + exists := newFake(t) + exists.handle("POST "+datasetsPath, 409, apiError(409, "duplicate", "Already Exists: Dataset test-project:fresh")) + exists.handle("GET "+datasetsPath+"/fresh", 200, created) + b2 := exists.open(t, backend.Config{}) + _, err = b2.CreateDataset(context.Background(), "fresh", backend.CreateDatasetOptions{}) + require.True(t, serr.Is(err, serr.AlreadyExists), "got %v", err) + again, err := b2.CreateDataset(context.Background(), "fresh", backend.CreateDatasetOptions{IfNotExists: true}) + require.NoError(t, err) + require.Equal(t, "fresh", again.Name) + + _, err = b.CreateDataset(context.Background(), "x", backend.CreateDatasetOptions{Location: "not a/location"}) + require.True(t, serr.Is(err, serr.InvalidArgument)) +} + +func TestDropDatasetCascadesOnlyWhenAskedAndMapsWhatBigQuerySays(t *testing.T) { + quiet(t) + f := newFake(t) + f.handleFunc("DELETE "+datasetsPath+"/full", func(w http.ResponseWriter, r *http.Request) { + if r.URL.Query().Get("deleteContents") != "true" { + writeJSON(w, 400, apiError(400, "resourceInUse", "Dataset test-project:full is still in use")) + return + } + w.WriteHeader(204) + }) + f.handle("DELETE "+datasetsPath+"/missing", 404, apiError(404, "notFound", "Not found: Dataset")) + b := f.open(t, backend.Config{}) + + err := b.DropDataset(context.Background(), "full", backend.DropDatasetOptions{}) + require.True(t, serr.Is(err, serr.PreconditionFailed), "a dataset that still holds tables needs cascade: %v", err) + require.NoError(t, b.DropDataset(context.Background(), "full", backend.DropDatasetOptions{Cascade: true})) + + require.True(t, serr.Is(b.DropDataset(context.Background(), "missing", backend.DropDatasetOptions{}), serr.NotFound)) + require.NoError(t, b.DropDataset(context.Background(), "missing", backend.DropDatasetOptions{IfExists: true})) + require.Len(t, f.seen("DELETE", datasetsPath+"/full"), 2, "the cascade is the only second attempt, and it is the caller's") +} + +func TestCreateTableSendsColumnsPartitioningAndClusteringThenReadsTheTableBack(t *testing.T) { + f := newFake(t) + f.handle("POST "+tablesPath, 200, map[string]any{"tableReference": tableRef("d", "t")}) + f.handle("GET "+tablePath, 200, map[string]any{ + "tableReference": tableRef("d", "t"), "type": "TABLE", + "schema": map[string]any{"fields": []map[string]any{field("id", "INTEGER", "REQUIRED")}}, + }) + b := f.open(t, backend.Config{}) + + got, err := b.CreateTable(context.Background(), backend.TableRef{Dataset: "d", Table: "t"}, backend.CreateTableOptions{ + Columns: []backend.Column{ + {Name: "id", Type: backend.TypeInt64}, + {Name: "at", Type: backend.TypeTimestampTZ, Nullable: true}, + {Name: "amount", Type: backend.TypeNumeric, Nullable: true, Precision: 10, Scale: 2}, + {Name: "big", Type: backend.TypeNumeric, Nullable: true, Precision: 50, Scale: 20}, + {Name: "tags", Type: backend.TypeArray, Fields: []backend.Column{{Name: "item", Type: backend.TypeString}}}, + {Name: "owner", Type: backend.TypeStruct, Nullable: true, Fields: []backend.Column{ + {Name: "name", Type: backend.TypeString, Nullable: true}, + }}, + }, + PartitionBy: []string{"at"}, + ClusterBy: []string{"id"}, + }) + require.NoError(t, err) + require.Equal(t, "id", got.Columns[0].Name, "the answer is the table as BigQuery now holds it") + + body := f.seen("POST", tablesPath)[0].Body + fields := body["schema"].(map[string]any)["fields"].([]any) + require.Len(t, fields, 6) + byName := map[string]map[string]any{} + for _, f := range fields { + byName[f.(map[string]any)["name"].(string)] = f.(map[string]any) + } + require.Equal(t, "INTEGER", byName["id"]["type"]) + require.Equal(t, "REQUIRED", byName["id"]["mode"]) + require.Equal(t, "TIMESTAMP", byName["at"]["type"]) + require.Equal(t, "NUMERIC", byName["amount"]["type"]) + require.Equal(t, "10", byName["amount"]["precision"]) + require.Equal(t, "BIGNUMERIC", byName["big"]["type"], "a precision NUMERIC cannot hold is BIGNUMERIC") + require.Equal(t, "REPEATED", byName["tags"]["mode"]) + require.Equal(t, "STRING", byName["tags"]["type"]) + require.Equal(t, "RECORD", byName["owner"]["type"]) + require.Equal(t, map[string]any{"type": "DAY", "field": "at"}, body["timePartitioning"]) + require.Equal(t, map[string]any{"fields": []any{"id"}}, body["clustering"]) +} + +func TestCreateTableIsIdempotentOnRequestAndRefusesWhatItCannotMake(t *testing.T) { + quiet(t) + f := newFake(t) + f.handle("GET "+tablePath, 200, map[string]any{"tableReference": tableRef("d", "t"), "type": "TABLE", + "schema": map[string]any{"fields": []map[string]any{field("id", "INTEGER", "NULLABLE")}}}) + f.handle("POST "+tablesPath, 409, apiError(409, "duplicate", "Already Exists: Table")) + b := f.open(t, backend.Config{}) + ref := backend.TableRef{Dataset: "d", Table: "t"} + cols := []backend.Column{{Name: "id", Type: backend.TypeInt64, Nullable: true}} + + got, err := b.CreateTable(context.Background(), ref, backend.CreateTableOptions{Columns: cols, IfNotExists: true}) + require.NoError(t, err) + require.Equal(t, "id", got.Columns[0].Name) + require.Empty(t, f.seen("POST", tablesPath), "an existing table is the answer; nothing is created") + + _, err = b.CreateTable(context.Background(), ref, backend.CreateTableOptions{Columns: cols}) + require.True(t, serr.Is(err, serr.AlreadyExists), "got %v", err) + + for name, tc := range map[string]struct { + opts backend.CreateTableOptions + want serr.Code + }{ + "no columns and no query": {backend.CreateTableOptions{}, serr.InvalidArgument}, + "columns and a query": {backend.CreateTableOptions{Columns: cols, AsSelect: "SELECT 1"}, serr.InvalidArgument}, + "two partition columns": {backend.CreateTableOptions{Columns: cols, PartitionBy: []string{"a", "b"}}, serr.Unsupported}, + "a column of unknown type": {backend.CreateTableOptions{Columns: []backend.Column{{Name: "x", Type: backend.TypeUnknown, NativeType: "RANGE"}}}, serr.Unsupported}, + "an array without elements": {backend.CreateTableOptions{Columns: []backend.Column{{Name: "x", Type: backend.TypeArray}}}, serr.InvalidArgument}, + } { + t.Run(name, func(t *testing.T) { + _, err := b.CreateTable(context.Background(), backend.TableRef{Dataset: "d", Table: "other"}, tc.opts) + require.True(t, serr.Is(err, tc.want), "got %v", err) + }) + } +} + +func TestCreateTableAsSelectRunsTheQueryIntoTheNewTable(t *testing.T) { + f := newFake(t) + f.handle("POST "+jobsPath, 200, jobResource("DONE", nil, nil)) + f.handle("GET "+jobPath, 200, jobResource("DONE", nil, nil)) + f.handle("GET "+resultsPath, 200, resultsResource([]map[string]any{field("id", "INTEGER", "NULLABLE")})) + f.handle("GET "+tablePath, 200, map[string]any{"tableReference": tableRef("d", "t"), "type": "TABLE", + "schema": map[string]any{"fields": []map[string]any{field("id", "INTEGER", "NULLABLE")}}}) + b := f.open(t, backend.Config{DefaultDataset: "src", MaxQueryBytes: 1 << 30}) + + got, err := b.CreateTable(context.Background(), backend.TableRef{Dataset: "d", Table: "t"}, backend.CreateTableOptions{ + AsSelect: "SELECT id FROM source", PartitionBy: []string{"_PARTITIONTIME"}, ClusterBy: []string{"id"}, + }) + require.NoError(t, err) + require.Equal(t, "id", got.Columns[0].Name) + + query := f.seen("POST", "/jobs")[0].Body["configuration"].(map[string]any)["query"].(map[string]any) + require.Equal(t, "SELECT id FROM source", query["query"]) + require.Equal(t, tableRef("d", "t"), query["destinationTable"]) + require.Equal(t, "WRITE_EMPTY", query["writeDisposition"], "a table that already holds rows is not overwritten") + require.Equal(t, "CREATE_IF_NEEDED", query["createDisposition"]) + require.Equal(t, map[string]any{"type": "DAY"}, query["timePartitioning"]) + require.Equal(t, map[string]any{"projectId": "test-project", "datasetId": "src"}, query["defaultDataset"]) + require.Equal(t, "1073741824", query["maximumBytesBilled"], "the server's byte cap binds a CTAS too") +} + +func TestDropTable(t *testing.T) { + quiet(t) + f := newFake(t) + f.handle("DELETE "+tablePath, 204, nil) + f.handle("DELETE "+tablesPath+"/gone", 404, apiError(404, "notFound", "Not found: Table")) + b := f.open(t, backend.Config{}) + + require.NoError(t, b.DropTable(context.Background(), backend.TableRef{Dataset: "d", Table: "t"}, backend.DropTableOptions{})) + require.True(t, serr.Is(b.DropTable(context.Background(), backend.TableRef{Dataset: "d", Table: "gone"}, backend.DropTableOptions{}), serr.NotFound)) + require.NoError(t, b.DropTable(context.Background(), backend.TableRef{Dataset: "d", Table: "gone"}, backend.DropTableOptions{IfExists: true})) +} + +// A name is expanded into a request path without escaping, so a dataset named +// "d/tables/t" would address a table: every operation refuses such names before +// sending anything. +func TestNamesCannotChangeWhatIsAddressed(t *testing.T) { + f := newFake(t) + b := f.open(t, backend.Config{}) + ctx := context.Background() + bad := []string{"d/tables/t", "d/../x", "d?x=1", "d#x", "d%2Ftables", "d.t", "", "d t"} + + for _, name := range bad { + ref := backend.TableRef{Dataset: name, Table: "t"} + _, err := b.GetTable(ctx, ref) + require.True(t, serr.Is(err, serr.InvalidArgument), "GetTable %q: %v", name, err) + _, err = b.CreateTable(ctx, ref, backend.CreateTableOptions{Columns: []backend.Column{{Name: "a", Type: backend.TypeInt64}}}) + require.True(t, serr.Is(err, serr.InvalidArgument), "CreateTable %q: %v", name, err) + require.True(t, serr.Is(b.DropTable(ctx, ref, backend.DropTableOptions{}), serr.InvalidArgument), "DropTable %q", name) + _, err = b.ListTables(ctx, backend.ListTablesOptions{Dataset: name}) + if name != "" { // an empty dataset means "the default", and there is none + require.True(t, serr.Is(err, serr.InvalidArgument), "ListTables %q: %v", name, err) + } + if name != "" { + _, err = b.CreateDataset(ctx, name, backend.CreateDatasetOptions{}) + require.True(t, serr.Is(err, serr.InvalidArgument), "CreateDataset %q: %v", name, err) + require.True(t, serr.Is(b.DropDataset(ctx, name, backend.DropDatasetOptions{Cascade: true}), serr.InvalidArgument), "DropDataset %q", name) + } + } + for _, name := range []string{"t/../x", "t?x", "t#", "t%2F", "t.u", "", "t\x00"} { + _, err := b.GetTable(ctx, backend.TableRef{Dataset: "d", Table: name}) + require.True(t, serr.Is(err, serr.InvalidArgument), "table %q: %v", name, err) + } + require.Empty(t, f.requests, "no such name reached BigQuery") +} + +func TestTableNamesMayHoldWhatBigQueryAllows(t *testing.T) { + f := newFake(t) + f.handle("GET /projects/test-project/datasets/d_1/tables/my-table%20two", 200, map[string]any{"tableReference": tableRef("d_1", "my-table two"), "type": "TABLE"}) + b := f.open(t, backend.Config{}) + got, err := b.GetTable(context.Background(), backend.TableRef{Dataset: "d_1", Table: "my-table two"}) + require.NoError(t, err) + require.Equal(t, "my-table two", got.Info.Ref.Table) +} + +func TestListTablesOfAMissingDatasetIsNotFound(t *testing.T) { + quiet(t) + f := newFake(t) + f.handle("GET "+tablesPath, 404, apiError(404, "notFound", "Not found: Dataset test-project:d")) + b := f.open(t, backend.Config{}) + + _, err := b.ListTables(context.Background(), backend.ListTablesOptions{Dataset: "d"}) + require.True(t, serr.Is(err, serr.NotFound), "got %v", err) +} diff --git a/internal/backend/bigquery/errors.go b/internal/backend/bigquery/errors.go new file mode 100644 index 0000000..95f01ba --- /dev/null +++ b/internal/backend/bigquery/errors.go @@ -0,0 +1,225 @@ +package bigquery + +import ( + "context" + "errors" + "log" + "net/http" + "strings" + + bq "cloud.google.com/go/bigquery" + "golang.org/x/oauth2" + "google.golang.org/api/googleapi" + + "github.com/codefly-dev/service-warehouse/internal/serr" +) + +// logf is where the vendor's own text goes: to the operator, never to a client. +// A test swaps it to keep its output quiet. +var logf = log.Printf + +// maxDetail bounds the diagnostic text of an invalid request that is passed on. +const maxDetail = 512 + +// failure is what a vendor error says about itself, read once so the +// classification below never touches an error type twice. +type failure struct { + // http is the HTTP status BigQuery answered with, or 0 when there was none. + http int + // reasons are BigQuery's machine-readable reason tokens, the first being the + // most specific. + reasons []string + // message is BigQuery's human-readable text. + message string +} + +// inspect reads the vendor error types this backend can meet. It reports false +// for an error that is none of them (a network failure, a context error). +func inspect(err error) (failure, bool) { + var f failure + found := false + + var api *googleapi.Error + if errors.As(err, &api) { + found = true + f.http = api.Code + f.message = api.Message + for _, item := range api.Errors { + if item.Reason != "" { + f.reasons = append(f.reasons, item.Reason) + } + } + } + + var byPointer *bq.Error + var byValue bq.Error + switch { + case errors.As(err, &byPointer): + found = true + f.reasons = append(f.reasons, byPointer.Reason) + if f.message == "" { + f.message = byPointer.Message + } + case errors.As(err, &byValue): + found = true + f.reasons = append(f.reasons, byValue.Reason) + if f.message == "" { + f.message = byValue.Message + } + } + + var multi bq.MultiError + if errors.As(err, &multi) { + for _, one := range multi { + if inner, ok := inspect(one); ok { + found = true + f.reasons = append(f.reasons, inner.reasons...) + if f.http == 0 { + f.http = inner.http + } + if f.message == "" { + f.message = inner.message + } + } + } + } + + var token *oauth2.RetrieveError + if errors.As(err, &token) { + // The credentials were refused when the server tried to sign in. + found = true + f.http = http.StatusUnauthorized + } + return f, found +} + +// reasonCodes maps BigQuery's reason tokens to the normalized codes. This is the +// one place vendor tokens are read; nothing past this file sees them. +var reasonCodes = map[string]serr.Code{ + "notFound": serr.NotFound, + "duplicate": serr.AlreadyExists, + "accessDenied": serr.PermissionDenied, + "forbidden": serr.PermissionDenied, + "billingNotEnabled": serr.PermissionDenied, + "invalidUser": serr.PermissionDenied, + "rateLimitExceeded": serr.Throttled, + "userRateLimitExceeded": serr.Throttled, + "quotaExceeded": serr.Throttled, + "bytesBilledLimitExceeded": serr.Throttled, + "billingTierLimitExceeded": serr.Throttled, + "resourcesExceeded": serr.Throttled, + "resourceInUse": serr.PreconditionFailed, + "invalid": serr.InvalidArgument, + "invalidQuery": serr.InvalidArgument, + "badRequest": serr.InvalidArgument, + "notImplemented": serr.Unsupported, + "timeout": serr.DeadlineExceeded, + "backendError": serr.Internal, + "internalError": serr.Internal, + "stopped": serr.Internal, +} + +// httpCodes is the fallback when BigQuery named no reason it is known by. +var httpCodes = map[int]serr.Code{ + http.StatusBadRequest: serr.InvalidArgument, + http.StatusUnauthorized: serr.PermissionDenied, + http.StatusForbidden: serr.PermissionDenied, + http.StatusNotFound: serr.NotFound, + http.StatusRequestTimeout: serr.DeadlineExceeded, + http.StatusConflict: serr.AlreadyExists, + http.StatusPreconditionFailed: serr.PreconditionFailed, + http.StatusTooManyRequests: serr.Throttled, + http.StatusNotImplemented: serr.Unsupported, + http.StatusGatewayTimeout: serr.DeadlineExceeded, + http.StatusInternalServerError: serr.Internal, + http.StatusBadGateway: serr.Internal, + http.StatusServiceUnavailable: serr.Internal, +} + +// codeFor is the normalized code of a vendor failure: the first reason BigQuery +// gave that is known, else the HTTP status, else Internal. +func codeFor(f failure) serr.Code { + for _, reason := range f.reasons { + if code, ok := reasonCodes[reason]; ok { + return code + } + } + if code, ok := httpCodes[f.http]; ok { + return code + } + return serr.Internal +} + +// phrases are what a client is told of a failure, written here and not taken +// from the vendor: the same failure reads the same on every backend. +var phrases = map[serr.Code]string{ + serr.NotFound: "the dataset, table or job does not exist", + serr.AlreadyExists: "the object already exists", + serr.PreconditionFailed: "a precondition of the operation was not met", + serr.Unsupported: "the warehouse does not offer this operation", + serr.PermissionDenied: "the server's credentials are not allowed to do this", + serr.Throttled: "the warehouse throttled the request or a quota or byte limit was exceeded", + serr.InvalidArgument: "the request was rejected as invalid", + serr.DeadlineExceeded: "the operation did not finish within its time limit", + serr.Internal: "the warehouse reported an internal failure", +} + +// describe is the message a client receives for code. An invalid request keeps +// the vendor's diagnosis, because that text describes the caller's own SQL or +// arguments ("Unrecognized name: x at [1:8]") and nothing else can; it never +// carries a reason token, a status or an identifier of the server's. Every +// other failure gets the fixed phrase: the vendor's text for a denial or an +// internal failure can name the server's own principals and resources, and goes +// to the operator's log instead. +func describe(code serr.Code, vendorMessage string) string { + if code == serr.InvalidArgument { + if text := strings.TrimSpace(vendorMessage); text != "" { + if len(text) > maxDetail { + text = text[:maxDetail] + } + return text + } + } + return phrases[code] +} + +// normalize is the normalized code of anything a call returned, and the message +// a client may be told of it. +func normalize(err error) (serr.Code, string) { + switch { + case errors.Is(err, context.DeadlineExceeded): + return serr.DeadlineExceeded, phrases[serr.DeadlineExceeded] + case errors.Is(err, context.Canceled): + return serr.Internal, "the operation was canceled" + } + if f, ok := inspect(err); ok { + code := codeFor(f) + return code, describe(code, f.message) + } + return serr.Internal, phrases[serr.Internal] +} + +// classify turns anything a call returned into a normalized serr error. An error +// already normalized passes through. The vendor's own error is logged when it is +// the operator's to read, and is never wrapped: Unwrap would hand its text to +// whoever formats the error. +func classify(op string, err error) error { + if err == nil { + return nil + } + var normalized *serr.Error + if errors.As(err, &normalized) { + return err + } + code, message := normalize(err) + switch code { + case serr.Internal, serr.PermissionDenied, serr.Throttled, serr.DeadlineExceeded, serr.Unsupported: + logf("bigquery: %s: %v", op, err) + } + return serr.New(code, op, message) +} + +// is reports whether err is the vendor failure that classifies as code. +func is(op string, err error, code serr.Code) bool { + return err != nil && serr.Is(classify(op, err), code) +} diff --git a/internal/backend/bigquery/errors_test.go b/internal/backend/bigquery/errors_test.go new file mode 100644 index 0000000..0a1e965 --- /dev/null +++ b/internal/backend/bigquery/errors_test.go @@ -0,0 +1,153 @@ +package bigquery + +import ( + "context" + "errors" + "fmt" + "net/http" + "net/url" + "strings" + "sync" + "testing" + + bq "cloud.google.com/go/bigquery" + "github.com/stretchr/testify/require" + "golang.org/x/oauth2" + "google.golang.org/api/googleapi" + + "github.com/codefly-dev/service-warehouse/internal/serr" +) + +func apiErr(status int, reasons ...string) *googleapi.Error { + e := &googleapi.Error{Code: status, Message: "VENDOR MESSAGE for " + http.StatusText(status)} + for _, r := range reasons { + e.Errors = append(e.Errors, googleapi.ErrorItem{Reason: r, Message: "VENDOR ITEM"}) + } + return e +} + +func TestClassifyMapsEveryBigQueryFailureToANormalizedCode(t *testing.T) { + quiet(t) + for name, tc := range map[string]struct { + err error + want serr.Code + }{ + "not found by reason": {apiErr(404, "notFound"), serr.NotFound}, + "not found by status": {apiErr(404), serr.NotFound}, + "duplicate": {apiErr(409, "duplicate"), serr.AlreadyExists}, + "conflict": {apiErr(409), serr.AlreadyExists}, + "access denied": {apiErr(403, "accessDenied"), serr.PermissionDenied}, + "forbidden": {apiErr(403, "forbidden"), serr.PermissionDenied}, + "billing not enabled": {apiErr(403, "billingNotEnabled"), serr.PermissionDenied}, + "unauthenticated": {apiErr(401), serr.PermissionDenied}, + "forbidden by status": {apiErr(403), serr.PermissionDenied}, + "rate limit": {apiErr(403, "rateLimitExceeded"), serr.Throttled}, + "user rate limit": {apiErr(403, "userRateLimitExceeded"), serr.Throttled}, + "quota": {apiErr(403, "quotaExceeded"), serr.Throttled}, + "bytes billed cap": {apiErr(400, "bytesBilledLimitExceeded"), serr.Throttled}, + "resources exceeded": {apiErr(400, "resourcesExceeded"), serr.Throttled}, + "too many requests": {apiErr(429), serr.Throttled}, + "dataset in use": {apiErr(400, "resourceInUse"), serr.PreconditionFailed}, + "precondition": {apiErr(412), serr.PreconditionFailed}, + "invalid query": {apiErr(400, "invalidQuery"), serr.InvalidArgument}, + "invalid": {apiErr(400, "invalid"), serr.InvalidArgument}, + "bad request": {apiErr(400), serr.InvalidArgument}, + "timeout reason": {apiErr(400, "timeout"), serr.DeadlineExceeded}, + "request timeout": {apiErr(408), serr.DeadlineExceeded}, + "gateway timeout": {apiErr(504), serr.DeadlineExceeded}, + "backend error": {apiErr(503, "backendError"), serr.Internal}, + "internal error": {apiErr(500, "internalError"), serr.Internal}, + "bad gateway": {apiErr(502), serr.Internal}, + "not implemented": {apiErr(501), serr.Unsupported}, + "not implemented reason": {apiErr(400, "notImplemented"), serr.Unsupported}, + "stopped job": {apiErr(400, "stopped"), serr.Internal}, + "reason beats status": {apiErr(400, "notFound"), serr.NotFound}, + "unknown reason by code": {apiErr(404, "somethingNew"), serr.NotFound}, + "first known reason wins": {apiErr(400, "somethingNew", "accessDenied"), serr.PermissionDenied}, + "job error by pointer": {&bq.Error{Reason: "invalidQuery", Message: "VENDOR"}, serr.InvalidArgument}, + "job error by value": {bq.Error{Reason: "accessDenied", Message: "VENDOR"}, serr.PermissionDenied}, + "multi error": {bq.MultiError{&bq.Error{Reason: "notFound"}}, serr.NotFound}, + "wrapped": {fmt.Errorf("while reading: %w", apiErr(404, "notFound")), serr.NotFound}, + "sign-in refused": {&oauth2.RetrieveError{Response: &http.Response{StatusCode: 400}}, serr.PermissionDenied}, + "deadline": {context.DeadlineExceeded, serr.DeadlineExceeded}, + "deadline in a transport": {&url.Error{Op: "Get", URL: "x", Err: context.DeadlineExceeded}, serr.DeadlineExceeded}, + "canceled": {context.Canceled, serr.Internal}, + "unknown": {errors.New("boom"), serr.Internal}, + } { + t.Run(name, func(t *testing.T) { + err := classify("Op", tc.err) + require.True(t, serr.Is(err, tc.want), "got %v", err) + }) + } + require.NoError(t, classify("Op", nil)) +} + +func TestClassifyLeavesANormalizedErrorAlone(t *testing.T) { + original := serr.New(serr.Unsupported, "Inner", "already normalized") + require.Same(t, error(original), classify("Outer", original)) + wrapped := fmt.Errorf("context: %w", original) + require.True(t, serr.Is(classify("Outer", wrapped), serr.Unsupported)) +} + +func TestWhatAClientIsToldIsWrittenHereNotTakenFromTheVendor(t *testing.T) { + logged := quiet(t) + + for name, err := range map[string]error{ + "denied": apiErr(403, "accessDenied"), + "internal": apiErr(500, "backendError"), + "throttle": apiErr(403, "quotaExceeded"), + "missing": apiErr(404, "notFound"), + "conflict": apiErr(409, "duplicate"), + "in use": apiErr(400, "resourceInUse"), + "unknown": errors.New("VENDOR SECRET svc@example.iam.gserviceaccount.com"), + } { + t.Run(name, func(t *testing.T) { + got := classify("Op", err) + text := got.Error() + require.NotContains(t, text, "VENDOR") + require.NotContains(t, text, "svc@") + var api *googleapi.Error + require.False(t, errors.As(got, &api), "the vendor's error is not wrapped, so nothing downstream can format it") + require.Equal(t, got, error(got.(*serr.Error))) + require.NotEmpty(t, phrases[serr.CodeOf(got)]) + }) + } + + // An invalid request keeps the vendor's diagnosis of the caller's own SQL, and + // only that: never its reason token or status. + got := classify("Query", apiErr(400, "invalidQuery")) + require.Contains(t, got.Error(), "VENDOR MESSAGE") + require.NotContains(t, got.Error(), "invalidQuery") + require.NotContains(t, got.Error(), "VENDOR ITEM") + + // Bounded, and never empty. + long := &googleapi.Error{Code: 400, Message: strings.Repeat("x", 10*maxDetail), Errors: []googleapi.ErrorItem{{Reason: "invalid"}}} + require.Less(t, len(classify("Query", long).Error()), 2*maxDetail) + empty := &googleapi.Error{Code: 400, Errors: []googleapi.ErrorItem{{Reason: "invalid"}}} + require.Contains(t, classify("Query", empty).Error(), phrases[serr.InvalidArgument]) + + // What is not for the client is for the operator. + require.NotEmpty(t, *logged) +} + +func TestOnlyWhatTheOperatorMustSeeIsLogged(t *testing.T) { + var mu sync.Mutex + var lines []string + previous := logf + logf = func(format string, args ...any) { + mu.Lock() + defer mu.Unlock() + lines = append(lines, fmt.Sprintf(format, args...)) + } + defer func() { logf = previous }() + + classify("Get", apiErr(404, "notFound")) + classify("Create", apiErr(409, "duplicate")) + classify("Query", apiErr(400, "invalidQuery")) + require.Empty(t, lines, "a caller's own mistake is not an operator's concern") + + classify("Query", apiErr(403, "accessDenied")) + classify("Query", apiErr(500, "internalError")) + require.Len(t, lines, 2) + require.Contains(t, lines[0], "VENDOR MESSAGE", "the vendor's text reaches the operator's log in full") +} diff --git a/internal/backend/bigquery/fake_test.go b/internal/backend/bigquery/fake_test.go new file mode 100644 index 0000000..c0dc96b --- /dev/null +++ b/internal/backend/bigquery/fake_test.go @@ -0,0 +1,128 @@ +package bigquery + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + + "github.com/stretchr/testify/require" + "google.golang.org/api/option" + + "github.com/codefly-dev/service-warehouse/internal/backend" +) + +// The tests below run the backend against a fake of BigQuery's REST API served +// from this process. It is not BigQuery: it answers what a test programs and +// records what the client library sent, so a test can assert the requests the +// backend builds and the way it reads the answers. Nothing leaves the machine. +// What only the real service can say is the integration test's to check. + +// recorded is one request the fake received. +type recorded struct { + Method string + Path string + Query map[string][]string + Body map[string]any +} + +type fake struct { + srv *httptest.Server + mux *http.ServeMux + + mu sync.Mutex + requests []recorded +} + +func newFake(t *testing.T) *fake { + t.Helper() + f := &fake{mux: http.NewServeMux()} + f.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + rec := recorded{Method: r.Method, Path: r.URL.Path, Query: r.URL.Query()} + if len(body) > 0 { + _ = json.Unmarshal(body, &rec.Body) + } + f.mu.Lock() + f.requests = append(f.requests, rec) + f.mu.Unlock() + f.mux.ServeHTTP(w, r) + })) + t.Cleanup(f.srv.Close) + return f +} + +// handle serves pattern ("POST /projects/{p}/jobs") with a JSON body. +func (f *fake) handle(pattern string, status int, body any) { + f.mux.HandleFunc(pattern, func(w http.ResponseWriter, _ *http.Request) { + writeJSON(w, status, body) + }) +} + +func (f *fake) handleFunc(pattern string, fn func(w http.ResponseWriter, r *http.Request)) { + f.mux.HandleFunc(pattern, fn) +} + +func writeJSON(w http.ResponseWriter, status int, body any) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + _ = json.NewEncoder(w).Encode(body) +} + +// apiError is the body BigQuery gives a failed request. +func apiError(status int, reason, message string) map[string]any { + return map[string]any{"error": map[string]any{ + "code": status, + "message": message, + "errors": []map[string]any{{"reason": reason, "message": message, "domain": "global"}}, + "status": "SOME_VENDOR_STATUS", + }} +} + +// seen returns the requests whose method and path match. +func (f *fake) seen(method, pathContains string) []recorded { + f.mu.Lock() + defer f.mu.Unlock() + var out []recorded + for _, r := range f.requests { + if r.Method == method && strings.Contains(r.Path, pathContains) { + out = append(out, r) + } + } + return out +} + +// open builds a Backend whose client talks to the fake. +func (f *fake) open(t *testing.T, cfg backend.Config) *Backend { + t.Helper() + if cfg.Database == "" { + cfg.Database = "test-project" + } + b, err := Open(context.Background(), cfg, + option.WithEndpoint(f.srv.URL), option.WithoutAuthentication()) + require.NoError(t, err) + t.Cleanup(func() { _ = b.Close() }) + return b +} + +// quiet keeps the operator log of a test run clean, and returns what was logged. +func quiet(t *testing.T) *[]string { + t.Helper() + var ( + mu sync.Mutex + logged []string + ) + previous := logf + logf = func(format string, args ...any) { + mu.Lock() + defer mu.Unlock() + logged = append(logged, fmt.Sprintf(format, args...)) + } + t.Cleanup(func() { logf = previous }) + return &logged +} diff --git a/internal/backend/bigquery/ident.go b/internal/backend/bigquery/ident.go new file mode 100644 index 0000000..4086d20 --- /dev/null +++ b/internal/backend/bigquery/ident.go @@ -0,0 +1,52 @@ +package bigquery + +import ( + "fmt" + "regexp" + "unicode/utf8" + + "github.com/codefly-dev/service-warehouse/internal/serr" +) + +// Names a caller supplies become parts of a request path, and the client library +// expands them without escaping reserved characters. A dataset named +// "d/tables/t" would then address a table, so every name is checked against what +// BigQuery itself allows before it is used: nothing a name contains can change +// what is addressed. + +// identifier is a name pattern with the longest name BigQuery accepts for it. +type identifier struct { + re *regexp.Regexp + max int +} + +func (i identifier) MatchString(s string) bool { + return s != "" && utf8.RuneCountInString(s) <= i.max && i.re.MatchString(s) +} + +var ( + // datasetID: letters, digits and underscores. + datasetID = identifier{regexp.MustCompile(`^[A-Za-z0-9_]+$`), 1024} + // tableID: Unicode letters, marks, numbers, connectors, dashes and spaces. + tableID = identifier{regexp.MustCompile(`^[\p{L}\p{M}\p{N}\p{Pc}\p{Pd} ]+$`), 1024} + // jobID: letters, digits, underscores and dashes. + jobID = identifier{regexp.MustCompile(`^[A-Za-z0-9_-]+$`), 1024} + // locationID: a region or multi-region such as US, EU or europe-west1. + locationID = identifier{regexp.MustCompile(`^[A-Za-z0-9-]+$`), 64} +) + +func checkDatasetID(op, name string) error { + if !datasetID.MatchString(name) { + return serr.New(serr.InvalidArgument, op, + fmt.Sprintf("%q is not a valid dataset name (letters, digits and underscores)", name)) + } + return nil +} + +func checkTableID(op, name string) error { + if !tableID.MatchString(name) { + return serr.New(serr.InvalidArgument, op, + fmt.Sprintf("%q is not a valid table name (letters, digits, underscores, dashes and spaces)", name)) + } + return nil +} diff --git a/internal/backend/bigquery/insert.go b/internal/backend/bigquery/insert.go new file mode 100644 index 0000000..094ce9c --- /dev/null +++ b/internal/backend/bigquery/insert.go @@ -0,0 +1,348 @@ +package bigquery + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "strings" + + bq "cloud.google.com/go/bigquery" + "github.com/apache/arrow-go/v18/arrow" + + "github.com/codefly-dev/service-warehouse/internal/backend" + "github.com/codefly-dev/service-warehouse/internal/serr" +) + +// maxRowsPerInsert is BigQuery's recommended streaming request size; a stream of +// more rows is sent as several requests. +const maxRowsPerInsert = 500 + +// maxRequestBytes bounds the rows one request carries, as BigQuery encodes them. +// BigQuery refuses a request over 10 MB and, with the same limit, a single row +// over it, and a row count says nothing of that: a few hundred rows with large +// text are many megabytes. The bound sits under the limit to leave room for the +// request's own envelope and for the difference between this estimate and the +// encoding. A row that alone passes it can never be sent and is refused as too +// large; the row is the problem, so retrying cannot help. +const maxRequestBytes = 8 << 20 + +// rowEnvelopeBytes is what a row adds to a request beyond its values: the keys +// and punctuation around them and its insert id. +const rowEnvelopeBytes = 64 + +// reasonInvalid is the reason BigQuery gives a streamed row whose own content it +// refuses. The other reasons a row carries (timeout, backendError, internalError, +// rateLimitExceeded, quotaExceeded, stopped) say nothing against the row. +const reasonInvalid = "invalid" + +// rowInserter is the streaming-insert seam; *bq.Inserter satisfies it. +type rowInserter interface { + Put(ctx context.Context, src any) error +} + +// insertRow is a row waiting for a request. +type insertRow struct { + index int64 + values map[string]bq.Value + size int +} + +// Save implements bq.ValueSaver. An empty insert id asks the client library for +// a random one, so no two rows are ever collapsed by BigQuery's best-effort +// de-duplication. +func (r insertRow) Save() (map[string]bq.Value, string, error) { return r.values, "", nil } + +// InsertRows appends the rows of every Arrow batch to a table, with streaming +// inserts, and reports what BigQuery refused row by row. +// +// What a failure says about a row decides how it is reported: +// +// - A row whose own content BigQuery refuses (it answers "invalid") is +// reported with RefusalInvalidValue. The rest of its request is still +// inserted: the inserter skips invalid rows, so the call is not atomic and a +// refused row does not shorten the batch. +// - A row that cannot be sent in any request, or whose values cannot be +// encoded, is reported with RefusalRowTooLarge or RefusalInvalidValue before +// anything of it is sent. +// - When the batch's columns disagree with the table's (a column the table +// lacks, a required column the batch lacks) every row is reported with +// RefusalSchemaMismatch and none is sent: the table, not any row, would have +// to change. The disagreement is found by comparing the batch's schema with +// the table's metadata, so no vendor text is read to decide it. +// - Anything else (a request that failed, a throttle, a quota, a timeout, a +// dead backend, a row reported with any other reason) fails the call. Rows of +// earlier requests may already be stored, so a caller that retries the call +// can store a row twice: delivery is at least once. +func (b *Backend) InsertRows(ctx context.Context, ref backend.TableRef, schema []byte, batches backend.BatchReader) (*backend.InsertResult, error) { + const op = "InsertRows" + ref, err := b.resolve(op, ref) + if err != nil { + return nil, err + } + stream, err := newBatchStream(schema, batches) + if err != nil { + return nil, err + } + defer stream.close() + + table := b.client.Dataset(ref.Dataset).Table(ref.Table) + md, err := table.Metadata(ctx) + if err != nil { + return nil, classify(op, err) + } + inserter := table.Inserter() + // Rows BigQuery finds invalid are skipped and reported back, so the valid + // rows of a request are stored and the invalid ones come back by index. + inserter.SkipInvalidRows = true + return insertAll(ctx, stream, md.Schema, inserter) +} + +// insertAll runs the stream through plan, send and the classification of what +// BigQuery answered. +func insertAll(ctx context.Context, stream *batchStream, table bq.Schema, inserter rowInserter) (*backend.InsertResult, error) { + const op = "InsertRows" + var ( + result backend.InsertResult + next int64 // the index of the next row of the stream + conflict string + checked bool + request []insertRow + reqBytes int + ) + flush := func() error { + if len(request) == 0 { + return nil + } + inserted, refused, err := send(ctx, inserter, request) + result.RowsInserted += int64(inserted) + result.Errors = append(result.Errors, refused...) + request, reqBytes = nil, 0 + return err + } + + for { + rec, err := stream.next() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + return nil, err + } + if !checked { + checked = true + if err := checkArrowSchema(rec.Schema()); err != nil { + rec.Release() + return nil, serr.New(serr.Unsupported, op, err.Error()) + } + if name, twice := repeatedColumn(rec.Schema()); twice { + // Two columns of one name would reach BigQuery as one, and a value + // would be dropped without a word. + rec.Release() + return nil, serr.New(serr.InvalidArgument, op, fmt.Sprintf("the batch has two columns named %q (BigQuery names are not case-sensitive)", name)) + } + conflict = schemaConflict(rec.Schema(), table) + } + rows := int(rec.NumRows()) + if conflict != "" { + for r := 0; r < rows; r++ { + result.Errors = append(result.Errors, backend.RowError{ + RowIndex: next + int64(r), Reason: backend.RefusalSchemaMismatch, Error: conflict, + }) + } + next += int64(rows) + rec.Release() + continue + } + for r := 0; r < rows; r++ { + row, refusal := encodeRow(rec, r, next+int64(r)) + if refusal != nil { + result.Errors = append(result.Errors, *refusal) + continue + } + if len(request) == maxRowsPerInsert || reqBytes+row.size > maxRequestBytes { + if err := flush(); err != nil { + rec.Release() + return nil, err + } + } + request = append(request, row) + reqBytes += row.size + } + next += int64(rows) + rec.Release() + } + if err := flush(); err != nil { + return nil, err + } + return &result, nil +} + +// encodeRow converts row r of rec to what a request carries. A row that cannot +// be sent comes back as the refusal to report instead. +func encodeRow(rec arrow.RecordBatch, r int, index int64) (insertRow, *backend.RowError) { + values := make(map[string]bq.Value, rec.NumCols()) + for c, field := range rec.Schema().Fields() { + v, err := jsonValue(field.Name, rec.Column(c), r) + if err != nil { + text := "the row has a value BigQuery cannot be given" + if ve, ok := asValueError(err); ok { + text = "the row has a value BigQuery cannot be given: " + ve.Error() + } + return insertRow{}, &backend.RowError{RowIndex: index, Reason: backend.RefusalInvalidValue, Error: text} + } + values[field.Name] = v + } + encoded, err := json.Marshal(values) + if err != nil { + return insertRow{}, &backend.RowError{ + RowIndex: index, Reason: backend.RefusalInvalidValue, Error: "the row's values cannot be encoded", + } + } + size := len(encoded) + rowEnvelopeBytes + if size > maxRequestBytes { + return insertRow{}, &backend.RowError{ + RowIndex: index, Reason: backend.RefusalRowTooLarge, + Error: fmt.Sprintf("the row is %d bytes and one request carries at most %d", size, maxRequestBytes), + } + } + return insertRow{index: index, values: values, size: size}, nil +} + +// send inserts one request and classifies the answer. It returns how many rows +// were stored and the rows BigQuery permanently refused; any other failure is +// the error, and nothing about it is said of a row. +func send(ctx context.Context, inserter rowInserter, request []insertRow) (int, []backend.RowError, error) { + savers := make([]bq.ValueSaver, len(request)) + for i, row := range request { + savers[i] = row + } + err := inserter.Put(ctx, savers) + if err == nil { + return len(request), nil, nil + } + var multi bq.PutMultiError + if !errors.As(err, &multi) || len(multi) == 0 { + return 0, nil, classify("InsertRows", err) + } + var refused []backend.RowError + for _, rowErr := range multi { + if rowErr.RowIndex < 0 || rowErr.RowIndex >= len(request) { + return 0, nil, serr.New(serr.Internal, "InsertRows", "BigQuery refused a row that is not in the request") + } + if !rowIsInvalid(rowErr) { + // A reason that says nothing against the row: the call failed, not the row. + return 0, nil, classify("InsertRows", rowErr.Errors) + } + refused = append(refused, backend.RowError{ + RowIndex: request[rowErr.RowIndex].index, + Reason: backend.RefusalInvalidValue, + Error: refusalText(rowErr), + }) + } + return len(request) - len(refused), refused, nil +} + +// rowIsInvalid reports whether BigQuery gave the row the reason "invalid". +func rowIsInvalid(row bq.RowInsertionError) bool { + for _, rowErr := range row.Errors { + var byPointer *bq.Error + var byValue bq.Error + switch { + case errors.As(rowErr, &byPointer): + if byPointer.Reason == reasonInvalid { + return true + } + case errors.As(rowErr, &byValue): + if byValue.Reason == reasonInvalid { + return true + } + } + } + return false +} + +// refusalText describes a refused row in words of this server's: the column +// BigQuery located the refusal at, never BigQuery's own message. +func refusalText(row bq.RowInsertionError) string { + for _, rowErr := range row.Errors { + var byPointer *bq.Error + var byValue bq.Error + location := "" + switch { + case errors.As(rowErr, &byPointer): + location = byPointer.Location + case errors.As(rowErr, &byValue): + location = byValue.Location + } + if location != "" { + return fmt.Sprintf("the warehouse refused the value of column %q", location) + } + } + return "the warehouse refused the row as invalid" +} + +// repeatedColumn reports the first column name the schema uses twice, comparing +// names the way BigQuery does. +func repeatedColumn(schema *arrow.Schema) (string, bool) { + seen := make(map[string]bool, schema.NumFields()) + for _, f := range schema.Fields() { + key := strings.ToLower(f.Name) + if seen[key] { + return f.Name, true + } + seen[key] = true + } + return "", false +} + +// schemaConflict says how a batch's columns disagree with the table's, or "" when +// they do not. It checks the names (BigQuery's are case-insensitive) and which +// columns are required, recursing into records; the types of the values are +// BigQuery's to judge, row by row. +func schemaConflict(batch *arrow.Schema, table bq.Schema) string { + return fieldsConflict("", batch.Fields(), table) +} + +func fieldsConflict(path string, batch []arrow.Field, table bq.Schema) string { + known := make(map[string]*bq.FieldSchema, len(table)) + for _, f := range table { + known[strings.ToLower(f.Name)] = f + } + carried := make(map[string]bool, len(batch)) + for _, f := range batch { + name := strings.ToLower(f.Name) + carried[name] = true + column, ok := known[name] + if !ok { + return fmt.Sprintf("the table has no column %q", path+f.Name) + } + if st, ok := recordFields(f.Type); ok && fieldType(column.Type) == bq.RecordFieldType { + if conflict := fieldsConflict(path+f.Name+".", st, column.Schema); conflict != "" { + return conflict + } + } + } + for _, f := range table { + if f.Required && !f.Repeated && strings.TrimSpace(f.DefaultValueExpression) == "" && !carried[strings.ToLower(f.Name)] { + return fmt.Sprintf("the table requires column %q and the batch has none", path+f.Name) + } + } + return "" +} + +// recordFields is the fields of a struct, or of a list of structs, which is how +// a repeated record arrives. +func recordFields(t arrow.DataType) ([]arrow.Field, bool) { + switch t := t.(type) { + case *arrow.StructType: + return t.Fields(), true + case *arrow.ListType: + return recordFields(t.Elem()) + case *arrow.LargeListType: + return recordFields(t.Elem()) + default: + return nil, false + } +} diff --git a/internal/backend/bigquery/insert_test.go b/internal/backend/bigquery/insert_test.go new file mode 100644 index 0000000..087e39c --- /dev/null +++ b/internal/backend/bigquery/insert_test.go @@ -0,0 +1,446 @@ +package bigquery + +import ( + "context" + "io" + "net/http" + "strings" + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/stretchr/testify/require" + + "github.com/codefly-dev/service-warehouse/internal/backend" + "github.com/codefly-dev/service-warehouse/internal/serr" +) + +// sliceReader hands a fixed list of payloads to InsertRows. +type sliceReader struct { + payloads [][]byte + next int +} + +func (r *sliceReader) Next() ([]byte, error) { + if r.next >= len(r.payloads) { + return nil, io.EOF + } + p := r.payloads[r.next] + r.next++ + return p, nil +} +func (r *sliceReader) Close() error { return nil } + +// batchOf builds an Arrow record batch and its framing: the header's schema +// message and the batch message. +func batchOf(t *testing.T, schema *arrow.Schema, fill func(b *array.RecordBuilder)) (schemaMsg, batch []byte) { + t.Helper() + mem := memory.NewGoAllocator() + builder := array.NewRecordBuilder(mem, schema) + defer builder.Release() + fill(builder) + rec := builder.NewRecordBatch() + defer rec.Release() + fr, err := newFramer(schema, mem) + require.NoError(t, err) + batch, err = fr.batch(rec) + require.NoError(t, err) + return fr.schemaMsg, batch +} + +var eventSchema = arrow.NewSchema([]arrow.Field{ + {Name: "id", Type: arrow.PrimitiveTypes.Int64}, + {Name: "name", Type: arrow.BinaryTypes.String, Nullable: true}, +}, nil) + +func eventBatch(t *testing.T, ids []int64, names []string) (schemaMsg, batch []byte) { + return batchOf(t, eventSchema, func(b *array.RecordBuilder) { + for i, id := range ids { + b.Field(0).(*array.Int64Builder).Append(id) + b.Field(1).(*array.StringBuilder).Append(names[i]) + } + }) +} + +// tableMetadata is what the fake answers a metadata read of the table with. +func tableMetadata(fields ...map[string]any) map[string]any { + return map[string]any{ + "kind": "bigquery#table", "type": "TABLE", + "tableReference": map[string]any{"projectId": "test-project", "datasetId": "d", "tableId": "t"}, + "schema": map[string]any{"fields": fields}, + } +} + +func field(name, typ, mode string) map[string]any { + return map[string]any{"name": name, "type": typ, "mode": mode} +} + +var eventTable = tableMetadata(field("id", "INTEGER", "REQUIRED"), field("name", "STRING", "NULLABLE")) + +const insertPath = "/projects/test-project/datasets/d/tables/t/insertAll" + +func insertRows(t *testing.T, b *Backend, header []byte, payloads ...[]byte) (*backend.InsertResult, error) { + t.Helper() + return b.InsertRows(context.Background(), backend.TableRef{Dataset: "d", Table: "t"}, header, &sliceReader{payloads: payloads}) +} + +func TestInsertRowsStoresRowsAndAsksBigQueryToSkipInvalidOnes(t *testing.T) { + f := newFake(t) + f.handle("GET /projects/test-project/datasets/d/tables/t", 200, eventTable) + f.handle("POST "+insertPath, 200, map[string]any{"kind": "bigquery#tableDataInsertAllResponse"}) + b := f.open(t, backend.Config{}) + + header, batch := eventBatch(t, []int64{1, 2, 3}, []string{"a", "b", "c"}) + res, err := insertRows(t, b, header, batch) + require.NoError(t, err) + require.EqualValues(t, 3, res.RowsInserted) + require.Empty(t, res.Errors) + + sent := f.seen("POST", "insertAll") + require.Len(t, sent, 1) + require.Equal(t, true, sent[0].Body["skipInvalidRows"]) + require.NotEqual(t, true, sent[0].Body["ignoreUnknownValues"], "an unknown value is a refusal, not something to drop") + rows := sent[0].Body["rows"].([]any) + require.Len(t, rows, 3) + first := rows[0].(map[string]any)["json"].(map[string]any) + require.Equal(t, map[string]any{"id": float64(1), "name": "a"}, first) +} + +func TestInsertRowsReportsTheRowsBigQueryRefusedByIndexAndKeepsTheRest(t *testing.T) { + f := newFake(t) + f.handle("GET /projects/test-project/datasets/d/tables/t", 200, eventTable) + f.handle("POST "+insertPath, 200, map[string]any{ + "kind": "bigquery#tableDataInsertAllResponse", + "insertErrors": []map[string]any{{ + "index": 1, + "errors": []map[string]any{{ + "reason": "invalid", "location": "id", + "message": "VENDOR TEXT must never reach a client: bad value", "debugInfo": "secret", + }}, + }}, + }) + b := f.open(t, backend.Config{}) + + header, batch := eventBatch(t, []int64{1, 2, 3}, []string{"a", "b", "c"}) + res, err := insertRows(t, b, header, batch) + require.NoError(t, err) + require.EqualValues(t, 2, res.RowsInserted) + require.Len(t, res.Errors, 1) + require.EqualValues(t, 1, res.Errors[0].RowIndex) + require.Equal(t, backend.RefusalInvalidValue, res.Errors[0].Reason) + require.Contains(t, res.Errors[0].Error, `"id"`) + require.NotContains(t, res.Errors[0].Error, "VENDOR") + require.NotContains(t, res.Errors[0].Error, "secret") +} + +func TestInsertRowsFailsTheCallWhenARowFailedForAReasonThatSaysNothingAgainstIt(t *testing.T) { + quiet(t) + for reason, want := range map[string]serr.Code{ + "backendError": serr.Internal, + "timeout": serr.DeadlineExceeded, + "rateLimitExceeded": serr.Throttled, + "quotaExceeded": serr.Throttled, + "stopped": serr.Internal, + } { + t.Run(reason, func(t *testing.T) { + f := newFake(t) + f.handle("GET /projects/test-project/datasets/d/tables/t", 200, eventTable) + f.handle("POST "+insertPath, 200, map[string]any{ + "insertErrors": []map[string]any{{"index": 0, "errors": []map[string]any{{"reason": reason, "message": "VENDOR"}}}}, + }) + b := f.open(t, backend.Config{}) + + header, batch := eventBatch(t, []int64{1, 2}, []string{"a", "b"}) + res, err := insertRows(t, b, header, batch) + require.Nil(t, res, "a failure that is not the row's is never reported as a refusal of it") + require.True(t, serr.Is(err, want), "got %v", err) + require.NotContains(t, err.Error(), "VENDOR") + }) + } +} + +func TestInsertRowsFailsTheCallWhenTheRequestFails(t *testing.T) { + quiet(t) + for name, tc := range map[string]struct { + status int + reason string + want serr.Code + }{ + // The client library retries a backendError, a rateLimitExceeded and any + // 5xx until the call's context ends, so those are classified in errors_test.go. + "quota": {403, "quotaExceeded", serr.Throttled}, + "denied": {403, "accessDenied", serr.PermissionDenied}, + "too large": {400, "invalid", serr.InvalidArgument}, + "missing": {404, "notFound", serr.NotFound}, + } { + t.Run(name, func(t *testing.T) { + f := newFake(t) + f.handle("GET /projects/test-project/datasets/d/tables/t", 200, eventTable) + f.handle("POST "+insertPath, tc.status, apiError(tc.status, tc.reason, "VENDOR MESSAGE")) + b := f.open(t, backend.Config{}) + + header, batch := eventBatch(t, []int64{1}, []string{"a"}) + _, err := insertRows(t, b, header, batch) + require.True(t, serr.Is(err, tc.want), "got %v", err) + }) + } +} + +func TestInsertRowsRefusesEveryRowWhenTheBatchDisagreesWithTheTable(t *testing.T) { + for name, table := range map[string]map[string]any{ + "a column the table lacks": tableMetadata(field("id", "INTEGER", "REQUIRED")), + "a required column the batch lacks": tableMetadata( + field("id", "INTEGER", "REQUIRED"), field("name", "STRING", "NULLABLE"), field("region", "STRING", "REQUIRED")), + } { + t.Run(name, func(t *testing.T) { + f := newFake(t) + f.handle("GET /projects/test-project/datasets/d/tables/t", 200, table) + f.handle("POST "+insertPath, 200, map[string]any{}) + b := f.open(t, backend.Config{}) + + header, batch := eventBatch(t, []int64{1, 2}, []string{"a", "b"}) + res, err := insertRows(t, b, header, batch) + require.NoError(t, err) + require.Zero(t, res.RowsInserted) + require.Len(t, res.Errors, 2) + for i, e := range res.Errors { + require.EqualValues(t, i, e.RowIndex) + require.Equal(t, backend.RefusalSchemaMismatch, e.Reason, "the table, not the row, is what disagrees") + } + require.Empty(t, f.seen("POST", "insertAll"), "nothing is sent for a batch the table cannot take") + }) + } +} + +func TestInsertRowsMatchesColumnNamesWithoutRegardToCase(t *testing.T) { + f := newFake(t) + f.handle("GET /projects/test-project/datasets/d/tables/t", 200, + tableMetadata(field("ID", "INTEGER", "REQUIRED"), field("Name", "STRING", "NULLABLE"))) + f.handle("POST "+insertPath, 200, map[string]any{}) + b := f.open(t, backend.Config{}) + + header, batch := eventBatch(t, []int64{1}, []string{"a"}) + res, err := insertRows(t, b, header, batch) + require.NoError(t, err) + require.EqualValues(t, 1, res.RowsInserted) +} + +func TestInsertRowsSplitsRequestsAndKeepsRowIndexesAcrossThem(t *testing.T) { + f := newFake(t) + f.handle("GET /projects/test-project/datasets/d/tables/t", 200, eventTable) + f.handle("POST "+insertPath, 200, map[string]any{}) + b := f.open(t, backend.Config{}) + + ids := make([]int64, maxRowsPerInsert+1) + names := make([]string, len(ids)) + for i := range ids { + ids[i], names[i] = int64(i), "n" + } + header, first := eventBatch(t, ids[:300], names[:300]) + _, second := eventBatch(t, ids[300:], names[300:]) + res, err := insertRows(t, b, header, first, second) + require.NoError(t, err) + require.EqualValues(t, maxRowsPerInsert+1, res.RowsInserted) + + sent := f.seen("POST", "insertAll") + require.Len(t, sent, 2, "501 rows are two requests") + require.Len(t, sent[0].Body["rows"], maxRowsPerInsert) + require.Len(t, sent[1].Body["rows"], 1) +} + +func TestInsertRowsIndexesRefusalsByTheirPlaceInTheWholeStream(t *testing.T) { + f := newFake(t) + f.handle("GET /projects/test-project/datasets/d/tables/t", 200, eventTable) + calls := 0 + f.handleFunc("POST "+insertPath, func(w http.ResponseWriter, r *http.Request) { + calls++ + resp := map[string]any{} + if calls == 2 { // the request that holds the stream's second half + resp["insertErrors"] = []map[string]any{{"index": 0, "errors": []map[string]any{{"reason": "invalid"}}}} + } + writeJSON(w, 200, resp) + }) + b := f.open(t, backend.Config{}) + + ids := make([]int64, maxRowsPerInsert+2) + names := make([]string, len(ids)) + for i := range ids { + ids[i], names[i] = int64(i), "n" + } + header, batch := eventBatch(t, ids, names) + res, err := insertRows(t, b, header, batch) + require.NoError(t, err) + require.Len(t, res.Errors, 1) + require.EqualValues(t, maxRowsPerInsert, res.Errors[0].RowIndex, "the first row of the second request is row 500 of the stream") + require.EqualValues(t, maxRowsPerInsert+1, res.RowsInserted) +} + +func TestInsertRowsRefusesARowTooLargeToSendAndStoresTheOthers(t *testing.T) { + f := newFake(t) + f.handle("GET /projects/test-project/datasets/d/tables/t", 200, eventTable) + f.handle("POST "+insertPath, 200, map[string]any{}) + b := f.open(t, backend.Config{}) + + huge := strings.Repeat("x", maxRequestBytes) + header, batch := eventBatch(t, []int64{1, 2, 3}, []string{"a", huge, "c"}) + res, err := insertRows(t, b, header, batch) + require.NoError(t, err) + require.EqualValues(t, 2, res.RowsInserted) + require.Len(t, res.Errors, 1) + require.EqualValues(t, 1, res.Errors[0].RowIndex) + require.Equal(t, backend.RefusalRowTooLarge, res.Errors[0].Reason) + + var sentIDs []float64 + for _, req := range f.seen("POST", "insertAll") { + for _, row := range req.Body["rows"].([]any) { + sentIDs = append(sentIDs, row.(map[string]any)["json"].(map[string]any)["id"].(float64)) + } + } + require.Equal(t, []float64{1, 3}, sentIDs) +} + +func TestInsertRowsSplitsARequestThatWouldPassTheByteLimit(t *testing.T) { + f := newFake(t) + f.handle("GET /projects/test-project/datasets/d/tables/t", 200, eventTable) + f.handle("POST "+insertPath, 200, map[string]any{}) + b := f.open(t, backend.Config{}) + + big := strings.Repeat("x", maxRequestBytes/2) + header, batch := eventBatch(t, []int64{1, 2, 3}, []string{big, big, big}) + res, err := insertRows(t, b, header, batch) + require.NoError(t, err) + require.EqualValues(t, 3, res.RowsInserted) + require.Len(t, f.seen("POST", "insertAll"), 3, "two such rows already pass the limit with their envelope") +} + +func TestInsertRowsRefusesAValueBigQueryCannotBeGiven(t *testing.T) { + f := newFake(t) + f.handle("GET /projects/test-project/datasets/d/tables/t", 200, eventTable) + f.handle("POST "+insertPath, 200, map[string]any{}) + b := f.open(t, backend.Config{}) + + header, batch := eventBatch(t, []int64{1, 2}, []string{"fine", "bad \xff text"}) + res, err := insertRows(t, b, header, batch) + require.NoError(t, err) + require.EqualValues(t, 1, res.RowsInserted) + require.Len(t, res.Errors, 1) + require.EqualValues(t, 1, res.Errors[0].RowIndex) + require.Equal(t, backend.RefusalInvalidValue, res.Errors[0].Reason) +} + +func TestInsertRowsReadsSelfDescribingBatchesWhenTheHeaderHasNoSchema(t *testing.T) { + f := newFake(t) + f.handle("GET /projects/test-project/datasets/d/tables/t", 200, eventTable) + f.handle("POST "+insertPath, 200, map[string]any{}) + b := f.open(t, backend.Config{}) + + header, batch := eventBatch(t, []int64{1, 2}, []string{"a", "b"}) + res, err := insertRows(t, b, nil, append(append([]byte{}, header...), batch...)) + require.NoError(t, err) + require.EqualValues(t, 2, res.RowsInserted) + + _, err = insertRows(t, b, nil, batch) + require.True(t, serr.Is(err, serr.InvalidArgument), "a batch with no schema anywhere cannot be read: %v", err) +} + +func TestInsertRowsAcceptsAHeaderSchemaThatEndsTheStream(t *testing.T) { + f := newFake(t) + f.handle("GET /projects/test-project/datasets/d/tables/t", 200, eventTable) + f.handle("POST "+insertPath, 200, map[string]any{}) + b := f.open(t, backend.Config{}) + + header, batch := eventBatch(t, []int64{1}, []string{"a"}) + res, err := insertRows(t, b, append(append([]byte{}, header...), eos...), batch) + require.NoError(t, err) + require.EqualValues(t, 1, res.RowsInserted) +} + +func TestInsertRowsRefusesMalformedAndMismatchedArrow(t *testing.T) { + f := newFake(t) + f.handle("GET /projects/test-project/datasets/d/tables/t", 200, eventTable) + f.handle("POST "+insertPath, 200, map[string]any{}) + b := f.open(t, backend.Config{}) + header, batch := eventBatch(t, []int64{1}, []string{"a"}) + + _, err := insertRows(t, b, []byte("not an arrow schema"), batch) + require.True(t, serr.Is(err, serr.InvalidArgument), "got %v", err) + + _, err = insertRows(t, b, header, []byte{0xff, 0xff, 0xff, 0xff, 0x10, 0, 0, 0, 1, 2, 3}) + require.True(t, serr.Is(err, serr.InvalidArgument), "got %v", err) + + other := arrow.NewSchema([]arrow.Field{{Name: "id", Type: arrow.PrimitiveTypes.Int64}}, nil) + otherHeader, otherBatch := batchOf(t, other, func(rb *array.RecordBuilder) { rb.Field(0).(*array.Int64Builder).Append(1) }) + _, err = insertRows(t, b, nil, append(append([]byte{}, header...), batch...), append(append([]byte{}, otherHeader...), otherBatch...)) + require.True(t, serr.Is(err, serr.InvalidArgument), "batches that disagree on their schema: %v", err) +} + +func TestInsertRowsIsUnsupportedForAnArrowTypeBigQueryHasNoEquivalentOf(t *testing.T) { + f := newFake(t) + f.handle("GET /projects/test-project/datasets/d/tables/t", 200, eventTable) + f.handle("POST "+insertPath, 200, map[string]any{}) + b := f.open(t, backend.Config{}) + + mapType := arrow.MapOf(arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int64) + schema := arrow.NewSchema([]arrow.Field{{Name: "id", Type: arrow.PrimitiveTypes.Int64}, {Name: "attrs", Type: mapType}}, nil) + header, batch := batchOf(t, schema, func(rb *array.RecordBuilder) { + rb.Field(0).(*array.Int64Builder).Append(1) + rb.Field(1).AppendNull() + }) + _, err := insertRows(t, b, header, batch) + require.True(t, serr.Is(err, serr.Unsupported), "got %v", err) + require.Empty(t, f.seen("POST", "insertAll")) +} + +func TestInsertRowsReportsAMissingTableAsNotFound(t *testing.T) { + quiet(t) + f := newFake(t) + f.handle("GET /projects/test-project/datasets/d/tables/t", 404, apiError(404, "notFound", "Not found: Table test-project:d.t")) + b := f.open(t, backend.Config{}) + + header, batch := eventBatch(t, []int64{1}, []string{"a"}) + _, err := insertRows(t, b, header, batch) + require.True(t, serr.Is(err, serr.NotFound), "got %v", err) + require.NotContains(t, err.Error(), "test-project", "the server's project is not the client's to read") +} + +func TestInsertRowsUsesTheDefaultDatasetAndRefusesBadNames(t *testing.T) { + f := newFake(t) + f.handle("GET /projects/test-project/datasets/dflt/tables/t", 200, eventTable) + f.handle("POST /projects/test-project/datasets/dflt/tables/t/insertAll", 200, map[string]any{}) + b := f.open(t, backend.Config{DefaultDataset: "dflt"}) + + header, batch := eventBatch(t, []int64{1}, []string{"a"}) + res, err := b.InsertRows(context.Background(), backend.TableRef{Table: "t"}, header, &sliceReader{payloads: [][]byte{batch}}) + require.NoError(t, err) + require.EqualValues(t, 1, res.RowsInserted) + + noDefault := newFake(t).open(t, backend.Config{}) + for _, ref := range []backend.TableRef{ + {Table: "t"}, // no dataset anywhere + {Dataset: "d/tables/x", Table: "t"}, + {Dataset: "d", Table: "t/../x"}, + {Dataset: "d", Table: ""}, + } { + _, err := noDefault.InsertRows(context.Background(), ref, header, &sliceReader{payloads: [][]byte{batch}}) + require.True(t, serr.Is(err, serr.InvalidArgument), "%+v: %v", ref, err) + } +} + +func TestInsertRowsRefusesABatchWithTwoColumnsOfOneName(t *testing.T) { + f := newFake(t) + f.handle("GET /projects/test-project/datasets/d/tables/t", 200, eventTable) + f.handle("POST "+insertPath, 200, map[string]any{}) + b := f.open(t, backend.Config{}) + + twice := arrow.NewSchema([]arrow.Field{ + {Name: "id", Type: arrow.PrimitiveTypes.Int64}, {Name: "ID", Type: arrow.PrimitiveTypes.Int64}, + }, nil) + header, batch := batchOf(t, twice, func(rb *array.RecordBuilder) { + rb.Field(0).(*array.Int64Builder).Append(1) + rb.Field(1).(*array.Int64Builder).Append(2) + }) + _, err := insertRows(t, b, header, batch) + require.True(t, serr.Is(err, serr.InvalidArgument), "got %v", err) + require.Empty(t, f.seen("POST", "insertAll"), "a value is never dropped silently") +} diff --git a/internal/backend/bigquery/integration_test.go b/internal/backend/bigquery/integration_test.go new file mode 100644 index 0000000..c4847b0 --- /dev/null +++ b/internal/backend/bigquery/integration_test.go @@ -0,0 +1,275 @@ +package bigquery + +import ( + "context" + "crypto/rand" + "encoding/hex" + "errors" + "io" + "math/big" + "os" + "testing" + "time" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/decimal128" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/stretchr/testify/require" + + "github.com/codefly-dev/service-warehouse/internal/backend" + "github.com/codefly-dev/service-warehouse/internal/serr" +) + +// This test runs the backend against a real BigQuery project. It is the only +// thing in this repository that can: everything else runs against a fake of the +// REST API, which can say what the backend sends and how it reads an answer, not +// what BigQuery does. +// +// It runs only when SWH_TEST_BIGQUERY_PROJECT names a project it may use, and is +// skipped otherwise. It creates one dataset of its own, named with a random +// suffix, puts everything it makes inside it, and deletes it when it ends; it +// never reads or changes anything else in the project. The project needs the +// permissions to create datasets, tables and query jobs, and to stream rows. +// +// SWH_TEST_BIGQUERY_PROJECT the project to use (required to run) +// SWH_TEST_BIGQUERY_LOCATION where to create the dataset (optional; the project's default otherwise) +// SWH_TEST_BIGQUERY_CREDENTIALS_FILE a service-account key file (optional; Application Default Credentials otherwise) +const ( + envProject = "SWH_TEST_BIGQUERY_PROJECT" + envLocation = "SWH_TEST_BIGQUERY_LOCATION" + envCredentials = "SWH_TEST_BIGQUERY_CREDENTIALS_FILE" +) + +func TestAgainstRealBigQuery(t *testing.T) { + project := os.Getenv(envProject) + if project == "" { + t.Skipf("skipped: set %s to a BigQuery project this test may create a scratch dataset in "+ + "(optionally %s and %s); no BigQuery was contacted", envProject, envLocation, envCredentials) + } + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute) + defer cancel() + + b, err := Open(ctx, backend.Config{ + Database: project, + Location: os.Getenv(envLocation), + CredentialsFile: os.Getenv(envCredentials), + }) + require.NoError(t, err) + defer func() { _ = b.Close() }() + + suffix := make([]byte, 6) + _, err = rand.Read(suffix) + require.NoError(t, err) + dataset := "swh_it_" + hex.EncodeToString(suffix) + t.Cleanup(func() { + // The dataset is this test's own, so it is dropped with everything in it. + cleanup, stop := context.WithTimeout(context.Background(), 2*time.Minute) + defer stop() + if err := b.DropDataset(cleanup, dataset, backend.DropDatasetOptions{Cascade: true, IfExists: true}); err != nil { + t.Errorf("could not delete the scratch dataset %s: %v", dataset, err) + } + }) + + t.Run("catalog", func(t *testing.T) { + created, err := b.CreateDataset(ctx, dataset, backend.CreateDatasetOptions{Description: "service-warehouse integration test"}) + require.NoError(t, err) + require.Equal(t, dataset, created.Name) + require.NotEmpty(t, created.Location) + + _, err = b.CreateDataset(ctx, dataset, backend.CreateDatasetOptions{}) + require.True(t, serr.Is(err, serr.AlreadyExists), "got %v", err) + again, err := b.CreateDataset(ctx, dataset, backend.CreateDatasetOptions{IfNotExists: true}) + require.NoError(t, err) + require.Equal(t, created.Name, again.Name) + + datasets, err := b.ListDatasets(ctx, backend.ListDatasetsOptions{Limit: 500}) + require.NoError(t, err) + names := map[string]bool{} + for _, d := range datasets.Datasets { + names[d.Name] = true + } + if datasets.NextPageToken == "" { + require.True(t, names[dataset], "the new dataset is listed") + } + + ref := backend.TableRef{Dataset: dataset, Table: "events"} + made, err := b.CreateTable(ctx, ref, backend.CreateTableOptions{ + Columns: []backend.Column{ + {Name: "id", Type: backend.TypeInt64}, + {Name: "name", Type: backend.TypeString, Nullable: true}, + {Name: "amount", Type: backend.TypeNumeric, Nullable: true, Precision: 10, Scale: 2}, + {Name: "at", Type: backend.TypeTimestampTZ, Nullable: true}, + {Name: "tags", Type: backend.TypeArray, Fields: []backend.Column{{Name: "item", Type: backend.TypeString}}}, + }, + PartitionBy: []string{"at"}, + ClusterBy: []string{"id"}, + }) + require.NoError(t, err) + require.Equal(t, []string{"at"}, made.PartitionBy) + require.Equal(t, []string{"id"}, made.ClusterBy) + require.Equal(t, backend.TypeInt64, made.Columns[0].Type) + require.False(t, made.Columns[0].Nullable) + require.Equal(t, "NUMERIC(10, 2)", made.Columns[2].NativeType) + require.Equal(t, backend.TypeArray, made.Columns[4].Type) + + _, err = b.CreateTable(ctx, ref, backend.CreateTableOptions{Columns: []backend.Column{{Name: "id", Type: backend.TypeInt64}}}) + require.True(t, serr.Is(err, serr.AlreadyExists), "got %v", err) + + tables, err := b.ListTables(ctx, backend.ListTablesOptions{Dataset: dataset}) + require.NoError(t, err) + require.Len(t, tables.Tables, 1) + require.Equal(t, backend.KindTable, tables.Tables[0].Kind) + + err = b.DropDataset(ctx, dataset, backend.DropDatasetOptions{}) + require.True(t, serr.Is(err, serr.PreconditionFailed), "a dataset with a table needs cascade: %v", err) + + _, err = b.GetTable(ctx, backend.TableRef{Dataset: dataset, Table: "nope"}) + require.True(t, serr.Is(err, serr.NotFound), "got %v", err) + }) + + t.Run("insert and query", func(t *testing.T) { + ref := backend.TableRef{Dataset: dataset, Table: "events"} + schema := arrow.NewSchema([]arrow.Field{ + {Name: "id", Type: arrow.PrimitiveTypes.Int64}, + {Name: "name", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "amount", Type: &arrow.Decimal128Type{Precision: 38, Scale: 9}, Nullable: true}, + {Name: "at", Type: &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: "UTC"}, Nullable: true}, + {Name: "tags", Type: arrow.ListOf(arrow.BinaryTypes.String)}, + }, nil) + mem := memory.NewGoAllocator() + builder := array.NewRecordBuilder(mem, schema) + defer builder.Release() + // BigQuery streams only into recent partitions, so the rows are dated now. + at := time.Now().UTC().Truncate(time.Hour) + for i := int64(1); i <= 3; i++ { + builder.Field(0).(*array.Int64Builder).Append(i) + builder.Field(1).(*array.StringBuilder).Append("event " + string(rune('a'+i-1))) + // Amounts are NUMERIC(38, 9): 10.00 and 20.00 fit the table's NUMERIC(10, 2); the + // third, 99,999,999,999, has eleven digits before the point and does not. + amount := decimal128.FromI64(i * 10_000_000_000) + if i == 3 { + amount = decimal128.FromBigInt(new(big.Int).Mul(big.NewInt(99_999_999_999), big.NewInt(1_000_000_000))) + } + builder.Field(2).(*array.Decimal128Builder).Append(amount) + builder.Field(3).(*array.TimestampBuilder).Append(arrow.Timestamp(at.Add(time.Duration(i) * time.Hour).UnixMicro())) + tags := builder.Field(4).(*array.ListBuilder) + tags.Append(true) + tags.ValueBuilder().(*array.StringBuilder).Append("t" + string(rune('0'+i))) + } + rec := builder.NewRecordBatch() + defer rec.Release() + fr, err := newFramer(schema, mem) + require.NoError(t, err) + batch, err := fr.batch(rec) + require.NoError(t, err) + + res, err := b.InsertRows(ctx, ref, fr.schemaMsg, &sliceReader{payloads: [][]byte{batch}}) + require.NoError(t, err) + require.EqualValues(t, 2, res.RowsInserted) + require.Len(t, res.Errors, 1) + require.EqualValues(t, 2, res.Errors[0].RowIndex) + require.Equal(t, backend.RefusalInvalidValue, res.Errors[0].Reason, "a value that does not fit its column is refused for that row alone") + + // A batch the table cannot take at all is refused as a schema mismatch. + mismatch := arrow.NewSchema([]arrow.Field{{Name: "id", Type: arrow.PrimitiveTypes.Int64}, {Name: "no_such_column", Type: arrow.BinaryTypes.String}}, nil) + mb := array.NewRecordBuilder(mem, mismatch) + defer mb.Release() + mb.Field(0).(*array.Int64Builder).Append(9) + mb.Field(1).(*array.StringBuilder).Append("x") + mrec := mb.NewRecordBatch() + defer mrec.Release() + mfr, err := newFramer(mismatch, mem) + require.NoError(t, err) + mbatch, err := mfr.batch(mrec) + require.NoError(t, err) + res, err = b.InsertRows(ctx, ref, mfr.schemaMsg, &sliceReader{payloads: [][]byte{mbatch}}) + require.NoError(t, err) + require.Zero(t, res.RowsInserted) + require.Len(t, res.Errors, 1) + require.Equal(t, backend.RefusalSchemaMismatch, res.Errors[0].Reason) + + _, err = b.InsertRows(ctx, backend.TableRef{Dataset: dataset, Table: "nope"}, fr.schemaMsg, &sliceReader{payloads: [][]byte{batch}}) + require.True(t, serr.Is(err, serr.NotFound), "got %v", err) + + // Streamed rows become visible to queries within moments, not instantly. + var result *backend.QueryResult + require.Eventually(t, func() bool { + r, err := b.Query(ctx, "SELECT id, name, amount, at, tags FROM `"+project+"."+dataset+".events` WHERE id >= @min ORDER BY id", backend.QueryOptions{ + Params: []backend.QueryParam{{Name: "min", Type: backend.TypeInt64, Value: "1"}}, + }) + if err != nil { + t.Logf("query: %v", err) + return false + } + if r.Job.Stats.RowsProduced != 2 { + _ = r.Batches.Close() + return false + } + result = r + return true + }, 2*time.Minute, 5*time.Second, "the two accepted rows are readable") + defer result.Batches.Close() + + require.NotEmpty(t, result.Job.ID) + require.Equal(t, backend.JobDone, result.Job.State) + _, records := decode(t, result) + require.Len(t, records, 1) + require.EqualValues(t, 2, records[0].NumRows()) + require.Equal(t, int64(1), records[0].Column(0).(*array.Int64).Value(0)) + require.Equal(t, "event b", records[0].Column(1).(*array.String).Value(1)) + + job, err := b.GetJob(ctx, result.Job.ID) + require.NoError(t, err) + require.Equal(t, backend.JobDone, job.State) + require.Equal(t, result.Job.ID, job.ID) + + // Dry run prices without running. + dry, err := b.Query(ctx, "SELECT id FROM `"+project+"."+dataset+".events`", backend.QueryOptions{DryRun: true}) + require.NoError(t, err) + require.Nil(t, dry.Batches) + require.Len(t, dry.Schema, 1) + + // A query over the byte cap is refused before it runs. + _, err = b.Query(ctx, "SELECT id FROM `"+project+"."+dataset+".events`", backend.QueryOptions{MaxBytesBilled: 1}) + require.True(t, serr.Is(err, serr.Throttled), "got %v", err) + + _, err = b.Query(ctx, "SELEC 1", backend.QueryOptions{}) + require.True(t, serr.Is(err, serr.InvalidArgument), "got %v", err) + _, err = b.Query(ctx, "SELECT * FROM `"+project+"."+dataset+".missing`", backend.QueryOptions{}) + require.True(t, serr.Is(err, serr.NotFound), "got %v", err) + }) + + t.Run("create table as select", func(t *testing.T) { + ref := backend.TableRef{Dataset: dataset, Table: "numbers"} + made, err := b.CreateTable(ctx, ref, backend.CreateTableOptions{AsSelect: "SELECT n, CAST(n AS STRING) AS label FROM UNNEST([1, 2, 3]) AS n"}) + require.NoError(t, err) + require.Len(t, made.Columns, 2) + require.EqualValues(t, 3, drain(t, b, "SELECT COUNT(*) FROM `"+project+"."+dataset+".numbers`")) + + require.NoError(t, b.DropTable(ctx, ref, backend.DropTableOptions{})) + require.True(t, serr.Is(b.DropTable(ctx, ref, backend.DropTableOptions{}), serr.NotFound)) + require.NoError(t, b.DropTable(ctx, ref, backend.DropTableOptions{IfExists: true})) + }) +} + +// drain runs a single-value query and returns its first cell as an integer. +func drain(t *testing.T, b *Backend, sql string) int64 { + t.Helper() + res, err := b.Query(context.Background(), sql, backend.QueryOptions{}) + require.NoError(t, err) + defer res.Batches.Close() + var value int64 + found := false + for { + raw, err := res.Batches.Next() + if errors.Is(err, io.EOF) { + break + } + require.NoError(t, err) + rec := readBack(t, res.ArrowSchema, raw) + value, found = rec.Column(0).(*array.Int64).Value(0), true + } + require.True(t, found, "the query returned a row") + return value +} diff --git a/internal/backend/bigquery/params.go b/internal/backend/bigquery/params.go new file mode 100644 index 0000000..490a9ee --- /dev/null +++ b/internal/backend/bigquery/params.go @@ -0,0 +1,238 @@ +package bigquery + +import ( + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "math" + "regexp" + "strconv" + "strings" + "time" + "unicode/utf8" + + bq "cloud.google.com/go/bigquery" + "cloud.google.com/go/civil" + + "github.com/codefly-dev/service-warehouse/internal/backend" + "github.com/codefly-dev/service-warehouse/internal/serr" +) + +// Parameters arrive as text plus a portable type, and are sent to BigQuery as +// that text with an explicit parameter type, so BigQuery parses and binds the +// value and nothing is ever spliced into the SQL. The text is checked here first, +// for what is cheap to check, so a malformed value is an invalid argument naming +// the parameter instead of a failed job. +// +// The SQL itself uses BigQuery's named-parameter syntax, @name. A parameter's +// name is given without the sigil; one leading @, : or $ is accepted and dropped. + +// errNotBindable marks a parameter type that has no scalar text form. +var errNotBindable = errors.New("not bindable from text") + +var paramName = regexp.MustCompile(`^[A-Za-z_][A-Za-z0-9_]*$`) + +// decimalText is a plain decimal number: a sign, digits, and optional fraction. +var decimalText = regexp.MustCompile(`^([+-]?)([0-9]+)(?:\.([0-9]+))?$`) + +// Layouts for text that is parsed here and then sent in BigQuery's own form. +const ( + timestampParamLayout = "2006-01-02 15:04:05.999999-07:00" + maxNumericInt = numericPrecision - numericScale + maxBigNumericInt = bigNumericPrecision - bigNumericScale +) + +// bindParams builds the query parameters of a request. +func bindParams(params []backend.QueryParam) ([]bq.QueryParameter, error) { + out := make([]bq.QueryParameter, 0, len(params)) + seen := make(map[string]bool, len(params)) + for _, p := range params { + name := p.Name + if name != "" && strings.ContainsAny(name[:1], "@:$") { + name = name[1:] + } + if !paramName.MatchString(name) { + return nil, serr.New(serr.InvalidArgument, "Query", + fmt.Sprintf("%q is not a valid parameter name (letters, digits and underscores, not starting with a digit)", p.Name)) + } + // BigQuery matches names case-insensitively. + key := strings.ToLower(name) + if seen[key] { + return nil, serr.New(serr.InvalidArgument, "Query", fmt.Sprintf("parameter %q is given twice", name)) + } + seen[key] = true + kind, text, err := paramValue(p) + if errors.Is(err, errNotBindable) { + return nil, serr.New(serr.Unsupported, "Query", + fmt.Sprintf("parameter %q: an array or struct cannot be bound from a text value", name)) + } + if err != nil { + return nil, serr.New(serr.InvalidArgument, "Query", fmt.Sprintf("parameter %q: %v", name, err)) + } + value := &bq.QueryParameterValue{Type: bq.StandardSQLDataType{TypeKind: kind}} + if p.IsNull { + // A typed NULL: the type is declared, the value is absent. + value.Value = bq.NullString{} + } else { + value.Value = text + } + out = append(out, bq.QueryParameter{Name: name, Value: value}) + } + return out, nil +} + +// paramValue returns BigQuery's type name for the parameter and its text, +// normalized. A null parameter returns only the type. +func paramValue(p backend.QueryParam) (kind, text string, err error) { + switch p.Type { + case backend.TypeBool: + kind = "BOOL" + case backend.TypeInt64: + kind = "INT64" + case backend.TypeFloat64: + kind = "FLOAT64" + case backend.TypeNumeric: + kind = "NUMERIC" + case backend.TypeString: + kind = "STRING" + case backend.TypeBytes: + kind = "BYTES" + case backend.TypeDate: + kind = "DATE" + case backend.TypeTime: + kind = "TIME" + case backend.TypeTimestamp: + kind = "DATETIME" + case backend.TypeTimestampTZ: + kind = "TIMESTAMP" + case backend.TypeInterval: + kind = "INTERVAL" + case backend.TypeJSON: + kind = "JSON" + case backend.TypeGeography: + kind = "GEOGRAPHY" + case backend.TypeArray, backend.TypeStruct: + return "", "", errNotBindable + default: + return "", "", fmt.Errorf("has no type; give it one of the scalar types") + } + if p.IsNull { + return kind, "", nil + } + v := p.Value + switch p.Type { + case backend.TypeBool: + b, err := strconv.ParseBool(v) + if err != nil { + return "", "", fmt.Errorf("%q is not a boolean", v) + } + return kind, strconv.FormatBool(b), nil + case backend.TypeInt64: + n, err := strconv.ParseInt(v, 10, 64) + if err != nil { + return "", "", fmt.Errorf("%q is not a 64-bit integer", v) + } + return kind, strconv.FormatInt(n, 10), nil + case backend.TypeFloat64: + f, err := strconv.ParseFloat(v, 64) + if err != nil { + return "", "", fmt.Errorf("%q is not a floating-point number", v) + } + switch { + case math.IsNaN(f): + return kind, "NaN", nil + case math.IsInf(f, 1): + return kind, "Infinity", nil + case math.IsInf(f, -1): + return kind, "-Infinity", nil + } + return kind, strconv.FormatFloat(f, 'g', -1, 64), nil + case backend.TypeNumeric: + return decimalParam(v) + case backend.TypeString: + if !utf8.ValidString(v) { + return "", "", fmt.Errorf("the text is not valid UTF-8") + } + return kind, v, nil + case backend.TypeBytes: + raw, err := base64.StdEncoding.DecodeString(v) + if err != nil { + return "", "", fmt.Errorf("bytes are given as standard base64 text") + } + return kind, base64.StdEncoding.EncodeToString(raw), nil + case backend.TypeDate: + d, err := civil.ParseDate(v) + if err != nil { + return "", "", fmt.Errorf("%q is not a date (YYYY-MM-DD)", v) + } + return kind, d.String(), nil + case backend.TypeTime: + t, err := civil.ParseTime(v) + if err != nil { + return "", "", fmt.Errorf("%q is not a time (HH:MM:SS[.ffffff])", v) + } + if t.Nanosecond%1000 != 0 { + return "", "", fmt.Errorf("BigQuery times have microsecond precision") + } + return kind, bq.CivilTimeString(t), nil + case backend.TypeTimestamp: + // civil reads only the T separator; BigQuery's own text uses a space. + dt, err := civil.ParseDateTime(strings.Replace(v, " ", "T", 1)) + if err != nil { + return "", "", fmt.Errorf("%q is not a wall-clock timestamp (YYYY-MM-DDTHH:MM:SS[.ffffff])", v) + } + if dt.Time.Nanosecond%1000 != 0 { + return "", "", fmt.Errorf("BigQuery timestamps have microsecond precision") + } + return kind, bq.CivilDateTimeString(dt), nil + case backend.TypeTimestampTZ: + t, err := time.Parse(time.RFC3339Nano, v) + if err != nil { + t, err = time.Parse("2006-01-02 15:04:05.999999999Z07:00", v) + } + if err != nil { + return "", "", fmt.Errorf("%q is not an instant (RFC 3339, with a zone)", v) + } + if t.Nanosecond()%1000 != 0 { + return "", "", fmt.Errorf("BigQuery timestamps have microsecond precision") + } + return kind, t.UTC().Format(timestampParamLayout), nil + case backend.TypeInterval: + iv, err := bq.ParseInterval(v) + if err != nil { + return "", "", fmt.Errorf("%q is not an interval (Y-M D H:M:S[.F])", v) + } + return kind, iv.String(), nil + case backend.TypeJSON: + if !json.Valid([]byte(v)) { + return "", "", fmt.Errorf("the text is not valid JSON") + } + return kind, v, nil + case backend.TypeGeography: + if strings.TrimSpace(v) == "" { + return "", "", fmt.Errorf("the geography is empty") + } + return kind, v, nil + } + return "", "", fmt.Errorf("unreachable parameter type") +} + +// decimalParam picks NUMERIC when the value fits it and BIGNUMERIC when only that +// fits, and refuses a value that fits neither: BigQuery would round it, and a +// bound value is never rounded. +func decimalParam(v string) (string, string, error) { + m := decimalText.FindStringSubmatch(v) + if m == nil { + return "", "", fmt.Errorf("%q is not a decimal number", v) + } + whole := strings.TrimLeft(m[2], "0") + fraction := m[3] + switch { + case len(whole) <= maxNumericInt && len(fraction) <= numericScale: + return "NUMERIC", strings.TrimPrefix(v, "+"), nil + case len(whole) <= maxBigNumericInt && len(fraction) <= bigNumericScale: + return "BIGNUMERIC", strings.TrimPrefix(v, "+"), nil + } + return "", "", fmt.Errorf("%q does not fit NUMERIC or BIGNUMERIC without rounding", v) +} diff --git a/internal/backend/bigquery/params_test.go b/internal/backend/bigquery/params_test.go new file mode 100644 index 0000000..c6c640b --- /dev/null +++ b/internal/backend/bigquery/params_test.go @@ -0,0 +1,122 @@ +package bigquery + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/codefly-dev/service-warehouse/internal/backend" + "github.com/codefly-dev/service-warehouse/internal/serr" +) + +func TestParameterValuesAreCheckedAndNormalized(t *testing.T) { + for name, tc := range map[string]struct { + param backend.QueryParam + kind string + text string + }{ + "bool": {backend.QueryParam{Type: backend.TypeBool, Value: "TRUE"}, "BOOL", "true"}, + "bool short": {backend.QueryParam{Type: backend.TypeBool, Value: "f"}, "BOOL", "false"}, + "int": {backend.QueryParam{Type: backend.TypeInt64, Value: "-9223372036854775808"}, "INT64", "-9223372036854775808"}, + "float": {backend.QueryParam{Type: backend.TypeFloat64, Value: "1.50"}, "FLOAT64", "1.5"}, + "float exponent": {backend.QueryParam{Type: backend.TypeFloat64, Value: "1e3"}, "FLOAT64", "1000"}, + "float nan": {backend.QueryParam{Type: backend.TypeFloat64, Value: "NaN"}, "FLOAT64", "NaN"}, + "float +inf": {backend.QueryParam{Type: backend.TypeFloat64, Value: "+Inf"}, "FLOAT64", "Infinity"}, + "float -inf": {backend.QueryParam{Type: backend.TypeFloat64, Value: "-Inf"}, "FLOAT64", "-Infinity"}, + "numeric": {backend.QueryParam{Type: backend.TypeNumeric, Value: "+12.345"}, "NUMERIC", "12.345"}, + "numeric at the edge": {backend.QueryParam{Type: backend.TypeNumeric, Value: strings.Repeat("9", 29) + "." + strings.Repeat("9", 9)}, "NUMERIC", strings.Repeat("9", 29) + "." + strings.Repeat("9", 9)}, + "numeric leading 0s": {backend.QueryParam{Type: backend.TypeNumeric, Value: "000" + strings.Repeat("9", 29)}, "NUMERIC", "000" + strings.Repeat("9", 29)}, + "numeric needing big": {backend.QueryParam{Type: backend.TypeNumeric, Value: "1." + strings.Repeat("1", 20)}, "BIGNUMERIC", "1." + strings.Repeat("1", 20)}, + "numeric wide": {backend.QueryParam{Type: backend.TypeNumeric, Value: strings.Repeat("9", 38)}, "BIGNUMERIC", strings.Repeat("9", 38)}, + "string": {backend.QueryParam{Type: backend.TypeString, Value: "it's a ' \" -- ; DROP"}, "STRING", "it's a ' \" -- ; DROP"}, + "empty string": {backend.QueryParam{Type: backend.TypeString, Value: ""}, "STRING", ""}, + "bytes": {backend.QueryParam{Type: backend.TypeBytes, Value: "AAEC/w=="}, "BYTES", "AAEC/w=="}, + "date": {backend.QueryParam{Type: backend.TypeDate, Value: "2024-02-29"}, "DATE", "2024-02-29"}, + "time": {backend.QueryParam{Type: backend.TypeTime, Value: "13:04:05.5"}, "TIME", "13:04:05.500000"}, + "wall clock": {backend.QueryParam{Type: backend.TypeTimestamp, Value: "2024-02-29 13:04:05"}, "DATETIME", "2024-02-29 13:04:05"}, + "wall clock T": {backend.QueryParam{Type: backend.TypeTimestamp, Value: "2024-02-29T13:04:05.25"}, "DATETIME", "2024-02-29 13:04:05.250000"}, + "instant": {backend.QueryParam{Type: backend.TypeTimestampTZ, Value: "2024-02-29T13:04:05.5+02:00"}, "TIMESTAMP", "2024-02-29 11:04:05.5+00:00"}, + "instant spaced": {backend.QueryParam{Type: backend.TypeTimestampTZ, Value: "2024-02-29 13:04:05Z"}, "TIMESTAMP", "2024-02-29 13:04:05+00:00"}, + "interval": {backend.QueryParam{Type: backend.TypeInterval, Value: "1-2 3 4:5:6"}, "INTERVAL", "1-2 3 4:5:6"}, + "json": {backend.QueryParam{Type: backend.TypeJSON, Value: `{"a":[1,2]}`}, "JSON", `{"a":[1,2]}`}, + "geography": {backend.QueryParam{Type: backend.TypeGeography, Value: "POINT(1 2)"}, "GEOGRAPHY", "POINT(1 2)"}, + "null keeps its type": {backend.QueryParam{Type: backend.TypeDate, IsNull: true, Value: "ignored"}, "DATE", ""}, + } { + t.Run(name, func(t *testing.T) { + kind, text, err := paramValue(tc.param) + require.NoError(t, err) + require.Equal(t, tc.kind, kind) + require.Equal(t, tc.text, text) + }) + } +} + +func TestMalformedParameterValuesAreInvalidArguments(t *testing.T) { + for name, p := range map[string]backend.QueryParam{ + "bool": {Type: backend.TypeBool, Value: "yes"}, + "int": {Type: backend.TypeInt64, Value: "1.5"}, + "int overflow": {Type: backend.TypeInt64, Value: "9223372036854775808"}, + "float": {Type: backend.TypeFloat64, Value: "one"}, + "numeric": {Type: backend.TypeNumeric, Value: "1e5"}, + "numeric rounding": {Type: backend.TypeNumeric, Value: "1." + strings.Repeat("1", 39)}, + "numeric too big": {Type: backend.TypeNumeric, Value: strings.Repeat("9", 40)}, + "string not UTF-8": {Type: backend.TypeString, Value: "a\xffb"}, + "bytes": {Type: backend.TypeBytes, Value: "not base64!"}, + "date": {Type: backend.TypeDate, Value: "29/02/2024"}, + "time": {Type: backend.TypeTime, Value: "25:00:00"}, + "time nanoseconds": {Type: backend.TypeTime, Value: "13:04:05.123456789"}, + "wall clock": {Type: backend.TypeTimestamp, Value: "yesterday"}, + "instant no zone": {Type: backend.TypeTimestampTZ, Value: "2024-02-29T13:04:05"}, + "instant ns": {Type: backend.TypeTimestampTZ, Value: "2024-02-29T13:04:05.123456789Z"}, + "interval": {Type: backend.TypeInterval, Value: "a while"}, + "json": {Type: backend.TypeJSON, Value: "{not json"}, + "geography": {Type: backend.TypeGeography, Value: " "}, + "no type": {Value: "1"}, + } { + t.Run(name, func(t *testing.T) { + p.Name = "p" + _, err := bindParams([]backend.QueryParam{p}) + require.True(t, serr.Is(err, serr.InvalidArgument), "got %v", err) + require.Contains(t, err.Error(), `"p"`, "the error names the parameter") + }) + } +} + +func TestArrayAndStructParametersAreUnsupportedNotMisbound(t *testing.T) { + for _, typ := range []backend.ColumnType{backend.TypeArray, backend.TypeStruct} { + _, err := bindParams([]backend.QueryParam{{Name: "p", Type: typ, Value: "[1,2]"}}) + require.True(t, serr.Is(err, serr.Unsupported), "got %v", err) + } +} + +func TestParameterNames(t *testing.T) { + got, err := bindParams([]backend.QueryParam{ + {Name: "plain", Type: backend.TypeInt64, Value: "1"}, + {Name: "@at", Type: backend.TypeInt64, Value: "2"}, + {Name: ":colon", Type: backend.TypeInt64, Value: "3"}, + {Name: "$dollar", Type: backend.TypeInt64, Value: "4"}, + {Name: "_u1", Type: backend.TypeInt64, Value: "5"}, + }) + require.NoError(t, err) + names := make([]string, len(got)) + for i, p := range got { + names[i] = p.Name + } + require.Equal(t, []string{"plain", "at", "colon", "dollar", "_u1"}, names) + + for _, bad := range []string{"", "@", "1a", "a b", "a-b", "a;b", "@@x", "x'", "é"} { + _, err := bindParams([]backend.QueryParam{{Name: bad, Type: backend.TypeInt64, Value: "1"}}) + require.True(t, serr.Is(err, serr.InvalidArgument), "%q: %v", bad, err) + } + _, err = bindParams([]backend.QueryParam{ + {Name: "Dup", Type: backend.TypeInt64, Value: "1"}, {Name: "dup", Type: backend.TypeInt64, Value: "2"}, + }) + require.True(t, serr.Is(err, serr.InvalidArgument), "BigQuery matches names case-insensitively: %v", err) +} + +func TestNoParametersIsNoParameters(t *testing.T) { + got, err := bindParams(nil) + require.NoError(t, err) + require.Empty(t, got) +} diff --git a/internal/backend/bigquery/query.go b/internal/backend/bigquery/query.go new file mode 100644 index 0000000..07e8ab9 --- /dev/null +++ b/internal/backend/bigquery/query.go @@ -0,0 +1,297 @@ +package bigquery + +import ( + "context" + "errors" + "fmt" + "io" + "strings" + "time" + + bq "cloud.google.com/go/bigquery" + "google.golang.org/api/iterator" + + "github.com/codefly-dev/service-warehouse/internal/backend" + "github.com/codefly-dev/service-warehouse/internal/serr" +) + +// cancelGrace bounds the best-effort cancellation of a job whose caller gave up. +const cancelGrace = 5 * time.Second + +// Job ids on the wire are "." when BigQuery reported a location: a +// job outside the US and EU multi-regions can only be found again with it, and a +// client holds the id, not the server's configuration. Neither part can contain a +// dot, so the split is unambiguous. An id with no dot is a BigQuery job id as it +// is, looked up in the configured location. +const jobIDSeparator = "." + +func encodeJobID(location, id string) string { + if location == "" || id == "" { + return id + } + return location + jobIDSeparator + id +} + +func (b *Backend) parseJobID(op, wire string) (id, location string, err error) { + id, location = wire, b.cfg.Location + if before, after, found := strings.Cut(wire, jobIDSeparator); found { + location, id = before, after + } + if !jobID.MatchString(id) || (location != "" && !locationID.MatchString(location)) { + return "", "", serr.New(serr.InvalidArgument, op, fmt.Sprintf("%q is not a job id", wire)) + } + return id, location, nil +} + +// limit is the effective limit of a request: the caller's own, never above the +// server's when it has one. Zero on both sides is no limit. +func limit[T ~int64 | ~int](requested, ceiling T) T { + switch { + case requested > 0 && ceiling > 0: + return min(requested, ceiling) + case requested > 0: + return requested + default: + return ceiling + } +} + +// Query runs SQL and streams its result as Arrow batches. The job is complete +// before this returns, because the header carries the result's schema, which +// BigQuery knows only then; for the same reason CancelJob cannot reach a query +// that is still running, and a caller that gives up cancels the call instead, +// which also cancels the job. +// +// A dry run only validates and prices the statement; a DML or DDL statement +// returns a result with no schema and no batches. +func (b *Backend) Query(ctx context.Context, sql string, opts backend.QueryOptions) (*backend.QueryResult, error) { + const op = "Query" + if strings.TrimSpace(sql) == "" { + return nil, serr.New(serr.InvalidArgument, op, "the SQL is empty") + } + params, err := bindParams(opts.Params) + if err != nil { + return nil, err + } + dataset := opts.DefaultDataset + if dataset == "" { + dataset = b.cfg.DefaultDataset + } + if dataset != "" { + if err := checkDatasetID(op, dataset); err != nil { + return nil, err + } + } + + timeout := limit(opts.Timeout, b.cfg.QueryTimeout) + var cancel context.CancelFunc + if timeout > 0 { + ctx, cancel = context.WithTimeout(ctx, timeout) + } else { + ctx, cancel = context.WithCancel(ctx) + } + // The context outlives this call when rows are still to be read, and is then + // cancelled when the reader closes. + handedOver := false + defer func() { + if !handedOver { + cancel() + } + }() + + q := b.client.Query(sql) + q.Parameters = params + if dataset != "" { + q.DefaultProjectID, q.DefaultDatasetID = b.project, dataset + } + q.DryRun = opts.DryRun + q.MaxBytesBilled = limit(opts.MaxBytesBilled, b.cfg.MaxQueryBytes) + if timeout > 0 { + q.JobTimeout = timeout + } + + job, err := q.Run(ctx) + if err != nil { + return nil, classify(op, err) + } + if opts.DryRun { + return b.dryRun(job) + } + + status, err := job.Wait(ctx) + if err != nil { + cancelAbandoned(ctx, job) + return nil, classify(op, err) + } + if err := status.Err(); err != nil { + return nil, classify(op, err) + } + it, err := job.Read(ctx) + if err != nil { + cancelAbandoned(ctx, job) + return nil, classify(op, err) + } + + result := &backend.QueryResult{Job: jobOut(job, status)} + result.Job.Stats.RowsProduced = int64(it.TotalRows) + if len(it.Schema) == 0 { + return result, nil // DML or DDL: nothing to read. + } + result.Schema = columnsFromSchema(it.Schema) + enc, err := newResultEncoder(result.Schema) + if err != nil { + return nil, err + } + result.ArrowSchema = enc.schemaMessage() + result.Batches = &batchReader{ctx: ctx, it: it, enc: enc, job: job, cancel: cancel} + handedOver = true + return result, nil +} + +// dryRun reports what BigQuery says a statement would cost and return. +func (b *Backend) dryRun(job *bq.Job) (*backend.QueryResult, error) { + status := job.LastStatus() + result := &backend.QueryResult{Job: jobOut(job, status)} + result.Job.State = backend.JobDone + if status != nil && status.Statistics != nil { + if qs, ok := status.Statistics.Details.(*bq.QueryStatistics); ok && len(qs.Schema) > 0 { + result.Schema = columnsFromSchema(qs.Schema) + enc, err := newResultEncoder(result.Schema) + if err != nil { + return nil, err + } + defer enc.release() + result.ArrowSchema = enc.schemaMessage() + } + } + return result, nil +} + +// cancelAbandoned asks BigQuery to stop a job whose caller stopped waiting, so a +// timeout or a hung-up client does not leave it running, and billing, behind. +func cancelAbandoned(ctx context.Context, job *bq.Job) { + if ctx.Err() == nil { + return + } + // The call's own context is done; the cancellation needs one that is not. + fresh, stop := context.WithTimeout(context.WithoutCancel(ctx), cancelGrace) + defer stop() + _ = job.Cancel(fresh) +} + +// batchReader turns the rows of a finished query into Arrow batches, a batch at +// a time, so a large result is never held in memory. +type batchReader struct { + ctx context.Context + it *bq.RowIterator + enc *resultEncoder + job *bq.Job + cancel context.CancelFunc + done bool +} + +func (r *batchReader) Next() ([]byte, error) { + for !r.done && !r.enc.full() { + var row []bq.Value + err := r.it.Next(&row) + if errors.Is(err, iterator.Done) { + r.done = true + break + } + if err != nil { + cancelAbandoned(r.ctx, r.job) + return nil, classify("Query", err) + } + if err := r.enc.add(row); err != nil { + return nil, err + } + } + if !r.enc.pending() { + return nil, io.EOF + } + return r.enc.flush() +} + +func (r *batchReader) Close() error { + r.enc.release() + r.cancel() + return nil +} + +// GetJob reports a job's state and accounting. +func (b *Backend) GetJob(ctx context.Context, id string) (*backend.Job, error) { + const op = "GetJob" + jobID, location, err := b.parseJobID(op, id) + if err != nil { + return nil, err + } + job, err := b.client.JobFromIDLocation(ctx, jobID, location) + if err != nil { + return nil, classify(op, err) + } + out := jobOut(job, job.LastStatus()) + return &out, nil +} + +// CancelJob asks BigQuery to cancel a job. It returns once the request is +// accepted; the job reports JobCancelled when the cancellation has taken effect. +func (b *Backend) CancelJob(ctx context.Context, id string) error { + const op = "CancelJob" + jobID, location, err := b.parseJobID(op, id) + if err != nil { + return err + } + job, err := b.client.JobFromIDLocation(ctx, jobID, location) + if err != nil { + return classify(op, err) + } + return classify(op, job.Cancel(ctx)) +} + +// jobOut is the portable view of a job and its accounting. The failure of a job +// is reported in this server's words, as every other failure is. +func jobOut(job *bq.Job, status *bq.JobStatus) backend.Job { + out := backend.Job{ID: encodeJobID(job.Location(), job.ID()), State: backend.JobDone} + if status == nil { + return out + } + switch status.State { + case bq.Pending: + out.State = backend.JobPending + case bq.Running: + out.State = backend.JobRunning + } + if err := status.Err(); err != nil { + out.State = backend.JobFailed + if f, ok := inspect(err); ok && containsReason(f, "stopped") { + out.State = backend.JobCancelled + } + _, out.Error = normalize(err) + } + if stats := status.Statistics; stats != nil { + out.Created, out.Started, out.Ended = stats.CreationTime, stats.StartTime, stats.EndTime + out.Stats.BytesScanned = stats.TotalBytesProcessed + out.Stats.SlotMillis = stats.TotalSlotDuration.Milliseconds() + if qs, ok := stats.Details.(*bq.QueryStatistics); ok { + out.Stats.BytesBilled = qs.TotalBytesBilled + out.Stats.CacheHit = qs.CacheHit + out.Stats.RowsAffected = qs.NumDMLAffectedRows + if qs.SlotMillis > 0 { + out.Stats.SlotMillis = qs.SlotMillis + } + if qs.TotalBytesProcessed > 0 { + out.Stats.BytesScanned = qs.TotalBytesProcessed + } + } + } + return out +} + +func containsReason(f failure, reason string) bool { + for _, r := range f.reasons { + if r == reason { + return true + } + } + return false +} diff --git a/internal/backend/bigquery/query_test.go b/internal/backend/bigquery/query_test.go new file mode 100644 index 0000000..5652f27 --- /dev/null +++ b/internal/backend/bigquery/query_test.go @@ -0,0 +1,447 @@ +package bigquery + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "net/http" + "testing" + "time" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/ipc" + "github.com/stretchr/testify/require" + + "github.com/codefly-dev/service-warehouse/internal/backend" + "github.com/codefly-dev/service-warehouse/internal/serr" +) + +const ( + jobsPath = "/projects/test-project/jobs" + resultsPath = "/projects/test-project/queries/job_1" + jobPath = "/projects/test-project/jobs/job_1" +) + +// jobResource is what BigQuery answers a job insert or a job read with. +func jobResource(state string, statistics, extra map[string]any) map[string]any { + job := map[string]any{ + "kind": "bigquery#job", + "jobReference": map[string]any{"projectId": "test-project", "jobId": "job_1", "location": "EU"}, + "status": map[string]any{"state": state}, + "statistics": map[string]any{ + "creationTime": "1700000000000", "startTime": "1700000001000", "endTime": "1700000002000", + }, + "configuration": map[string]any{"query": map[string]any{"query": "SELECT 1", "useLegacySql": false}}, + } + for k, v := range statistics { + job["statistics"].(map[string]any)[k] = v + } + for k, v := range extra { + job[k] = v + } + return job +} + +func resultsResource(fields []map[string]any, rows ...[]any) map[string]any { + out := make([]map[string]any, len(rows)) + for i, row := range rows { + cells := make([]map[string]any, len(row)) + for j, v := range row { + cells[j] = map[string]any{"v": v} + } + out[i] = map[string]any{"f": cells} + } + return map[string]any{ + "kind": "bigquery#getQueryResultsResponse", "jobComplete": true, + "jobReference": map[string]any{"projectId": "test-project", "jobId": "job_1", "location": "EU"}, + "schema": map[string]any{"fields": fields}, + "totalRows": fmt.Sprint(len(rows)), + "rows": out, + } +} + +// queryFake serves one successful query job over the fake. +func queryFake(t *testing.T, results map[string]any, stats map[string]any) *fake { + t.Helper() + f := newFake(t) + f.handle("POST "+jobsPath, 200, jobResource("DONE", stats, nil)) + f.handle("GET "+jobPath, 200, jobResource("DONE", stats, nil)) + f.handle("GET "+resultsPath, 200, results) + return f +} + +// decode reads a result's batches back with a plain Arrow reader, exactly as a +// client does: the header's schema message followed by each batch message. +func decode(t *testing.T, res *backend.QueryResult) (schema *arrow.Schema, records []arrow.RecordBatch) { + t.Helper() + require.NotEmpty(t, res.ArrowSchema) + require.NotNil(t, res.Batches) + for { + batch, err := res.Batches.Next() + if errors.Is(err, io.EOF) { + break + } + require.NoError(t, err) + reader, err := ipc.NewReader(io.MultiReader(bytes.NewReader(res.ArrowSchema), bytes.NewReader(batch))) + require.NoError(t, err) + schema = reader.Schema() + for reader.Next() { + rec := reader.RecordBatch() + rec.Retain() + records = append(records, rec) + } + require.NoError(t, reader.Err()) + reader.Release() + } + t.Cleanup(func() { + for _, rec := range records { + rec.Release() + } + }) + return schema, records +} + +func TestQueryStreamsTheResultAsArrowAndBindsParameters(t *testing.T) { + fields := []map[string]any{ + field("id", "INTEGER", "NULLABLE"), + field("name", "STRING", "NULLABLE"), + field("score", "FLOAT", "NULLABLE"), + } + f := queryFake(t, + resultsResource(fields, []any{"1", "ada", "1.5"}, []any{"2", nil, "2.5"}), + map[string]any{"totalBytesProcessed": "2048", "query": map[string]any{ + "totalBytesBilled": "10485760", "cacheHit": true, "totalBytesProcessed": "2048", + }}) + b := f.open(t, backend.Config{Location: "EU", DefaultDataset: "dflt"}) + + res, err := b.Query(context.Background(), "SELECT id, name, score FROM t WHERE id > @min AND name = @who", backend.QueryOptions{ + Params: []backend.QueryParam{ + {Name: "min", Type: backend.TypeInt64, Value: "0"}, + {Name: "@who", Type: backend.TypeString, IsNull: true}, + }, + MaxBytesBilled: 5000, + Timeout: 90 * time.Second, + }) + require.NoError(t, err) + defer res.Batches.Close() + + require.Equal(t, "EU.job_1", res.Job.ID, "the job id carries the location it ran in") + require.Equal(t, backend.JobDone, res.Job.State) + require.EqualValues(t, 2048, res.Job.Stats.BytesScanned) + require.EqualValues(t, 10485760, res.Job.Stats.BytesBilled) + require.True(t, res.Job.Stats.CacheHit) + require.EqualValues(t, 2, res.Job.Stats.RowsProduced) + + require.Equal(t, []backend.Column{ + {Name: "id", Type: backend.TypeInt64, Nullable: true, NativeType: "INT64"}, + {Name: "name", Type: backend.TypeString, Nullable: true, NativeType: "STRING"}, + {Name: "score", Type: backend.TypeFloat64, Nullable: true, NativeType: "FLOAT64"}, + }, res.Schema) + + schema, records := decode(t, res) + require.Equal(t, "id", schema.Field(0).Name) + require.Len(t, records, 1) + require.EqualValues(t, 2, records[0].NumRows()) + require.Equal(t, int64(2), records[0].Column(0).(*array.Int64).Value(1)) + require.True(t, records[0].Column(1).IsNull(1)) + require.Equal(t, 1.5, records[0].Column(2).(*array.Float64).Value(0)) + + insert := f.seen("POST", "/jobs") + require.Len(t, insert, 1) + query := insert[0].Body["configuration"].(map[string]any)["query"].(map[string]any) + require.Equal(t, false, query["useLegacySql"]) + require.Equal(t, "5000", query["maximumBytesBilled"]) + require.Equal(t, "90000", insert[0].Body["configuration"].(map[string]any)["jobTimeoutMs"]) + require.Equal(t, map[string]any{"projectId": "test-project", "datasetId": "dflt"}, query["defaultDataset"]) + require.Equal(t, []any{ + map[string]any{ + "name": "min", "parameterType": map[string]any{"type": "INT64"}, + "parameterValue": map[string]any{"value": "0"}, + }, + map[string]any{ + "name": "who", "parameterType": map[string]any{"type": "STRING"}, + "parameterValue": map[string]any{"value": nil}, + }, + }, query["queryParameters"], "a null is a typed parameter whose value is null, not the empty string") + require.Equal(t, "EU", insert[0].Body["jobReference"].(map[string]any)["location"], "the job is placed in the configured location") +} + +func TestQuerySplitsALargeResultIntoBatchesThatEachDecode(t *testing.T) { + rows := make([][]any, 2*batchMaxRows+5) + for i := range rows { + rows[i] = []any{fmt.Sprint(i)} + } + f := queryFake(t, resultsResource([]map[string]any{field("id", "INTEGER", "NULLABLE")}, rows...), nil) + b := f.open(t, backend.Config{}) + + res, err := b.Query(context.Background(), "SELECT id FROM t", backend.QueryOptions{}) + require.NoError(t, err) + defer res.Batches.Close() + _, records := decode(t, res) + require.Len(t, records, 3) + total := int64(0) + for _, rec := range records { + total += rec.NumRows() + } + require.EqualValues(t, len(rows), total) + require.EqualValues(t, 2*batchMaxRows, records[2].Column(0).(*array.Int64).Value(0)) + _, err = res.Batches.Next() + require.ErrorIs(t, err, io.EOF, "a finished result keeps answering end of stream") +} + +func TestQueryOfAnEmptyResultHasASchemaAndNoBatches(t *testing.T) { + f := queryFake(t, resultsResource([]map[string]any{field("id", "INTEGER", "NULLABLE")}), nil) + b := f.open(t, backend.Config{}) + + res, err := b.Query(context.Background(), "SELECT id FROM t WHERE FALSE", backend.QueryOptions{}) + require.NoError(t, err) + defer res.Batches.Close() + require.Len(t, res.Schema, 1) + require.NotEmpty(t, res.ArrowSchema) + _, err = res.Batches.Next() + require.ErrorIs(t, err, io.EOF) +} + +func TestQueryDryRunPricesTheStatementWithoutReadingAnything(t *testing.T) { + f := newFake(t) + f.handle("POST "+jobsPath, 200, jobResource("DONE", map[string]any{ + "totalBytesProcessed": "4096", + "query": map[string]any{ + "totalBytesProcessed": "4096", + "schema": map[string]any{"fields": []map[string]any{field("n", "INTEGER", "NULLABLE")}}, + }, + }, nil)) + b := f.open(t, backend.Config{}) + + res, err := b.Query(context.Background(), "SELECT COUNT(*) AS n FROM t", backend.QueryOptions{DryRun: true}) + require.NoError(t, err) + require.Nil(t, res.Batches, "a dry run has a header and no batches") + require.EqualValues(t, 4096, res.Job.Stats.BytesScanned) + require.Equal(t, backend.JobDone, res.Job.State) + require.Len(t, res.Schema, 1) + require.NotEmpty(t, res.ArrowSchema) + + insert := f.seen("POST", "/jobs") + require.Len(t, insert, 1) + require.Equal(t, true, insert[0].Body["configuration"].(map[string]any)["dryRun"]) + require.Empty(t, f.seen("GET", "/queries/"), "nothing is fetched for a dry run") +} + +func TestQueryOfADMLStatementReturnsRowsAffectedAndNoResult(t *testing.T) { + f := newFake(t) + stats := map[string]any{"query": map[string]any{"statementType": "UPDATE", "numDmlAffectedRows": "7"}} + f.handle("POST "+jobsPath, 200, jobResource("DONE", stats, nil)) + f.handle("GET "+jobPath, 200, jobResource("DONE", stats, nil)) + f.handle("GET "+resultsPath, 200, resultsResource(nil)) + b := f.open(t, backend.Config{}) + + res, err := b.Query(context.Background(), "UPDATE t SET x = 1 WHERE TRUE", backend.QueryOptions{}) + require.NoError(t, err) + require.Nil(t, res.Batches) + require.Empty(t, res.Schema) + require.Empty(t, res.ArrowSchema) + require.EqualValues(t, 7, res.Job.Stats.RowsAffected) +} + +func TestQueryOfAColumnWithNoPortableTypeIsUnsupportedNotALookalike(t *testing.T) { + rangeField := map[string]any{"name": "span", "type": "RANGE", "mode": "NULLABLE", "rangeElementType": map[string]any{"type": "DATE"}} + f := queryFake(t, resultsResource([]map[string]any{rangeField}), nil) + b := f.open(t, backend.Config{}) + + res, err := b.Query(context.Background(), "SELECT span FROM t", backend.QueryOptions{}) + require.Nil(t, res) + require.True(t, serr.Is(err, serr.Unsupported), "got %v", err) + require.Contains(t, err.Error(), "span") +} + +func TestQueryReportsEveryFailureInNormalizedTerms(t *testing.T) { + logged := quiet(t) + for name, tc := range map[string]struct { + status int + reason string + message string + want serr.Code + }{ + "invalid sql": {400, "invalidQuery", "Syntax error: Unexpected identifier \"foo\" at [1:8]", serr.InvalidArgument}, + "no such table": {404, "notFound", "Not found: Table test-project:d.nope", serr.NotFound}, + "denied": {403, "accessDenied", "Access Denied: User svc@test-project.iam.gserviceaccount.com does not have permission", serr.PermissionDenied}, + "over the cap": {400, "bytesBilledLimitExceeded", "Query exceeded limit for bytes billed: 1000. 10485760 or higher required.", serr.Throttled}, + "quota": {403, "quotaExceeded", "Quota exceeded: Your project exceeded quota", serr.Throttled}, + } { + t.Run(name, func(t *testing.T) { + f := newFake(t) + f.handle("POST "+jobsPath, tc.status, apiError(tc.status, tc.reason, tc.message)) + b := f.open(t, backend.Config{}) + + _, err := b.Query(context.Background(), "SELECT 1", backend.QueryOptions{}) + require.True(t, serr.Is(err, tc.want), "got %v", err) + require.NotContains(t, err.Error(), tc.reason, "a reason token is the vendor's") + require.NotContains(t, err.Error(), "SOME_VENDOR_STATUS") + if tc.want == serr.InvalidArgument { + require.Contains(t, err.Error(), "Syntax error", "the diagnosis of the caller's own SQL is kept") + } else { + require.NotContains(t, err.Error(), tc.message, "the vendor's text is the operator's, not the client's") + require.NotContains(t, err.Error(), "svc@") + } + }) + } + require.NotEmpty(t, *logged, "what is not a client's to read is logged for the operator") +} + +func TestQueryReportsAJobThatFailsWhileWaiting(t *testing.T) { + quiet(t) + f := newFake(t) + f.handle("POST "+jobsPath, 200, jobResource("RUNNING", nil, nil)) + f.handle("GET "+resultsPath, 400, apiError(400, "bytesBilledLimitExceeded", "Query exceeded limit for bytes billed")) + b := f.open(t, backend.Config{}) + + _, err := b.Query(context.Background(), "SELECT * FROM big", backend.QueryOptions{}) + require.True(t, serr.Is(err, serr.Throttled), "got %v", err) +} + +func TestQueryHonorsTheTighterOfTheCallersAndTheServersLimits(t *testing.T) { + for name, tc := range map[string]struct { + serverBytes, askBytes int64 + serverTime, askTime time.Duration + wantBytes, wantTimeMS string + }{ + "the caller asks for more than the server allows": {1000, 5000, time.Minute, time.Hour, "1000", "60000"}, + "the caller asks for less": {1000, 500, time.Minute, time.Second, "500", "1000"}, + "the caller asks for nothing": {1000, 0, time.Minute, 0, "1000", "60000"}, + "the server has no limit": {0, 700, 0, 2 * time.Second, "700", "2000"}, + "neither has a limit": {0, 0, 0, 0, "", ""}, + } { + t.Run(name, func(t *testing.T) { + f := queryFake(t, resultsResource([]map[string]any{field("id", "INTEGER", "NULLABLE")}, []any{"1"}), nil) + b := f.open(t, backend.Config{MaxQueryBytes: tc.serverBytes, QueryTimeout: tc.serverTime}) + res, err := b.Query(context.Background(), "SELECT 1", backend.QueryOptions{MaxBytesBilled: tc.askBytes, Timeout: tc.askTime}) + require.NoError(t, err) + _ = res.Batches.Close() + + config := f.seen("POST", "/jobs")[0].Body["configuration"].(map[string]any) + query := config["query"].(map[string]any) + if tc.wantBytes == "" { + require.NotContains(t, query, "maximumBytesBilled") + require.NotContains(t, config, "jobTimeoutMs") + return + } + require.Equal(t, tc.wantBytes, query["maximumBytesBilled"]) + require.Equal(t, tc.wantTimeMS, config["jobTimeoutMs"]) + }) + } +} + +func TestQueryThatOutlivesItsTimeLimitIsCancelledAndReportedAsDeadlineExceeded(t *testing.T) { + quiet(t) + f := newFake(t) + f.handle("POST "+jobsPath, 200, jobResource("RUNNING", nil, nil)) + f.handleFunc("GET "+resultsPath, func(w http.ResponseWriter, r *http.Request) { + writeJSON(w, 200, map[string]any{"kind": "bigquery#getQueryResultsResponse", "jobComplete": false}) + }) + f.handle("POST "+jobPath+"/cancel", 200, map[string]any{"kind": "bigquery#jobCancelResponse", "job": jobResource("RUNNING", nil, nil)}) + b := f.open(t, backend.Config{}) + + _, err := b.Query(context.Background(), "SELECT slow()", backend.QueryOptions{Timeout: 300 * time.Millisecond}) + require.True(t, serr.Is(err, serr.DeadlineExceeded), "got %v", err) + require.Len(t, f.seen("POST", "/cancel"), 1, "the abandoned job is cancelled so it stops billing") +} + +func TestQueryRefusesWhatIsNotAQueryBeforeAskingBigQuery(t *testing.T) { + f := newFake(t) + b := f.open(t, backend.Config{}) + for name, tc := range map[string]struct { + sql string + opts backend.QueryOptions + want serr.Code + }{ + "empty sql": {" ", backend.QueryOptions{}, serr.InvalidArgument}, + "bad parameter": {"SELECT 1", backend.QueryOptions{Params: []backend.QueryParam{{Name: "n", Type: backend.TypeInt64, Value: "x"}}}, serr.InvalidArgument}, + "array parameter": {"SELECT 1", backend.QueryOptions{Params: []backend.QueryParam{{Name: "n", Type: backend.TypeArray, Value: "[1]"}}}, serr.Unsupported}, + "bad default": {"SELECT 1", backend.QueryOptions{DefaultDataset: "a/b"}, serr.InvalidArgument}, + "untyped parameter": {"SELECT 1", backend.QueryOptions{Params: []backend.QueryParam{{Name: "n", Value: "1"}}}, serr.InvalidArgument}, + } { + t.Run(name, func(t *testing.T) { + _, err := b.Query(context.Background(), tc.sql, tc.opts) + require.True(t, serr.Is(err, tc.want), "got %v", err) + }) + } + require.Empty(t, f.requests, "nothing was sent to BigQuery") +} + +func TestGetJobReportsStateAndAccounting(t *testing.T) { + quiet(t) + for name, tc := range map[string]struct { + state string + extra map[string]any + want backend.JobState + fail bool + }{ + "running": {"RUNNING", nil, backend.JobRunning, false}, + "pending": {"PENDING", nil, backend.JobPending, false}, + "done": {"DONE", nil, backend.JobDone, false}, + "failed": {"DONE", map[string]any{"status": map[string]any{"state": "DONE", "errorResult": map[string]any{ + "reason": "invalidQuery", "message": "Syntax error: bad sql"}}}, backend.JobFailed, true}, + "cancelled": {"DONE", map[string]any{"status": map[string]any{"state": "DONE", "errorResult": map[string]any{ + "reason": "stopped", "message": "Job execution was cancelled: User requested cancellation"}}}, backend.JobCancelled, true}, + } { + t.Run(name, func(t *testing.T) { + f := newFake(t) + f.handle("GET "+jobPath, 200, jobResource(tc.state, map[string]any{ + "totalBytesProcessed": "100", + "totalSlotMs": "250", + "query": map[string]any{"totalBytesBilled": "10485760", "totalBytesProcessed": "100", "numDmlAffectedRows": "3"}, + }, tc.extra)) + b := f.open(t, backend.Config{}) + + job, err := b.GetJob(context.Background(), "EU.job_1") + require.NoError(t, err) + require.Equal(t, "EU.job_1", job.ID) + require.Equal(t, tc.want, job.State) + require.EqualValues(t, 100, job.Stats.BytesScanned) + require.EqualValues(t, 10485760, job.Stats.BytesBilled) + require.EqualValues(t, 3, job.Stats.RowsAffected) + require.Equal(t, time.UnixMilli(1700000000000).UTC(), job.Created.UTC()) + if tc.fail { + require.NotEmpty(t, job.Error) + require.NotContains(t, job.Error, "invalidQuery") + require.NotContains(t, job.Error, "stopped") + } else { + require.Empty(t, job.Error) + } + require.Equal(t, "EU", f.seen("GET", "/jobs/job_1")[0].Query["location"][0], "the location in the id is the one asked for") + }) + } +} + +func TestJobIDsAreCheckedAndFallBackToTheConfiguredLocation(t *testing.T) { + quiet(t) + f := newFake(t) + f.handle("GET "+jobPath, 200, jobResource("DONE", nil, nil)) + f.handle("POST "+jobPath+"/cancel", 200, map[string]any{"job": jobResource("DONE", nil, nil)}) + f.handle("GET /projects/test-project/jobs/missing", 404, apiError(404, "notFound", "Not found: Job")) + b := f.open(t, backend.Config{Location: "asia-northeast1"}) + + _, err := b.GetJob(context.Background(), "job_1") + require.NoError(t, err) + require.Equal(t, "asia-northeast1", f.seen("GET", "/jobs/job_1")[0].Query["location"][0]) + + require.NoError(t, b.CancelJob(context.Background(), "EU.job_1")) + require.Equal(t, "EU", f.seen("POST", "/cancel")[0].Query["location"][0]) + + for _, bad := range []string{"", "a/b", "EU.a/b", "../x", "EU.job_1.extra", "e u.job"} { + _, err := b.GetJob(context.Background(), bad) + require.True(t, serr.Is(err, serr.InvalidArgument), "%q: %v", bad, err) + require.True(t, serr.Is(b.CancelJob(context.Background(), bad), serr.InvalidArgument), "%q", bad) + } + _, err = b.GetJob(context.Background(), "missing") + require.True(t, serr.Is(err, serr.NotFound), "got %v", err) +} + +func TestQueryIdentifiesAJobWithoutALocationByItsBareID(t *testing.T) { + require.Equal(t, "job_1", encodeJobID("", "job_1")) + require.Equal(t, "EU.job_1", encodeJobID("EU", "job_1")) + require.Equal(t, "", encodeJobID("EU", ""), "a dry run has no job to name") +} diff --git a/internal/backend/bigquery/types.go b/internal/backend/bigquery/types.go new file mode 100644 index 0000000..9a1de41 --- /dev/null +++ b/internal/backend/bigquery/types.go @@ -0,0 +1,362 @@ +package bigquery + +import ( + "fmt" + "strings" + + bq "cloud.google.com/go/bigquery" + "github.com/apache/arrow-go/v18/arrow" + + "github.com/codefly-dev/service-warehouse/internal/backend" + "github.com/codefly-dev/service-warehouse/internal/serr" +) + +// This file is the type mapping: BigQuery's field schema to the portable +// Column, and the portable Column to an Arrow type. Both directions are total +// over what is mapped and refuse the rest with serr.Unsupported; nothing falls +// into a type that merely looks right. + +// Defaults of the two decimal types when a column states no precision. +const ( + numericPrecision, numericScale = 38, 9 + bigNumericPrecision, bigNumericScale = 76, 38 +) + +// itemName names the single element column of an ARRAY, which is Arrow's own +// name for it. +const itemName = "item" + +// fieldType normalizes the standard-SQL spellings BigQuery also accepts to the +// ones its schemas report. +func fieldType(t bq.FieldType) bq.FieldType { + switch strings.ToUpper(string(t)) { + case "INT64": + return bq.IntegerFieldType + case "FLOAT64": + return bq.FloatFieldType + case "BOOL": + return bq.BooleanFieldType + case "STRUCT": + return bq.RecordFieldType + default: + return bq.FieldType(strings.ToUpper(string(t))) + } +} + +// portableType is the portable bucket of a BigQuery type. A type with no bucket +// is TypeUnknown and keeps its spelling in Column.NativeType. +func portableType(t bq.FieldType) backend.ColumnType { + switch fieldType(t) { + case bq.BooleanFieldType: + return backend.TypeBool + case bq.IntegerFieldType: + return backend.TypeInt64 + case bq.FloatFieldType: + return backend.TypeFloat64 + case bq.NumericFieldType, bq.BigNumericFieldType: + return backend.TypeNumeric + case bq.StringFieldType: + return backend.TypeString + case bq.BytesFieldType: + return backend.TypeBytes + case bq.DateFieldType: + return backend.TypeDate + case bq.TimeFieldType: + return backend.TypeTime + case bq.DateTimeFieldType: + return backend.TypeTimestamp + case bq.TimestampFieldType: + return backend.TypeTimestampTZ + case bq.IntervalFieldType: + return backend.TypeInterval + case bq.JSONFieldType: + return backend.TypeJSON + case bq.RecordFieldType: + return backend.TypeStruct + case bq.GeographyFieldType: + return backend.TypeGeography + default: + return backend.TypeUnknown + } +} + +// nativeType is BigQuery's own standard-SQL spelling of a field's type, with its +// parameters: NUMERIC(10, 2), STRING(32), ARRAY>. +func nativeType(f *bq.FieldSchema) string { + base := baseNativeType(f) + if f.Repeated { + return "ARRAY<" + base + ">" + } + return base +} + +func baseNativeType(f *bq.FieldSchema) string { + switch t := fieldType(f.Type); t { + case bq.IntegerFieldType: + return "INT64" + case bq.FloatFieldType: + return "FLOAT64" + case bq.BooleanFieldType: + return "BOOL" + case bq.RecordFieldType: + parts := make([]string, len(f.Schema)) + for i, sub := range f.Schema { + parts[i] = sub.Name + " " + nativeType(sub) + } + return "STRUCT<" + strings.Join(parts, ", ") + ">" + case bq.NumericFieldType, bq.BigNumericFieldType: + switch { + case f.Precision > 0 && f.Scale > 0: + return fmt.Sprintf("%s(%d, %d)", t, f.Precision, f.Scale) + case f.Precision > 0: + return fmt.Sprintf("%s(%d)", t, f.Precision) + } + return string(t) + case bq.StringFieldType, bq.BytesFieldType: + if f.MaxLength > 0 { + return fmt.Sprintf("%s(%d)", t, f.MaxLength) + } + return string(t) + case bq.RangeFieldType: + if f.RangeElementType != nil { + return "RANGE<" + string(fieldType(f.RangeElementType.Type)) + ">" + } + return string(t) + default: + return string(t) + } +} + +// columnFromField is the portable description of a BigQuery field. A REPEATED +// field is an ARRAY whose single sub-column is the element; a RECORD is a STRUCT +// whose sub-columns are its fields. +func columnFromField(f *bq.FieldSchema) backend.Column { + element := backend.Column{ + Name: f.Name, + Type: portableType(f.Type), + Nullable: !f.Required, + NativeType: baseNativeType(f), + Description: f.Description, + } + switch fieldType(f.Type) { + case bq.NumericFieldType: + element.Precision, element.Scale = decimalParameters(f, numericPrecision, numericScale) + case bq.BigNumericFieldType: + element.Precision, element.Scale = decimalParameters(f, bigNumericPrecision, bigNumericScale) + case bq.RecordFieldType: + element.Fields = columnsFromSchema(f.Schema) + } + if !f.Repeated { + return element + } + // BigQuery has no NULL array and no NULL element: an absent array is empty. + element.Name, element.Nullable, element.Description = itemName, false, "" + return backend.Column{ + Name: f.Name, + Type: backend.TypeArray, + Nullable: false, + NativeType: nativeType(f), + Description: f.Description, + Fields: []backend.Column{element}, + } +} + +func decimalParameters(f *bq.FieldSchema, defaultPrecision, defaultScale int32) (int32, int32) { + if f.Precision > 0 { + return int32(f.Precision), int32(f.Scale) + } + return defaultPrecision, defaultScale +} + +func columnsFromSchema(s bq.Schema) []backend.Column { + if s == nil { + return nil + } + out := make([]backend.Column, len(s)) + for i, f := range s { + out[i] = columnFromField(f) + } + return out +} + +// fieldFromColumn is the BigQuery field a portable column creates. A column of +// a type BigQuery has no equivalent for is Unsupported, never a lookalike. +func fieldFromColumn(c backend.Column) (*bq.FieldSchema, error) { + if strings.TrimSpace(c.Name) == "" { + return nil, serr.New(serr.InvalidArgument, "CreateTable", "a column has no name") + } + if c.Type == backend.TypeArray { + if len(c.Fields) != 1 { + return nil, serr.New(serr.InvalidArgument, "CreateTable", + fmt.Sprintf("column %q is an ARRAY and must carry exactly one element column", c.Name)) + } + element := c.Fields[0] + if element.Type == backend.TypeArray { + return nil, serr.New(serr.Unsupported, "CreateTable", + fmt.Sprintf("column %q: BigQuery does not allow an array of arrays", c.Name)) + } + element.Name = c.Name + f, err := fieldFromColumn(element) + if err != nil { + return nil, err + } + f.Repeated, f.Required = true, false + f.Description = c.Description + return f, nil + } + + f := &bq.FieldSchema{Name: c.Name, Description: c.Description, Required: !c.Nullable} + switch c.Type { + case backend.TypeBool: + f.Type = bq.BooleanFieldType + case backend.TypeInt64: + f.Type = bq.IntegerFieldType + case backend.TypeFloat64: + f.Type = bq.FloatFieldType + case backend.TypeNumeric: + f.Type = bq.NumericFieldType + if c.Precision > numericPrecision || c.Scale > numericScale || + strings.HasPrefix(strings.ToUpper(c.NativeType), "BIGNUMERIC") { + f.Type = bq.BigNumericFieldType + } + // A precision and scale equal to the type's own bounds is the unparameterized + // type, which is how BigQuery reports it back. + defaultPrecision, defaultScale := int32(numericPrecision), int32(numericScale) + if f.Type == bq.BigNumericFieldType { + defaultPrecision, defaultScale = bigNumericPrecision, bigNumericScale + } + if c.Precision > 0 && !(c.Precision == defaultPrecision && c.Scale == defaultScale) { + f.Precision, f.Scale = int64(c.Precision), int64(c.Scale) + } + case backend.TypeString: + f.Type = bq.StringFieldType + case backend.TypeBytes: + f.Type = bq.BytesFieldType + case backend.TypeDate: + f.Type = bq.DateFieldType + case backend.TypeTime: + f.Type = bq.TimeFieldType + case backend.TypeTimestamp: + f.Type = bq.DateTimeFieldType + case backend.TypeTimestampTZ: + f.Type = bq.TimestampFieldType + case backend.TypeInterval: + f.Type = bq.IntervalFieldType + case backend.TypeJSON: + f.Type = bq.JSONFieldType + case backend.TypeGeography: + f.Type = bq.GeographyFieldType + case backend.TypeStruct: + if len(c.Fields) == 0 { + return nil, serr.New(serr.InvalidArgument, "CreateTable", + fmt.Sprintf("column %q is a STRUCT and must carry its fields", c.Name)) + } + f.Type = bq.RecordFieldType + for _, sub := range c.Fields { + subField, err := fieldFromColumn(sub) + if err != nil { + return nil, err + } + f.Schema = append(f.Schema, subField) + } + default: + return nil, serr.New(serr.Unsupported, "CreateTable", + fmt.Sprintf("column %q has a type with no BigQuery equivalent (native type %q)", c.Name, c.NativeType)) + } + return f, nil +} + +func schemaFromColumns(cols []backend.Column) (bq.Schema, error) { + out := make(bq.Schema, 0, len(cols)) + for _, c := range cols { + f, err := fieldFromColumn(c) + if err != nil { + return nil, err + } + out = append(out, f) + } + return out, nil +} + +// arrowField is the Arrow field a portable column becomes in a result. +func arrowField(c backend.Column) (arrow.Field, error) { + t, err := arrowType(c) + if err != nil { + return arrow.Field{}, err + } + return arrow.Field{Name: c.Name, Type: t, Nullable: c.Nullable}, nil +} + +// arrowType is the Arrow type of a portable column. A column whose type has no +// portable bucket (a BigQuery RANGE, say) is Unsupported: a result is never +// sent with a column the client could mistake for another type. +func arrowType(c backend.Column) (arrow.DataType, error) { + switch c.Type { + case backend.TypeBool: + return arrow.FixedWidthTypes.Boolean, nil + case backend.TypeInt64: + return arrow.PrimitiveTypes.Int64, nil + case backend.TypeFloat64: + return arrow.PrimitiveTypes.Float64, nil + case backend.TypeNumeric: + precision, scale := c.Precision, c.Scale + if precision == 0 { + precision, scale = numericPrecision, numericScale + if strings.HasPrefix(strings.ToUpper(c.NativeType), "BIGNUMERIC") { + precision, scale = bigNumericPrecision, bigNumericScale + } + } + if precision > numericPrecision { + return &arrow.Decimal256Type{Precision: precision, Scale: scale}, nil + } + return &arrow.Decimal128Type{Precision: precision, Scale: scale}, nil + case backend.TypeString, backend.TypeJSON, backend.TypeGeography: + return arrow.BinaryTypes.String, nil + case backend.TypeBytes: + return arrow.BinaryTypes.Binary, nil + case backend.TypeDate: + return arrow.FixedWidthTypes.Date32, nil + case backend.TypeTime: + return arrow.FixedWidthTypes.Time64us, nil + case backend.TypeTimestamp: + return &arrow.TimestampType{Unit: arrow.Microsecond}, nil + case backend.TypeTimestampTZ: + return &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: "UTC"}, nil + case backend.TypeInterval: + return arrow.FixedWidthTypes.MonthDayNanoInterval, nil + case backend.TypeArray: + if len(c.Fields) != 1 { + return nil, serr.New(serr.Internal, "arrow", fmt.Sprintf("array column %q has no element type", c.Name)) + } + element, err := arrowField(c.Fields[0]) + if err != nil { + return nil, err + } + return arrow.ListOfField(element), nil + case backend.TypeStruct: + fields := make([]arrow.Field, len(c.Fields)) + for i, sub := range c.Fields { + f, err := arrowField(sub) + if err != nil { + return nil, err + } + fields[i] = f + } + return arrow.StructOf(fields...), nil + default: + return nil, serr.New(serr.Unsupported, "arrow", + fmt.Sprintf("column %q (native type %q) has no portable type and cannot be sent as Arrow", c.Name, c.NativeType)) + } +} + +func arrowSchema(cols []backend.Column) (*arrow.Schema, error) { + fields := make([]arrow.Field, len(cols)) + for i, c := range cols { + f, err := arrowField(c) + if err != nil { + return nil, err + } + fields[i] = f + } + return arrow.NewSchema(fields, nil), nil +} diff --git a/internal/backend/bigquery/types_test.go b/internal/backend/bigquery/types_test.go new file mode 100644 index 0000000..19b49c8 --- /dev/null +++ b/internal/backend/bigquery/types_test.go @@ -0,0 +1,160 @@ +package bigquery + +import ( + "testing" + + bq "cloud.google.com/go/bigquery" + "github.com/stretchr/testify/require" + + "github.com/codefly-dev/service-warehouse/internal/backend" + "github.com/codefly-dev/service-warehouse/internal/serr" +) + +func TestPortableTypeOfEveryBigQueryType(t *testing.T) { + for in, want := range map[bq.FieldType]backend.ColumnType{ + bq.BooleanFieldType: backend.TypeBool, "BOOL": backend.TypeBool, + bq.IntegerFieldType: backend.TypeInt64, "INT64": backend.TypeInt64, + bq.FloatFieldType: backend.TypeFloat64, "FLOAT64": backend.TypeFloat64, + bq.NumericFieldType: backend.TypeNumeric, bq.BigNumericFieldType: backend.TypeNumeric, + bq.StringFieldType: backend.TypeString, bq.BytesFieldType: backend.TypeBytes, + bq.DateFieldType: backend.TypeDate, bq.TimeFieldType: backend.TypeTime, + // The portable TIMESTAMP is wall-clock; the zoned instant is TIMESTAMPTZ. + bq.DateTimeFieldType: backend.TypeTimestamp, bq.TimestampFieldType: backend.TypeTimestampTZ, + bq.IntervalFieldType: backend.TypeInterval, bq.JSONFieldType: backend.TypeJSON, + bq.RecordFieldType: backend.TypeStruct, "STRUCT": backend.TypeStruct, + bq.GeographyFieldType: backend.TypeGeography, + bq.RangeFieldType: backend.TypeUnknown, "SOMETHING_NEW": backend.TypeUnknown, + } { + require.Equal(t, want, portableType(in), string(in)) + } +} + +func TestNativeTypeKeepsBigQuerysSpelling(t *testing.T) { + for name, tc := range map[string]struct { + field *bq.FieldSchema + want string + }{ + "integer": {&bq.FieldSchema{Type: bq.IntegerFieldType}, "INT64"}, + "float": {&bq.FieldSchema{Type: bq.FloatFieldType}, "FLOAT64"}, + "boolean": {&bq.FieldSchema{Type: bq.BooleanFieldType}, "BOOL"}, + "numeric": {&bq.FieldSchema{Type: bq.NumericFieldType}, "NUMERIC"}, + "numeric with p,s": {&bq.FieldSchema{Type: bq.NumericFieldType, Precision: 10, Scale: 2}, "NUMERIC(10, 2)"}, + "numeric with p": {&bq.FieldSchema{Type: bq.NumericFieldType, Precision: 10}, "NUMERIC(10)"}, + "bignumeric": {&bq.FieldSchema{Type: bq.BigNumericFieldType, Precision: 40, Scale: 10}, "BIGNUMERIC(40, 10)"}, + "sized string": {&bq.FieldSchema{Type: bq.StringFieldType, MaxLength: 32}, "STRING(32)"}, + "sized bytes": {&bq.FieldSchema{Type: bq.BytesFieldType, MaxLength: 8}, "BYTES(8)"}, + "array": {&bq.FieldSchema{Type: bq.DateFieldType, Repeated: true}, "ARRAY"}, + "range": {&bq.FieldSchema{Type: bq.RangeFieldType, RangeElementType: &bq.RangeElementType{Type: bq.DateTimeFieldType}}, "RANGE"}, + "struct": {&bq.FieldSchema{Type: bq.RecordFieldType, Schema: bq.Schema{{Name: "a", Type: bq.IntegerFieldType}, {Name: "b", Type: bq.StringFieldType, Repeated: true}}}, "STRUCT>"}, + "array of structs": {&bq.FieldSchema{Type: bq.RecordFieldType, Repeated: true, Schema: bq.Schema{{Name: "a", Type: bq.IntegerFieldType}}}, "ARRAY>"}, + "a type not yet seen": {&bq.FieldSchema{Type: "NEW_TYPE"}, "NEW_TYPE"}, + } { + t.Run(name, func(t *testing.T) { + require.Equal(t, tc.want, nativeType(tc.field)) + }) + } +} + +func TestColumnsRoundTripThroughBigQueryFields(t *testing.T) { + cols := []backend.Column{ + {Name: "b", Type: backend.TypeBool}, + {Name: "i", Type: backend.TypeInt64, Nullable: true, Description: "an int"}, + {Name: "f", Type: backend.TypeFloat64, Nullable: true}, + {Name: "n", Type: backend.TypeNumeric, Nullable: true, Precision: 10, Scale: 2}, + {Name: "big", Type: backend.TypeNumeric, Nullable: true, Precision: 76, Scale: 38, NativeType: "BIGNUMERIC"}, + {Name: "s", Type: backend.TypeString, Nullable: true}, + {Name: "y", Type: backend.TypeBytes, Nullable: true}, + {Name: "d", Type: backend.TypeDate, Nullable: true}, + {Name: "t", Type: backend.TypeTime, Nullable: true}, + {Name: "wall", Type: backend.TypeTimestamp, Nullable: true}, + {Name: "at", Type: backend.TypeTimestampTZ, Nullable: true}, + {Name: "iv", Type: backend.TypeInterval, Nullable: true}, + {Name: "j", Type: backend.TypeJSON, Nullable: true}, + {Name: "g", Type: backend.TypeGeography, Nullable: true}, + {Name: "tags", Type: backend.TypeArray, Fields: []backend.Column{{Name: "item", Type: backend.TypeString}}}, + {Name: "owner", Type: backend.TypeStruct, Nullable: true, Fields: []backend.Column{ + {Name: "name", Type: backend.TypeString, Nullable: true}, + {Name: "ids", Type: backend.TypeArray, Fields: []backend.Column{{Name: "item", Type: backend.TypeInt64}}}, + }}, + {Name: "lines", Type: backend.TypeArray, Fields: []backend.Column{{Name: "item", Type: backend.TypeStruct, Fields: []backend.Column{ + {Name: "sku", Type: backend.TypeString, Nullable: true}, + }}}}, + } + schema, err := schemaFromColumns(cols) + require.NoError(t, err) + + // Back to portable columns: the types, nullability and shape survive. The + // native spelling is what BigQuery reports, so it is not part of the round + // trip's claim. + back := columnsFromSchema(schema) + require.Len(t, back, len(cols)) + for i := range cols { + assertSameShape(t, cols[i], back[i]) + } +} + +func assertSameShape(t *testing.T, want, got backend.Column) { + t.Helper() + require.Equal(t, want.Name, got.Name) + require.Equal(t, want.Type, got.Type, want.Name) + require.Equal(t, want.Nullable, got.Nullable, want.Name) + require.Equal(t, want.Description, got.Description, want.Name) + if want.Precision != 0 { + require.Equal(t, want.Precision, got.Precision, want.Name) + require.Equal(t, want.Scale, got.Scale, want.Name) + } + require.Len(t, got.Fields, len(want.Fields), want.Name) + for i := range want.Fields { + assertSameShape(t, want.Fields[i], got.Fields[i]) + } +} + +func TestNumericChoosesTheWiderTypeOnlyWhenTheNarrowOneCannotHoldIt(t *testing.T) { + for name, tc := range map[string]struct { + col backend.Column + wantType bq.FieldType + wantP, wS int64 + }{ + "unparameterized": {backend.Column{Name: "n", Type: backend.TypeNumeric}, bq.NumericFieldType, 0, 0}, + "its own bounds": {backend.Column{Name: "n", Type: backend.TypeNumeric, Precision: 38, Scale: 9}, bq.NumericFieldType, 0, 0}, + "parameterized": {backend.Column{Name: "n", Type: backend.TypeNumeric, Precision: 10, Scale: 2}, bq.NumericFieldType, 10, 2}, + "scale past 9": {backend.Column{Name: "n", Type: backend.TypeNumeric, Precision: 30, Scale: 12}, bq.BigNumericFieldType, 30, 12}, + "precision past 38": {backend.Column{Name: "n", Type: backend.TypeNumeric, Precision: 50, Scale: 5}, bq.BigNumericFieldType, 50, 5}, + "its own spelling": {backend.Column{Name: "n", Type: backend.TypeNumeric, NativeType: "BIGNUMERIC"}, bq.BigNumericFieldType, 0, 0}, + "bignumeric default": {backend.Column{Name: "n", Type: backend.TypeNumeric, Precision: 76, Scale: 38, NativeType: "BIGNUMERIC"}, bq.BigNumericFieldType, 0, 0}, + } { + t.Run(name, func(t *testing.T) { + f, err := fieldFromColumn(tc.col) + require.NoError(t, err) + require.Equal(t, tc.wantType, f.Type) + require.Equal(t, tc.wantP, f.Precision) + require.Equal(t, tc.wS, f.Scale) + }) + } +} + +func TestColumnsBigQueryCannotHoldAreRefusedAndNeverBecomeSomethingElse(t *testing.T) { + for name, tc := range map[string]struct { + col backend.Column + want serr.Code + }{ + "unknown type": {backend.Column{Name: "x", Type: backend.TypeUnknown, NativeType: "RANGE"}, serr.Unsupported}, + "array of arrays": {backend.Column{Name: "x", Type: backend.TypeArray, Fields: []backend.Column{{Name: "item", Type: backend.TypeArray, Fields: []backend.Column{{Type: backend.TypeInt64}}}}}, serr.Unsupported}, + "array w/o element": {backend.Column{Name: "x", Type: backend.TypeArray}, serr.InvalidArgument}, + "struct w/o fields": {backend.Column{Name: "x", Type: backend.TypeStruct}, serr.InvalidArgument}, + "no name": {backend.Column{Type: backend.TypeInt64}, serr.InvalidArgument}, + "nested unknown": {backend.Column{Name: "x", Type: backend.TypeStruct, Fields: []backend.Column{{Name: "y", Type: backend.TypeUnknown}}}, serr.Unsupported}, + } { + t.Run(name, func(t *testing.T) { + _, err := fieldFromColumn(tc.col) + require.True(t, serr.Is(err, tc.want), "got %v", err) + }) + } +} + +func TestAResultColumnWithNoPortableTypeCannotBecomeArrow(t *testing.T) { + _, err := arrowSchema([]backend.Column{{Name: "span", Type: backend.TypeUnknown, NativeType: "RANGE"}}) + require.True(t, serr.Is(err, serr.Unsupported)) + require.Contains(t, err.Error(), "span") + require.Contains(t, err.Error(), "RANGE") +} diff --git a/internal/backend/bigquery/values.go b/internal/backend/bigquery/values.go new file mode 100644 index 0000000..004f9a6 --- /dev/null +++ b/internal/backend/bigquery/values.go @@ -0,0 +1,254 @@ +package bigquery + +import ( + "encoding/base64" + "errors" + "fmt" + "math" + "unicode/utf8" + + bq "cloud.google.com/go/bigquery" + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" +) + +// This file turns one cell of an Arrow batch into the JSON value a BigQuery +// streaming insert sends for it. Every value comes out as a JSON-native Go value +// (string, int64, float64, bool, nil, a slice or a map), so what the insert sends +// does not depend on how the client library would have encoded a Go type. +// +// A cell that cannot be sent is a valueError: the row is refused as invalid. A +// column whose Arrow type has no BigQuery equivalent is not a cell's fault and is +// caught for the whole stream by checkArrowSchema before any row is read. + +// valueError is a cell BigQuery cannot be given, with the column it is in. +type valueError struct { + column string + reason string +} + +func (e *valueError) Error() string { return fmt.Sprintf("column %q: %s", e.column, e.reason) } + +// Timestamp layouts BigQuery accepts in a JSON row, at its microsecond precision. +const ( + timestampLayoutTZ = "2006-01-02T15:04:05.000000Z" + timestampLayoutNoTZ = "2006-01-02T15:04:05.000000" + dateLayout = "2006-01-02" + microsPerSecond = int64(1_000_000) + nanosPerMicro = int64(1_000) + microsPerDay = 24 * 3600 * microsPerSecond +) + +// checkArrowSchema refuses, before anything is read, an Arrow schema with a type +// that has no BigQuery equivalent. That is a property of the request, not of a +// row, so it fails the call. +func checkArrowSchema(schema *arrow.Schema) error { + for _, f := range schema.Fields() { + if err := checkArrowType(f.Name, f.Type); err != nil { + return err + } + } + return nil +} + +func checkArrowType(name string, t arrow.DataType) error { + switch t := t.(type) { + case *arrow.BooleanType, *arrow.Int8Type, *arrow.Int16Type, *arrow.Int32Type, *arrow.Int64Type, + *arrow.Uint8Type, *arrow.Uint16Type, *arrow.Uint32Type, *arrow.Uint64Type, + *arrow.Float32Type, *arrow.Float64Type, + *arrow.StringType, *arrow.LargeStringType, + *arrow.BinaryType, *arrow.LargeBinaryType, *arrow.FixedSizeBinaryType, + *arrow.Date32Type, *arrow.Date64Type, *arrow.Time32Type, *arrow.Time64Type, + *arrow.TimestampType, *arrow.Decimal128Type, *arrow.Decimal256Type, + *arrow.MonthDayNanoIntervalType: + return nil + case *arrow.ListType: + return checkArrowType(name, t.Elem()) + case *arrow.LargeListType: + return checkArrowType(name, t.Elem()) + case *arrow.StructType: + for _, f := range t.Fields() { + if err := checkArrowType(name+"."+f.Name, f.Type); err != nil { + return err + } + } + return nil + default: + return &arrowTypeError{column: name, typeName: t.String()} + } +} + +// arrowTypeError is an Arrow type BigQuery has no equivalent of. +type arrowTypeError struct{ column, typeName string } + +func (e *arrowTypeError) Error() string { + return fmt.Sprintf("column %q has the Arrow type %s, which BigQuery has no equivalent of", e.column, e.typeName) +} + +// jsonValue is the JSON value of arr[i] for column name. +func jsonValue(name string, arr arrow.Array, i int) (any, error) { + if arr.IsNull(i) { + return nil, nil + } + bad := func(format string, args ...any) (any, error) { + return nil, &valueError{column: name, reason: fmt.Sprintf(format, args...)} + } + switch a := arr.(type) { + case *array.Boolean: + return a.Value(i), nil + case *array.Int8: + return int64(a.Value(i)), nil + case *array.Int16: + return int64(a.Value(i)), nil + case *array.Int32: + return int64(a.Value(i)), nil + case *array.Int64: + return a.Value(i), nil + case *array.Uint8: + return int64(a.Value(i)), nil + case *array.Uint16: + return int64(a.Value(i)), nil + case *array.Uint32: + return int64(a.Value(i)), nil + case *array.Uint64: + if a.Value(i) > uint64(math.MaxInt64) { + return bad("the integer %d does not fit a 64-bit signed integer", a.Value(i)) + } + return int64(a.Value(i)), nil + case *array.Float32: + return floatJSON(float64(a.Value(i))), nil + case *array.Float64: + return floatJSON(a.Value(i)), nil + case *array.String: + return textJSON(name, a.Value(i)) + case *array.LargeString: + return textJSON(name, a.Value(i)) + case *array.Binary: + return base64.StdEncoding.EncodeToString(a.Value(i)), nil + case *array.LargeBinary: + return base64.StdEncoding.EncodeToString(a.Value(i)), nil + case *array.FixedSizeBinary: + return base64.StdEncoding.EncodeToString(a.Value(i)), nil + case *array.Date32: + return a.Value(i).ToTime().Format(dateLayout), nil + case *array.Date64: + return a.Value(i).ToTime().Format(dateLayout), nil + case *array.Time32: + unit := a.DataType().(*arrow.Time32Type).Unit + return timeOfDay(name, int64(a.Value(i)), unit) + case *array.Time64: + unit := a.DataType().(*arrow.Time64Type).Unit + return timeOfDay(name, int64(a.Value(i)), unit) + case *array.Timestamp: + t := a.DataType().(*arrow.TimestampType) + if t.Unit == arrow.Nanosecond && int64(a.Value(i))%nanosPerMicro != 0 { + return bad("the timestamp has sub-microsecond precision, which BigQuery cannot store") + } + instant := a.Value(i).ToTime(t.Unit) + if t.TimeZone != "" { + return instant.UTC().Format(timestampLayoutTZ), nil + } + return instant.UTC().Format(timestampLayoutNoTZ), nil + case *array.Decimal128: + return a.Value(i).ToString(a.DataType().(*arrow.Decimal128Type).Scale), nil + case *array.Decimal256: + return a.Value(i).ToString(a.DataType().(*arrow.Decimal256Type).Scale), nil + case *array.MonthDayNanoInterval: + v := a.Value(i) + const nanosPerHour, nanosPerMinute, nanosPerSecond = int64(3_600_000_000_000), int64(60_000_000_000), int64(1_000_000_000) + n := v.Nanoseconds + interval := bq.IntervalValue{ + Years: v.Months / 12, + Months: v.Months % 12, + Days: v.Days, + Hours: int32(n / nanosPerHour), + Minutes: int32(n % nanosPerHour / nanosPerMinute), + Seconds: int32(n % nanosPerMinute / nanosPerSecond), + SubSecondNanos: int32(n % nanosPerSecond), + } + return interval.String(), nil + case *array.List: + start, end := a.ValueOffsets(i) + return listJSON(name, a.ListValues(), int(start), int(end)) + case *array.LargeList: + start, end := a.ValueOffsets(i) + return listJSON(name, a.ListValues(), int(start), int(end)) + case *array.Struct: + st := a.DataType().(*arrow.StructType) + out := make(map[string]any, st.NumFields()) + for k := 0; k < st.NumFields(); k++ { + v, err := jsonValue(name+"."+st.Field(k).Name, a.Field(k), i) + if err != nil { + return nil, err + } + out[st.Field(k).Name] = v + } + return out, nil + default: + return nil, &arrowTypeError{column: name, typeName: arr.DataType().String()} + } +} + +func listJSON(name string, values arrow.Array, start, end int) (any, error) { + out := make([]any, 0, end-start) + for k := start; k < end; k++ { + v, err := jsonValue(name, values, k) + if err != nil { + return nil, err + } + out = append(out, v) + } + return out, nil +} + +// floatJSON keeps a float a number, except the three values JSON has no number +// for, which BigQuery reads as these strings. +func floatJSON(v float64) any { + switch { + case math.IsNaN(v): + return "NaN" + case math.IsInf(v, 1): + return "Infinity" + case math.IsInf(v, -1): + return "-Infinity" + default: + return v + } +} + +// textJSON refuses text BigQuery would refuse: a STRING must be valid UTF-8. +func textJSON(name, s string) (any, error) { + if !utf8.ValidString(s) { + return nil, &valueError{column: name, reason: "the text is not valid UTF-8"} + } + return s, nil +} + +// timeOfDay formats a time of day as BigQuery's TIME text, at microseconds. +func timeOfDay(name string, v int64, unit arrow.TimeUnit) (any, error) { + var micros int64 + switch unit { + case arrow.Second: + micros = v * microsPerSecond + case arrow.Millisecond: + micros = v * 1000 + case arrow.Microsecond: + micros = v + default: + if v%nanosPerMicro != 0 { + return nil, &valueError{column: name, reason: "the time has sub-microsecond precision, which BigQuery cannot store"} + } + micros = v / nanosPerMicro + } + if micros < 0 || micros >= microsPerDay { + return nil, &valueError{column: name, reason: "the time of day is outside 00:00:00 to 23:59:59.999999"} + } + seconds := micros / microsPerSecond + return fmt.Sprintf("%02d:%02d:%02d.%06d", seconds/3600, seconds%3600/60, seconds%60, micros%microsPerSecond), nil +} + +// asValueError reports whether err is a cell's fault. +func asValueError(err error) (*valueError, bool) { + var ve *valueError + return ve, errors.As(err, &ve) +} From 6ac88d49f1902d7a761278f627e6807c8ab6421c Mon Sep 17 00:00:00 2001 From: crisfiorini-oss <301773486+crisfiorini-oss@users.noreply.github.com> Date: Thu, 8 Oct 2026 16:29:08 -0400 Subject: [PATCH 05/14] docs: state what the bigquery backend serves, and how to run its test The backend matrix marked bigquery as planned and said mem was the only wired backend. It now lists what each wired backend serves, per RPC, and the BigQuery notes a client needs and cannot guess: the parameter syntax and value forms, the Arrow framing (one Schema message, then RecordBatch messages), the type mapping, what a listing costs, and InsertRows semantics: the per-row refusal reasons, the stream-wide row index, and at-least-once delivery after a failed call. The env table says which existing variables mean what to BigQuery and that settings it never reads are refused. AGENTS.md said there was no fake for a cloud backend, which is no longer true of bigquery, so its "say what you did not verify" rule now says what the fake shows and what only a real project does. Co-Authored-By: Claude Sonnet 5.5 --- AGENTS.md | 12 +++-- README.md | 135 ++++++++++++++++++++++++++++++++++++++++++++++-------- 2 files changed, 125 insertions(+), 22 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 3e8d740..e0d6bee 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -88,10 +88,14 @@ preferences; each is a rule an agent broke at real cost. database would have "worked" and taught the wrong thing; the fix was the guard. 6. **Say what you did not verify.** Unverified is not the same as working. No - cloud backend can be exercised from this repo — there is no BigQuery, - Snowflake, Redshift, or ClickHouse in CI and no fake for one — so a change - touching a cloud path is a change whose behaviour is unproven, and the PR - says so. + cloud backend is exercised against its real warehouse from CI — there is no + BigQuery, Snowflake, Redshift, or ClickHouse there. The `bigquery` backend is + unit-tested against a fake of BigQuery's REST API + (`internal/backend/bigquery/fake_test.go`), which shows what the backend + sends and how it reads an answer, not what BigQuery does; its integration + test runs only when `SWH_TEST_BIGQUERY_PROJECT` is set. A change touching a + cloud path is therefore unproven against the real service unless that test + was run, and the PR says which. 7. **Prove reuse before you build.** The default answer is that the platform already has it and your change fills a gap or adds a helper on it, so the PR diff --git a/README.md b/README.md index 7c2af44..959886a 100644 --- a/README.md +++ b/README.md @@ -43,8 +43,10 @@ BigQuery Storage / Snowflake / ADBC converged on). Anything backend-specific | `Native` | escape hatch for backend-specific verbs | Errors are normalized to gRPC status codes (`NotFound`, `AlreadyExists`, -`FailedPrecondition`, `Unimplemented`, `ResourceExhausted`, …) so clients never -parse a vendor SQLSTATE or reason string. +`FailedPrecondition`, `Unimplemented`, `ResourceExhausted`, `DeadlineExceeded`, +…) so clients never parse a vendor SQLSTATE or reason string. The text of a +failure is written by this server; the one piece of vendor text a client can +read is the diagnosis of its own invalid SQL (`InvalidArgument`). ## Results are Arrow @@ -65,17 +67,97 @@ schema, or a ClickHouse database — and a **table** is the leaf. | Kind | Status | Notes | |------|--------|-------| | `mem` | **working** | in-memory catalog + DDL, no query engine — the zero-dep default for tests | +| `bigquery` | **working** | catalog, DDL (incl. CTAS), queries with named parameters and a dry run, job get/cancel, streaming `InsertRows`; no `Load`, `Unload` or `Native` | | `duckdb` | skeleton | intended embedded local engine (the "MinIO of warehouses") | -| `bigquery` | planned | | | `snowflake` | planned | | | `redshift` | planned | | | `clickhouse` | planned | | -Today `mem` is the only fully-wired backend: it serves the whole catalog/DDL -plane and returns `Unimplemented` for query and data-movement, so a client -introspects `Capabilities` and never discovers the gap by a wrong answer. The -next milestone is `duckdb` (SQL execution + Arrow encoding), which makes the -"develop on DuckDB, ship on BigQuery" story real. +What each wired backend serves, per RPC. Anything marked `Unimplemented` returns +that gRPC code and says so in `Capabilities`, so a client introspects first and +never discovers a gap by a wrong answer. + +| RPC | `mem` | `bigquery` | +|-----|-------|------------| +| `Query` | `Unimplemented` | yes — named parameters, dry run, byte cap, timeout | +| `GetJob` / `CancelJob` | `NotFound` (it runs no jobs) | yes | +| `ListDatasets` / `ListTables` / `GetTable` | yes | yes | +| `CreateDataset` / `DropDataset` | yes | yes | +| `CreateTable` / `DropTable` | yes (CTAS: `Unimplemented`) | yes, including CTAS | +| `InsertRows` | `Unimplemented` | yes — streaming inserts, per-row refusals | +| `Load` / `Unload` | `Unimplemented` | `Unimplemented` | +| `Native` | `Unimplemented` | `Unimplemented` (no verbs) | + +The next milestone is `duckdb` (SQL execution + Arrow encoding), which makes the +"develop on DuckDB, ship on BigQuery" story real; until then the zero-dependency +way to run the server is `SWH_BACKEND=mem`. + +### The BigQuery backend + +Credentials are Application Default Credentials, or the service-account key +named by `SWH_CREDENTIALS_FILE`; nothing else is read. The server is bound to +the project in `SWH_DATABASE`: jobs run and are billed in it, and `dataset` names +are datasets of it. What the credentials must be allowed to do follows the RPCs a +deployment uses: run query jobs and read data (`Query`, CTAS), read table +metadata (`GetTable`, `ListTables`, and `InsertRows`, which reads it to compare +the batch with the table), create and delete datasets and tables, and stream rows +into a table. + +**Queries.** `sql` uses BigQuery's named parameters, `@name`; a parameter's `name` +is given without the sigil. `value` is text, checked here and bound by BigQuery +with the type you give, never spliced into the SQL: booleans as `true`/`false`, +integers and floats as decimal text, `NUMERIC` as a plain decimal (it becomes +`BIGNUMERIC` when only that holds it, and a value that would be rounded is +refused), bytes as standard base64, dates as `YYYY-MM-DD`, `TIME` and wall-clock +timestamps without a zone, instants as RFC 3339 with a zone, intervals as +`Y-M D H:M:S[.F]`, JSON and geography (WKT) as text. Array and struct parameters +are `Unimplemented`; sub-microsecond precision is refused, not truncated. The +query has finished by the time its header is sent (the schema is known only +then), so `CancelJob` cannot reach a running query: cancel the call, which also +cancels the job. The job `id` is opaque; hand it back to `GetJob`/`CancelJob`. +`max_bytes_billed` and `timeout_ms` are tightened by `SWH_MAX_QUERY_BYTES` and +`SWH_QUERY_TIMEOUT`: the smaller of the two applies. A `DML`/`DDL` statement +returns a header with an empty schema and `rows_affected`. + +**Results.** The header's `arrow_ipc_schema` is one Arrow IPC *Schema message*, and +each `arrow_batch` is one IPC *RecordBatch message* with no schema, up to 8192 +rows or about 1 MiB of text and bytes. The schema followed by the batches is a +valid Arrow IPC stream, so a client concatenates them and reads with any Arrow +reader. Types map as `BOOL`→bool, `INT64`→int64, `FLOAT64`→float64, +`NUMERIC`→decimal128 (`BIGNUMERIC`→decimal256), `STRING`/`JSON`/`GEOGRAPHY`→utf8, +`BYTES`→binary, `DATE`→date32, `TIME`→time64[us], `DATETIME`→timestamp[us], +`TIMESTAMP`→timestamp[us, UTC], `INTERVAL`→month-day-nano interval, +`ARRAY`→list, `STRUCT`→struct. A column with no portable bucket (BigQuery +`RANGE`) is `UNKNOWN` in the catalog with its spelling in `native_type`, and a +query that returns one is `Unimplemented`. + +**Catalog.** Listings carry BigQuery's names only, so each entry's detail is a +metadata request of its own; pages default to 50 entries and never exceed 500. +`partition_by` is one column at day granularity (`_PARTITIONTIME` for ingestion +time); more than one is `Unimplemented`. Names are checked against what BigQuery +allows before they are used. + +**`InsertRows`.** Arrow comes in the same framing: the header's schema message +and one RecordBatch message per batch, or batches that each carry their own +schema. Rows are sent as streaming inserts of at most 500 rows and about 8 MiB. +The call is not atomic: BigQuery is asked to skip invalid rows, so valid rows are +stored and refused ones come back as `errors`, each with its `row_index` (counted +from 0 across every batch of the call) and a `reason`: + +| `RowRefusal` | meaning | set by `bigquery` when | +|--------------|---------|------------------------| +| `INVALID_VALUE` | a value cannot be stored in its column | BigQuery refuses the row's content, or a value cannot be encoded (text that is not UTF-8, an out-of-range time) | +| `ROW_TOO_LARGE` | the row alone is bigger than a row may be | the row passes 8 MiB as BigQuery encodes it | +| `SCHEMA_MISMATCH` | the table, not the row, disagrees | the batch has a column the table lacks, or lacks a required column; every row of the batch is refused and none is sent | +| `UNSPECIFIED` | refused, no reason stated | not set by `bigquery` | + +A failure that says nothing against a row (a throttle, a quota, a timeout, a dead +backend) fails the call with a status code instead and is never a row refusal; +rows of earlier requests of that call may already be stored, so a retry can store +a row twice — delivery is at least once. A schema mismatch is found by comparing +the batch's columns with the table's metadata, so no vendor text decides it; a +table altered while the call runs can surface as `INVALID_VALUE`. An Arrow type +BigQuery has no equivalent of (map, duration, dictionary) is `Unimplemented`. ## Configuration (env) @@ -84,18 +166,23 @@ next milestone is `duckdb` (SQL execution + Arrow encoding), which makes the | `SWH_LISTEN` | `:9465` | gRPC listen address | | `SWH_BACKEND` | `duckdb` | `mem` \| `duckdb` \| `bigquery` \| `snowflake` \| `redshift` \| `clickhouse` | | `SWH_DATABASE` | — | required for cloud backends (BQ project / DB name); DuckDB path | -| `SWH_DATASET` | — | default namespace for unqualified names | -| `SWH_LOCATION` | — | region for datasets/jobs (BigQuery) | -| `SWH_DSN` | — | backend-native connection string (overrides discrete fields) | -| `SWH_HOST` / `SWH_PORT` | — | SQL backend host/port | -| `SWH_USER` / `SWH_PASSWORD` | — | credentials | -| `SWH_ACCOUNT` | — | Snowflake account identifier | +| `SWH_DATASET` | — | default namespace for unqualified names (for BigQuery also the default dataset of a query's SQL) | +| `SWH_LOCATION` | — | region for datasets/jobs (BigQuery: where jobs run and where a dataset is created when the request names none) | +| `SWH_DSN` | — | backend-native connection string (overrides discrete fields); `bigquery` has none and refuses it | +| `SWH_HOST` / `SWH_PORT` | — | SQL backend host/port; `bigquery` refuses them | +| `SWH_USER` / `SWH_PASSWORD` | — | credentials; `bigquery` refuses them | +| `SWH_ACCOUNT` | — | Snowflake account identifier; `bigquery` refuses it | | `SWH_CREDENTIALS_FILE` | — | BigQuery service-account JSON (else ADC) | -| `SWH_MAX_QUERY_BYTES` | `0` | per-query scan cap (0 = backend default) | -| `SWH_QUERY_TIMEOUT` | `0` | per-query timeout (Go duration; 0 = backend default) | +| `SWH_MAX_QUERY_BYTES` | `0` | per-query scan cap (0 = backend default); for BigQuery the ceiling of `maximumBytesBilled`, which a request may only lower | +| `SWH_QUERY_TIMEOUT` | `0` | per-query timeout (Go duration; 0 = backend default); a request may only shorten it | | `SWH_AUTH_TOKEN` | — | shared secret every caller must present; required unless `SWH_ALLOW_ANONYMOUS=true` | | `SWH_ALLOW_ANONYMOUS` | `false` | `true` accepts callers with no token; refused together with `SWH_AUTH_TOKEN` | +The `bigquery` backend reads `SWH_DATABASE` (required), `SWH_DATASET`, +`SWH_LOCATION`, `SWH_CREDENTIALS_FILE`, `SWH_MAX_QUERY_BYTES` and +`SWH_QUERY_TIMEOUT`, and adds no variable of its own. A setting it never reads is +a startup error, not something silently ignored. + ## Authentication The server holds the warehouse credentials and runs SQL against the database it @@ -341,7 +428,8 @@ codefly generate proto --proto ./proto --output ./gen go build ./... go vet ./... -go test -race ./... # unit tests (mem backend + bufconn server) and the agent's +go test -race ./... # unit tests: mem, a bufconn server, the agent, and the + # bigquery backend against a fake of BigQuery's REST API ``` The release-gate regression tests use `git` and `jq` (also present in CI). @@ -369,6 +457,17 @@ SWH_GATEWAY_IMAGE=service-warehouse:e2e go test -v -tags e2e -count=1 . With `SWH_GATEWAY_IMAGE` unset they skip, and a skipped run still prints `ok`: read `-v` for `--- PASS`. +The `bigquery` unit tests contact nothing: they run the real client against a +fake of BigQuery's REST API served from the test process, so they check what the +backend sends and how it reads an answer, not what BigQuery does. The one test +that talks to BigQuery, `TestAgainstRealBigQuery` in +[`internal/backend/bigquery/integration_test.go`](internal/backend/bigquery/integration_test.go), +is skipped unless `SWH_TEST_BIGQUERY_PROJECT` names a project it may use. It +creates a scratch dataset of its own with a random name, works only inside it, and +deletes it. `SWH_TEST_BIGQUERY_LOCATION` and `SWH_TEST_BIGQUERY_CREDENTIALS_FILE` +are optional (Application Default Credentials otherwise). CI runs no BigQuery, so +that test has not run in CI. + `buf` is never run on the host. `codefly generate proto` runs it inside the versioned proto companion image (`ghcr.io/codefly-dev/proto`), so the plugin versions and the `goimports` pass are the image's, fixed by its tag — two @@ -390,7 +489,7 @@ companion's — the drift gate, without putting `buf` in CI. proto/codefly/warehouse/v0/ the uniform API proto/buf.{yaml,gen.yaml} buf config, read by the proto companion gen/ generated gRPC stubs -internal/backend/ Backend interface + mem, duckdb (+ cloud backends) +internal/backend/ Backend interface + mem, duckdb, bigquery (+ more cloud backends) internal/server/ gRPC Warehouse implementation + proto↔backend mapping internal/config/ env configuration internal/serr/ normalized error model From e5bb78bf88d31d5fbd68677b3921f88e7cd175ea Mon Sep 17 00:00:00 2001 From: crisfiorini-oss <301773486+crisfiorini-oss@users.noreply.github.com> Date: Thu, 8 Oct 2026 16:34:46 -0400 Subject: [PATCH 06/14] feat(arrowipc): frame the Arrow schema and batch messages once, for every backend A Query header carries the result schema as an Arrow IPC schema message, each message after it carries one record-batch message, and InsertRows takes the same two parts. Joined, they are a valid Arrow IPC stream, but arrow-go writes and reads whole streams, so the wire form needs its own code: a writer that hands back the schema and each batch as separate messages, and a reader that joins a header schema and a source of batch messages back into a stream. Every backend that produces or ingests Arrow needs exactly this, so it lives in its own package instead of in the first backend that wants it. Dictionary-encoded fields are refused by the encoder: the wire carries no dictionary messages. Tested against arrow-go's own reader: schema plus batches round-trip, the first batch message does not repeat the schema, a self-describing stream (no header schema) reads, an empty stream reads, a source error surfaces as itself. Co-Authored-By: Claude Sonnet 5.5 --- internal/arrowipc/arrowipc.go | 182 +++++++++++++++++++++++++++++ internal/arrowipc/arrowipc_test.go | 145 +++++++++++++++++++++++ 2 files changed, 327 insertions(+) create mode 100644 internal/arrowipc/arrowipc.go create mode 100644 internal/arrowipc/arrowipc_test.go diff --git a/internal/arrowipc/arrowipc.go b/internal/arrowipc/arrowipc.go new file mode 100644 index 0000000..29b5e1e --- /dev/null +++ b/internal/arrowipc/arrowipc.go @@ -0,0 +1,182 @@ +// Package arrowipc is the one place that knows how Arrow travels on the +// Warehouse wire. A result header carries the schema as an Arrow IPC schema +// message and every following message carries one Arrow IPC record-batch +// message; both are the "encapsulated" form (continuation marker, metadata +// length, flatbuffer, body), so the schema message followed by the batch +// messages is a valid Arrow IPC stream and any Arrow reader consumes it. +// +// The package is backend-neutral: a backend that produces Arrow records uses +// Encoder, and a backend that ingests them uses NewReader, so no backend +// re-derives the framing. +package arrowipc + +import ( + "bytes" + "errors" + "fmt" + "io" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/ipc" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +// endOfStream is the Arrow IPC stream terminator: the continuation marker +// followed by a zero metadata length. +var endOfStream = [8]byte{0xFF, 0xFF, 0xFF, 0xFF, 0, 0, 0, 0} + +// messageSink is an ipc.PayloadWriter that keeps each message as its own byte +// slice instead of flattening them into one stream. +type messageSink struct { + messages [][]byte +} + +func (s *messageSink) Start() error { return nil } +func (s *messageSink) Close() error { return nil } + +func (s *messageSink) WritePayload(p ipc.Payload) error { + var buf bytes.Buffer + if _, err := p.WritePayload(&buf); err != nil { + return err + } + s.messages = append(s.messages, buf.Bytes()) + return nil +} + +func (s *messageSink) take() [][]byte { + out := s.messages + s.messages = nil + return out +} + +func rejectDictionaries(schema *arrow.Schema) error { + for _, f := range schema.Fields() { + if hasDictionary(f.Type) { + return fmt.Errorf("arrowipc: dictionary-encoded field %q is not part of the wire contract", f.Name) + } + } + return nil +} + +func hasDictionary(t arrow.DataType) bool { + switch v := t.(type) { + case *arrow.DictionaryType: + return true + case *arrow.ListType: + return hasDictionary(v.Elem()) + case *arrow.LargeListType: + return hasDictionary(v.Elem()) + case *arrow.FixedSizeListType: + return hasDictionary(v.Elem()) + case *arrow.StructType: + for _, f := range v.Fields() { + if hasDictionary(f.Type) { + return true + } + } + } + return false +} + +// SchemaMessage returns the Arrow IPC schema message for schema. +func SchemaMessage(schema *arrow.Schema) ([]byte, error) { + if err := rejectDictionaries(schema); err != nil { + return nil, err + } + sink := &messageSink{} + w := ipc.NewWriterWithPayloadWriter(sink, ipc.WithSchema(schema), ipc.WithAllocator(memory.DefaultAllocator)) + if err := w.Close(); err != nil { + return nil, fmt.Errorf("arrowipc: encode schema: %w", err) + } + msgs := sink.take() + if len(msgs) != 1 { + return nil, fmt.Errorf("arrowipc: expected one schema message, got %d", len(msgs)) + } + return msgs[0], nil +} + +// Encoder serializes record batches of one schema as record-batch messages. +type Encoder struct { + sink *messageSink + w *ipc.Writer + // schemaSeen is false until the writer has emitted the schema message that + // precedes the first batch; that message belongs in the header, not here. + schemaSeen bool +} + +// NewEncoder returns an Encoder for batches of schema. +func NewEncoder(schema *arrow.Schema) (*Encoder, error) { + if err := rejectDictionaries(schema); err != nil { + return nil, err + } + sink := &messageSink{} + return &Encoder{ + sink: sink, + w: ipc.NewWriterWithPayloadWriter(sink, ipc.WithSchema(schema), ipc.WithAllocator(memory.DefaultAllocator)), + }, nil +} + +// Encode returns the record-batch message for rec. +func (e *Encoder) Encode(rec arrow.RecordBatch) ([]byte, error) { + if err := e.w.Write(rec); err != nil { + return nil, fmt.Errorf("arrowipc: encode batch: %w", err) + } + msgs := e.sink.take() + if !e.schemaSeen { + // The writer emits the schema ahead of the first batch. + if len(msgs) == 0 { + return nil, errors.New("arrowipc: writer produced no messages") + } + msgs = msgs[1:] + e.schemaSeen = true + } + if len(msgs) != 1 { + return nil, fmt.Errorf("arrowipc: expected one batch message, got %d", len(msgs)) + } + return msgs[0], nil +} + +// Close releases the encoder. +func (e *Encoder) Close() error { return e.w.Close() } + +// Source yields the messages that follow a schema, one per Next call, and +// io.EOF when there are no more. backend.BatchReader satisfies it. +type Source interface { + Next() ([]byte, error) +} + +// NewReader reads the records of a stream whose schema message is schema and +// whose record-batch messages come from src. schema may be empty when src's +// first message is self-describing (a schema message followed by batches). +func NewReader(schema []byte, src Source) (*ipc.Reader, error) { + return ipc.NewReader(&messageStream{pending: schema, src: src}, ipc.WithAllocator(memory.DefaultAllocator)) +} + +// messageStream presents a header schema plus a Source of messages as one +// contiguous Arrow IPC stream, closed with the end-of-stream marker. +type messageStream struct { + pending []byte + src Source + done bool +} + +func (m *messageStream) Read(p []byte) (int, error) { + for len(m.pending) == 0 { + if m.done { + return 0, io.EOF + } + next, err := m.src.Next() + switch { + case errors.Is(err, io.EOF): + m.done = true + m.pending = endOfStream[:] + case err != nil: + return 0, err + default: + m.pending = next + } + } + n := copy(p, m.pending) + m.pending = m.pending[n:] + return n, nil +} diff --git a/internal/arrowipc/arrowipc_test.go b/internal/arrowipc/arrowipc_test.go new file mode 100644 index 0000000..8783f73 --- /dev/null +++ b/internal/arrowipc/arrowipc_test.go @@ -0,0 +1,145 @@ +package arrowipc + +import ( + "errors" + "io" + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/stretchr/testify/require" +) + +type msgs struct { + list [][]byte + err error +} + +func (m *msgs) Next() ([]byte, error) { + if len(m.list) == 0 { + if m.err != nil { + return nil, m.err + } + return nil, io.EOF + } + next := m.list[0] + m.list = m.list[1:] + return next, nil +} + +var testSchema = arrow.NewSchema([]arrow.Field{ + {Name: "id", Type: arrow.PrimitiveTypes.Int64}, + {Name: "name", Type: arrow.BinaryTypes.String, Nullable: true}, +}, nil) + +func batch(ids []int64, names []string) arrow.RecordBatch { + b := array.NewRecordBuilder(memory.DefaultAllocator, testSchema) + defer b.Release() + b.Field(0).(*array.Int64Builder).AppendValues(ids, nil) + b.Field(1).(*array.StringBuilder).AppendValues(names, nil) + return b.NewRecordBatch() +} + +func TestSchemaThenBatchesIsAStream(t *testing.T) { + enc, err := NewEncoder(testSchema) + require.NoError(t, err) + defer enc.Close() + + first, second := batch([]int64{1, 2}, []string{"a", "b"}), batch([]int64{3}, []string{"c"}) + defer first.Release() + defer second.Release() + m1, err := enc.Encode(first) + require.NoError(t, err) + m2, err := enc.Encode(second) + require.NoError(t, err) + + header, err := SchemaMessage(testSchema) + require.NoError(t, err) + + rd, err := NewReader(header, &msgs{list: [][]byte{m1, m2}}) + require.NoError(t, err) + defer rd.Release() + require.True(t, rd.Schema().Equal(testSchema)) + + var ids []int64 + for rd.Next() { + col := rd.RecordBatch().Column(0).(*array.Int64) + for i := 0; i < col.Len(); i++ { + ids = append(ids, col.Value(i)) + } + } + require.NoError(t, rd.Err()) + require.Equal(t, []int64{1, 2, 3}, ids) +} + +// The first batch message must not carry the schema: it is in the header, and a +// reader that is handed it twice would see a second schema where a batch belongs. +func TestBatchMessagesDoNotCarryTheSchema(t *testing.T) { + enc, err := NewEncoder(testSchema) + require.NoError(t, err) + defer enc.Close() + rec := batch([]int64{1}, []string{"a"}) + defer rec.Release() + m, err := enc.Encode(rec) + require.NoError(t, err) + + // A batch message alone is not a stream: with no schema ahead of it, it + // cannot be read. + _, err = NewReader(nil, &msgs{list: [][]byte{m}}) + require.Error(t, err) +} + +func TestSelfDescribingStream(t *testing.T) { + enc, err := NewEncoder(testSchema) + require.NoError(t, err) + defer enc.Close() + rec := batch([]int64{7}, []string{"x"}) + defer rec.Release() + m, err := enc.Encode(rec) + require.NoError(t, err) + header, err := SchemaMessage(testSchema) + require.NoError(t, err) + + // With no header schema, the first message carries it. + rd, err := NewReader(nil, &msgs{list: [][]byte{append(append([]byte(nil), header...), m...)}}) + require.NoError(t, err) + defer rd.Release() + require.True(t, rd.Next()) + require.Equal(t, int64(7), rd.RecordBatch().Column(0).(*array.Int64).Value(0)) +} + +func TestNoBatchesIsAnEmptyStream(t *testing.T) { + header, err := SchemaMessage(testSchema) + require.NoError(t, err) + rd, err := NewReader(header, &msgs{}) + require.NoError(t, err) + defer rd.Release() + require.False(t, rd.Next()) + require.NoError(t, rd.Err()) +} + +func TestSourceErrorSurfaces(t *testing.T) { + boom := errors.New("transport failed") + header, err := SchemaMessage(testSchema) + require.NoError(t, err) + rd, err := NewReader(header, &msgs{err: boom}) + if err == nil { + defer rd.Release() + for rd.Next() { + } + err = rd.Err() + } + require.ErrorIs(t, err, boom) +} + +func TestDictionaryEncodingIsRefused(t *testing.T) { + schema := arrow.NewSchema([]arrow.Field{{ + Name: "c", + Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.BinaryTypes.String}, + }}, nil) + _, err := SchemaMessage(schema) + require.Error(t, err) + _, err = NewEncoder(schema) + require.Error(t, err) +} From 213249a7ff91485d5d90cc5410a24016e73e622c Mon Sep 17 00:00:00 2001 From: crisfiorini-oss <301773486+crisfiorini-oss@users.noreply.github.com> Date: Thu, 8 Oct 2026 16:44:31 -0400 Subject: [PATCH 07/14] feat(arrowipc): tell a client's bad Arrow from a source that failed, and read a header that ends the stream Two backends read InsertRows through this package, and each needs the same three things from it, which the first version left to the caller. A header schema that arrives terminated by the end-of-stream marker was read as the end of the stream: the batches after it were never reached and every row was dropped without an error. The marker is now removed from the header before it is read. A failure to read was one undifferentiated error, so a caller could not return the transport's own error as itself and report only the client's malformed bytes as an invalid argument. NewReader now returns a Reader whose errors are the source's own error when the source failed, and wrap ErrMalformed when the bytes did. The Reader has the read methods of arrow-go's ipc.Reader, so it is a drop-in for a caller that only reads. An empty message is refused as malformed. Skipping it would hide a client that sent a batch it never filled in. Tested: a terminated header still yields its rows; each kind of bad input (not a schema, a batch cut short, bytes that are not a message, a second schema where a batch belongs, an empty message, nothing at all) is ErrMalformed; a failed source is not. Co-Authored-By: Claude Sonnet 5.5 --- internal/arrowipc/arrowipc.go | 93 +++++++++++++++++++++++++++++- internal/arrowipc/arrowipc_test.go | 87 +++++++++++++++++++++++++++- 2 files changed, 176 insertions(+), 4 deletions(-) diff --git a/internal/arrowipc/arrowipc.go b/internal/arrowipc/arrowipc.go index 29b5e1e..29f5fa3 100644 --- a/internal/arrowipc/arrowipc.go +++ b/internal/arrowipc/arrowipc.go @@ -145,11 +145,96 @@ type Source interface { Next() ([]byte, error) } +// ErrMalformed marks a failure to read the bytes as an Arrow IPC stream: not a +// schema where one belongs, a batch that does not match it, a message cut short, +// an empty message. It is the client's data that is wrong. A Source that fails to +// deliver the bytes at all is not malformed data, and its error comes back as +// itself, so a caller maps ErrMalformed to an invalid-argument and returns +// anything else unchanged. +var ErrMalformed = errors.New("malformed Arrow IPC stream") + +// errEmptyMessage is what an empty message is: not a batch, and not skippable. +var errEmptyMessage = errors.New("an empty message is not an Arrow IPC message") + +// Reader reads the records of an Arrow IPC stream that arrives as wire messages. +// It has the read side of an arrow-go ipc.Reader, with errors classified: see +// ErrMalformed. +type Reader struct { + r *ipc.Reader + src *watchedSource + err error +} + // NewReader reads the records of a stream whose schema message is schema and // whose record-batch messages come from src. schema may be empty when src's -// first message is self-describing (a schema message followed by batches). -func NewReader(schema []byte, src Source) (*ipc.Reader, error) { - return ipc.NewReader(&messageStream{pending: schema, src: src}, ipc.WithAllocator(memory.DefaultAllocator)) +// first message is self-describing (a schema message followed by batches). A +// schema that arrives terminated by the end-of-stream marker is read as the +// schema it is: left on, the marker would end the stream before the first batch +// and every row would be silently dropped. +// +// It reads the schema before it returns, so a first message that is not one is +// reported here. +func NewReader(schema []byte, src Source) (*Reader, error) { + w := &watchedSource{src: src} + r, err := ipc.NewReader( + &messageStream{pending: bytes.TrimSuffix(schema, endOfStream[:]), src: w}, + ipc.WithAllocator(memory.DefaultAllocator), + ) + if err != nil { + return nil, w.classify(err) + } + return &Reader{r: r, src: w}, nil +} + +// Schema is the schema every record has. +func (r *Reader) Schema() *arrow.Schema { return r.r.Schema() } + +// Next reads the next record, and reports false at the end of the stream or on +// an error; Err says which. The record is valid until the next call to Next. +func (r *Reader) Next() bool { + if r.err != nil { + return false + } + if r.r.Next() { + return true + } + if err := r.r.Err(); err != nil { + r.err = r.src.classify(err) + } + return false +} + +// RecordBatch is the record Next read. +func (r *Reader) RecordBatch() arrow.RecordBatch { return r.r.RecordBatch() } + +// Err is nil at a clean end of stream. +func (r *Reader) Err() error { return r.err } + +// Release releases the reader and its current record. +func (r *Reader) Release() { r.r.Release() } + +// watchedSource remembers why its Source failed, so the failure is returned as +// itself and not as a complaint about Arrow. +type watchedSource struct { + src Source + err error +} + +func (w *watchedSource) Next() ([]byte, error) { + msg, err := w.src.Next() + if err != nil && !errors.Is(err, io.EOF) && w.err == nil { + w.err = err + } + return msg, err +} + +// classify is the source's own error when the source failed, and ErrMalformed +// around err when the bytes it delivered are what failed to read. +func (w *watchedSource) classify(err error) error { + if w.err != nil { + return w.err + } + return fmt.Errorf("%w: %w", ErrMalformed, err) } // messageStream presents a header schema plus a Source of messages as one @@ -172,6 +257,8 @@ func (m *messageStream) Read(p []byte) (int, error) { m.pending = endOfStream[:] case err != nil: return 0, err + case len(next) == 0: + return 0, errEmptyMessage default: m.pending = next } diff --git a/internal/arrowipc/arrowipc_test.go b/internal/arrowipc/arrowipc_test.go index 8783f73..eb12b4a 100644 --- a/internal/arrowipc/arrowipc_test.go +++ b/internal/arrowipc/arrowipc_test.go @@ -87,7 +87,7 @@ func TestBatchMessagesDoNotCarryTheSchema(t *testing.T) { // A batch message alone is not a stream: with no schema ahead of it, it // cannot be read. _, err = NewReader(nil, &msgs{list: [][]byte{m}}) - require.Error(t, err) + require.ErrorIs(t, err, ErrMalformed) } func TestSelfDescribingStream(t *testing.T) { @@ -131,6 +131,7 @@ func TestSourceErrorSurfaces(t *testing.T) { err = rd.Err() } require.ErrorIs(t, err, boom) + require.NotErrorIs(t, err, ErrMalformed, "a source that failed is not data that is wrong") } func TestDictionaryEncodingIsRefused(t *testing.T) { @@ -143,3 +144,87 @@ func TestDictionaryEncodingIsRefused(t *testing.T) { _, err = NewEncoder(schema) require.Error(t, err) } + +// A header schema that is terminated by the end-of-stream marker is still the +// schema: read as a stream, the marker would end it before the first batch and +// the rows would vanish without an error. +func TestHeaderSchemaThatEndsTheStreamIsStillTheSchema(t *testing.T) { + enc, err := NewEncoder(testSchema) + require.NoError(t, err) + defer enc.Close() + rec := batch([]int64{1, 2}, []string{"a", "b"}) + defer rec.Release() + m, err := enc.Encode(rec) + require.NoError(t, err) + header, err := SchemaMessage(testSchema) + require.NoError(t, err) + + rd, err := NewReader(append(append([]byte(nil), header...), endOfStream[:]...), &msgs{list: [][]byte{m}}) + require.NoError(t, err) + defer rd.Release() + var rows int64 + for rd.Next() { + rows += rd.RecordBatch().NumRows() + } + require.NoError(t, rd.Err()) + require.EqualValues(t, 2, rows) +} + +func TestMalformedBytesAreErrMalformed(t *testing.T) { + enc, err := NewEncoder(testSchema) + require.NoError(t, err) + defer enc.Close() + rec := batch([]int64{1}, []string{"a"}) + defer rec.Release() + m, err := enc.Encode(rec) + require.NoError(t, err) + header, err := SchemaMessage(testSchema) + require.NoError(t, err) + + for name, tc := range map[string]struct { + header []byte + list [][]byte + }{ + "a header that is not a schema": {[]byte{0xff, 0xff, 0xff, 0xff, 4, 0, 0, 0, 1, 2, 3, 4}, [][]byte{m}}, + "a batch cut short": {header, [][]byte{m[:len(m)/2]}}, + "a message that is not a message": {header, [][]byte{{0xff, 0xff, 0xff, 0xff, 0x10, 0, 0, 0, 1, 2, 3}}}, + "a second schema where a batch belongs": {header, [][]byte{append(append([]byte(nil), header...), m...)}}, + "an empty message": {header, [][]byte{{}}}, + "no schema and no message": {nil, nil}, + } { + t.Run(name, func(t *testing.T) { + rd, err := NewReader(tc.header, &msgs{list: tc.list}) + if err == nil { + defer rd.Release() + for rd.Next() { + } + err = rd.Err() + } + require.ErrorIs(t, err, ErrMalformed) + }) + } +} + +// A message that follows an empty one is not read past it: skipping would hide a +// client that sent a batch it never filled in. +func TestEmptyMessageAfterBatchesStopsTheStream(t *testing.T) { + enc, err := NewEncoder(testSchema) + require.NoError(t, err) + defer enc.Close() + rec := batch([]int64{1}, []string{"a"}) + defer rec.Release() + m, err := enc.Encode(rec) + require.NoError(t, err) + header, err := SchemaMessage(testSchema) + require.NoError(t, err) + + rd, err := NewReader(header, &msgs{list: [][]byte{m, {}, m}}) + require.NoError(t, err) + defer rd.Release() + var rows int64 + for rd.Next() { + rows += rd.RecordBatch().NumRows() + } + require.EqualValues(t, 1, rows) + require.ErrorIs(t, rd.Err(), ErrMalformed) +} From d789c26a262ad83b9adac800dde4cb2f5bca049c Mon Sep 17 00:00:00 2001 From: crisfiorini-oss <301773486+crisfiorini-oss@users.noreply.github.com> Date: Thu, 8 Oct 2026 16:46:23 -0400 Subject: [PATCH 08/14] refactor(bigquery): frame Arrow through arrowipc, and drop the backend's own copy The bigquery backend carried its own Arrow IPC framing, a writer that cut the schema and each batch out of a stream and a reader that rejoined them, and internal/arrowipc now does exactly that for every backend. Two copies of how Arrow travels on the wire is one too many: the next backend would have picked whichever it read first. Query writes its header schema with arrowipc.SchemaMessage and each batch with an arrowipc.Encoder. InsertRows reads with arrowipc.NewReader and returns a malformed stream as InvalidArgument and a source that failed as itself, which is what the backend's own reader did. One thing narrows, to what the proto already says. InsertRows used to also accept a header schema together with batches that each repeat their own schema; the contract is the header's schema and then bare batches, or an empty header and a first message that carries the schema, and that is what is read now. A second schema where a batch belongs is InvalidArgument. The framing tests that belonged to the copy are gone with it; arrowipc's own tests cover the same cases, and the backend's insert tests still drive the whole path, now with a test that a failed transport comes back as itself. Co-Authored-By: Claude Sonnet 5.5 --- README.md | 13 +- internal/backend/bigquery/arrow.go | 208 +++--------------- internal/backend/bigquery/arrow_test.go | 82 ------- internal/backend/bigquery/insert.go | 40 ++-- internal/backend/bigquery/insert_test.go | 47 +++- internal/backend/bigquery/integration_test.go | 21 +- 6 files changed, 110 insertions(+), 301 deletions(-) diff --git a/README.md b/README.md index 959886a..899b26a 100644 --- a/README.md +++ b/README.md @@ -138,11 +138,13 @@ time); more than one is `Unimplemented`. Names are checked against what BigQuery allows before they are used. **`InsertRows`.** Arrow comes in the same framing: the header's schema message -and one RecordBatch message per batch, or batches that each carry their own -schema. Rows are sent as streaming inserts of at most 500 rows and about 8 MiB. -The call is not atomic: BigQuery is asked to skip invalid rows, so valid rows are -stored and refused ones come back as `errors`, each with its `row_index` (counted -from 0 across every batch of the call) and a `reason`: +and one RecordBatch message per batch, or, with no schema in the header, a first +message that is a schema message followed by its batch. A message that is not +Arrow (cut short, empty, a second schema where a batch belongs) is +`InvalidArgument`. Rows are sent as streaming inserts of at most 500 rows and +about 8 MiB. The call is not atomic: BigQuery is asked to skip invalid rows, so +valid rows are stored and refused ones come back as `errors`, each with its +`row_index` (counted from 0 across every batch of the call) and a `reason`: | `RowRefusal` | meaning | set by `bigquery` when | |--------------|---------|------------------------| @@ -490,6 +492,7 @@ proto/codefly/warehouse/v0/ the uniform API proto/buf.{yaml,gen.yaml} buf config, read by the proto companion gen/ generated gRPC stubs internal/backend/ Backend interface + mem, duckdb, bigquery (+ more cloud backends) +internal/arrowipc/ the Arrow IPC framing every backend shares (schema message, batch messages) internal/server/ gRPC Warehouse implementation + proto↔backend mapping internal/config/ env configuration internal/serr/ normalized error model diff --git a/internal/backend/bigquery/arrow.go b/internal/backend/bigquery/arrow.go index 78b84e3..e03a8f3 100644 --- a/internal/backend/bigquery/arrow.go +++ b/internal/backend/bigquery/arrow.go @@ -1,10 +1,7 @@ package bigquery import ( - "bytes" - "errors" "fmt" - "io" "math/big" "time" @@ -14,9 +11,9 @@ import ( "github.com/apache/arrow-go/v18/arrow/array" "github.com/apache/arrow-go/v18/arrow/decimal128" "github.com/apache/arrow-go/v18/arrow/decimal256" - "github.com/apache/arrow-go/v18/arrow/ipc" "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/codefly-dev/service-warehouse/internal/arrowipc" "github.com/codefly-dev/service-warehouse/internal/backend" "github.com/codefly-dev/service-warehouse/internal/serr" ) @@ -25,11 +22,9 @@ import ( // result values into Arrow-IPC batches for Query, and Arrow-IPC batches into // the values a streaming insert sends for InsertRows. // -// Framing, which is the contract a client reads against: the header carries one -// encapsulated IPC Schema message, and every batch is one encapsulated IPC -// RecordBatch message without the schema. The schema followed by the batches is -// a valid Arrow IPC stream (an end-of-stream marker is optional), which is how -// BigQuery's own Storage Read API frames Arrow too. +// The framing, which is the contract a client reads against, is not here: the +// header's schema and each batch as an IPC message are written and read by +// internal/arrowipc, the one place that knows how Arrow travels on the wire. const ( // batchMaxRows and batchTargetBytes bound one result batch, so a wide table @@ -39,61 +34,12 @@ const ( batchTargetBytes = 1 << 20 ) -// eos is the 8-byte end-of-stream marker an IPC stream writer appends: the -// continuation indicator followed by a zero metadata length. -var eos = []byte{0xff, 0xff, 0xff, 0xff, 0, 0, 0, 0} - -// framer writes the schema message once and each record batch as a message of -// its own, by running one IPC stream writer and cutting what it emits. -type framer struct { - schema *arrow.Schema - schemaMsg []byte - buf bytes.Buffer - writer *ipc.Writer - primed bool -} - -func newFramer(schema *arrow.Schema, mem memory.Allocator) (*framer, error) { - // A stream with no batches is the schema message and the end marker. - var only bytes.Buffer - empty := ipc.NewWriter(&only, ipc.WithSchema(schema), ipc.WithAllocator(mem)) - if err := empty.Close(); err != nil { - return nil, serr.Wrap(serr.Internal, "arrow", fmt.Errorf("encode the schema: %w", err)) - } - encoded := only.Bytes() - if !bytes.HasSuffix(encoded, eos) { - return nil, serr.New(serr.Internal, "arrow", "the IPC writer's stream does not end with the end-of-stream marker") - } - f := &framer{schema: schema, schemaMsg: append([]byte(nil), encoded[:len(encoded)-len(eos)]...)} - f.writer = ipc.NewWriter(&f.buf, ipc.WithSchema(schema), ipc.WithAllocator(mem)) - return f, nil -} - -// batch returns rec as one RecordBatch message. -func (f *framer) batch(rec arrow.RecordBatch) ([]byte, error) { - f.buf.Reset() - if err := f.writer.Write(rec); err != nil { - return nil, serr.Wrap(serr.Internal, "arrow", fmt.Errorf("encode a batch: %w", err)) - } - out := f.buf.Bytes() - if !f.primed { - // The first write also carries the schema message, which the header - // already sent. - if !bytes.HasPrefix(out, f.schemaMsg) { - return nil, serr.New(serr.Internal, "arrow", "the IPC writer's first batch does not start with the schema message") - } - out = out[len(f.schemaMsg):] - f.primed = true - } - return append([]byte(nil), out...), nil -} - // resultEncoder turns BigQuery rows into Arrow batches. type resultEncoder struct { cols []backend.Column - mem memory.Allocator + header []byte // the schema as an IPC message builder *array.RecordBuilder - framer *framer + encoder *arrowipc.Encoder rows int size int } @@ -103,21 +49,24 @@ func newResultEncoder(cols []backend.Column) (*resultEncoder, error) { if err != nil { return nil, err } - mem := memory.NewGoAllocator() - fr, err := newFramer(schema, mem) + header, err := arrowipc.SchemaMessage(schema) if err != nil { - return nil, err + return nil, serr.Wrap(serr.Internal, "arrow", err) + } + fr, err := arrowipc.NewEncoder(schema) + if err != nil { + return nil, serr.Wrap(serr.Internal, "arrow", err) } return &resultEncoder{ cols: cols, - mem: mem, - builder: array.NewRecordBuilder(mem, schema), - framer: fr, + header: header, + builder: array.NewRecordBuilder(memory.NewGoAllocator(), schema), + encoder: fr, }, nil } // schemaMessage is the IPC Schema message for the result header. -func (e *resultEncoder) schemaMessage() []byte { return e.framer.schemaMsg } +func (e *resultEncoder) schemaMessage() []byte { return e.header } // add appends one row, which must carry one value per column. func (e *resultEncoder) add(row []bq.Value) error { @@ -143,10 +92,17 @@ func (e *resultEncoder) flush() ([]byte, error) { rec := e.builder.NewRecordBatch() defer rec.Release() e.rows, e.size = 0, 0 - return e.framer.batch(rec) + out, err := e.encoder.Encode(rec) + if err != nil { + return nil, serr.Wrap(serr.Internal, "arrow", err) + } + return out, nil } -func (e *resultEncoder) release() { e.builder.Release() } +func (e *resultEncoder) release() { + e.builder.Release() + _ = e.encoder.Close() +} // valueSize is a cheap estimate of a value's weight, enough to bound a batch. func valueSize(v bq.Value) int { @@ -333,117 +289,3 @@ func appendDecimal(b array.Builder, c backend.Column, r *big.Rat) error { } return nil } - -// batchStream reads the Arrow batches of an InsertRows call. The header's schema -// message, when there is one, completes each batch; a batch that carries its own -// schema is read as the stream it is. -type batchStream struct { - mem memory.Allocator - header []byte // the header's Schema message with any end marker removed; nil when absent - src backend.BatchReader - - schema *arrow.Schema // the schema every record must have, known after the first - cur *ipc.Reader -} - -func newBatchStream(header []byte, src backend.BatchReader) (*batchStream, error) { - s := &batchStream{mem: memory.NewGoAllocator(), src: src} - if len(header) == 0 { - return s, nil - } - s.header = bytes.TrimSuffix(header, eos) - schema, err := s.readSchema(bytes.NewReader(s.header)) - if err != nil { - return nil, err - } - s.schema = schema - return s, nil -} - -func invalidBatch(err error) error { - return serr.Wrap(serr.InvalidArgument, "InsertRows", fmt.Errorf("the Arrow data is malformed: %w", err)) -} - -// guard turns a panic inside the Arrow reader, which is fed bytes from the -// network, into an invalid-argument error. -func guard(err *error) { - if r := recover(); r != nil { - *err = invalidBatch(fmt.Errorf("%v", r)) - } -} - -func (s *batchStream) readSchema(r io.Reader) (schema *arrow.Schema, err error) { - defer guard(&err) - reader, err := ipc.NewReader(r, ipc.WithAllocator(s.mem)) - if err != nil { - return nil, invalidBatch(err) - } - defer reader.Release() - return reader.Schema(), nil -} - -// carriesSchema reports whether an IPC payload begins with a Schema message. -func carriesSchema(payload []byte) bool { - messages := ipc.NewMessageReader(bytes.NewReader(payload)) - defer messages.Release() - message, err := messages.Message() - if err != nil { - return false - } - defer message.Release() - return message.Type() == ipc.MessageSchema -} - -// next returns the next record batch, or io.EOF when the stream is exhausted. The -// caller releases the record. -func (s *batchStream) next() (rec arrow.RecordBatch, err error) { - defer guard(&err) - for { - if s.cur != nil { - if s.cur.Next() { - rec := s.cur.RecordBatch() - if s.schema == nil { - s.schema = rec.Schema() - } else if !s.schema.Equal(rec.Schema()) { - return nil, serr.New(serr.InvalidArgument, "InsertRows", "a batch's schema differs from the schema of the batches before it") - } - rec.Retain() - return rec, nil - } - err := s.cur.Err() - s.cur.Release() - s.cur = nil - if err != nil { - return nil, invalidBatch(err) - } - } - payload, err := s.src.Next() - if errors.Is(err, io.EOF) { - return nil, io.EOF - } - if err != nil { - return nil, err - } - var stream io.Reader - switch { - case carriesSchema(payload): - stream = bytes.NewReader(payload) - case s.header != nil: - stream = io.MultiReader(bytes.NewReader(s.header), bytes.NewReader(payload)) - default: - return nil, serr.New(serr.InvalidArgument, "InsertRows", "a batch carries no schema and the header gave none") - } - reader, err := ipc.NewReader(stream, ipc.WithAllocator(s.mem)) - if err != nil { - return nil, invalidBatch(err) - } - s.cur = reader - } -} - -func (s *batchStream) close() { - if s.cur != nil { - s.cur.Release() - s.cur = nil - } -} diff --git a/internal/backend/bigquery/arrow_test.go b/internal/backend/bigquery/arrow_test.go index 143f0c6..4cbcd97 100644 --- a/internal/backend/bigquery/arrow_test.go +++ b/internal/backend/bigquery/arrow_test.go @@ -2,7 +2,6 @@ package bigquery import ( "bytes" - "errors" "io" "math" "math/big" @@ -183,87 +182,6 @@ func TestResultEncoderRefusesAValueOfTheWrongType(t *testing.T) { require.Error(t, enc.add([]bq.Value{}), "a row with too few values") } -func TestFramingIsTheSchemaMessageThenEachBatchMessage(t *testing.T) { - schema := arrow.NewSchema([]arrow.Field{{Name: "id", Type: arrow.PrimitiveTypes.Int64}}, nil) - mem := memory.NewGoAllocator() - fr, err := newFramer(schema, mem) - require.NoError(t, err) - require.True(t, carriesSchema(fr.schemaMsg)) - require.False(t, bytes.HasSuffix(fr.schemaMsg, eos), "the schema message is not terminated") - - build := func(vals ...int64) []byte { - b := array.NewRecordBuilder(mem, schema) - defer b.Release() - for _, v := range vals { - b.Field(0).(*array.Int64Builder).Append(v) - } - rec := b.NewRecordBatch() - defer rec.Release() - out, err := fr.batch(rec) - require.NoError(t, err) - return out - } - first, second := build(1, 2), build(3) - require.False(t, carriesSchema(first), "a batch does not repeat the schema") - require.False(t, carriesSchema(second)) - require.False(t, bytes.HasPrefix(second, fr.schemaMsg)) - - // The schema and the batches together are one valid Arrow IPC stream, which - // is what lets a client hand them to any Arrow reader. - stream := io.MultiReader(bytes.NewReader(fr.schemaMsg), bytes.NewReader(first), bytes.NewReader(second), bytes.NewReader(eos)) - reader, err := ipc.NewReader(stream) - require.NoError(t, err) - defer reader.Release() - var seen []int64 - for reader.Next() { - col := reader.RecordBatch().Column(0).(*array.Int64) - for i := 0; i < col.Len(); i++ { - seen = append(seen, col.Value(i)) - } - } - require.NoError(t, reader.Err()) - require.Equal(t, []int64{1, 2, 3}, seen) -} - -func TestBatchStreamReadsHeaderSchemaWithBatchesAndSelfDescribingStreams(t *testing.T) { - header, batch := eventBatch(t, []int64{1, 2}, []string{"a", "b"}) - - read := func(header []byte, payloads ...[]byte) (rows int64, err error) { - s, err := newBatchStream(header, &sliceReader{payloads: payloads}) - if err != nil { - return 0, err - } - defer s.close() - for { - rec, err := s.next() - if errors.Is(err, io.EOF) { - return rows, nil - } - if err != nil { - return rows, err - } - rows += rec.NumRows() - rec.Release() - } - } - - n, err := read(header, batch, batch) - require.NoError(t, err) - require.EqualValues(t, 4, n) - - n, err = read(nil, append(append([]byte{}, header...), batch...), batch[:0:0]) - require.Error(t, err, "an empty payload is not a batch") - _ = n - - n, err = read(nil, append(append(append([]byte{}, header...), batch...), eos...)) - require.NoError(t, err) - require.EqualValues(t, 2, n) - - n, err = read(header) - require.NoError(t, err) - require.Zero(t, n, "no batches is no rows") -} - func TestJSONValuesOfTheCellsAStreamingInsertSends(t *testing.T) { mem := memory.NewGoAllocator() build := func(dt arrow.DataType, fill func(b array.Builder)) arrow.Array { diff --git a/internal/backend/bigquery/insert.go b/internal/backend/bigquery/insert.go index 094ce9c..8e0f427 100644 --- a/internal/backend/bigquery/insert.go +++ b/internal/backend/bigquery/insert.go @@ -5,12 +5,12 @@ import ( "encoding/json" "errors" "fmt" - "io" "strings" bq "cloud.google.com/go/bigquery" "github.com/apache/arrow-go/v18/arrow" + "github.com/codefly-dev/service-warehouse/internal/arrowipc" "github.com/codefly-dev/service-warehouse/internal/backend" "github.com/codefly-dev/service-warehouse/internal/serr" ) @@ -81,11 +81,11 @@ func (b *Backend) InsertRows(ctx context.Context, ref backend.TableRef, schema [ if err != nil { return nil, err } - stream, err := newBatchStream(schema, batches) + stream, err := arrowipc.NewReader(schema, batches) if err != nil { - return nil, err + return nil, streamError(err) } - defer stream.close() + defer stream.Release() table := b.client.Dataset(ref.Dataset).Table(ref.Table) md, err := table.Metadata(ctx) @@ -99,9 +99,19 @@ func (b *Backend) InsertRows(ctx context.Context, ref backend.TableRef, schema [ return insertAll(ctx, stream, md.Schema, inserter) } +// streamError is the error to return for a failure to read the call's Arrow: the +// client's malformed data is an invalid argument, and anything else, such as the +// transport failing, is returned as itself. +func streamError(err error) error { + if errors.Is(err, arrowipc.ErrMalformed) { + return serr.Wrap(serr.InvalidArgument, "InsertRows", fmt.Errorf("the Arrow data is malformed: %w", err)) + } + return err +} + // insertAll runs the stream through plan, send and the classification of what // BigQuery answered. -func insertAll(ctx context.Context, stream *batchStream, table bq.Schema, inserter rowInserter) (*backend.InsertResult, error) { +func insertAll(ctx context.Context, stream *arrowipc.Reader, table bq.Schema, inserter rowInserter) (*backend.InsertResult, error) { const op = "InsertRows" var ( result backend.InsertResult @@ -122,24 +132,18 @@ func insertAll(ctx context.Context, stream *batchStream, table bq.Schema, insert return err } - for { - rec, err := stream.next() - if errors.Is(err, io.EOF) { - break - } - if err != nil { - return nil, err - } + // The stream owns each record it hands out until the next call, so none is + // released here. + for stream.Next() { + rec := stream.RecordBatch() if !checked { checked = true if err := checkArrowSchema(rec.Schema()); err != nil { - rec.Release() return nil, serr.New(serr.Unsupported, op, err.Error()) } if name, twice := repeatedColumn(rec.Schema()); twice { // Two columns of one name would reach BigQuery as one, and a value // would be dropped without a word. - rec.Release() return nil, serr.New(serr.InvalidArgument, op, fmt.Sprintf("the batch has two columns named %q (BigQuery names are not case-sensitive)", name)) } conflict = schemaConflict(rec.Schema(), table) @@ -152,7 +156,6 @@ func insertAll(ctx context.Context, stream *batchStream, table bq.Schema, insert }) } next += int64(rows) - rec.Release() continue } for r := 0; r < rows; r++ { @@ -163,7 +166,6 @@ func insertAll(ctx context.Context, stream *batchStream, table bq.Schema, insert } if len(request) == maxRowsPerInsert || reqBytes+row.size > maxRequestBytes { if err := flush(); err != nil { - rec.Release() return nil, err } } @@ -171,7 +173,9 @@ func insertAll(ctx context.Context, stream *batchStream, table bq.Schema, insert reqBytes += row.size } next += int64(rows) - rec.Release() + } + if err := stream.Err(); err != nil { + return nil, streamError(err) } if err := flush(); err != nil { return nil, err diff --git a/internal/backend/bigquery/insert_test.go b/internal/backend/bigquery/insert_test.go index 087e39c..5cde053 100644 --- a/internal/backend/bigquery/insert_test.go +++ b/internal/backend/bigquery/insert_test.go @@ -12,6 +12,7 @@ import ( "github.com/apache/arrow-go/v18/arrow/memory" "github.com/stretchr/testify/require" + "github.com/codefly-dev/service-warehouse/internal/arrowipc" "github.com/codefly-dev/service-warehouse/internal/backend" "github.com/codefly-dev/service-warehouse/internal/serr" ) @@ -36,17 +37,19 @@ func (r *sliceReader) Close() error { return nil } // message and the batch message. func batchOf(t *testing.T, schema *arrow.Schema, fill func(b *array.RecordBuilder)) (schemaMsg, batch []byte) { t.Helper() - mem := memory.NewGoAllocator() - builder := array.NewRecordBuilder(mem, schema) + builder := array.NewRecordBuilder(memory.NewGoAllocator(), schema) defer builder.Release() fill(builder) rec := builder.NewRecordBatch() defer rec.Release() - fr, err := newFramer(schema, mem) + schemaMsg, err := arrowipc.SchemaMessage(schema) require.NoError(t, err) - batch, err = fr.batch(rec) + enc, err := arrowipc.NewEncoder(schema) require.NoError(t, err) - return fr.schemaMsg, batch + defer enc.Close() + batch, err = enc.Encode(rec) + require.NoError(t, err) + return schemaMsg, batch } var eventSchema = arrow.NewSchema([]arrow.Field{ @@ -351,7 +354,8 @@ func TestInsertRowsAcceptsAHeaderSchemaThatEndsTheStream(t *testing.T) { b := f.open(t, backend.Config{}) header, batch := eventBatch(t, []int64{1}, []string{"a"}) - res, err := insertRows(t, b, append(append([]byte{}, header...), eos...), batch) + endOfStream := []byte{0xff, 0xff, 0xff, 0xff, 0, 0, 0, 0} + res, err := insertRows(t, b, append(append([]byte{}, header...), endOfStream...), batch) require.NoError(t, err) require.EqualValues(t, 1, res.RowsInserted) } @@ -375,6 +379,37 @@ func TestInsertRowsRefusesMalformedAndMismatchedArrow(t *testing.T) { require.True(t, serr.Is(err, serr.InvalidArgument), "batches that disagree on their schema: %v", err) } +// A transport that fails mid-stream is the call's failure, not the client's +// malformed Arrow: it is returned as itself and never as an invalid argument. +func TestInsertRowsReturnsAFailedSourceAsItself(t *testing.T) { + f := newFake(t) + f.handle("GET /projects/test-project/datasets/d/tables/t", 200, eventTable) + f.handle("POST "+insertPath, 200, map[string]any{}) + b := f.open(t, backend.Config{}) + header, batch := eventBatch(t, []int64{1}, []string{"a"}) + + lost := serr.New(serr.Throttled, "recv", "the client went away") + _, err := b.InsertRows(context.Background(), backend.TableRef{Dataset: "d", Table: "t"}, header, + &failingReader{payloads: [][]byte{batch}, err: lost}) + require.Same(t, lost, err) +} + +// failingReader hands out its payloads and then fails instead of ending. +type failingReader struct { + payloads [][]byte + err error +} + +func (r *failingReader) Next() ([]byte, error) { + if len(r.payloads) == 0 { + return nil, r.err + } + p := r.payloads[0] + r.payloads = r.payloads[1:] + return p, nil +} +func (r *failingReader) Close() error { return nil } + func TestInsertRowsIsUnsupportedForAnArrowTypeBigQueryHasNoEquivalentOf(t *testing.T) { f := newFake(t) f.handle("GET /projects/test-project/datasets/d/tables/t", 200, eventTable) diff --git a/internal/backend/bigquery/integration_test.go b/internal/backend/bigquery/integration_test.go index c4847b0..5d79ead 100644 --- a/internal/backend/bigquery/integration_test.go +++ b/internal/backend/bigquery/integration_test.go @@ -17,6 +17,7 @@ import ( "github.com/apache/arrow-go/v18/arrow/memory" "github.com/stretchr/testify/require" + "github.com/codefly-dev/service-warehouse/internal/arrowipc" "github.com/codefly-dev/service-warehouse/internal/backend" "github.com/codefly-dev/service-warehouse/internal/serr" ) @@ -159,12 +160,15 @@ func TestAgainstRealBigQuery(t *testing.T) { } rec := builder.NewRecordBatch() defer rec.Release() - fr, err := newFramer(schema, mem) + header, err := arrowipc.SchemaMessage(schema) require.NoError(t, err) - batch, err := fr.batch(rec) + enc, err := arrowipc.NewEncoder(schema) + require.NoError(t, err) + defer enc.Close() + batch, err := enc.Encode(rec) require.NoError(t, err) - res, err := b.InsertRows(ctx, ref, fr.schemaMsg, &sliceReader{payloads: [][]byte{batch}}) + res, err := b.InsertRows(ctx, ref, header, &sliceReader{payloads: [][]byte{batch}}) require.NoError(t, err) require.EqualValues(t, 2, res.RowsInserted) require.Len(t, res.Errors, 1) @@ -179,17 +183,20 @@ func TestAgainstRealBigQuery(t *testing.T) { mb.Field(1).(*array.StringBuilder).Append("x") mrec := mb.NewRecordBatch() defer mrec.Release() - mfr, err := newFramer(mismatch, mem) + mheader, err := arrowipc.SchemaMessage(mismatch) + require.NoError(t, err) + menc, err := arrowipc.NewEncoder(mismatch) require.NoError(t, err) - mbatch, err := mfr.batch(mrec) + defer menc.Close() + mbatch, err := menc.Encode(mrec) require.NoError(t, err) - res, err = b.InsertRows(ctx, ref, mfr.schemaMsg, &sliceReader{payloads: [][]byte{mbatch}}) + res, err = b.InsertRows(ctx, ref, mheader, &sliceReader{payloads: [][]byte{mbatch}}) require.NoError(t, err) require.Zero(t, res.RowsInserted) require.Len(t, res.Errors, 1) require.Equal(t, backend.RefusalSchemaMismatch, res.Errors[0].Reason) - _, err = b.InsertRows(ctx, backend.TableRef{Dataset: dataset, Table: "nope"}, fr.schemaMsg, &sliceReader{payloads: [][]byte{batch}}) + _, err = b.InsertRows(ctx, backend.TableRef{Dataset: dataset, Table: "nope"}, header, &sliceReader{payloads: [][]byte{batch}}) require.True(t, serr.Is(err, serr.NotFound), "got %v", err) // Streamed rows become visible to queries within moments, not instantly. From 0f9b31022da083f993fa4101b53d8e12a03a11a6 Mon Sep 17 00:00:00 2001 From: crisfiorini-oss <301773486+crisfiorini-oss@users.noreply.github.com> Date: Thu, 8 Oct 2026 18:26:47 -0400 Subject: [PATCH 09/14] fix(arrowipc): refuse an Arrow message that states a size over the limit The Arrow a client sends to InsertRows states its own sizes. arrow-go documents size limits for its message reader, but ipc.NewReader builds its message reader without applying them, so a message's stated metadata and body sizes were taken at face value. arrowipc.NewReader now builds the message reader with the limits and hands it to ipc.NewReaderFromMessageReader: metadata over 1 MiB (a constant; a schema or a buffer layout is a few bytes a column) and a body over SWH_MAX_ARROW_MESSAGE_BYTES (default 64 MiB, positive, malformed or zero rejected at startup like the other limits) are refused as ErrMalformed, which the backends return as InvalidArgument. A bound that is not positive means the default, never no bound, because arrow-go reads it as unbounded. The backends pass the configured value through backend.Config. A test fails if any other file in the repository, test files included, builds an Arrow reader directly; the two bigquery tests that did now go through arrowipc. Tested: a batch message whose metadata states a body of 16 TiB is refused with under 8 MiB allocated, by default, with a bound of zero and with a negative one; metadata that states 2 GiB, and a stream of text, likewise; a body over a small configured bound is refused and one under it read; bigquery InsertRows with a small SWH_MAX_ARROW_MESSAGE_BYTES returns InvalidArgument and sends nothing. Co-Authored-By: Claude Sonnet 5.5 --- README.md | 12 ++- internal/arrowipc/arrowipc.go | 46 ++++++++- internal/arrowipc/construction_test.go | 98 +++++++++++++++++++ internal/arrowipc/limits_test.go | 117 +++++++++++++++++++++++ internal/backend/backend.go | 4 + internal/backend/bigquery/arrow_test.go | 6 +- internal/backend/bigquery/insert.go | 2 +- internal/backend/bigquery/insert_test.go | 19 ++++ internal/backend/bigquery/query_test.go | 5 +- internal/config/config.go | 12 +++ internal/config/config_test.go | 17 +++- 11 files changed, 321 insertions(+), 17 deletions(-) create mode 100644 internal/arrowipc/construction_test.go create mode 100644 internal/arrowipc/limits_test.go diff --git a/README.md b/README.md index 899b26a..ccf0ec2 100644 --- a/README.md +++ b/README.md @@ -140,8 +140,9 @@ allows before they are used. **`InsertRows`.** Arrow comes in the same framing: the header's schema message and one RecordBatch message per batch, or, with no schema in the header, a first message that is a schema message followed by its batch. A message that is not -Arrow (cut short, empty, a second schema where a batch belongs) is -`InvalidArgument`. Rows are sent as streaming inserts of at most 500 rows and +Arrow (cut short, empty, a second schema where a batch belongs) or that states a +size over a limit (metadata over 1 MiB, a body over `SWH_MAX_ARROW_MESSAGE_BYTES`) +is `InvalidArgument`. Rows are sent as streaming inserts of at most 500 rows and about 8 MiB. The call is not atomic: BigQuery is asked to skip invalid rows, so valid rows are stored and refused ones come back as `errors`, each with its `row_index` (counted from 0 across every batch of the call) and a `reason`: @@ -179,11 +180,12 @@ BigQuery has no equivalent of (map, duration, dictionary) is `Unimplemented`. | `SWH_QUERY_TIMEOUT` | `0` | per-query timeout (Go duration; 0 = backend default); a request may only shorten it | | `SWH_AUTH_TOKEN` | — | shared secret every caller must present; required unless `SWH_ALLOW_ANONYMOUS=true` | | `SWH_ALLOW_ANONYMOUS` | `false` | `true` accepts callers with no token; refused together with `SWH_AUTH_TOKEN` | +| `SWH_MAX_ARROW_MESSAGE_BYTES` | `67108864` (64 MiB) | largest body of one Arrow IPC message a client may send (`InsertRows`); a larger one is `InvalidArgument`, refused before any of it is allocated. Must be positive. grpc-go's own 4 MiB receive limit is checked first | The `bigquery` backend reads `SWH_DATABASE` (required), `SWH_DATASET`, -`SWH_LOCATION`, `SWH_CREDENTIALS_FILE`, `SWH_MAX_QUERY_BYTES` and -`SWH_QUERY_TIMEOUT`, and adds no variable of its own. A setting it never reads is -a startup error, not something silently ignored. +`SWH_LOCATION`, `SWH_CREDENTIALS_FILE`, `SWH_MAX_QUERY_BYTES`, +`SWH_QUERY_TIMEOUT` and `SWH_MAX_ARROW_MESSAGE_BYTES`, and adds no variable of its +own. A setting it never reads is a startup error, not something silently ignored. ## Authentication diff --git a/internal/arrowipc/arrowipc.go b/internal/arrowipc/arrowipc.go index 29f5fa3..330fe3f 100644 --- a/internal/arrowipc/arrowipc.go +++ b/internal/arrowipc/arrowipc.go @@ -165,6 +165,33 @@ type Reader struct { err error } +const ( + // DefaultMaxMessageBytes is the largest message body a reader accepts when + // the caller names no limit: what SWH_MAX_ARROW_MESSAGE_BYTES defaults to. + DefaultMaxMessageBytes int64 = 64 << 20 + + // maxMetadataBytes is the largest message metadata a reader accepts. The + // metadata is a schema or a batch's buffer layout, a few bytes per column; a + // megabyte is thousands of columns. + maxMetadataBytes int64 = 1 << 20 +) + +// Option configures a Reader. +type Option func(*readerConfig) + +type readerConfig struct{ maxMessageBytes int64 } + +// WithMaxMessageBytes bounds the body of one message. A message over the bound is +// refused as ErrMalformed before any of it is allocated; a value that is not +// positive leaves the default, never "no bound". +func WithMaxMessageBytes(n int64) Option { + return func(c *readerConfig) { + if n > 0 { + c.maxMessageBytes = n + } + } +} + // NewReader reads the records of a stream whose schema message is schema and // whose record-batch messages come from src. schema may be empty when src's // first message is self-describing (a schema message followed by batches). A @@ -172,14 +199,27 @@ type Reader struct { // schema it is: left on, the marker would end the stream before the first batch // and every row would be silently dropped. // +// The bytes come from a client, so the sizes a message states are not taken on +// trust. arrow-go's ipc.NewReader does not apply the size limits that its message +// reader has, so every reader is built here, from a message reader that does, and +// nowhere else (a test holds the repository to that). +// // It reads the schema before it returns, so a first message that is not one is // reported here. -func NewReader(schema []byte, src Source) (*Reader, error) { +func NewReader(schema []byte, src Source, opts ...Option) (*Reader, error) { + cfg := readerConfig{maxMessageBytes: DefaultMaxMessageBytes} + for _, opt := range opts { + opt(&cfg) + } w := &watchedSource{src: src} - r, err := ipc.NewReader( + mem := memory.DefaultAllocator + messages := ipc.NewMessageReader( &messageStream{pending: bytes.TrimSuffix(schema, endOfStream[:]), src: w}, - ipc.WithAllocator(memory.DefaultAllocator), + ipc.WithAllocator(mem), + ipc.WithMetadataSizeLimit(maxMetadataBytes), + ipc.WithBodySizeLimit(cfg.maxMessageBytes), ) + r, err := ipc.NewReaderFromMessageReader(messages, ipc.WithAllocator(mem)) if err != nil { return nil, w.classify(err) } diff --git a/internal/arrowipc/construction_test.go b/internal/arrowipc/construction_test.go new file mode 100644 index 0000000..080e530 --- /dev/null +++ b/internal/arrowipc/construction_test.go @@ -0,0 +1,98 @@ +package arrowipc + +import ( + "go/ast" + "go/parser" + "go/token" + "io/fs" + "path/filepath" + "strconv" + "testing" + + "github.com/stretchr/testify/require" +) + +const arrowIPCImport = "github.com/apache/arrow-go/v18/arrow/ipc" + +// readerConstructors are the arrow-go calls that build something which reads +// Arrow IPC bytes. All of them are given the size a message states for itself +// unless the caller says otherwise, and the bytes of InsertRows come from a +// client. +var readerConstructors = map[string]bool{ + "NewReader": true, + "NewReaderFromMessageReader": true, + "NewMessageReader": true, + "NewFileReader": true, + "NewFileReaderFromMemory": true, +} + +// directReaders lists the arrow-go reader constructors a source file calls. +func directReaders(t *testing.T, name string, src any) []string { + t.Helper() + file, err := parser.ParseFile(token.NewFileSet(), name, src, 0) + require.NoError(t, err, name) + + local := "" + for _, imp := range file.Imports { + path, err := strconv.Unquote(imp.Path.Value) + require.NoError(t, err) + if path != arrowIPCImport { + continue + } + local = "ipc" + if imp.Name != nil { + local = imp.Name.Name + } + } + if local == "" { + return nil + } + var found []string + ast.Inspect(file, func(n ast.Node) bool { + sel, ok := n.(*ast.SelectorExpr) + if !ok { + return true + } + if id, ok := sel.X.(*ast.Ident); ok && id.Name == local && readerConstructors[sel.Sel.Name] { + found = append(found, sel.Sel.Name) + } + return true + }) + return found +} + +func TestTheDetectorSeesADirectReader(t *testing.T) { + src := `package x +import arrowipc "github.com/apache/arrow-go/v18/arrow/ipc" +func f() { _, _ = arrowipc.NewReader(nil) }` + require.Equal(t, []string{"NewReader"}, directReaders(t, "x.go", src)) + require.Empty(t, directReaders(t, "y.go", `package y +import "github.com/apache/arrow-go/v18/arrow/ipc" +func f() { _ = ipc.NewWriter }`), "writers are not readers") +} + +// A reader built anywhere but here does not have the size limits. This fails the +// day any other file, test files included, builds one. +func TestNoReaderIsBuiltOutsideThisPackage(t *testing.T) { + root := filepath.Join("..", "..") + err := filepath.WalkDir(root, func(path string, d fs.DirEntry, err error) error { + if err != nil { + return err + } + rel, _ := filepath.Rel(root, path) + if d.IsDir() { + if rel == "gen" || d.Name() == ".git" { + return filepath.SkipDir + } + return nil + } + if filepath.Ext(path) != ".go" || rel == filepath.Join("internal", "arrowipc", "arrowipc.go") { + return nil + } + for _, call := range directReaders(t, path, nil) { + t.Errorf("%s calls ipc.%s: build the reader with arrowipc.NewReader, which applies the message size limits", rel, call) + } + return nil + }) + require.NoError(t, err) +} diff --git a/internal/arrowipc/limits_test.go b/internal/arrowipc/limits_test.go new file mode 100644 index 0000000..96ade80 --- /dev/null +++ b/internal/arrowipc/limits_test.go @@ -0,0 +1,117 @@ +package arrowipc + +import ( + "bytes" + "encoding/binary" + "runtime" + "testing" + + "github.com/stretchr/testify/require" +) + +// allocated is how many bytes f allocated: what the sizes a message states must +// not be able to choose. +func allocated(f func()) uint64 { + var before, after runtime.MemStats + runtime.GC() + runtime.ReadMemStats(&before) + f() + runtime.ReadMemStats(&after) + return after.TotalAlloc - before.TotalAlloc +} + +// boundedAlloc is far above what refusing a message costs and far below what any +// message the tests state would cost to read. +const boundedAlloc = 8 << 20 + +func readAll(header []byte, src Source, opts ...Option) error { + rd, err := NewReader(header, src, opts...) + if err != nil { + return err + } + defer rd.Release() + for rd.Next() { + } + return rd.Err() +} + +// oversizedBatch is a valid batch message, a few hundred bytes, whose metadata +// states a body of 16 TiB: the body bound refuses it before any of it is read. +func oversizedBatch(t *testing.T) []byte { + t.Helper() + enc, err := NewEncoder(testSchema) + require.NoError(t, err) + defer enc.Close() + rec := batch([]int64{1, 2}, []string{"a", "b"}) + defer rec.Release() + m, err := enc.Encode(rec) + require.NoError(t, err) + + // The body length is the one 8-byte value in the metadata that equals the + // bytes after it. + metaLen := int(binary.LittleEndian.Uint32(m[4:8])) + meta := m[8 : 8+metaLen] + bodyLen := make([]byte, 8) + binary.LittleEndian.PutUint64(bodyLen, uint64(len(m)-8-metaLen)) + at := bytes.Index(meta, bodyLen) + require.GreaterOrEqual(t, at, 0, "the metadata states the body length") + require.Equal(t, -1, bytes.Index(meta[at+1:], bodyLen), "only once") + + oversized := append([]byte(nil), m...) + binary.LittleEndian.PutUint64(oversized[8+at:], 1<<44) + require.Less(t, len(oversized), 512) + return oversized +} + +func TestAMessageThatClaimsAHugeBodyIsRefusedBeforeItIsAllocated(t *testing.T) { + header, err := SchemaMessage(testSchema) + require.NoError(t, err) + oversized := oversizedBatch(t) + + // A bound that is not positive is the default bound, never no bound. + for name, opts := range map[string][]Option{ + "by default": nil, + "with a bound of zero": {WithMaxMessageBytes(0)}, + "with a negative bound": {WithMaxMessageBytes(-1)}, + "with a bound it exceeds": {WithMaxMessageBytes(1 << 20)}, + } { + t.Run(name, func(t *testing.T) { + var got error + grew := allocated(func() { got = readAll(header, &msgs{list: [][]byte{oversized}}, opts...) }) + require.ErrorIs(t, got, ErrMalformed) + require.Less(t, grew, uint64(boundedAlloc)) + }) + } +} + +// The metadata is a schema or a buffer layout, so its size is bounded too, and the +// first four bytes of a stream that is not Arrow at all are read as the length. +func TestMetadataThatClaimsToBeHugeIsRefusedBeforeItIsAllocated(t *testing.T) { + for name, header := range map[string][]byte{ + "a length of 2 GiB": {0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x7f}, + "text, read as length": []byte("not an arrow schema"), + } { + t.Run(name, func(t *testing.T) { + var got error + grew := allocated(func() { got = readAll(header, &msgs{}) }) + require.ErrorIs(t, got, ErrMalformed) + require.Less(t, grew, uint64(boundedAlloc)) + }) + } +} + +func TestTheBodyBoundIsTheCallersToSet(t *testing.T) { + enc, err := NewEncoder(testSchema) + require.NoError(t, err) + defer enc.Close() + rec := batch([]int64{1, 2}, []string{"a", "b"}) + defer rec.Release() + m, err := enc.Encode(rec) + require.NoError(t, err) + header, err := SchemaMessage(testSchema) + require.NoError(t, err) + + require.ErrorIs(t, readAll(header, &msgs{list: [][]byte{m}}, WithMaxMessageBytes(16)), ErrMalformed, + "a batch whose body is over the bound is refused") + require.NoError(t, readAll(header, &msgs{list: [][]byte{m}}, WithMaxMessageBytes(1<<10))) +} diff --git a/internal/backend/backend.go b/internal/backend/backend.go index 2734226..0c73cf3 100644 --- a/internal/backend/backend.go +++ b/internal/backend/backend.go @@ -45,6 +45,10 @@ type Config struct { MaxQueryBytes int64 // QueryTimeout bounds a single query; 0 means the backend default. QueryTimeout time.Duration + // MaxArrowMessageBytes bounds the body of one Arrow IPC message a client + // sends (InsertRows); a larger one is refused as invalid before any of it is + // read. 0 means the default of internal/arrowipc, never no bound. + MaxArrowMessageBytes int64 } // ColumnType is the portable logical type — the honest intersection across the diff --git a/internal/backend/bigquery/arrow_test.go b/internal/backend/bigquery/arrow_test.go index 4cbcd97..d0eea4c 100644 --- a/internal/backend/bigquery/arrow_test.go +++ b/internal/backend/bigquery/arrow_test.go @@ -1,8 +1,6 @@ package bigquery import ( - "bytes" - "io" "math" "math/big" "testing" @@ -13,10 +11,10 @@ import ( "github.com/apache/arrow-go/v18/arrow" "github.com/apache/arrow-go/v18/arrow/array" "github.com/apache/arrow-go/v18/arrow/decimal128" - "github.com/apache/arrow-go/v18/arrow/ipc" "github.com/apache/arrow-go/v18/arrow/memory" "github.com/stretchr/testify/require" + "github.com/codefly-dev/service-warehouse/internal/arrowipc" "github.com/codefly-dev/service-warehouse/internal/serr" ) @@ -87,7 +85,7 @@ func everyType() (bq.Schema, [][]bq.Value, [][]any) { func readBack(t *testing.T, schemaMsg []byte, batch []byte) arrow.RecordBatch { t.Helper() - reader, err := ipc.NewReader(io.MultiReader(bytes.NewReader(schemaMsg), bytes.NewReader(batch))) + reader, err := arrowipc.NewReader(schemaMsg, &sliceReader{payloads: [][]byte{batch}}) require.NoError(t, err) t.Cleanup(reader.Release) require.True(t, reader.Next(), "the batch holds a record: %v", reader.Err()) diff --git a/internal/backend/bigquery/insert.go b/internal/backend/bigquery/insert.go index 8e0f427..e86bf75 100644 --- a/internal/backend/bigquery/insert.go +++ b/internal/backend/bigquery/insert.go @@ -81,7 +81,7 @@ func (b *Backend) InsertRows(ctx context.Context, ref backend.TableRef, schema [ if err != nil { return nil, err } - stream, err := arrowipc.NewReader(schema, batches) + stream, err := arrowipc.NewReader(schema, batches, arrowipc.WithMaxMessageBytes(b.cfg.MaxArrowMessageBytes)) if err != nil { return nil, streamError(err) } diff --git a/internal/backend/bigquery/insert_test.go b/internal/backend/bigquery/insert_test.go index 5cde053..65f9a21 100644 --- a/internal/backend/bigquery/insert_test.go +++ b/internal/backend/bigquery/insert_test.go @@ -410,6 +410,25 @@ func (r *failingReader) Next() ([]byte, error) { } func (r *failingReader) Close() error { return nil } +// SWH_MAX_ARROW_MESSAGE_BYTES reaches the reader: a message whose body is over it +// is the client's invalid argument, and nothing of the call is sent to BigQuery. +func TestInsertRowsRefusesAMessageOverTheConfiguredBound(t *testing.T) { + f := newFake(t) + f.handle("GET /projects/test-project/datasets/d/tables/t", 200, eventTable) + f.handle("POST "+insertPath, 200, map[string]any{}) + header, batch := eventBatch(t, []int64{1, 2, 3}, []string{"a", "b", "c"}) + + small := f.open(t, backend.Config{MaxArrowMessageBytes: 16}) + _, err := insertRows(t, small, header, batch) + require.True(t, serr.Is(err, serr.InvalidArgument), "got %v", err) + require.Empty(t, f.seen("POST", "insertAll")) + + roomy := f.open(t, backend.Config{MaxArrowMessageBytes: 1 << 20}) + res, err := insertRows(t, roomy, header, batch) + require.NoError(t, err) + require.EqualValues(t, 3, res.RowsInserted) +} + func TestInsertRowsIsUnsupportedForAnArrowTypeBigQueryHasNoEquivalentOf(t *testing.T) { f := newFake(t) f.handle("GET /projects/test-project/datasets/d/tables/t", 200, eventTable) diff --git a/internal/backend/bigquery/query_test.go b/internal/backend/bigquery/query_test.go index 5652f27..fff76c6 100644 --- a/internal/backend/bigquery/query_test.go +++ b/internal/backend/bigquery/query_test.go @@ -1,7 +1,6 @@ package bigquery import ( - "bytes" "context" "errors" "fmt" @@ -12,9 +11,9 @@ import ( "github.com/apache/arrow-go/v18/arrow" "github.com/apache/arrow-go/v18/arrow/array" - "github.com/apache/arrow-go/v18/arrow/ipc" "github.com/stretchr/testify/require" + "github.com/codefly-dev/service-warehouse/internal/arrowipc" "github.com/codefly-dev/service-warehouse/internal/backend" "github.com/codefly-dev/service-warehouse/internal/serr" ) @@ -85,7 +84,7 @@ func decode(t *testing.T, res *backend.QueryResult) (schema *arrow.Schema, recor break } require.NoError(t, err) - reader, err := ipc.NewReader(io.MultiReader(bytes.NewReader(res.ArrowSchema), bytes.NewReader(batch))) + reader, err := arrowipc.NewReader(res.ArrowSchema, &sliceReader{payloads: [][]byte{batch}}) require.NoError(t, err) schema = reader.Schema() for reader.Next() { diff --git a/internal/config/config.go b/internal/config/config.go index 47983e4..b38073c 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -13,6 +13,7 @@ import ( "time" "github.com/codefly-dev/core/callertoken" + "github.com/codefly-dev/service-warehouse/internal/arrowipc" "github.com/codefly-dev/service-warehouse/internal/backend" ) @@ -42,6 +43,15 @@ func FromEnv() (Config, error) { if err != nil { return Config{}, err } + // Unlike the two above, this limit has no "0 means the default": a zero that + // reached the Arrow reader would be read as no bound, so it is refused here. + maxArrowMessageBytes, err := envInt("SWH_MAX_ARROW_MESSAGE_BYTES", arrowipc.DefaultMaxMessageBytes) + if err != nil { + return Config{}, err + } + if maxArrowMessageBytes < 1 { + return Config{}, fmt.Errorf("SWH_MAX_ARROW_MESSAGE_BYTES: must be positive, got %d", maxArrowMessageBytes) + } // A negative limit is not a smaller limit: both fields spend 0 as "use the // backend default", so a negative value would reach the backend as no limit // at all — the opposite of what an operator setting one intends. @@ -68,6 +78,8 @@ func FromEnv() (Config, error) { CredentialsFile: os.Getenv("SWH_CREDENTIALS_FILE"), MaxQueryBytes: maxQueryBytes, QueryTimeout: queryTimeout, + + MaxArrowMessageBytes: maxArrowMessageBytes, }, AuthToken: strings.TrimSpace(os.Getenv("SWH_AUTH_TOKEN")), } diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 092be3c..47e2f8f 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -46,7 +46,7 @@ func clearEnv(t *testing.T) { for _, k := range []string{ "SWH_LISTEN", "SWH_BACKEND", "SWH_DATABASE", "SWH_DATASET", "SWH_LOCATION", "SWH_DSN", "SWH_HOST", "SWH_PORT", "SWH_USER", "SWH_PASSWORD", "SWH_ACCOUNT", - "SWH_CREDENTIALS_FILE", "SWH_MAX_QUERY_BYTES", "SWH_QUERY_TIMEOUT", + "SWH_CREDENTIALS_FILE", "SWH_MAX_QUERY_BYTES", "SWH_QUERY_TIMEOUT", "SWH_MAX_ARROW_MESSAGE_BYTES", "SWH_AUTH_TOKEN", "SWH_ALLOW_ANONYMOUS", } { t.Setenv(k, "") @@ -128,6 +128,11 @@ func TestMalformedLimitsAreRejected(t *testing.T) { {"SWH_QUERY_TIMEOUT", "5 minutes", "not a duration"}, {"SWH_QUERY_TIMEOUT", "-30s", "must not be negative"}, {"SWH_PORT", "abc", "not an integer"}, + {"SWH_MAX_ARROW_MESSAGE_BYTES", "64MiB", "not an integer"}, + {"SWH_MAX_ARROW_MESSAGE_BYTES", "-1", "must be positive"}, + // Zero is refused rather than read as "the default": a zero that reached + // the Arrow reader would be no bound at all. + {"SWH_MAX_ARROW_MESSAGE_BYTES", "0", "must be positive"}, } { t.Run(tc.key+"="+tc.value, func(t *testing.T) { clearEnv(t) @@ -154,6 +159,16 @@ func TestDefaultsAndParsing(t *testing.T) { require.Equal(t, int64(4096), cfg.Backend.MaxQueryBytes) require.Equal(t, 90*time.Second, cfg.Backend.QueryTimeout) require.Equal(t, 5439, cfg.Backend.Port) + require.Equal(t, int64(64<<20), cfg.Backend.MaxArrowMessageBytes, "the Arrow message bound is on by default") +} + +func TestArrowMessageBoundIsConfigurable(t *testing.T) { + clearEnv(t) + t.Setenv("SWH_MAX_ARROW_MESSAGE_BYTES", "1048576") + + cfg, err := FromEnv() + require.NoError(t, err) + require.Equal(t, int64(1<<20), cfg.Backend.MaxArrowMessageBytes) } // Unset limits keep meaning "use the backend default"; rejecting malformed From 4aaa64c3facb023eabbd4b76c5cd8bd836bf9e84 Mon Sep 17 00:00:00 2001 From: crisfiorini-oss <301773486+crisfiorini-oss@users.noreply.github.com> Date: Thu, 8 Oct 2026 18:26:49 -0400 Subject: [PATCH 10/14] fix(arrowipc): bound a schema message's stated counts and offsets by its own bytes before decoding it STOPGAP: a defensive check, to be removed if the decoder validates its own input. A schema message states counts (fields, children, metadata entries, type ids), lengths and offsets, and the decoder sizes what it builds from them. The message size limits do not bound those, so arrowipc now checks them against the message's own bytes first. Before arrow-go sees the first message of a stream, and only when it is a schema, arrowipc walks its flatbuffer with the public flatbuffers table API: the schema's fields, each field's children, every custom metadata vector, a union's type ids and a timestamp's time zone, which is what the decoder reads from a schema. It refuses as ErrMalformed any vector that states more elements than the bytes after it can hold (at 4 bytes an element, the smallest there is), any offset that points outside the message or to or before itself (which also rules out a cycle), and any walk that would visit more structure than the message has bytes, counting a table that is referenced more than once at each reference. A schema that is a tree costs no more than its size, so a producer that writes an ordinary schema never reaches the last rule. What passes is bounded by the 1 MiB metadata cap, and so is what the decoder builds from it. Nothing else changes for a message that is not a schema. Tested: a nested schema (struct and list children, custom metadata on the schema and its fields, a union, a timestamp with a zone), a 4,000-column schema with metadata on every field and an empty one all still read and equal the schema that was sent; thirteen schemas, each the nested one with a count, a length or an offset changed to state more than the bytes hold, are refused with under 8 MiB allocated, in the header and in a first message that carries the schema; a 400 KB message whose 50,000 fields share 50,000 children is refused; setting any single byte of the nested schema message to any of five values never costs 8 MiB. The fuzz target FuzzSchemaReader runs its seeds in the ordinary suite; 5.5 million executions of `go test -fuzz` in a container found no input over the bound. Not changed: a message may still make the reader allocate the body it states, up to SWH_MAX_ARROW_MESSAGE_BYTES, before it finds the bytes are not there. Co-Authored-By: Claude Sonnet 5.5 --- README.md | 5 +- go.mod | 2 +- internal/arrowipc/arrowipc.go | 19 +- internal/arrowipc/limits_test.go | 1 - internal/arrowipc/schemacheck.go | 266 +++++++++++++++++++++++++ internal/arrowipc/schemacheck_test.go | 276 ++++++++++++++++++++++++++ 6 files changed, 564 insertions(+), 5 deletions(-) create mode 100644 internal/arrowipc/schemacheck.go create mode 100644 internal/arrowipc/schemacheck_test.go diff --git a/README.md b/README.md index ccf0ec2..14a252a 100644 --- a/README.md +++ b/README.md @@ -140,8 +140,9 @@ allows before they are used. **`InsertRows`.** Arrow comes in the same framing: the header's schema message and one RecordBatch message per batch, or, with no schema in the header, a first message that is a schema message followed by its batch. A message that is not -Arrow (cut short, empty, a second schema where a batch belongs) or that states a -size over a limit (metadata over 1 MiB, a body over `SWH_MAX_ARROW_MESSAGE_BYTES`) +Arrow (cut short, empty, a second schema where a batch belongs), that states a +size over a limit (metadata over 1 MiB, a body over `SWH_MAX_ARROW_MESSAGE_BYTES`), +or whose schema states more fields, children or metadata than its bytes can hold, is `InvalidArgument`. Rows are sent as streaming inserts of at most 500 rows and about 8 MiB. The call is not atomic: BigQuery is asked to skip invalid rows, so valid rows are stored and refused ones come back as `errors`, each with its diff --git a/go.mod b/go.mod index 7ec7d4c..892273d 100644 --- a/go.mod +++ b/go.mod @@ -7,6 +7,7 @@ require ( cloud.google.com/go/bigquery v1.85.0 github.com/apache/arrow-go/v18 v18.8.0 github.com/codefly-dev/core v0.19.4 + github.com/google/flatbuffers v25.12.19+incompatible github.com/stretchr/testify v1.12.1 golang.org/x/oauth2 v0.36.0 golang.org/x/sync v0.23.0 @@ -74,7 +75,6 @@ require ( github.com/go-viper/mapstructure/v2 v2.5.0 // indirect github.com/goccy/go-json v0.10.6 // indirect github.com/gofrs/flock v0.13.1 // indirect - github.com/google/flatbuffers v25.12.19+incompatible // indirect github.com/google/gnostic-models v0.7.1 // indirect github.com/google/go-github/v89 v89.0.0 // indirect github.com/google/go-querystring v1.2.0 // indirect diff --git a/internal/arrowipc/arrowipc.go b/internal/arrowipc/arrowipc.go index 330fe3f..e39914b 100644 --- a/internal/arrowipc/arrowipc.go +++ b/internal/arrowipc/arrowipc.go @@ -8,6 +8,13 @@ // The package is backend-neutral: a backend that produces Arrow records uses // Encoder, and a backend that ingests them uses NewReader, so no backend // re-derives the framing. +// +// The bytes NewReader reads come from a client, so the sizes, counts and offsets +// they state are not taken on trust. NewReader applies size limits to every +// message. A schema message's stated counts and offsets are also checked against +// its own bytes before the decoder sizes anything from them: a STOPGAP, labelled +// as one, in schemacheck.go, which can be removed if the decoder validates its own +// input. package arrowipc import ( @@ -204,6 +211,9 @@ func WithMaxMessageBytes(n int64) Option { // reader has, so every reader is built here, from a message reader that does, and // nowhere else (a test holds the repository to that). // +// A schema message is also walked before it is decoded, and refused if its stated +// counts and offsets exceed its own bytes: a stopgap, see schemacheck.go. +// // It reads the schema before it returns, so a first message that is not one is // reported here. func NewReader(schema []byte, src Source, opts ...Option) (*Reader, error) { @@ -213,8 +223,12 @@ func NewReader(schema []byte, src Source, opts ...Option) (*Reader, error) { } w := &watchedSource{src: src} mem := memory.DefaultAllocator + stream, err := vetFirstMessage(&messageStream{pending: bytes.TrimSuffix(schema, endOfStream[:]), src: w}) + if err != nil { + return nil, w.classify(err) + } messages := ipc.NewMessageReader( - &messageStream{pending: bytes.TrimSuffix(schema, endOfStream[:]), src: w}, + stream, ipc.WithAllocator(mem), ipc.WithMetadataSizeLimit(maxMetadataBytes), ipc.WithBodySizeLimit(cfg.maxMessageBytes), @@ -274,6 +288,9 @@ func (w *watchedSource) classify(err error) error { if w.err != nil { return w.err } + if errors.Is(err, ErrMalformed) { + return err + } return fmt.Errorf("%w: %w", ErrMalformed, err) } diff --git a/internal/arrowipc/limits_test.go b/internal/arrowipc/limits_test.go index 96ade80..a58306e 100644 --- a/internal/arrowipc/limits_test.go +++ b/internal/arrowipc/limits_test.go @@ -13,7 +13,6 @@ import ( // not be able to choose. func allocated(f func()) uint64 { var before, after runtime.MemStats - runtime.GC() runtime.ReadMemStats(&before) f() runtime.ReadMemStats(&after) diff --git a/internal/arrowipc/schemacheck.go b/internal/arrowipc/schemacheck.go new file mode 100644 index 0000000..5a54ca9 --- /dev/null +++ b/internal/arrowipc/schemacheck.go @@ -0,0 +1,266 @@ +package arrowipc + +import ( + "bytes" + "encoding/binary" + "fmt" + "io" + + flatbuffers "github.com/google/flatbuffers/go" +) + +// STOPGAP: a defensive check on a schema message, before it is decoded. It is not +// part of the wire contract, and it can be removed if the decoder validates its own +// input. +// +// A schema message states counts (fields, children, metadata entries, type ids), +// lengths and offsets, and the decoder sizes what it builds from them. They are +// checked here against the message's own bytes first. Before the decoder sees the +// first message of a stream, and only when it is a schema, this walks its +// flatbuffer with the public flatbuffers table API and refuses it, as ErrMalformed, +// when +// +// - a vector states more elements than the bytes after it can hold, at the +// smallest element size there is (4 bytes: an offset or an int32); +// - an offset points outside the message, or to or before itself, which also +// rules out a cycle; +// - or the walk would visit more structure than the message has bytes: every +// table costs 4, every vector 4 plus 4 for each element, every string 4 plus +// its length, and a table that is referenced more than once is paid for at +// each reference. A message that is a tree costs no more than its own size, +// so a producer that writes an ordinary schema never reaches this. +// +// The walk covers what the decoder reads from a schema (fields, their children, +// every custom metadata vector, a union's type ids, a timestamp's time zone), so +// what it admits is bounded by the metadata cap, 1 MiB, and so is what the +// decoder builds from it. The slot numbers are those of Schema.fbs in the Arrow +// format and do not change: a field is only ever added at the end of a table. + +// vtable offsets of the fields read below: 4 plus 2 a slot. +const ( + slotMessageHeaderType = 6 + slotMessageHeader = 8 + + slotSchemaFields = 6 + slotSchemaMetadata = 8 + + slotFieldName = 4 + slotFieldTypeType = 8 + slotFieldType = 10 + slotFieldChildren = 14 + slotFieldMetadata = 16 + + slotKeyValueKey = 4 + slotKeyValueValue = 6 + + slotUnionTypeIDs = 6 + + slotTimestampZone = 6 + + messageHeaderSchema = 1 // MessageHeader.Schema + typeTimestamp = 10 + typeUnion = 14 +) + +// shape is a walk of one flatbuffer. +type shape struct { + buf []byte + budget uint64 // bytes of structure the walk may still visit +} + +type shapeFault string + +func (s *shape) fail(format string, args ...any) { panic(shapeFault(fmt.Sprintf(format, args...))) } + +// spend pays for structure the walk visits. +func (s *shape) spend(n uint64) { + if n > s.budget { + s.fail("it states more structure than it has bytes") + } + s.budget -= n +} + +// need checks that n bytes at pos are inside the message. pos and n are below +// 1<<33, so the sum cannot wrap. +func (s *shape) need(pos, n uint64) { + if pos+n > uint64(len(s.buf)) { + s.fail("an offset or a length points outside the message") + } +} + +// indirect follows the offset stored at pos, which must lead strictly forward to +// somewhere with at least a word to read. +func (s *shape) indirect(pos uint64) uint64 { + s.need(pos, 4) + delta := uint64(flatbuffers.GetUOffsetT(s.buf[pos:])) + if delta == 0 { + s.fail("an offset points to itself") + } + dest := pos + delta + s.need(dest, 4) + return dest +} + +// table is the table at pos. Its vtable has to be inside the message too, which +// is all Table.Offset needs to read safely. +func (s *shape) table(pos uint64) flatbuffers.Table { + s.spend(4) + s.need(pos, 4) + t := flatbuffers.Table{Bytes: s.buf, Pos: flatbuffers.UOffsetT(pos)} + vtable := int64(pos) - int64(t.GetSOffsetT(t.Pos)) + if vtable < 0 { + s.fail("a table's vtable is before the message") + } + s.need(uint64(vtable), 4) + size := uint64(t.GetVOffsetT(flatbuffers.UOffsetT(vtable))) + if size < 4 { + s.fail("a vtable is smaller than its own header") + } + s.need(uint64(vtable), size) + return t +} + +// field is where the vtable slot of t is stored in t, and whether it is set. +func (s *shape) field(t flatbuffers.Table, slot flatbuffers.VOffsetT) (uint64, bool) { + o := t.Offset(slot) + if o == 0 { + return 0, false + } + return uint64(t.Pos) + uint64(o), true +} + +func (s *shape) byteField(t flatbuffers.Table, slot flatbuffers.VOffsetT) byte { + pos, ok := s.field(t, slot) + if !ok { + return 0 + } + s.need(pos, 1) + return s.buf[pos] +} + +func (s *shape) tableField(t flatbuffers.Table, slot flatbuffers.VOffsetT) (flatbuffers.Table, bool) { + pos, ok := s.field(t, slot) + if !ok { + return flatbuffers.Table{}, false + } + return s.table(s.indirect(pos)), true +} + +// vector is the position of the first element of the vector in a slot of t and +// how many it states, after checking that the bytes hold that many. +func (s *shape) vector(t flatbuffers.Table, slot flatbuffers.VOffsetT) (first, n uint64) { + pos, ok := s.field(t, slot) + if !ok { + return 0, 0 + } + at := s.indirect(pos) + n = uint64(flatbuffers.GetUOffsetT(s.buf[at:])) + first = at + 4 + s.need(first, 4*n) + s.spend(4 + 4*n) + return first, n +} + +func (s *shape) str(t flatbuffers.Table, slot flatbuffers.VOffsetT) { + pos, ok := s.field(t, slot) + if !ok { + return + } + at := s.indirect(pos) + n := uint64(flatbuffers.GetUOffsetT(s.buf[at:])) + s.need(at+4, n) + s.spend(4 + n) +} + +func (s *shape) walkSchema(t flatbuffers.Table) { + first, n := s.vector(t, slotSchemaFields) + for i := uint64(0); i < n; i++ { + s.walkField(s.table(s.indirect(first + 4*i))) + } + s.walkKeyValues(t, slotSchemaMetadata) +} + +func (s *shape) walkField(t flatbuffers.Table) { + s.str(t, slotFieldName) + switch s.byteField(t, slotFieldTypeType) { + case typeUnion: + if u, ok := s.tableField(t, slotFieldType); ok { + s.vector(u, slotUnionTypeIDs) + } + case typeTimestamp: + if ts, ok := s.tableField(t, slotFieldType); ok { + s.str(ts, slotTimestampZone) + } + } + first, n := s.vector(t, slotFieldChildren) + for i := uint64(0); i < n; i++ { + s.walkField(s.table(s.indirect(first + 4*i))) + } + s.walkKeyValues(t, slotFieldMetadata) +} + +func (s *shape) walkKeyValues(t flatbuffers.Table, slot flatbuffers.VOffsetT) { + first, n := s.vector(t, slot) + for i := uint64(0); i < n; i++ { + kv := s.table(s.indirect(first + 4*i)) + s.str(kv, slotKeyValueKey) + s.str(kv, slotKeyValueValue) + } +} + +// checkSchemaMetadata refuses a message's metadata, which is its flatbuffer, if it +// is a schema that states more than its bytes can hold; see the note at the top +// of the file. Anything that is not a schema is not its concern: arrow-go refuses +// it, or reads a record batch, whose counts it does not size slices from. +func checkSchemaMetadata(buf []byte) (err error) { + defer func() { + if r := recover(); r != nil { + // A fault of ours, or an out-of-range read of the table API on bytes + // that were checked less closely than the ones it reads: either way + // the bytes are not a schema arrow-go can be trusted with. + err = fmt.Errorf("%w: schema metadata: %v", ErrMalformed, r) + } + }() + s := &shape{buf: buf, budget: uint64(len(buf))} + s.need(0, 4) + msg := s.table(uint64(flatbuffers.GetUOffsetT(buf))) + if s.byteField(msg, slotMessageHeaderType) != messageHeaderSchema { + return nil + } + if schema, ok := s.tableField(msg, slotMessageHeader); ok { + s.walkSchema(schema) + } + return nil +} + +// vetFirstMessage reads the first message off stream, framing and metadata, and +// refuses it with checkSchemaMetadata when it is a schema. It returns a reader +// that replays what it read ahead of the rest of the stream, so arrow-go reads +// the same bytes. A message whose framing is not one arrow-go would accept is +// replayed unjudged: arrow-go refuses it itself, by its own limits. +func vetFirstMessage(stream io.Reader) (io.Reader, error) { + var prefix [8]byte + if _, err := io.ReadFull(stream, prefix[:4]); err != nil { + return nil, err + } + word := binary.LittleEndian.Uint32(prefix[:4]) + prefixLen, length := 4, int64(int32(word)) // a stream from before the continuation marker + if word == 0xFFFFFFFF { + if _, err := io.ReadFull(stream, prefix[4:]); err != nil { + return nil, err + } + prefixLen, length = 8, int64(int32(binary.LittleEndian.Uint32(prefix[4:]))) + } + replay := bytes.NewReader(prefix[:prefixLen]) + if length < 4 || length > maxMetadataBytes { + return io.MultiReader(replay, stream), nil + } + metadata := make([]byte, length) + if _, err := io.ReadFull(stream, metadata); err != nil { + return nil, err + } + if err := checkSchemaMetadata(metadata); err != nil { + return nil, err + } + return io.MultiReader(replay, bytes.NewReader(metadata), stream), nil +} diff --git a/internal/arrowipc/schemacheck_test.go b/internal/arrowipc/schemacheck_test.go new file mode 100644 index 0000000..30440a8 --- /dev/null +++ b/internal/arrowipc/schemacheck_test.go @@ -0,0 +1,276 @@ +package arrowipc + +import ( + "encoding/binary" + "fmt" + "testing" + + "github.com/apache/arrow-go/v18/arrow" + flatbuffers "github.com/google/flatbuffers/go" + "github.com/stretchr/testify/require" +) + +// nestedSchema has everything the walk follows: struct and list children, +// custom metadata on the schema and on fields, a union with type ids and a +// timestamp with a time zone. +func nestedSchema() *arrow.Schema { + meta := func(kv ...string) arrow.Metadata { + var keys, vals []string + for i := 0; i < len(kv); i += 2 { + keys, vals = append(keys, kv[i]), append(vals, kv[i+1]) + } + return arrow.NewMetadata(keys, vals) + } + union := arrow.SparseUnionOf( + []arrow.Field{{Name: "i", Type: arrow.PrimitiveTypes.Int32}, {Name: "s", Type: arrow.BinaryTypes.String}}, + []arrow.UnionTypeCode{3, 5}, + ) + m := meta("origin", "tests", "ARROW:note", "kept") + return arrow.NewSchema([]arrow.Field{ + {Name: "id", Type: arrow.PrimitiveTypes.Int64}, + {Name: "s", Type: arrow.StructOf( + arrow.Field{Name: "x", Type: arrow.PrimitiveTypes.Int32, Metadata: meta("unit", "m")}, + arrow.Field{Name: "y", Type: arrow.ListOf(arrow.BinaryTypes.String), Nullable: true}, + ), Metadata: meta("a", "1", "b", "2")}, + {Name: "ts", Type: &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: "Europe/Paris"}, Metadata: meta("k", "v")}, + {Name: "u", Type: union}, + }, &m) +} + +func schemaOf(t testing.TB, s *arrow.Schema) []byte { + t.Helper() + m, err := SchemaMessage(s) + require.NoError(t, err) + return m +} + +// nav walks a message's flatbuffer to the words a test patches. +type nav struct { + t testing.TB + buf []byte // the flatbuffer, after the 8-byte framing +} + +func (n nav) root() flatbuffers.Table { + return flatbuffers.Table{Bytes: n.buf, Pos: flatbuffers.GetUOffsetT(n.buf)} +} + +// union is the table in the union slot of t. +func (n nav) union(t flatbuffers.Table, slot flatbuffers.VOffsetT) flatbuffers.Table { + o := t.Offset(slot) + require.NotZero(n.t, o, "slot %d is set", slot) + var out flatbuffers.Table + t.Union(&out, flatbuffers.UOffsetT(o)) + return out +} + +// target is where the offset stored in slot of t leads. +func (n nav) target(t flatbuffers.Table, slot flatbuffers.VOffsetT) int { + o := t.Offset(slot) + require.NotZero(n.t, o, "slot %d is set", slot) + at := int(t.Pos) + int(o) + return at + int(flatbuffers.GetUOffsetT(n.buf[at:])) +} + +// slotAt is where slot of t is stored. +func (n nav) slotAt(t flatbuffers.Table, slot flatbuffers.VOffsetT) int { + o := t.Offset(slot) + require.NotZero(n.t, o, "slot %d is set", slot) + return int(t.Pos) + int(o) +} + +// element is the i-th table of the vector in slot of t. +func (n nav) element(t flatbuffers.Table, slot flatbuffers.VOffsetT, i int) flatbuffers.Table { + ep := n.target(t, slot) + 4 + 4*i + return flatbuffers.Table{Bytes: n.buf, Pos: flatbuffers.UOffsetT(ep + int(flatbuffers.GetUOffsetT(n.buf[ep:])))} +} + +// schema and field are the tables the cases below reach into. +func (n nav) schema() flatbuffers.Table { return n.union(n.root(), slotMessageHeader) } + +func put32(buf []byte, at int, v uint32) { binary.LittleEndian.PutUint32(buf[at:], v) } + +// inflatedSchemas are the nested schema's message with one word changed to state +// more than the bytes hold: a count, a length, an offset. +func inflatedSchemas(t testing.TB) map[string][]byte { + t.Helper() + cases := map[string]func(n nav, buf []byte){ + "four billion fields": func(n nav, buf []byte) { put32(buf, n.target(n.schema(), slotSchemaFields), 0xFFFFFFFF) }, + "2^30 fields": func(n nav, buf []byte) { put32(buf, n.target(n.schema(), slotSchemaFields), 1<<30) }, + "more fields than bytes": func(n nav, buf []byte) { put32(buf, n.target(n.schema(), slotSchemaFields), 1<<18) }, + "a billion schema metadata": func(n nav, buf []byte) { put32(buf, n.target(n.schema(), slotSchemaMetadata), 1<<30) }, + "a billion children": func(n nav, buf []byte) { + put32(buf, n.target(n.element(n.schema(), slotSchemaFields, 1), slotFieldChildren), 1<<30) + }, + "a billion field metadata": func(n nav, buf []byte) { + put32(buf, n.target(n.element(n.schema(), slotSchemaFields, 1), slotFieldMetadata), 1<<30) + }, + "a billion union type ids": func(n nav, buf []byte) { + u := n.union(n.element(n.schema(), slotSchemaFields, 3), slotFieldType) + put32(buf, n.target(u, slotUnionTypeIDs), 1<<30) + }, + "a name longer than the message": func(n nav, buf []byte) { + put32(buf, n.target(n.element(n.schema(), slotSchemaFields, 0), slotFieldName), 1<<30) + }, + "a time zone longer than the message": func(n nav, buf []byte) { + ts := n.union(n.element(n.schema(), slotSchemaFields, 2), slotFieldType) + put32(buf, n.target(ts, slotTimestampZone), 1<<30) + }, + "a metadata key longer than the message": func(n nav, buf []byte) { + kv := n.element(n.schema(), slotSchemaMetadata, 0) + put32(buf, n.target(kv, slotKeyValueKey), 1<<30) + }, + "a fields vector that starts outside": func(n nav, buf []byte) { + put32(buf, n.slotAt(n.schema(), slotSchemaFields), 1<<30) + }, + "a field that is its own offset": func(n nav, buf []byte) { + ep := n.target(n.schema(), slotSchemaFields) + 4 + put32(buf, ep, 0) + }, + "a field that points backwards": func(n nav, buf []byte) { + ep := n.target(n.schema(), slotSchemaFields) + 4 + put32(buf, ep, 0xFFFFFFF0) + }, + } + out := map[string][]byte{} + for name, patch := range cases { + m := schemaOf(t, nestedSchema()) + patch(nav{t: t, buf: m[8:]}, m[8:]) + out[name] = m + } + return out +} + +func TestASchemaThatStatesMoreThanItsBytesHoldIsRefusedBeforeItIsDecoded(t *testing.T) { + for name, m := range inflatedSchemas(t) { + t.Run(name, func(t *testing.T) { + var got error + grew := allocated(func() { got = readAll(m, &msgs{}) }) + require.ErrorIs(t, got, ErrMalformed) + require.Less(t, grew, uint64(boundedAlloc)) + }) + } +} + +// A table that is referenced many times is walked, and paid for, at each +// reference: a message of 400 KB whose 50,000 fields each have 50,000 children, +// all of them one table, is refused. +func TestASchemaThatPointsManyReferencesAtOneTableIsRefused(t *testing.T) { + const n = 50_000 + b := flatbuffers.NewBuilder(1 << 20) + name := b.CreateString("x") + b.StartObject(7) + b.PrependUOffsetTSlot(0, name, 0) + leaf := b.EndObject() + + vector := func(of flatbuffers.UOffsetT) flatbuffers.UOffsetT { + b.StartVector(4, n, 4) + for i := 0; i < n; i++ { + b.PrependUOffsetT(of) + } + return b.EndVector(n) + } + children := vector(leaf) + b.StartObject(7) + b.PrependUOffsetTSlot(0, name, 0) + b.PrependUOffsetTSlot(5, children, 0) + mid := b.EndObject() + fields := vector(mid) + + b.StartObject(4) + b.PrependUOffsetTSlot(1, fields, 0) + schema := b.EndObject() + b.StartObject(5) + b.PrependByteSlot(1, messageHeaderSchema, 0) + b.PrependUOffsetTSlot(2, schema, 0) + b.Finish(b.EndObject()) + flat := b.FinishedBytes() + require.Less(t, len(flat), int(maxMetadataBytes), "it fits the metadata cap, so only the walk can refuse it") + + m := append([]byte{0xFF, 0xFF, 0xFF, 0xFF, 0, 0, 0, 0}, flat...) + put32(m, 4, uint32(len(flat))) + var got error + grew := allocated(func() { got = readAll(m, &msgs{}) }) + require.ErrorIs(t, got, ErrMalformed) + require.Less(t, grew, uint64(boundedAlloc)) +} + +// The walk is no stricter than arrow-go on a schema that is what it says it is. +func TestSchemasThatAreWhatTheySayPassTheWalk(t *testing.T) { + wide := make([]arrow.Field, 4000) + for i := range wide { + wide[i] = arrow.Field{Name: fmt.Sprintf("column_%04d", i), Type: arrow.PrimitiveTypes.Int64, Nullable: true, + Metadata: arrow.NewMetadata([]string{"i"}, []string{fmt.Sprint(i)})} + } + for name, schema := range map[string]*arrow.Schema{ + "nested, with every vector the walk follows": nestedSchema(), + "wide, with metadata on every field": arrow.NewSchema(wide, nil), + "empty": arrow.NewSchema(nil, nil), + } { + t.Run(name, func(t *testing.T) { + rd, err := NewReader(schemaOf(t, schema), &msgs{}) + require.NoError(t, err) + defer rd.Release() + require.True(t, rd.Schema().Equal(schema), "got %s want %s", rd.Schema(), schema) + if schema.HasMetadata() { + require.Equal(t, schema.Metadata().Keys(), rd.Schema().Metadata().Keys()) + require.Equal(t, schema.Metadata().Values(), rd.Schema().Metadata().Values()) + } + for i, f := range schema.Fields() { + require.Equal(t, f.Metadata.Keys(), rd.Schema().Field(i).Metadata.Keys()) + } + }) + } +} + +// A schema carried in the first message of a stream with no header is vetted +// like one in the header. +func TestAFirstMessageThatCarriesTheSchemaIsVettedToo(t *testing.T) { + for name, m := range inflatedSchemas(t) { + t.Run(name, func(t *testing.T) { + var got error + grew := allocated(func() { got = readAll(nil, &msgs{list: [][]byte{m}}) }) + require.ErrorIs(t, got, ErrMalformed) + require.Less(t, grew, uint64(boundedAlloc)) + }) + } +} + +// Whatever one byte of a schema message is made, it is read, or refused, with +// bounded memory. +func TestNoSingleByteMakesASchemaMessageCostMemory(t *testing.T) { + m := schemaOf(t, nestedSchema()) + for at := 0; at < len(m); at++ { + for _, v := range []byte{0x00, 0x01, 0x7F, 0x80, 0xFF} { + if m[at] == v { + continue + } + mutated := append([]byte(nil), m...) + mutated[at] = v + grew := allocated(func() { _ = readAll(mutated, &msgs{}) }) + require.Less(t, grew, uint64(boundedAlloc), "byte %d set to %#x", at, v) + } + } +} + +// FuzzSchemaReader runs its seeds in the ordinary suite, and `go test -fuzz +// FuzzSchemaReader ./internal/arrowipc` explores from them. The body bound is +// small here: a message that states a body is allocated up to the bound before +// it is read, which is the bound's job and not what this looks for. +func FuzzSchemaReader(f *testing.F) { + for _, s := range []*arrow.Schema{nestedSchema(), testSchema, arrow.NewSchema(nil, nil)} { + m, err := SchemaMessage(s) + if err != nil { + f.Fatal(err) + } + f.Add(m) + } + for _, m := range inflatedSchemas(f) { + f.Add(m) + } + f.Fuzz(func(t *testing.T, data []byte) { + grew := allocated(func() { _ = readAll(data, &msgs{}, WithMaxMessageBytes(1<<20)) }) + if grew > boundedAlloc { + t.Fatalf("%d bytes of input allocated %d", len(data), grew) + } + }) +} From 2869574360ae14604c75700225ff563edc84790a Mon Sep 17 00:00:00 2001 From: crisfiorini-oss <301773486+crisfiorini-oss@users.noreply.github.com> Date: Thu, 8 Oct 2026 18:26:51 -0400 Subject: [PATCH 11/14] fix(config): default SWH_MAX_ARROW_MESSAGE_BYTES to grpc-go's 4 MiB receive limit An Arrow message reaches InsertRows inside one gRPC message and cannot be larger than the message carrying it, so a bound above the gRPC receive limit refuses nothing a client can send. The default is now 4 MiB, grpc-go's own; the knob stays, for a deployment that raises the gRPC server's limit. The two bigquery tests that send a batch of more than 4 MiB straight to the backend, to exercise its own 8 MiB row limit, raise the bound for themselves. Co-Authored-By: Claude Sonnet 5.5 --- README.md | 2 +- internal/arrowipc/arrowipc.go | 7 +++++-- internal/backend/bigquery/insert_test.go | 10 ++++++++-- internal/config/config_test.go | 2 +- 4 files changed, 15 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 14a252a..21e2aac 100644 --- a/README.md +++ b/README.md @@ -181,7 +181,7 @@ BigQuery has no equivalent of (map, duration, dictionary) is `Unimplemented`. | `SWH_QUERY_TIMEOUT` | `0` | per-query timeout (Go duration; 0 = backend default); a request may only shorten it | | `SWH_AUTH_TOKEN` | — | shared secret every caller must present; required unless `SWH_ALLOW_ANONYMOUS=true` | | `SWH_ALLOW_ANONYMOUS` | `false` | `true` accepts callers with no token; refused together with `SWH_AUTH_TOKEN` | -| `SWH_MAX_ARROW_MESSAGE_BYTES` | `67108864` (64 MiB) | largest body of one Arrow IPC message a client may send (`InsertRows`); a larger one is `InvalidArgument`, refused before any of it is allocated. Must be positive. grpc-go's own 4 MiB receive limit is checked first | +| `SWH_MAX_ARROW_MESSAGE_BYTES` | `4194304` (4 MiB) | largest body of one Arrow IPC message a client may send (`InsertRows`); a larger one is `InvalidArgument`, refused before any of it is allocated. Must be positive. The default is grpc-go's own receive limit, which a message cannot exceed anyway: raising this above it has no effect until the gRPC server's limit is raised too | The `bigquery` backend reads `SWH_DATABASE` (required), `SWH_DATASET`, `SWH_LOCATION`, `SWH_CREDENTIALS_FILE`, `SWH_MAX_QUERY_BYTES`, diff --git a/internal/arrowipc/arrowipc.go b/internal/arrowipc/arrowipc.go index e39914b..4f77d32 100644 --- a/internal/arrowipc/arrowipc.go +++ b/internal/arrowipc/arrowipc.go @@ -174,8 +174,11 @@ type Reader struct { const ( // DefaultMaxMessageBytes is the largest message body a reader accepts when - // the caller names no limit: what SWH_MAX_ARROW_MESSAGE_BYTES defaults to. - DefaultMaxMessageBytes int64 = 64 << 20 + // the caller names no limit: what SWH_MAX_ARROW_MESSAGE_BYTES defaults to. It + // is grpc-go's default receive limit, because a message arrives inside one + // gRPC message and cannot be larger than the one carrying it, so a larger + // bound refuses nothing a client can send. + DefaultMaxMessageBytes int64 = 4 << 20 // maxMetadataBytes is the largest message metadata a reader accepts. The // metadata is a schema or a batch's buffer layout, a few bytes per column; a diff --git a/internal/backend/bigquery/insert_test.go b/internal/backend/bigquery/insert_test.go index 65f9a21..8c940e1 100644 --- a/internal/backend/bigquery/insert_test.go +++ b/internal/backend/bigquery/insert_test.go @@ -283,7 +283,10 @@ func TestInsertRowsRefusesARowTooLargeToSendAndStoresTheOthers(t *testing.T) { f := newFake(t) f.handle("GET /projects/test-project/datasets/d/tables/t", 200, eventTable) f.handle("POST "+insertPath, 200, map[string]any{}) - b := f.open(t, backend.Config{}) + // A batch this size is a message of more than the default bound, which the + // gRPC server's receive limit would refuse first; the bound is raised so the + // backend's own row limit is what is under test. + b := f.open(t, backend.Config{MaxArrowMessageBytes: 64 << 20}) huge := strings.Repeat("x", maxRequestBytes) header, batch := eventBatch(t, []int64{1, 2, 3}, []string{"a", huge, "c"}) @@ -307,7 +310,10 @@ func TestInsertRowsSplitsARequestThatWouldPassTheByteLimit(t *testing.T) { f := newFake(t) f.handle("GET /projects/test-project/datasets/d/tables/t", 200, eventTable) f.handle("POST "+insertPath, 200, map[string]any{}) - b := f.open(t, backend.Config{}) + // A batch this size is a message of more than the default bound, which the + // gRPC server's receive limit would refuse first; the bound is raised so the + // backend's own row limit is what is under test. + b := f.open(t, backend.Config{MaxArrowMessageBytes: 64 << 20}) big := strings.Repeat("x", maxRequestBytes/2) header, batch := eventBatch(t, []int64{1, 2, 3}, []string{big, big, big}) diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 47e2f8f..c7a5581 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -159,7 +159,7 @@ func TestDefaultsAndParsing(t *testing.T) { require.Equal(t, int64(4096), cfg.Backend.MaxQueryBytes) require.Equal(t, 90*time.Second, cfg.Backend.QueryTimeout) require.Equal(t, 5439, cfg.Backend.Port) - require.Equal(t, int64(64<<20), cfg.Backend.MaxArrowMessageBytes, "the Arrow message bound is on by default") + require.Equal(t, int64(4<<20), cfg.Backend.MaxArrowMessageBytes, "the Arrow message bound is on by default") } func TestArrowMessageBoundIsConfigurable(t *testing.T) { From 1a58363784c89057bfb57f8d1bdd503241840a57 Mon Sep 17 00:00:00 2001 From: crisfiorini-oss <301773486+crisfiorini-oss@users.noreply.github.com> Date: Thu, 15 Oct 2026 15:17:05 +0000 Subject: [PATCH 12/14] chore(proto): reproduce bindings through pinned companion formatter --- gen/codefly/warehouse/v0/warehouse.pb.go | 5 +++-- gen/codefly/warehouse/v0/warehouse_grpc.pb.go | 1 + 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/gen/codefly/warehouse/v0/warehouse.pb.go b/gen/codefly/warehouse/v0/warehouse.pb.go index c67bb8f..bc0991b 100644 --- a/gen/codefly/warehouse/v0/warehouse.pb.go +++ b/gen/codefly/warehouse/v0/warehouse.pb.go @@ -15,11 +15,12 @@ package warehousev0 import ( - protoreflect "google.golang.org/protobuf/reflect/protoreflect" - protoimpl "google.golang.org/protobuf/runtime/protoimpl" reflect "reflect" sync "sync" unsafe "unsafe" + + protoreflect "google.golang.org/protobuf/reflect/protoreflect" + protoimpl "google.golang.org/protobuf/runtime/protoimpl" ) const ( diff --git a/gen/codefly/warehouse/v0/warehouse_grpc.pb.go b/gen/codefly/warehouse/v0/warehouse_grpc.pb.go index 4462838..97dbaab 100644 --- a/gen/codefly/warehouse/v0/warehouse_grpc.pb.go +++ b/gen/codefly/warehouse/v0/warehouse_grpc.pb.go @@ -16,6 +16,7 @@ package warehousev0 import ( context "context" + grpc "google.golang.org/grpc" codes "google.golang.org/grpc/codes" status "google.golang.org/grpc/status" From 5fca15bbf773dd8f3eaa72de5e4a6be3be157ea9 Mon Sep 17 00:00:00 2001 From: crisfiorini-oss <301773486+crisfiorini-oss@users.noreply.github.com> Date: Thu, 15 Oct 2026 15:17:05 +0000 Subject: [PATCH 13/14] docs(bigquery): record reuse, scope exception and Scan dependency --- README.md | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/README.md b/README.md index 21e2aac..34713ae 100644 --- a/README.md +++ b/README.md @@ -163,6 +163,14 @@ the batch's columns with the table's metadata, so no vendor text decides it; a table altered while the call runs can surface as `INVALID_VALUE`. An Arrow type BigQuery has no equivalent of (map, duration, dictionary) is `Unimplemented`. +**Implementation decisions (2026-10-15).** + +| Decision | Evidence and outcome | +|----------|----------------------| +| Build or reuse | Searched this repository and the shared `codefly-dev/core` and `codefly-dev/sdk-go` implementations: the existing BigQuery backend in [#14](https://github.com/codefly-dev/service-warehouse/pull/14) is the warehouse implementation; the shared libraries supply caller authentication and agent/runtime support, with no other BigQuery backend. Retain and extend the existing backend and `internal/arrowipc`. Go's standard library has no BigQuery or Arrow implementation; reuse Google's `cloud.google.com/go/bigquery` client (its query, inserter and per-row error implementations and tests) and Apache Arrow's IPC implementation, with the portable contract and bounded framing owned here. | +| Scope exception | On 2026-10-15 the user approved retaining #14's existing implementation, which already exceeds the 2,000 hand-written-line stop threshold in [#23](https://github.com/codefly-dev/service-warehouse/pull/23). This exception retains the documented catalog, DDL, query, job and streaming-insert behavior; it does not authorize additional features. | +| Feature order | Finish and merge #14 first, then the DuckDB engine, then restricted `Scan`, then export the shared Arrow helpers. Open each dependent PR only after its base has merged. `Scan` is absent from this version's contract; the backend's query path uses REST results and creates query jobs. | + ## Configuration (env) | Var | Default | Notes | @@ -473,6 +481,13 @@ deletes it. `SWH_TEST_BIGQUERY_LOCATION` and `SWH_TEST_BIGQUERY_CREDENTIALS_FILE are optional (Application Default Credentials otherwise). CI runs no BigQuery, so that test has not run in CI. +The current live test proves streamed inserts are visible to `Query`, when run. +It cannot prove visibility through the Storage Read API because `Scan` is a +later feature. The `Scan` implementation must extend this live test to read a +streamed row through a read session, accept `IN` and timestamp comparisons in +the session's row restriction, and verify that those filters select the expected +rows. A REST fake or a successful `Query` is not evidence for that requirement. + `buf` is never run on the host. `codefly generate proto` runs it inside the versioned proto companion image (`ghcr.io/codefly-dev/proto`), so the plugin versions and the `goimports` pass are the image's, fixed by its tag — two From b0420dbe2e72c7291eb235ba2e0c8d5de01cf5ea Mon Sep 17 00:00:00 2001 From: crisfiorini-oss Date: Wed, 21 Oct 2026 03:01:29 +0000 Subject: [PATCH 14/14] chore: mark generated warehouse bindings for review --- .gitattributes | 2 ++ 1 file changed, 2 insertions(+) create mode 100644 .gitattributes diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..296ff3a --- /dev/null +++ b/.gitattributes @@ -0,0 +1,2 @@ +# Protocol bindings are generated by the proto owner. Keep authored source visible. +gen/** linguist-generated=true