Repository: google/exposure-notifications-server
Pull request #964 (2020-09-08 23:38:12)

Add support for partial revisions
Fixes GH-893

**Release Note**

```release-note
Add configurable support for partial revisions
```

/assign @mikehelmick 

google-oss-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/google/exposure-notifications-server/pull/964#" title="Author self-approved">sethvargo</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=google%2Fexposure-notifications-server).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/google/exposure-notifications-server/blob/main/OWNERS)~~ [sethvargo]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->
google-oss-robot: @sethvargo: The following test **failed**, say `/retest` to rerun all failed tests:

Test name | Commit | Details | Rerun command
--- | --- | --- | ---
pull-en-server-release-unit | d868db0bf0dd184086f20dae4378f740e594b72f | [link](https://oss-prow.knative.dev/view/gs/oss-prow/pr-logs/pull/google_exposure-notifications-server/964/pull-en-server-release-unit/1303478057822064640) | `/test pull-en-server-release-unit`

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository. I understand the commands that are listed [here](https://go.k8s.io/bot-commands).
</details>
<!-- test report -->
google-oss-robot: @sethvargo: The following test **failed**, say `/retest` to rerun all failed tests:

Test name | Commit | Details | Rerun command
--- | --- | --- | ---
pull-en-server-release-unit | a239f9ecb1489cb2331e47eda9a1dd5370c738f8 | [link](https://oss-prow.knative.dev/view/gs/oss-prow/pr-logs/pull/google_exposure-notifications-server/964/pull-en-server-release-unit/1304086944866635776) | `/test pull-en-server-release-unit`

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository. I understand the commands that are listed [here](https://go.k8s.io/bot-commands).
</details>
<!-- test report -->
google-oss-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/google/exposure-notifications-server/pull/964#pullrequestreview-486327371" title="Approved">mikehelmick</a>*, *<a href="https://github.com/google/exposure-notifications-server/pull/964#" title="Author self-approved">sethvargo</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=google%2Fexposure-notifications-server).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/google/exposure-notifications-server/blob/main/OWNERS)~~ [mikehelmick,sethvargo]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->

Patch:
diff --git a/internal/federationin/federationin.go b/internal/federationin/federationin.go
index 712eba4..80baf58 100644
--- a/internal/federationin/federationin.go
+++ b/internal/federationin/federationin.go
@@ -57,7 +57,7 @@ var (
 
 type (
 	fetchFn               func(context.Context, *pb.FederationFetchRequest, ...grpc.CallOption) (*pb.FederationFetchResponse, error)
-	insertExposuresFn     func(context.Context, []*publishmodel.Exposure, *pb.RevisionTokenData, bool) (int, error)
+	insertExposuresFn     func(context.Context, *publishdb.InsertAndReviseExposuresRequest) (*publishdb.InsertAndReviseExposuresResponse, error)
 	startFederationSyncFn func(context.Context, *model.FederationInQuery, time.Time) (int64, database.FinalizeSyncFn, error)
 )
 
@@ -292,24 +292,28 @@ func pull(ctx context.Context, metrics metrics.Exporter, opts *pullOptions) (err
 					})
 
 					if len(exposures) == fetchBatchSize {
-						n := 0
-						if n, err = opts.deps.insertExposures(ctx, exposures, nil, false); err != nil {
+						resp, err := opts.deps.insertExposures(ctx, &publishdb.InsertAndReviseExposuresRequest{
+							Incoming: exposures,
+						})
+						if err != nil {
 							metrics.WriteInt("federation-pull-inserts", false, len(exposures))
 							return fmt.Errorf("inserting %d exposures: %w", len(exposures), err)
 						}
-						total += n
+						total += int(resp.Inserted)
 						exposures = nil // Start a new batch.
 					}
 				}
 			}
 		}
 		if len(exposures) > 0 {
-			n := 0
-			if n, err = opts.deps.insertExposures(ctx, exposures, nil, false); err != nil {
+			resp, err := opts.deps.insertExposures(ctx, &publishdb.InsertAndReviseExposuresRequest{
+				Incoming: exposures,
+			})
+			if err != nil {
 				metrics.WriteInt("federation-pull-inserts", false, len(exposures))
 				return fmt.Errorf("inserting %d exposures: %w", len(exposures), err)
 			}
-			total += n
+			total += int(resp.Inserted)
 		}
 
 		partial = response.PartialResponse
diff --git a/internal/generate/handler.go b/internal/generate/handler.go
index 7590e76..c57a2e8 100644
--- a/internal/generate/handler.go
+++ b/internal/generate/handler.go
@@ -156,7 +156,10 @@ func (h *generateHandler) generate(ctx context.Context, regions []string) error
 			return fmt.Errorf("failed to transform generated exposures: %w", err)
 		}
 
-		n, err := h.database.InsertAndReviseExposures(ctx, exposures, nil, true)
+		n, err := h.database.InsertAndReviseExposures(ctx, &publishdb.InsertAndReviseExposuresRequest{
+			Incoming:     exposures,
+			RequireToken: true,
+		})
 		if err != nil {
 			return fmt.Errorf("failed to write exposure record: %w", err)
 		}
@@ -186,7 +189,12 @@ func (h *generateHandler) generate(ctx context.Context, regions []string) error
 				})
 			}
 
-			n, err := h.database.InsertAndReviseExposures(ctx, exposures, &token, true)
+			// Bypass revision token enforcement on generated data.
+			n, err := h.database.InsertAndReviseExposures(ctx, &publishdb.InsertAndReviseExposuresRequest{
+				Incoming:     exposures,
+				Token:        &token,
+				RequireToken: true,
+			})
 			if err != nil {
 				return fmt.Errorf("failed to revise exposure record: %w", err)
 			}
diff --git a/internal/publish/config.go b/internal/publish/config.go
index 5c8d274..1f17ff9 100644
--- a/internal/publish/config.go
+++ b/internal/publish/config.go
@@ -61,6 +61,12 @@ type Config struct {
 
 	RevisionKeyCacheDuration time.Duration `env:"REVISION_KEY_CACHE_DURATION, default=1m"`
 
+	// AllowPartialRevisions permits uploading multiple exposure keys with a
+	// revision token where only a subset of the keys are in the token. In that
+	// case, only the incoming exposure keys that match the revision token are
+	// uploaded and the remainder are discarded.
+	AllowPartialRevisions bool `env:"ALLOW_PARTIAL_REVISIONS, default=false"`
+
 	// API Versions.
 	EnableV1Alpha1API bool `env:"ENABLE_V1ALPHA1_API, default=true"`
 
diff --git a/internal/publish/database/exposure.go b/internal/publish/database/exposure.go
index 186ec54..bf83372 100644
--- a/internal/publish/database/exposure.go
+++ b/internal/publish/database/exposure.go
@@ -367,17 +367,62 @@ func executeReviseExposure(ctx context.Context, tx pgx.Tx, stmtName string, exp
 	return nil
 }
 
-// InsertAndReviseExposures transactionally revises and inserts a set of keys as necessary.
-func (db *PublishDB) InsertAndReviseExposures(ctx context.Context, incoming []*model.Exposure, token *pb.RevisionTokenData, tokenRequired bool) (int, error) {
-	logger := logging.FromContext(ctx)
-	updated := 0
-	err := db.db.InTx(ctx, pgx.ReadCommitted, func(tx pgx.Tx) error {
+// InsertAndReviseExposuresRequest is used as input to InsertAndReviseExposures.
+type InsertAndReviseExposuresRequest struct {
+	Incoming []*model.Exposure
+	Token    *pb.RevisionTokenData
+
+	// RequireToken requires that the request supply a revision token to re-upload
+	// existing keys.
+	RequireToken bool
+
+	// AllowPartialRevisions allows revising a subset of exposures if other
+	// exposures are included that are not part of the revision token. This exists
+	// to support roaming scenarios. This is only used if RequireToken is true.
+	AllowPartialRevisions bool
+}
+
+// InsertAndReviseExposuresResponse is the response from an
+// InsertAndReviseExposures call.
+type InsertAndReviseExposuresResponse struct {
+	// Inserted is the number of new exposures that were inserted into the
+	// database.
+	Inserted uint64
+
+	// Revised is the number of exposures that matched an existing TEK and were
+	// subsequently revised.
+	Revised uint64
+
+	// Dropped is the number of exposures that were not inserted or updated. This
+	// could be because they weren't present in the revision token, etc.
+	Dropped uint64
+
+	// Exposures is the actual exposures that were inserted or updated in this
+	// call.
+	Exposures []*model.Exposure
+}
+
+// InsertAndReviseExposures transactionally revises and inserts a set of keys as
+// necessary.
+func (db *PublishDB) InsertAndReviseExposures(ctx context.Context, req *InsertAndReviseExposuresRequest) (*InsertAndReviseExposuresResponse, error) {
+	logger := logging.FromContext(ctx).Named("InsertAndReviseExposures")
+
+	if req == nil {
+		return nil, fmt.Errorf("missing request")
+	}
+
+	// Maintain a record of the number of exposures inserted and updated, and a
+	// record of the exposures that were actually inserted/updated after merge
+	// logic.
+	var resp InsertAndReviseExposuresResponse
+
+	if err := db.db.InTx(ctx, pgx.ReadCommitted, func(tx pgx.Tx) error {
 		// Build the base64-encoded list of keys - this is needed so we can lookup
 		// the keys in the database. Also build a lookup map by key for validation
 		// later.
-		b64keys := make([]string, len(incoming))
-		incomingMap := make(map[string]*model.Exposure, len(incoming))
-		for i, v := range incoming {
+		b64keys := make([]string, len(req.Incoming))
+		incomingMap := make(map[string]*model.Exposure, len(req.Incoming))
+		for i, v := range req.Incoming {
 			b64keys[i] = v.ExposureKeyBase64()
 			incomingMap[v.ExposureKeyBase64()] = v
 		}
@@ -397,18 +442,18 @@ func (db *PublishDB) InsertAndReviseExposures(ctx context.Context, incoming []*m
 		// are being revised.
 		if len(existing) > 0 {
 			// Check if a revision token is required.
-			if token == nil {
+			if req.Token == nil {
 				logger.Errorw("attempted to revise keys, but revision token is missing")
-				if tokenRequired {
+				if req.RequireToken {
 					return ErrNoRevisionToken
 				}
 			}
 
 			// Build a map of allowed revisions for validation and comparison.
 			allowedRevisions := make(map[string]*pb.RevisableKey)
-			if token != nil {
+			if req.Token != nil {
 				// Special handling for the allow bypass scenario where no token was presented.
-				for _, v := range token.RevisableKeys {
+				for _, v := range req.Token.RevisableKeys {
 					b := base64.StdEncoding.EncodeToString(v.TemporaryExposureKey)
 					allowedRevisions[b] = v
 				}
@@ -424,7 +469,7 @@ func (db *PublishDB) InsertAndReviseExposures(ctx context.Context, incoming []*m
 							"existing_number", ex.IntervalNumber,
 							"incoming_count", in.IntervalCount,
 							"incoming_number", in.IntervalNumber)
-						if tokenRequired {
+						if req.RequireToken {
 							return ErrIncomingMetadataMismatch
 						}
 					}
@@ -438,28 +483,63 @@ func (db *PublishDB) InsertAndReviseExposures(ctx context.Context, incoming []*m
 							"existing_number", ex.IntervalNumber,
 							"incoming_count", rk.IntervalCount,
 							"incoming_number", rk.IntervalNumber)
-						if tokenRequired {
+						if req.RequireToken {
 							return ErrRevisionTokenMetadataMismatch
 						}
 					}
 				} else {
-					// User sent in an existing key they they do not have the token for.
-					logger.Errorw("cannot revise key: not in revision token")
-					if tokenRequired {
-						return ErrExistingKeyNotInToken
+					// The user sent an existing key for which they do not have a revision
+					// token. There's a plausible scenario with roaming where this user
+					// has changed apps in the middle of a roaming period, and thus this
+					// could be legitimate.
+					//
+					// Suppose a user was in California, using the California app. They
+					// feel ill, do a video call and get a "likely" diagnosis. They decide
+					// to drive to Montana to quarantine, downloading the Montana app.
+					// Once in Montanan, their symptoms worsen and they get a clinical
+					// diagnosis as "positive", marking as such in the app. Since the
+					// Montana app does not have a revision token for the keys that were
+					// previously uploaded when the user was in California, some of the
+					// keys cannot be revised.
+					if req.RequireToken {
+						if req.AllowPartialRevisions {
+							logger.Warnw("skipping key: not in revision token, but " +
+								"partial revision is permitted")
+							delete(incomingMap, k)
+						} else {
+							logger.Errorw("cannot revise key: not in revision token")
+							return ErrExistingKeyNotInToken
+						}
 					}
 				}
 			}
 		}
 
+		// Calculate the number of dropped responses.
+		resp.Dropped = uint64(len(req.Incoming) - len(incomingMap))
+
 		// If we got this far, the revision token is valid for this request, not
-		// required, or bypassed.
+		// required, or bypassed. It's possible that keys were given, but none of
+		// those keys matched the revision token keys and partial responses were
+		// allowed. In that case, stop here.
+		if len(incomingMap) == 0 {
+			return nil
+		}
+
+		// Build the true incoming list from the map. It's possible items were removed
+		// from the map if they did not exist in the revision token and partial
+		// responses were allowed.
+		incoming := make([]*model.Exposure, 0, len(incomingMap))
+		for _, v := range incomingMap {
+			incoming = append(incoming, v)
+		}
 
 		// Run through the merge logic.
 		exposures, err := model.ReviseKeys(ctx, existing, incoming)
 		if err != nil {
 			return fmt.Errorf("unable to revise keys: %w", err)
 		}
+		resp.Exposures = exposures
 
 		// Prepare the insert and update statements.
 		insertStmt, err := prepareInsertExposure(ctx, tx)
@@ -479,20 +559,20 @@ func (db *PublishDB) InsertAndReviseExposures(ctx context.Context, incoming []*m
 				if err := executeInsertExposure(ctx, tx, insertStmt, exp); err != nil {
 					return err
 				}
-				updated++
+				resp.Inserted++
 			} else {
 				if err := executeReviseExposure(ctx, tx, updateStmt, exp); err != nil {
 					return err
 				}
-				updated++
+				resp.Revised++
 			}
 		}
 		return nil
-	})
-	if err != nil {
-		updated = 0
+	}); err != nil {
+		return nil, err
 	}
-	return updated, err
+
+	return &resp, nil
 }
 
 // DeleteExposure deletes exposure
diff --git a/internal/publish/model/exposure_model.go b/internal/publish/model/exposure_model.go
index 9f82a10..e638d18 100644
--- a/internal/publish/model/exposure_model.go
+++ b/internal/publish/model/exposure_model.go
@@ -330,7 +330,7 @@ func TransformExposureKey(exposureKey verifyapi.ExposureKey, appPackageName stri
 }
 
 // ReviseKeys takes a set of existing keys, and a list of keys currently being uploaded.
-// Only keys that need to be revised or are being created fir the first time
+// Only keys that need to be revised or are being created for the first time
 // are returned in the output set.
 func ReviseKeys(ctx context.Context, existing map[string]*Exposure, incoming []*Exposure) ([]*Exposure, error) {
 	//logger := logging.FromContext(ctx)
diff --git a/internal/publish/publish.go b/internal/publish/publish.go
index 3a693d2..7a7dda9 100644
--- a/internal/publish/publish.go
+++ b/internal/publish/publish.go
@@ -121,8 +121,9 @@ type PublishHandler struct {
 type response struct {
 	status      int
 	pubResponse *verifyapi.PublishResponse
-	metric      string
-	count       int // metricCount
+
+	// metrics is a function to execute to publish metrics
+	metrics func()
 }
 
 func (r *response) padResponse(c *Config) error {
@@ -173,7 +174,7 @@ func newVersionBridge(regions []string) *versionBridge {
 
 // process runs the publish business logic over a "v1" version of the publish request
 // and knows how to join in data from previous versions (the provided versionBridge)
-func (h *PublishHandler) process(ctx context.Context, data *verifyapi.Publish, bridge *versionBridge) response {
+func (h *PublishHandler) process(ctx context.Context, data *verifyapi.Publish, bridge *versionBridge) *response {
 	ctx, span := trace.StartSpan(ctx, "(*publish.PublishHandler).process")
 	defer span.End()
 
@@ -188,14 +189,15 @@ func (h *PublishHandler) process(ctx context.Context, data *verifyapi.Publish, b
 		if errors.Is(err, authorizedapp.ErrAppNotFound) {
 			message := fmt.Sprintf("unauthorized health authority: %v", data.HealthAuthorityID)
 			span.SetStatus(trace.Status{Code: trace.StatusCodeInternal, Message: message})
-			return response{
+			return &response{
 				status: http.StatusUnauthorized,
 				pubResponse: &verifyapi.PublishResponse{
 					ErrorMessage: message,
 					Code:         verifyapi.ErrorUnknownHealthAuthorityID,
 				},
-				metric: "publish-health-authority-not-authorized",
-				count:  1,
+				metrics: func() {
+					metrics.WriteInt("publish-health-authority-not-authorized", true, 1)
+				},
 			}
 		}
 
@@ -206,14 +208,15 @@ func (h *PublishHandler) process(ctx context.Context, data *verifyapi.Publish, b
 		// and no other data from the request.
 		message := fmt.Sprintf("error loading health authority config: %v", err)
 		span.SetStatus(trace.Status{Code: trace.StatusCodeInternal, Message: message})
-		return response{
+		return &response{
 			status: http.StatusNotFound,
 			pubResponse: &verifyapi.PublishResponse{
 				ErrorMessage: message,
 				Code:         verifyapi.ErrorUnableToLoadHealthAuthority,
 			},
-			metric: "publish-error-loading-authorizedapp",
-			count:  1,
+			metrics: func() {
+				metrics.WriteInt("publish-error-loading-authorizedapp", true, 1)
+			},
 		}
 	}
 
@@ -228,13 +231,14 @@ func (h *PublishHandler) process(ctx context.Context, data *verifyapi.Publish, b
 				err := fmt.Errorf("app %v tried to write to unauthorized region %v", appConfig.AppPackageName, r)
 				message := fmt.Sprintf("verifying allowed regions: %v", err)
 				span.SetStatus(trace.Status{Code: trace.StatusCodePermissionDenied, Message: message})
-				return response{
+				return &response{
 					status: http.StatusUnauthorized,
 					pubResponse: &verifyapi.PublishResponse{
 						ErrorMessage: message, // Error code omitted, since this isn't in the v1 path.
 					},
-					metric: "publish-region-not-authorized",
-					count:  1,
+					metrics: func() {
+						metrics.WriteInt("publish-region-not-authorized", true, 1)
+					},
 				}
 			}
 		}
@@ -256,14 +260,15 @@ func (h *PublishHandler) process(ctx context.Context, data *verifyapi.Publish, b
 		logger.Errorf("No regions present in request or configured for healthAuthorityID: %v", data.HealthAuthorityID)
 		message := fmt.Sprintf("unknown health authority regions for %v", data.HealthAuthorityID)
 		span.SetStatus(trace.Status{Code: trace.StatusCodePermissionDenied, Message: message})
-		return response{
+		return &response{
 			status: http.StatusInternalServerError,
 			pubResponse: &verifyapi.PublishResponse{
 				ErrorMessage: message,
 				Code:         verifyapi.ErrorHealthAuthorityMissingRegionConfiguration,
 			},
-			metric: "publish-region-not-specified",
-			count:  1,
+			metrics: func() {
+				metrics.WriteInt("publish-region-not-specified", true, 1)
+			},
 		}
 	}
 
@@ -281,13 +286,16 @@ func (h *PublishHandler) process(ctx context.Context, data *verifyapi.Publish, b
 				logger.Errorw(message, "error", err)
 			}
 			span.SetStatus(trace.Status{Code: trace.StatusCodeInvalidArgument, Message: message})
-			return response{
+			return &response{
 				status: http.StatusUnauthorized,
 				pubResponse: &verifyapi.PublishResponse{
 					ErrorMessage: message,
 					Code:         verifyapi.ErrorVerificationCertificateInvalid,
 				},
-				metric: "publish-bad-verification", count: 1}
+				metrics: func() {
+					metrics.WriteInt("publish-bad-verification", true, 1)
+				},
+			}
 		}
 	}
 
@@ -317,16 +325,25 @@ func (h *PublishHandler) process(ctx context.Context, data *verifyapi.Publish, b
 		message := fmt.Sprintf("unable to read request data: %v", transformError)
 		logger.Error(message)
 		span.SetStatus(trace.Status{Code: trace.StatusCodeInvalidArgument, Message: message})
-		return response{
+		return &response{
 			status: http.StatusBadRequest,
 			pubResponse: &verifyapi.PublishResponse{
 				ErrorMessage: message,
 				Code:         verifyapi.ErrorBadRequest,
 			},
-			metric: "publish-transform-fail", count: 1}
+			metrics: func() {
+				metrics.WriteInt("publish-transform-fail", true, 1)
+			},
+		}
 	}
 
-	n, err := h.database.InsertAndReviseExposures(ctx, exposures, token, !appConfig.BypassRevisionToken)
+	resp, err := h.database.InsertAndReviseExposures(ctx, &database.InsertAndReviseExposuresRequest{
+		Incoming: exposures,
+		Token:    token,
+
+		RequireToken:          !appConfig.BypassRevisionToken,
+		AllowPartialRevisions: h.config.AllowPartialRevisions,
+	})
 	if err != nil {
 		status := http.StatusBadRequest
 		var logMessage, errorMessage, errorCode string
@@ -348,13 +365,16 @@ func (h *PublishHandler) process(ctx context.Context, data *verifyapi.Publish, b
 		}
 		logger.Error(logMessage)
 		span.SetStatus(trace.Status{Code: trace.StatusCodeInternal, Message: logMessage})
-		return response{
+		return &response{
 			status: status,
 			pubResponse: &verifyapi.PublishResponse{
 				ErrorMessage: errorMessage,
 				Code:         errorCode,
 			},
-			metric: metric, count: 1}
+			metrics: func() {
+				metrics.WriteInt(metric, true, 1)
+			},
+		}
 	}
 
 	// Build the new revision token. Union of existing token take + new exposures.
@@ -368,7 +388,7 @@ func (h *PublishHandler) process(ctx context.Context, data *verifyapi.Publish, b
 			}
 		}
 	}
-	newToken, err := h.tokenManager.MakeRevisionToken(ctx, &keep, exposures, h.tokenAAD)
+	newToken, err := h.tokenManager.MakeRevisionToken(ctx, &keep, resp.Exposures, h.tokenAAD)
 	if err != nil {
 		// In this case, an empty revision token is returned. The rest of the request
 		// was handled correctly.
@@ -376,13 +396,17 @@ func (h *PublishHandler) process(ctx context.Context, data *verifyapi.Publish, b
 		newToken = make([]byte, 0)
 	}
 
-	message := fmt.Sprintf("Inserted %d exposures.", n)
-	span.AddAttributes(trace.Int64Attribute("inserted_exposures", int64(n)))
-	logger.Info(message)
+	span.AddAttributes(trace.Int64Attribute("exposures_inserted", int64(resp.Inserted)))
+	span.AddAttributes(trace.Int64Attribute("exposures_revised", int64(resp.Revised)))
+	span.AddAttributes(trace.Int64Attribute("exposures_dropped", int64(resp.Dropped)))
+	logger.Infow("published exposures",
+		"inserted", resp.Inserted,
+		"updated", resp.Revised,
+		"dropped", resp.Dropped)
 
 	publishResponse := verifyapi.PublishResponse{
 		RevisionToken:     base64.StdEncoding.EncodeToString(newToken),
-		InsertedExposures: n,
+		InsertedExposures: int(resp.Inserted),
 	}
 	// If there was a partial failure on transform, add that information back into the success response.
 	if transformError != nil {
@@ -390,10 +414,13 @@ func (h *PublishHandler) process(ctx context.Context, data *verifyapi.Publish, b
 		publishResponse.ErrorMessage = transformError.Error()
 	}
 
-	return response{
+	return &response{
 		status:      http.StatusOK,
 		pubResponse: &publishResponse,
-		metric:      "publish-exposures-written",
-		count:       n,
+		metrics: func() {
+			metrics.WriteInt("publish-exposures-inserted", true, int(resp.Inserted))
+			metrics.WriteInt("publish-exposures-revised", true, int(resp.Revised))
+			metrics.WriteInt("publish-exposures-dropped", true, int(resp.Dropped))
+		},
 	}
 }
diff --git a/internal/publish/publish_v1.go b/internal/publish/publish_v1.go
index ebc7738..dd4a4d6 100644
--- a/internal/publish/publish_v1.go
+++ b/internal/publish/publish_v1.go
@@ -28,10 +28,12 @@ import (
 	"github.com/mikehelmick/go-chaff"
 )
 
-func (h *PublishHandler) handleRequest(w http.ResponseWriter, r *http.Request) response {
+func (h *PublishHandler) handleRequest(w http.ResponseWriter, r *http.Request) *response {
 	ctx, span := trace.StartSpan(r.Context(), "(*publish.PublishHandler).handleRequest")
 	defer span.End()
 
+	metrics := h.serverenv.MetricsExporter(ctx)
+
 	w.Header().Set(HeaderAPIVersion, "v1")
 
 	var data verifyapi.Publish
@@ -43,14 +45,15 @@ func (h *PublishHandler) handleRequest(w http.ResponseWriter, r *http.Request) r
 		if code == http.StatusInternalServerError {
 			errorCode = verifyapi.ErrorInternalError
 		}
-		return response{
+		return &response{
 			status: code,
 			pubResponse: &verifyapi.PublishResponse{
 				ErrorMessage: message,
 				Code:         errorCode,
 			},
-			metric: "publish-bad-json",
-			count:  1,
+			metrics: func() {
+				metrics.WriteInt("publish-bad-json", true, 1)
+			},
 		}
 	}
 
@@ -71,8 +74,8 @@ func (h *PublishHandler) Handle() http.Handler {
 				logging.FromContext(ctx).Errorw("failed to padd response", "error", err)
 			}
 
-			if response.metric != "" {
-				metrics.WriteInt(response.metric, true, response.count)
+			if response.metrics != nil {
+				response.metrics()
 			}
 
 			data, err := json.Marshal(response.pubResponse)
diff --git a/internal/publish/publish_v1alpha1.go b/internal/publish/publish_v1alpha1.go
index f02ab84..b14f65b 100644
--- a/internal/publish/publish_v1alpha1.go
+++ b/internal/publish/publish_v1alpha1.go
@@ -29,10 +29,12 @@ import (
 	"github.com/mikehelmick/go-chaff"
 )
 
-func (h *PublishHandler) handleV1Apha1Request(w http.ResponseWriter, r *http.Request) response {
+func (h *PublishHandler) handleV1Apha1Request(w http.ResponseWriter, r *http.Request) *response {
 	ctx, span := trace.StartSpan(r.Context(), "(*publish.PublishHandler).handleRequest")
 	defer span.End()
 
+	metrics := h.serverenv.MetricsExporter(ctx)
+
 	w.Header().Set(HeaderAPIVersion, "v1alpha")
 
 	var data v1alpha1.Publish
@@ -40,10 +42,13 @@ func (h *PublishHandler) handleV1Apha1Request(w http.ResponseWriter, r *http.Req
 	if err != nil {
 		message := fmt.Sprintf("error unmarshaling API call, code: %v: %v", code, err)
 		span.SetStatus(trace.Status{Code: trace.StatusCodeInternal, Message: message})
-		return response{
+		return &response{
 			status:      code,
 			pubResponse: &verifyapi.PublishResponse{ErrorMessage: message}, // will be down-converted in ServeHTTP
-			metric:      "publish-bad-json", count: 1}
+			metrics: func() {
+				metrics.WriteInt("publish-bad-json", true, 1)
+			},
+		}
 	}
 
 	// Upconvert the exposure key records.
@@ -94,8 +99,8 @@ func (h *PublishHandler) HandleV1Alpha1() http.Handler {
 				Padding:           response.pubResponse.Padding,
 			}
 
-			if response.metric != "" {
-				metrics.WriteInt(response.metric, true, response.count)
+			if response.metrics != nil {
+				response.metrics()
 			}
 
 			data, err := json.Marshal(&alpha1Response)


Test patch:
diff --git a/internal/export/database/export_test.go b/internal/export/database/export_test.go
index 35bc4d9..9415c4f 100644
--- a/internal/export/database/export_test.go
+++ b/internal/export/database/export_test.go
@@ -488,7 +488,10 @@ func TestTravelerKeys(t *testing.T) {
 	}
 
 	// Add the keys to the database.
-	if _, err := publishdb.New(testDB).InsertAndReviseExposures(ctx, []*publishmodel.Exposure{trav, notTrav}, nil, true); err != nil {
+	if _, err := publishdb.New(testDB).InsertAndReviseExposures(ctx, &publishdb.InsertAndReviseExposuresRequest{
+		Incoming:     []*publishmodel.Exposure{trav, notTrav},
+		RequireToken: true,
+	}); err != nil {
 		t.Fatal(err)
 	}
 
@@ -585,7 +588,10 @@ func TestKeysInBatch(t *testing.T) {
 	}
 
 	// Add the keys to the database.
-	if _, err := publishdb.New(testDB).InsertAndReviseExposures(ctx, []*publishmodel.Exposure{sek, eek}, nil, true); err != nil {
+	if _, err := publishdb.New(testDB).InsertAndReviseExposures(ctx, &publishdb.InsertAndReviseExposuresRequest{
+		Incoming:     []*publishmodel.Exposure{sek, eek},
+		RequireToken: true,
+	}); err != nil {
 		t.Fatal(err)
 	}
 
diff --git a/internal/federationin/federationin_test.go b/internal/federationin/federationin_test.go
index 0829bf4..2e1fef1 100644
--- a/internal/federationin/federationin_test.go
+++ b/internal/federationin/federationin_test.go
@@ -21,6 +21,7 @@ import (
 
 	"github.com/google/exposure-notifications-server/internal/federationin/database"
 	"github.com/google/exposure-notifications-server/internal/federationin/model"
+	publishdb "github.com/google/exposure-notifications-server/internal/publish/database"
 	publishmodel "github.com/google/exposure-notifications-server/internal/publish/model"
 
 	"github.com/google/exposure-notifications-server/internal/metrics"
@@ -70,9 +71,12 @@ type publishDB struct {
 	exposures []*publishmodel.Exposure
 }
 
-func (idb *publishDB) insertExposures(ctx context.Context, exposures []*publishmodel.Exposure, rt *pb.RevisionTokenData, enforce bool) (int, error) {
-	idb.exposures = append(idb.exposures, exposures...)
-	return len(exposures), nil
+func (idb *publishDB) insertExposures(ctx context.Context, req *publishdb.InsertAndReviseExposuresRequest) (*publishdb.InsertAndReviseExposuresResponse, error) {
+	idb.exposures = append(idb.exposures, req.Incoming...)
+	return &publishdb.InsertAndReviseExposuresResponse{
+		Exposures: req.Incoming,
+		Inserted:  uint64(len(req.Incoming)),
+	}, nil
 }
 
 // syncDB mocks the database, recording start and complete invocations for a sync record.
diff --git a/internal/performance/export_test.go b/internal/performance/export_test.go
index 5962600..e56d40a 100644
--- a/internal/performance/export_test.go
+++ b/internal/performance/export_test.go
@@ -152,13 +152,12 @@ func TestExport(t *testing.T) {
 			m.ExposureKey, _ = base64util.DecodeString(newKey.Key)
 			revisedExposures = append(revisedExposures, &m)
 		}
-		updated, err := publishdb.New(db).InsertAndReviseExposures(ctx, revisedExposures,
-			nil, false)
+		resp, err := publishdb.New(db).InsertAndReviseExposures(ctx, &publishdb.InsertAndReviseExposuresRequest{Incoming: revisedExposures})
 		if err != nil {
 			t.Fatal(err)
 		}
-		if updated != keysPerPublish {
-			t.Fatalf("Want updated: %d, got %d", keysPerPublish, updated)
+		if got, want := int(resp.Revised), keysPerPublish; got != want {
+			t.Fatalf("Want revised: %d, got %d", want, got)
 		}
 	}
 
diff --git a/internal/publish/database/exposure_test.go b/internal/publish/database/exposure_test.go
index e566ca6..d87497b 100644
--- a/internal/publish/database/exposure_test.go
+++ b/internal/publish/database/exposure_test.go
@@ -20,6 +20,7 @@ import (
 	"encoding/base64"
 	"errors"
 	"fmt"
+	"reflect"
 	"sort"
 	"strings"
 	"testing"
@@ -69,7 +70,10 @@ func TestReadExposures(t *testing.T) {
 			Traveler:        true,
 		},
 	}
-	if _, err := testPublishDB.InsertAndReviseExposures(ctx, exposures, nil, true); err != nil {
+	if _, err := testPublishDB.InsertAndReviseExposures(ctx, &InsertAndReviseExposuresRequest{
+		Incoming:     exposures,
+		RequireToken: true,
+	}); err != nil {
 		t.Fatal(err)
 	}
 
@@ -184,7 +188,10 @@ func TestExposures(t *testing.T) {
 			LocalProvenance: false,
 		},
 	}
-	if _, err := testPublishDB.InsertAndReviseExposures(ctx, exposures, nil, true); err != nil {
+	if _, err := testPublishDB.InsertAndReviseExposures(ctx, &InsertAndReviseExposuresRequest{
+		Incoming:     exposures,
+		RequireToken: true,
+	}); err != nil {
 		t.Fatal(err)
 	}
 
@@ -294,206 +301,416 @@ func testRandomTEK(tb testing.TB) []byte {
 	return b
 }
 
-func TestReviseExposures(t *testing.T) {
-	t.Parallel()
-
-	ctx := context.Background()
-	testDB := database.NewTestDatabase(t)
-	pubDB := New(testDB)
-
-	createdAt := time.Now().UTC().Add(-2 * time.Hour).Truncate(time.Hour)
-	revisedAt := time.Now().UTC().Truncate(time.Hour)
+func testExposure(tb testing.TB) *model.Exposure {
+	tb.Helper()
 
-	existingTEK := testRandomTEK(t)
-	existingExp := &model.Exposure{
-		ExposureKey:      existingTEK,
+	exposure := &model.Exposure{
+		ExposureKey:      testRandomTEK(tb),
 		TransmissionRisk: verifyapi.TransmissionRiskClinical,
 		AppPackageName:   "foo.bar",
 		Regions:          []string{"US", "CA", "MX"},
 		IntervalNumber:   100,
 		IntervalCount:    144,
-		CreatedAt:        createdAt,
+		CreatedAt:        time.Now().UTC().Add(-2 * time.Hour).Truncate(time.Hour),
 		LocalProvenance:  true,
 		ReportType:       verifyapi.ReportTypeClinical,
 	}
-	existingExp.SetDaysSinceSymptomOnset(0)
-	if n, err := pubDB.InsertAndReviseExposures(ctx, []*model.Exposure{existingExp}, nil, true); err != nil {
-		t.Fatalf("error inserting exposures: %v", err)
-	} else if n != 1 {
-		t.Fatalf("wrong number of changed exposures, want: 1, got: %v", n)
-	}
+	exposure.SetDaysSinceSymptomOnset(0)
+	return exposure
+}
+
+func TestReviseExposures(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	testDB := database.NewTestDatabase(t)
+	pubDB := New(testDB)
 
 	// Attempt to revise without a revision token - this should fail.
-	{
-		revisedExp := &model.Exposure{}
-		*revisedExp = *existingExp
-		revisedExp.TransmissionRisk = verifyapi.TransmissionRiskConfirmedStandard
-		revisedExp.CreatedAt = revisedAt
-		revisedExp.ReportType = verifyapi.ReportTypeConfirmed
+	t.Run("revise_without_token", func(t *testing.T) {
+		t.Parallel()
 
-		if _, err := pubDB.InsertAndReviseExposures(ctx, []*model.Exposure{revisedExp}, nil, true); !errors.Is(err, ErrNoRevisionToken) {
-			t.Fatalf("expected %#v to be %#v", err, ErrNoRevisionToken)
+		// Insert the exposure
+		exposure := testExposure(t)
+		insertResp, err := pubDB.InsertAndReviseExposures(ctx, &InsertAndReviseExposuresRequest{
+			Incoming: []*model.Exposure{exposure},
+		})
+		if err != nil {
+			t.Fatal(err)
 		}
-	}
+		if got, want := int(insertResp.Inserted), 1; got != want {
+			t.Fatalf("expected %d to be %d", got, want)
+		}
+
+		// Change and revise
+		exposure.TransmissionRisk = verifyapi.TransmissionRiskConfirmedStandard
+		exposure.ReportType = verifyapi.ReportTypeConfirmed
+
+		if _, err := pubDB.InsertAndReviseExposures(ctx, &InsertAndReviseExposuresRequest{
+			Incoming:     []*model.Exposure{exposure},
+			RequireToken: true,
+		}); !errors.Is(err, ErrNoRevisionToken) {
+			t.Errorf("expected %#v to be %#v", err, ErrNoRevisionToken)
+		}
+	})
 
 	// Attempt to revise with a revision token that doesn't include the TEK - this
 	// should fail.
-	{
-		revisedExp := &model.Exposure{}
-		*revisedExp = *existingExp
-		revisedExp.TransmissionRisk = verifyapi.TransmissionRiskConfirmedStandard
-		revisedExp.CreatedAt = revisedAt
-		revisedExp.ReportType = verifyapi.ReportTypeConfirmed
+	t.Run("revise_with_token_without_matching_tek", func(t *testing.T) {
+		t.Parallel()
+
+		// Insert the exposure
+		exposure := testExposure(t)
+		insertResp, err := pubDB.InsertAndReviseExposures(ctx, &InsertAndReviseExposuresRequest{
+			Incoming: []*model.Exposure{exposure},
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got, want := int(insertResp.Inserted), 1; got != want {
+			t.Fatalf("expected %d to be %d", got, want)
+		}
+
+		// Change and revise
+		exposure.TransmissionRisk = verifyapi.TransmissionRiskConfirmedStandard
+		exposure.ReportType = verifyapi.ReportTypeConfirmed
 
 		token := &pb.RevisionTokenData{
 			RevisableKeys: []*pb.RevisableKey{
 				{
 					TemporaryExposureKey: testRandomTEK(t), // TEK mismatch
-					IntervalNumber:       100,
-					IntervalCount:        144,
+					IntervalNumber:       exposure.IntervalNumber,
+					IntervalCount:        exposure.IntervalCount,
 				},
 			},
 		}
 
-		if _, err := pubDB.InsertAndReviseExposures(ctx, []*model.Exposure{revisedExp}, token, true); !errors.Is(err, ErrExistingKeyNotInToken) {
+		if _, err := pubDB.InsertAndReviseExposures(ctx, &InsertAndReviseExposuresRequest{
+			Incoming:     []*model.Exposure{exposure},
+			Token:        token,
+			RequireToken: true,
+		}); !errors.Is(err, ErrExistingKeyNotInToken) {
 			t.Fatalf("expected %#v to be %#v", err, ErrExistingKeyNotInToken)
 		}
-	}
+	})
 
 	// Attempt to revise with a revision token that has the wrong metadata - this
 	// should fail.
-	{
-		revisedExp := &model.Exposure{}
-		*revisedExp = *existingExp
-		revisedExp.TransmissionRisk = verifyapi.TransmissionRiskConfirmedStandard
-		revisedExp.CreatedAt = revisedAt
-		revisedExp.ReportType = verifyapi.ReportTypeConfirmed
+	t.Run("revise_with_token_wrong_metadata", func(t *testing.T) {
+		t.Parallel()
+
+		// Insert the exposure
+		exposure := testExposure(t)
+		insertResp, err := pubDB.InsertAndReviseExposures(ctx, &InsertAndReviseExposuresRequest{
+			Incoming: []*model.Exposure{exposure},
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got, want := int(insertResp.Inserted), 1; got != want {
+			t.Fatalf("expected %d to be %d", got, want)
+		}
 
 		token := &pb.RevisionTokenData{
 			RevisableKeys: []*pb.RevisableKey{
 				{
-					TemporaryExposureKey: existingTEK,
-					IntervalNumber:       101, // Changed, should cause failure
-					IntervalCount:        144,
+					TemporaryExposureKey: exposure.ExposureKey,
+					IntervalNumber:       exposure.IntervalNumber + 1, // Changed, should cause failure
+					IntervalCount:        exposure.IntervalCount,
 				},
 			},
 		}
 
-		if _, err := pubDB.InsertAndReviseExposures(ctx, []*model.Exposure{revisedExp}, token, true); !errors.Is(err, ErrRevisionTokenMetadataMismatch) {
+		if _, err := pubDB.InsertAndReviseExposures(ctx, &InsertAndReviseExposuresRequest{
+			Incoming:     []*model.Exposure{exposure},
+			Token:        token,
+			RequireToken: true,
+		}); !errors.Is(err, ErrRevisionTokenMetadataMismatch) {
 			t.Fatalf("expected %#v to be %#v", err, ErrRevisionTokenMetadataMismatch)
 		}
-	}
+	})
 
 	// Attempt to modify metadata on incoming keys - this should fail.
-	{
-		revisedExp := &model.Exposure{}
-		*revisedExp = *existingExp
-		revisedExp.IntervalCount = 124 // Changed, should make request invalid
+	t.Run("revise_metadata_incoming_keys", func(t *testing.T) {
+		t.Parallel()
+
+		// Insert the exposure
+		exposure := testExposure(t)
+		insertResp, err := pubDB.InsertAndReviseExposures(ctx, &InsertAndReviseExposuresRequest{
+			Incoming: []*model.Exposure{exposure},
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got, want := int(insertResp.Inserted), 1; got != want {
+			t.Fatalf("expected %d to be %d", got, want)
+		}
+
+		// Change metadata - should make request invalid
+		exposure.IntervalCount = 124
 
 		token := &pb.RevisionTokenData{
 			RevisableKeys: []*pb.RevisableKey{
 				{
-					TemporaryExposureKey: existingTEK,
-					IntervalNumber:       101, // Changed, should cause failure
-					IntervalCount:        144,
+					TemporaryExposureKey: exposure.ExposureKey,
+					IntervalNumber:       exposure.IntervalNumber,
+					IntervalCount:        exposure.IntervalCount,
 				},
 			},
 		}
 
-		if _, err := pubDB.InsertAndReviseExposures(ctx, []*model.Exposure{revisedExp}, token, true); !errors.Is(err, ErrIncomingMetadataMismatch) {
+		if _, err := pubDB.InsertAndReviseExposures(ctx, &InsertAndReviseExposuresRequest{
+			Incoming:     []*model.Exposure{exposure},
+			Token:        token,
+			RequireToken: true,
+		}); !errors.Is(err, ErrIncomingMetadataMismatch) {
 			t.Fatalf("expected %#v to be %#v", err, ErrIncomingMetadataMismatch)
 		}
-	}
+	})
 
-	// Modify the existing on in place.
-	revisedExp := &model.Exposure{
-		ExposureKey:      existingTEK,
-		TransmissionRisk: verifyapi.TransmissionRiskConfirmedStandard,
-		AppPackageName:   "foo.bar",
-		Regions:          []string{"US", "CA", "MX"},
-		IntervalNumber:   100,
-		IntervalCount:    144,
-		CreatedAt:        revisedAt,
-		LocalProvenance:  true,
-		ReportType:       verifyapi.ReportTypeConfirmed,
-	}
-	// Add a new TEK to insert.
-	newTEK := []byte{1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 5}
-	newExp := &model.Exposure{
-		ExposureKey:      newTEK,
-		TransmissionRisk: verifyapi.TransmissionRiskClinical,
-		AppPackageName:   "foo.bar",
-		Regions:          []string{"US", "CA", "MX"},
-		IntervalNumber:   244,
-		IntervalCount:    144,
-		CreatedAt:        revisedAt,
-		LocalProvenance:  true,
-		ReportType:       verifyapi.ReportTypeConfirmed,
-	}
-	newExp.SetDaysSinceSymptomOnset(1)
-
-	revisions := []*model.Exposure{revisedExp, newExp}
-	// In the first pass - try to revise the key without the necessary revision token.
-	if _, err := pubDB.InsertAndReviseExposures(ctx, revisions, nil, true); err == nil {
-		t.Fatalf("expected error revising without token data")
-	} else if !errors.Is(err, ErrNoRevisionToken) {
-		t.Fatalf("wrong error, want: '%v' got: %v", ErrNoRevisionToken, err)
-	}
+	// Attempt to revise where a subset of the TEKs match the revision token and
+	// partial revisions are permitted. This should succeed.
+	t.Run("revise_subset_of_keys_partial", func(t *testing.T) {
+		t.Parallel()
 
-	// Revision token that allows revision of the revised key.
-	var token pb.RevisionTokenData
-	token.RevisableKeys = append(token.RevisableKeys, &pb.RevisableKey{
-		TemporaryExposureKey: existingTEK,
-		IntervalNumber:       100,
-		IntervalCount:        144,
-	})
+		// Insert the exposure
+		exposure1 := testExposure(t)
+		insertResp1, err := pubDB.InsertAndReviseExposures(ctx, &InsertAndReviseExposuresRequest{
+			Incoming: []*model.Exposure{exposure1},
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got, want := int(insertResp1.Inserted), 1; got != want {
+			t.Fatalf("expected %d to be %d", got, want)
+		}
 
-	if n, err := pubDB.InsertAndReviseExposures(ctx, revisions, &token, true); err != nil {
-		t.Fatalf("error revising exposures: %v", err)
-	} else if n != 2 {
-		t.Fatalf("wrong number of changed exposures, want: 2, got: %v", n)
-	}
+		// Change the diagnosis
+		exposure1.ReportType = verifyapi.ReportTypeConfirmed
 
-	// Read back and compare.
-	expectedKeys := []string{existingExp.ExposureKeyBase64(), newExp.ExposureKeyBase64()}
-	var got map[string]*model.Exposure
-	err := testDB.InTx(ctx, pgx.ReadCommitted, func(tx pgx.Tx) error {
-		var err error
-		got, err = pubDB.ReadExposures(ctx, tx, expectedKeys)
-		return err
+		// Create new exposure which is not included in the revision token, but
+		// should also not return an error due to partial revisions being permitted.
+		exposure2 := testExposure(t)
+		insertResp2, err := pubDB.InsertAndReviseExposures(ctx, &InsertAndReviseExposuresRequest{
+			Incoming: []*model.Exposure{exposure2},
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got, want := int(insertResp2.Inserted), 1; got != want {
+			t.Fatalf("expected %d to be %d", got, want)
+		}
+
+		// Create the revision token, only including the first token.
+		token := &pb.RevisionTokenData{
+			RevisableKeys: []*pb.RevisableKey{
+				{
+					TemporaryExposureKey: exposure1.ExposureKey,
+					IntervalNumber:       exposure1.IntervalNumber,
+					IntervalCount:        exposure1.IntervalCount,
+				},
+			},
+		}
+
+		// Attempt to revise
+		resp, err := pubDB.InsertAndReviseExposures(ctx, &InsertAndReviseExposuresRequest{
+			Incoming:              []*model.Exposure{exposure1, exposure2},
+			Token:                 token,
+			RequireToken:          true,
+			AllowPartialRevisions: true,
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got, want := int(resp.Inserted), 0; got != want {
+			t.Errorf("expected %d to be %d", got, want)
+		}
+		if got, want := int(resp.Revised), 1; got != want {
+			t.Errorf("expected %d to be %d", got, want)
+		}
+		if got, want := int(resp.Dropped), 1; got != want {
+			t.Errorf("expected %d to be %d", got, want)
+		}
 	})
-	if err != nil {
-		t.Fatalf("error reading exposures: %v", err)
-	}
 
-	if l := len(got); l != 2 {
-		t.Fatalf("didn't read back both exposures, got: %v", l)
-	}
+	t.Run("revise_already_revised", func(t *testing.T) {
+		t.Parallel()
+
+		// Insert the exposure
+		exposure := testExposure(t)
+		insertResp, err := pubDB.InsertAndReviseExposures(ctx, &InsertAndReviseExposuresRequest{
+			Incoming: []*model.Exposure{exposure},
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got, want := int(insertResp.Inserted), 1; got != want {
+			t.Fatalf("expected %d to be %d", got, want)
+		}
 
-	want := []*model.Exposure{existingExp, newExp}
-	// need to modify the exitingExp to match what we expect to get back.
-	existingExp.SetRevisedTransmissionRisk(verifyapi.TransmissionRiskConfirmedStandard)
-	existingExp.SetRevisedAt(revisedAt)
-	existingExp.SetRevisedReportType(verifyapi.ReportTypeConfirmed)
+		// Revise the exposure
+		exposure.ReportType = verifyapi.ReportTypeConfirmed
 
-	for i := 0; i <= 1; i++ {
-		if diff := cmp.Diff(want[i], got[want[i].ExposureKeyBase64()], ignoreUnexportedExposure); diff != "" {
-			t.Errorf("mismatch (-want, +got):\n%s", diff)
+		token := &pb.RevisionTokenData{
+			RevisableKeys: []*pb.RevisableKey{
+				{
+					TemporaryExposureKey: exposure.ExposureKey,
+					IntervalNumber:       exposure.IntervalNumber,
+					IntervalCount:        exposure.IntervalCount,
+				},
+			},
+		}
+
+		reviseResp, err := pubDB.InsertAndReviseExposures(ctx, &InsertAndReviseExposuresRequest{
+			Incoming:     []*model.Exposure{exposure},
+			Token:        token,
+			RequireToken: true,
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got, want := int(reviseResp.Revised), 1; got != want {
+			t.Errorf("expected %d to be %d: %#v", got, want, reviseResp)
 		}
-	}
 
-	token.RevisableKeys = append(token.RevisableKeys, &pb.RevisableKey{
-		TemporaryExposureKey: newTEK,
-		IntervalNumber:       244,
-		IntervalCount:        144,
+		// Attempt to revise the same key again - this should fail.
+		if _, err := pubDB.InsertAndReviseExposures(ctx, &InsertAndReviseExposuresRequest{
+			Incoming:     []*model.Exposure{exposure},
+			Token:        token,
+			RequireToken: true,
+		}); !errors.Is(err, model.ErrorKeyAlreadyRevised) {
+			t.Errorf("expected %#v to be %#v", err, model.ErrorKeyAlreadyRevised)
+		}
 	})
 
-	// Attempt to revise the already revised key.
-	if _, err := pubDB.InsertAndReviseExposures(ctx, []*model.Exposure{revisedExp}, &token, true); err == nil {
-		t.Fatalf("expected error on revising already revised key, got nil")
-	} else if !strings.Contains(err.Error(), "key has already been revised ") {
-		t.Fatalf("wrong error, want 'invalid key revision request', got: %v", err)
-	}
+	t.Run("revise_partial_new_flow", func(t *testing.T) {
+		t.Parallel()
+
+		// Insert 2 exposures
+		exposure1 := testExposure(t)
+		exposure2 := testExposure(t)
+		insertResp, err := pubDB.InsertAndReviseExposures(ctx, &InsertAndReviseExposuresRequest{
+			Incoming: []*model.Exposure{exposure1, exposure2},
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got, want := int(insertResp.Inserted), 2; got != want {
+			t.Fatalf("expected %d to be %d", got, want)
+		}
+
+		// Change the first exposure
+		exposure1.ReportType = verifyapi.ReportTypeConfirmed
+
+		// Change the second exposure
+		exposure2.ReportType = verifyapi.ReportTypeConfirmed
+
+		// Create another new exposure to insert
+		exposure3 := testExposure(t)
+		exposure3.SetDaysSinceSymptomOnset(1)
+		exposure3.IntervalNumber = 244
+		exposure3.IntervalCount = 144
+		exposure3.ReportType = verifyapi.ReportTypeConfirmed
+
+		// Create a revision token that includes exposure1, but not exposure2 or
+		// exposure3.
+		token := &pb.RevisionTokenData{
+			RevisableKeys: []*pb.RevisableKey{
+				{
+					TemporaryExposureKey: exposure1.ExposureKey,
+					IntervalNumber:       exposure1.IntervalNumber,
+					IntervalCount:        exposure1.IntervalCount,
+				},
+			},
+		}
+
+		// Exposure2 should be dropped since it exists and is not in the database,
+		// but exposure3 should be inserted because its new.
+		resp, err := pubDB.InsertAndReviseExposures(ctx, &InsertAndReviseExposuresRequest{
+			Incoming:              []*model.Exposure{exposure1, exposure2, exposure3},
+			Token:                 token,
+			RequireToken:          true,
+			AllowPartialRevisions: true,
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got, want := int(resp.Inserted), 1; got != want {
+			// exposure3
+			t.Errorf("expected inserted %d to be %d", got, want)
+		}
+		if got, want := int(resp.Revised), 1; got != want {
+			// exposure1
+			t.Errorf("expected revised %d to be %d", got, want)
+		}
+		if got, want := int(resp.Dropped), 1; got != want {
+			// exposure2
+			t.Errorf("expected dropped %d to be %d", got, want)
+		}
+		if got, want := len(resp.Exposures), 2; got != want {
+			t.Errorf("expected exposures %d to be %d", got, want)
+		}
+
+		// Ensure the returned exposures don't include exposure 2 which should have
+		// been dropped. This will ensure it's not included in the revision token as
+		// well.
+		gotExposures := make([]string, len(resp.Exposures))
+		for i, e := range resp.Exposures {
+			gotExposures[i] = e.ExposureKeyBase64()
+		}
+		sort.Strings(gotExposures)
+		wantExposures := []string{
+			exposure1.ExposureKeyBase64(),
+			exposure3.ExposureKeyBase64(),
+		}
+		sort.Strings(wantExposures)
+		if got, want := gotExposures, wantExposures; !reflect.DeepEqual(got, want) {
+			t.Errorf("expected %#v to be %#v", got, want)
+		}
+
+		// Read back and ensure they actually went in the database.
+		expectedTEKs := []string{
+			exposure1.ExposureKeyBase64(),
+			exposure2.ExposureKeyBase64(),
+			exposure3.ExposureKeyBase64(),
+		}
+		var got map[string]*model.Exposure
+		if err := testDB.InTx(ctx, pgx.ReadCommitted, func(tx pgx.Tx) error {
+			var err error
+			got, err = pubDB.ReadExposures(ctx, tx, expectedTEKs)
+			return err
+		}); err != nil {
+			t.Fatal(err)
+		}
+
+		// Ensure exposure1 was revised.
+		gotExposure1, ok := got[exposure1.ExposureKeyBase64()]
+		if !ok {
+			t.Fatal("exposure1 is missing in response")
+		}
+		if got, want := gotExposure1.RevisedReportType, verifyapi.ReportTypeConfirmed; got == nil || *got != want {
+			t.Errorf("expected %#v to be %#v", got, want)
+		}
+
+		// Ensure exposure2 was NOT revised (it was not in the revision token and
+		// should have been dropped).
+		gotExposure2, ok := got[exposure2.ExposureKeyBase64()]
+		if !ok {
+			t.Fatal("exposure2 is missing in response")
+		}
+		if got, want := gotExposure2.ReportType, verifyapi.ReportTypeClinical; got != want {
+			t.Errorf("expected %#v to be %#v", got, want)
+		}
+
+		// Ensure exposure3 was inserted.
+		gotExposure3, ok := got[exposure3.ExposureKeyBase64()]
+		if !ok {
+			t.Fatal("exposure3 is missing in response")
+		}
+		if got, want := gotExposure3.IntervalNumber, exposure3.IntervalNumber; got != want {
+			t.Errorf("expected %#v to be %#v", got, want)
+		}
+	})
 }
 
 func TestIterateExposuresCursor(t *testing.T) {
@@ -521,7 +738,10 @@ func TestIterateExposuresCursor(t *testing.T) {
 			Regions:        []string{"MX", "CA"},
 		},
 	}
-	if _, err := testPublishDB.InsertAndReviseExposures(ctx, exposures, nil, true); err != nil {
+	if _, err := testPublishDB.InsertAndReviseExposures(ctx, &InsertAndReviseExposuresRequest{
+		Incoming:     exposures,
+		RequireToken: true,
+	}); err != nil {
 		t.Fatal(err)
 	}
 	// Iterate over them, canceling the context in the middle.