feat: Support google.protobuf.Timestamp values for the storage write API

internal/contentdata/repository.go

@@ -398,23 +398,7 @@ func (r *Repository) AddTableData(ctx context.Context, tx *connection.Tx, projec
 		values := make([]interface{}, 0, len(table.Columns))
 
 		for _, column := range columns {
-			if value, found := data[column.Name]; found {
-				isTimestampColumn := column.Type == types.TIMESTAMP
-				inputString, isInputString := value.(string)
-
-				if isInputString && isTimestampColumn {
-					parsedTimestamp, err := zetasqlite.TimeFromTimestampValue(inputString)
-					// If we could parse the timestamp, use it when inserting, otherwise fallback to the supplied value
-					if err == nil {
-						values = append(values, parsedTimestamp)
-						continue
-					}
-				}
-
-				values = append(values, value)
-			} else {
-				values = append(values, nil)
-			}
+			values = append(values, columnValue(data, column))
 		}
 
 		if _, err := stmt.ExecContext(ctx, values...); err != nil {
@@ -425,6 +409,28 @@ func (r *Repository) AddTableData(ctx context.Context, tx *connection.Tx, projec
 	return nil
 }
 
+func columnValue(data map[string]interface{}, column *types.Column) interface{} {
+	value, found := data[column.Name]
+	if !found {
+		return nil
+	}
+
+	switch column.Type {
+	case types.TIMESTAMP:
+		s, ok := value.(string)
+		if !ok {
+			return value
+		}
+		if t, err := zetasqlite.TimeFromTimestampValue(s); err == nil {
+			return t
+		}
+		// Fallback to the supplied value when it can't be parsed
+		return value
+	default:
+		return value
+	}
+}
+
 func (r *Repository) DeleteTables(ctx context.Context, tx *connection.Tx, projectID, datasetID string, tableIDs []string) error {
 	tx.SetProjectAndDataset(projectID, datasetID)
 	if err := tx.ContentRepoMode(); err != nil {

internal/types/types.go

@@ -2,7 +2,6 @@ package types
 
 import (
 	"fmt"
-	"time"
 
 	"github.com/apache/arrow/go/v10/arrow/array"
 	"github.com/goccy/bigquery-emulator/types"
@@ -192,19 +191,31 @@ func Format(schema *bigqueryv2.TableSchema, rows []*TableRow, useInt64Timestamp
 	for _, row := range rows {
 		cells := make([]*TableCell, 0, len(row.F))
 		for colIdx, cell := range row.F {
-			if schema.Fields[colIdx].Type == "TIMESTAMP" && cell.V != nil {
-				t, _ := zetasqlite.TimeFromTimestampValue(cell.V.(string))
-				microsec := t.UnixNano() / int64(time.Microsecond)
-				cells = append(cells, &TableCell{
-					V: fmt.Sprint(microsec),
-				})
-			} else {
-				cells = append(cells, cell)
-			}
+			cells = append(cells, formatCell(schema.Fields[colIdx], cell))
 		}
 		formattedRows = append(formattedRows, &TableRow{
 			F: cells,
 		})
 	}
 	return formattedRows
 }
+
+// formatCell formats timestamp cells as microseconds to match what the
+// client libraries expect
+func formatCell(schema *bigqueryv2.TableFieldSchema, cell *TableCell) *TableCell {
+	switch schema.Type {
+	case "TIMESTAMP":
+		s, ok := cell.V.(string)
+		if !ok {
+			return cell
+		}
+		t, _ := zetasqlite.TimeFromTimestampValue(s)
+
+		microsec := t.UnixMicro()
+		return &TableCell{
+			V: fmt.Sprint(microsec),
+		}
+	default:
+		return cell
+	}
+}

server/storage_handler.go

@@ -524,7 +524,6 @@ func (s *storageWriteServer) appendRows(req *storagepb.AppendRowsRequest, msgDes
 		status.rows = append(status.rows, data...)
 	}
 	return s.sendResult(stream, streamName, offset+int64(len(rows)))
-
 }
 
 func (s *storageWriteServer) sendResult(stream storagepb.BigQueryWrite_AppendRowsServer, streamName string, offset int64) error {
@@ -598,9 +597,43 @@ func (s *storageWriteServer) decodeProtoReflectValue(f protoreflect.FieldDescrip
 		}
 		return ret, nil
 	}
+
+	// The BigQuery SDK sends dynamic, well known types with underscore separators.
+	// They're also prefixed with a scope, so we have to check the suffix.
+	//
+	// BigQuery supports timestamps being int64 and [timestamppb.Timestamp]:
+	// https://cloud.google.com/bigquery/docs/supported-data-types
+	var fullName string
+	if f.Message() != nil {
+		fullName = string(f.Message().FullName())
+	}
+	if strings.HasSuffix(fullName, "google_protobuf_Timestamp") || strings.HasSuffix(fullName, "google.protobuf.Timestamp") {
+		return decodeTimestamp(v.Message().Interface())
+	}
 	return s.decodeProtoReflectValueFromKind(f.Kind(), v)
 }
 
+// decodeTimestamp unwraps a [timestamppb.Timestamp] wire-compatible message into the
+// underlying timestamp.
+//
+// The message may be a [dynamicpb.Message] sent to us via the storage write API, so we
+// need a round-trip encode/decode for conversion.
+func decodeTimestamp(msg proto.Message) (interface{}, error) {
+	b, err := proto.Marshal(msg)
+	if err != nil {
+		return nil, fmt.Errorf("encoding timestamppb.Timestamp: %w", err)
+	}
+	ts := new(timestamppb.Timestamp)
+	if err := proto.Unmarshal(b, ts); err != nil {
+		return nil, fmt.Errorf("decoding timestamppb.Timestamp: %w", err)
+	}
+	conv := ts.AsTime()
+	if conv.IsZero() {
+		return time.Time{}, nil
+	}
+	return conv.Truncate(time.Microsecond), nil
+}
+
 func (s *storageWriteServer) decodeProtoReflectValueFromKind(kind protoreflect.Kind, v protoreflect.Value) (interface{}, error) {
 	if !v.IsValid() {
 		return nil, nil

server/storage_test.go

@@ -31,7 +31,14 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/encoding/protojson"
 	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/reflect/protodesc"
+	"google.golang.org/protobuf/reflect/protoreflect"
+	"google.golang.org/protobuf/types/descriptorpb"
+	"google.golang.org/protobuf/types/dynamicpb"
+	"google.golang.org/protobuf/types/known/timestamppb"
+	"google.golang.org/protobuf/types/known/wrapperspb"
 
 	"github.com/goccy/bigquery-emulator/types"
 )
@@ -315,7 +322,6 @@ func processAvro(t *testing.T, ctx context.Context, schema string, ch <-chan *st
 			undecoded := rows.GetAvroRows().GetSerializedBinaryRows()
 			for len(undecoded) > 0 {
 				datum, remainingBytes, err := codec.NativeFromBinary(undecoded)
-
 				if err != nil {
 					if err == io.EOF {
 						break
@@ -721,3 +727,191 @@ func generateExampleMessages(numMessages int) ([][]byte, error) {
 	}
 	return msgs, nil
 }
+
+func TestStorageWriteDynamicDescriptor(t *testing.T) {
+	const (
+		projectID = "test"
+		datasetID = "test"
+		tableID   = "sample"
+	)
+
+	// Ensure that the zero value and populated timestamps round-trip the way we expect
+	testTimestamps := []time.Time{{}, time.Now()}
+	for _, expected := range testTimestamps {
+		ctx := context.Background()
+		bqServer, err := server.New(server.TempStorage)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if err := bqServer.Load(
+			server.StructSource(
+				types.NewProject(
+					projectID,
+					types.NewDataset(
+						datasetID,
+						types.NewTable(
+							tableID,
+							[]*types.Column{
+								types.NewColumn("timestamp", types.TIMESTAMP),
+								types.NewColumn("msg", types.STRING),
+							},
+							nil,
+						),
+					),
+				),
+			),
+		); err != nil {
+			t.Fatal(err)
+		}
+		testServer := bqServer.TestServer()
+		defer func() {
+			testServer.Close()
+			bqServer.Close()
+		}()
+		opts, err := testServer.GRPCClientOptions(ctx)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		client, err := managedwriter.NewClient(ctx, projectID, opts...)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer client.Close()
+
+		fullTableName := managedwriter.TableParentFromParts(
+			projectID,
+			datasetID,
+			tableID,
+		)
+
+		msgDesc, protoDesc := dynamicProtoDescriptors(t)
+		stream, err := client.NewManagedStream(
+			ctx,
+			managedwriter.WithDestinationTable(fullTableName),
+			managedwriter.WithType(managedwriter.DefaultStream),
+			managedwriter.WithSchemaDescriptor(protoDesc),
+		)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		bqc, err := bigquery.NewClient(
+			ctx,
+			projectID,
+			option.WithEndpoint(testServer.URL),
+			option.WithoutAuthentication(),
+		)
+		if err != nil {
+			t.Fatal(err)
+		}
+		t.Cleanup(func() {
+			_ = bqc.Close()
+		})
+
+		payload := fmt.Appendf(nil,
+			`{"timestamp": "%s", "msg": "hello"}`,
+			expected.Format(time.RFC3339Nano),
+		)
+		msg := dynamicpb.NewMessage(msgDesc)
+		if err := protojson.Unmarshal(payload, msg); err != nil {
+			t.Fatal(err)
+		}
+
+		row, err := proto.Marshal(msg)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		res, err := stream.AppendRows(ctx, [][]byte{row})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if _, err := res.GetResult(ctx); err != nil {
+			t.Fatal(err)
+		}
+
+		// Ensure the value retrieved is the value written
+		table := bqc.Dataset(datasetID).Table(tableID)
+		it := table.Read(ctx)
+
+		got := map[string]bigquery.Value{}
+		if err := it.Next(&got); err != nil {
+			t.Fatal(err)
+		}
+
+		ts, ok := got["timestamp"]
+		if !ok {
+			t.Fatalf("expected 'timestamp' in map, got %+v", got)
+		}
+		gotTime, ok := ts.(time.Time)
+		if !ok {
+			t.Fatalf("expected 'time.Time', got %T", ts)
+		}
+		if !gotTime.Equal(expected) {
+			t.Fatalf("expected %v, got %v", expected, gotTime)
+		}
+	}
+}
+
+// dynamicProtoDescriptors creates a protobuf at runtime, returning the message and
+// type descriptors.
+//
+// This is specifically needed to verify sending [timestamppb.Timestamp] values to the
+// storage write API for feature parity with BigQuery.
+func dynamicProtoDescriptors(t *testing.T) (protoreflect.MessageDescriptor, *descriptorpb.DescriptorProto) {
+	t.Helper()
+
+	scope := "test"
+	dp := &descriptorpb.DescriptorProto{
+		Name: proto.String(scope),
+		Field: []*descriptorpb.FieldDescriptorProto{
+			{
+				Name:     proto.String("timestamp"),
+				Number:   proto.Int32(1),
+				Type:     descriptorpb.FieldDescriptorProto_TYPE_MESSAGE.Enum(),
+				TypeName: proto.String(".google.protobuf.Timestamp"),
+				Label:    descriptorpb.FieldDescriptorProto_LABEL_REQUIRED.Enum(),
+			},
+			{
+				Name:   proto.String("msg"),
+				Number: proto.Int32(2),
+				Type:   descriptorpb.FieldDescriptorProto_TYPE_STRING.Enum(),
+				Label:  descriptorpb.FieldDescriptorProto_LABEL_REQUIRED.Enum(),
+			},
+		},
+	}
+	fdp := &descriptorpb.FileDescriptorProto{
+		MessageType: []*descriptorpb.DescriptorProto{dp},
+		Name:        proto.String(scope + ".proto"),
+		Syntax:      proto.String("proto2"),
+		Dependency: []string{
+			"google/protobuf/wrappers.proto",
+			"google/protobuf/timestamp.proto",
+		},
+	}
+
+	fdpList := []*descriptorpb.FileDescriptorProto{
+		fdp,
+		protodesc.ToFileDescriptorProto(wrapperspb.File_google_protobuf_wrappers_proto),
+		protodesc.ToFileDescriptorProto(timestamppb.File_google_protobuf_timestamp_proto),
+	}
+	fds := &descriptorpb.FileDescriptorSet{File: fdpList}
+
+	files, err := protodesc.NewFiles(fds)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	found, err := files.FindDescriptorByName(protoreflect.FullName(scope))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	messageDescriptor := found.(protoreflect.MessageDescriptor)
+	protoDescriptor, err := adapt.NormalizeDescriptor(messageDescriptor)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return messageDescriptor, protoDescriptor
+}