[grpc] Add basic bytestream QueryWriteStatus support

We don't support partial writes, so QueryWriteStatus will either
return 0 bytes committed/incomplete, or all bytes committed/complete.

server/grpc_bytestream.go

@@ -580,7 +580,21 @@ func (s *grpcServer) Write(srv bytestream.ByteStream_WriteServer) error {
 	return nil
 }
 
-func (s *grpcServer) QueryWriteStatus(context.Context, *bytestream.QueryWriteStatusRequest) (*bytestream.QueryWriteStatusResponse, error) {
-	return nil, status.Error(codes.Unimplemented,
-		"QueryWriteStatus is not implemented")
+func (s *grpcServer) QueryWriteStatus(ctx context.Context, req *bytestream.QueryWriteStatusRequest) (*bytestream.QueryWriteStatusResponse, error) {
+
+	hash, size, _, err := s.parseWriteResource(req.ResourceName)
+	if err != nil {
+		return nil, err
+	}
+
+	// We don't support partial writes, so the status will either be fully written
+	// and complete, or 0 written and incomplete.
+
+	exists, _ := s.cache.Contains(ctx, cache.CAS, hash, size)
+
+	if !exists {
+		return &bytestream.QueryWriteStatusResponse{CommittedSize: 0, Complete: false}, nil
+	}
+
+	return &bytestream.QueryWriteStatusResponse{CommittedSize: size, Complete: true}, nil
 }

server/grpc_test.go

@@ -1135,6 +1135,108 @@ func TestGrpcByteStreamSkippedWrite(t *testing.T) {
 	}
 }
 
+func TestGrpcByteStreamQueryWriteStatus(t *testing.T) {
+
+	testBlob, testBlobHash := testutils.RandomDataAndHash(123)
+	testBlobDigest := pb.Digest{
+		Hash:      testBlobHash,
+		SizeBytes: int64(len(testBlob)),
+	}
+
+	instance := "unusedInstance"
+
+	resourceName := fmt.Sprintf(
+		"%s/uploads/%s/blobs/%s/%d",
+		instance,
+		uuid.New().String(),
+		testBlobDigest.Hash,
+		len(testBlob),
+	)
+
+	req := &bytestream.QueryWriteStatusRequest{ResourceName: resourceName}
+
+	resp, err := bsClient.QueryWriteStatus(context.Background(), req)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Confirm that the blob does not exist, and we get "0 bytes committed, incomplete".
+
+	if resp == nil {
+		t.Fatal("Expected non-nil *bytestream.QueryWriteStatusResponse")
+	}
+	if resp.CommittedSize != 0 {
+		t.Fatalf("Expected CommittedSize == 0, got: %d", resp.CommittedSize)
+	}
+	if resp.Complete {
+		t.Fatal("Expected incomplete response")
+	}
+
+	// Write the blob.
+	bswc, err := bsClient.Write(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	err = bswc.Send(&bytestream.WriteRequest{
+		ResourceName: resourceName,
+		FinishWrite:  true,
+		Data:         testBlob,
+		WriteOffset:  0,
+	})
+	if err != nil && err == io.EOF {
+		t.Fatal(err)
+	}
+
+	_, err = bswc.CloseAndRecv()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Confirm that the blob now exists, and we get "len(testBlob) bytes committed, complete".
+
+	resp, err = bsClient.QueryWriteStatus(context.Background(), req)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if resp == nil {
+		t.Fatal("Expected non-nil *bytestream.QueryWriteStatusResponse")
+	}
+	if resp.CommittedSize != int64(len(testBlob)) {
+		t.Fatalf("Expected CommittedSize == %d, got: %d", len(testBlob), resp.CommittedSize)
+	}
+	if !resp.Complete {
+		t.Fatal("Expected complete response")
+	}
+
+	// Check the empty blob special case, which should always exist.
+
+	resourceName = fmt.Sprintf(
+		"%s/uploads/%s/blobs/%s/0",
+		instance,
+		uuid.New().String(),
+		emptySha256,
+	)
+
+	req = &bytestream.QueryWriteStatusRequest{ResourceName: resourceName}
+
+	resp, err = bsClient.QueryWriteStatus(context.Background(), req)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if resp == nil {
+		t.Fatal("Expected non-nil *bytestream.QueryWriteStatusResponse")
+	}
+	if resp.CommittedSize != 0 {
+		t.Fatalf("Expected CommittedSize == 0, got: %d", resp.CommittedSize)
+	}
+	if !resp.Complete {
+		t.Fatal("Expected complete response")
+	}
+}
+
 func TestGrpcCasBasics(t *testing.T) {
 
 	testBlob, testBlobHash := testutils.RandomDataAndHash(256)