Add OnlyNonTravelers to exports. (#1072)

Fixes #1054

internal/admin/exports/form.go

@@ -16,6 +16,7 @@
 package exports
 
 import (
+	"errors"
 	"fmt"
 	"sort"
 	"strings"
@@ -25,10 +26,13 @@ import (
 	"github.com/google/exposure-notifications-server/internal/export/model"
 )
 
+var ErrCannotSetBothTravelers = errors.New("cannot have both 'include travelers', and 'only non-travelers' set")
+
 type formData struct {
 	OutputRegion     string        `form:"OutputRegion"`
 	InputRegions     string        `form:"InputRegions"`
 	IncludeTravelers bool          `form:"IncludeTravelers"`
+	OnlyNonTravelers bool          `form:"OnlyNonTravelers"`
 	ExcludeRegions   string        `form:"ExcludeRegions"`
 	BucketName       string        `form:"BucketName"`
 	FilenameRoot     string        `form:"FilenameRoot"`
@@ -63,13 +67,17 @@ func (f *formData) PopulateExportConfig(ec *model.ExportConfig) error {
 	if err != nil {
 		return err
 	}
+	if f.IncludeTravelers && f.OnlyNonTravelers {
+		return ErrCannotSetBothTravelers
+	}
 
 	ec.BucketName = strings.TrimSpace(f.BucketName)
 	ec.FilenameRoot = strings.TrimSpace(f.FilenameRoot)
 	ec.Period = f.Period
 	ec.OutputRegion = strings.TrimSpace(f.OutputRegion)
 	ec.InputRegions = splitRegions(f.InputRegions)
 	ec.IncludeTravelers = f.IncludeTravelers
+	ec.OnlyNonTravelers = f.OnlyNonTravelers
 	ec.ExcludeRegions = splitRegions(f.ExcludeRegions)
 	ec.From = from
 	ec.Thru = thru

internal/export/batcher.go

@@ -132,6 +132,8 @@ func (s *Server) maybeCreateBatches(ctx context.Context, ec *model.ExportConfig,
 			OutputRegion:     ec.OutputRegion,
 			InputRegions:     ec.InputRegions,
 			IncludeTravelers: ec.IncludeTravelers,
+			OnlyNonTravelers: ec.OnlyNonTravelers,
+			ExcludeRegions:   ec.ExcludeRegions,
 			Status:           model.ExportBatchOpen,
 			SignatureInfoIDs: infoIds,
 		})

internal/export/database/export.go

@@ -52,13 +52,13 @@ func (db *ExportDB) AddExportConfig(ctx context.Context, ec *model.ExportConfig)
 		row := tx.QueryRow(ctx, `
 			INSERT INTO
 				ExportConfig
-				(bucket_name, filename_root, period_seconds, output_region, from_timestamp, thru_timestamp, signature_info_ids, input_regions, include_travelers, exclude_regions)
+				(bucket_name, filename_root, period_seconds, output_region, from_timestamp, thru_timestamp, signature_info_ids, input_regions, include_travelers, exclude_regions, only_non_travelers)
 			VALUES
-				($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
+				($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
 			RETURNING config_id
 		`, ec.BucketName, ec.FilenameRoot, int(ec.Period.Seconds()), ec.OutputRegion,
 			ec.From, thru, ec.SignatureInfoIDs, ec.InputRegions, ec.IncludeTravelers,
-			ec.ExcludeRegions)
+			ec.ExcludeRegions, ec.OnlyNonTravelers)
 
 		if err := row.Scan(&ec.ConfigID); err != nil {
 			return fmt.Errorf("fetching config_id: %w", err)
@@ -79,11 +79,11 @@ func (db *ExportDB) UpdateExportConfig(ctx context.Context, ec *model.ExportConf
 			UPDATE
 				ExportConfig
 			SET
-				bucket_name = $1, filename_root = $2, period_seconds = $3, output_region = $4, from_timestamp = $5, thru_timestamp = $6, signature_info_ids = $7, input_regions = $8, include_travelers = $9, exclude_regions = $10
-			WHERE config_id = $11
+				bucket_name = $1, filename_root = $2, period_seconds = $3, output_region = $4, from_timestamp = $5, thru_timestamp = $6, signature_info_ids = $7, input_regions = $8, include_travelers = $9, exclude_regions = $10, only_non_travelers = $11
+			WHERE config_id = $12
 		`, ec.BucketName, ec.FilenameRoot, int(ec.Period.Seconds()), ec.OutputRegion,
 			ec.From, thru, ec.SignatureInfoIDs, ec.InputRegions, ec.IncludeTravelers,
-			ec.ExcludeRegions, ec.ConfigID)
+			ec.ExcludeRegions, ec.OnlyNonTravelers, ec.ConfigID)
 		if err != nil {
 			return fmt.Errorf("updating signatureinfo: %w", err)
 		}
@@ -99,7 +99,7 @@ func (db *ExportDB) GetExportConfig(ctx context.Context, id int64) (*model.Expor
 	if err := db.db.InTx(ctx, pgx.ReadCommitted, func(tx pgx.Tx) error {
 		row := tx.QueryRow(ctx, `
 			SELECT
-				config_id, bucket_name, filename_root, period_seconds, output_region, from_timestamp, thru_timestamp, signature_info_ids, input_regions, include_travelers, exclude_regions
+				config_id, bucket_name, filename_root, period_seconds, output_region, from_timestamp, thru_timestamp, signature_info_ids, input_regions, include_travelers, exclude_regions, only_non_travelers
 			FROM
 				ExportConfig
 			WHERE
@@ -125,7 +125,7 @@ func (db *ExportDB) GetAllExportConfigs(ctx context.Context) ([]*model.ExportCon
 	if err := db.db.InTx(ctx, pgx.ReadCommitted, func(tx pgx.Tx) error {
 		rows, err := tx.Query(ctx, `
 			SELECT
-				config_id, bucket_name, filename_root, period_seconds, output_region, from_timestamp, thru_timestamp, signature_info_ids, input_regions, include_travelers, exclude_regions
+				config_id, bucket_name, filename_root, period_seconds, output_region, from_timestamp, thru_timestamp, signature_info_ids, input_regions, include_travelers, exclude_regions, only_non_travelers
 			FROM
 				ExportConfig
 			ORDER BY config_id
@@ -162,7 +162,7 @@ func (db *ExportDB) IterateExportConfigs(ctx context.Context, t time.Time, f fun
 	if err := db.db.InTx(ctx, pgx.ReadCommitted, func(tx pgx.Tx) error {
 		rows, err := tx.Query(ctx, `
 			SELECT
-				config_id, bucket_name, filename_root, period_seconds, output_region, from_timestamp, thru_timestamp, signature_info_ids, input_regions, include_travelers, exclude_regions
+				config_id, bucket_name, filename_root, period_seconds, output_region, from_timestamp, thru_timestamp, signature_info_ids, input_regions, include_travelers, exclude_regions, only_non_travelers
 			FROM
 				ExportConfig
 			WHERE
@@ -204,7 +204,7 @@ func scanOneExportConfig(row pgx.Row) (*model.ExportConfig, error) {
 		periodSeconds int
 		thru          *time.Time
 	)
-	if err := row.Scan(&m.ConfigID, &m.BucketName, &m.FilenameRoot, &periodSeconds, &outputRegion, &m.From, &thru, &m.SignatureInfoIDs, &m.InputRegions, &m.IncludeTravelers, &m.ExcludeRegions); err != nil {
+	if err := row.Scan(&m.ConfigID, &m.BucketName, &m.FilenameRoot, &periodSeconds, &outputRegion, &m.From, &thru, &m.SignatureInfoIDs, &m.InputRegions, &m.IncludeTravelers, &m.ExcludeRegions, &m.OnlyNonTravelers); err != nil {
 		return nil, err
 	}
 
@@ -470,17 +470,17 @@ func (db *ExportDB) AddExportBatches(ctx context.Context, batches []*model.Expor
 		_, err := tx.Prepare(ctx, stmtName, `
 			INSERT INTO
 				ExportBatch
-				(config_id, bucket_name, filename_root, start_timestamp, end_timestamp, output_region, status, signature_info_ids, input_regions, include_travelers, exclude_regions)
+				(config_id, bucket_name, filename_root, start_timestamp, end_timestamp, output_region, status, signature_info_ids, input_regions, include_travelers, exclude_regions, only_non_travelers)
 			VALUES
-				($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
+				($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
 		`)
 		if err != nil {
 			return err
 		}
 
 		for _, eb := range batches {
 			if _, err := tx.Exec(ctx, stmtName,
-				eb.ConfigID, eb.BucketName, eb.FilenameRoot, eb.StartTimestamp, eb.EndTimestamp, eb.OutputRegion, eb.Status, eb.SignatureInfoIDs, eb.InputRegions, eb.IncludeTravelers, eb.ExcludeRegions); err != nil {
+				eb.ConfigID, eb.BucketName, eb.FilenameRoot, eb.StartTimestamp, eb.EndTimestamp, eb.OutputRegion, eb.Status, eb.SignatureInfoIDs, eb.InputRegions, eb.IncludeTravelers, eb.ExcludeRegions, eb.OnlyNonTravelers); err != nil {
 				return err
 			}
 		}
@@ -609,7 +609,7 @@ type queryRowFn func(ctx context.Context, query string, args ...interface{}) pgx
 func lookupExportBatch(ctx context.Context, batchID int64, queryRow queryRowFn) (*model.ExportBatch, error) {
 	row := queryRow(ctx, `
 		SELECT
-			batch_id, config_id, bucket_name, filename_root, start_timestamp, end_timestamp, output_region, status, lease_expires, signature_info_ids, input_regions, include_travelers, exclude_regions
+			batch_id, config_id, bucket_name, filename_root, start_timestamp, end_timestamp, output_region, status, lease_expires, signature_info_ids, input_regions, include_travelers, exclude_regions, only_non_travelers
 		FROM
 			ExportBatch
 		WHERE
@@ -619,7 +619,7 @@ func lookupExportBatch(ctx context.Context, batchID int64, queryRow queryRowFn)
 
 	var expires *time.Time
 	eb := model.ExportBatch{}
-	if err := row.Scan(&eb.BatchID, &eb.ConfigID, &eb.BucketName, &eb.FilenameRoot, &eb.StartTimestamp, &eb.EndTimestamp, &eb.OutputRegion, &eb.Status, &expires, &eb.SignatureInfoIDs, &eb.InputRegions, &eb.IncludeTravelers, &eb.ExcludeRegions); err != nil {
+	if err := row.Scan(&eb.BatchID, &eb.ConfigID, &eb.BucketName, &eb.FilenameRoot, &eb.StartTimestamp, &eb.EndTimestamp, &eb.OutputRegion, &eb.Status, &expires, &eb.SignatureInfoIDs, &eb.InputRegions, &eb.IncludeTravelers, &eb.ExcludeRegions, &eb.OnlyNonTravelers); err != nil {
 		if err == pgx.ErrNoRows {
 			return nil, database.ErrNotFound
 		}
@@ -643,6 +643,7 @@ func (db *ExportDB) FinalizeBatch(ctx context.Context, eb *model.ExportBatch, fi
 				OutputRegion:     eb.OutputRegion,
 				InputRegions:     eb.InputRegions,
 				IncludeTravelers: eb.IncludeTravelers,
+				OnlyNonTravelers: eb.OnlyNonTravelers,
 				ExcludeRegions:   eb.ExcludeRegions,
 				BatchNum:         i + 1,
 				BatchSize:        batchSize,
@@ -732,15 +733,15 @@ func (db *ExportDB) LookupExportFile(ctx context.Context, filename string) (*mod
 	if err := db.db.InTx(ctx, pgx.ReadCommitted, func(tx pgx.Tx) error {
 		row := tx.QueryRow(ctx, `
 			SELECT
-				bucket_name, filename, batch_id, output_region, batch_num, batch_size, status, input_regions, include_travelers, exclude_regions
+				bucket_name, filename, batch_id, output_region, batch_num, batch_size, status, input_regions, include_travelers, exclude_regions, only_non_travelers
 			FROM
 				ExportFile
 			WHERE
 				filename = $1
 			LIMIT 1
 		`, filename)
 
-		if err := row.Scan(&file.BucketName, &file.Filename, &file.BatchID, &file.OutputRegion, &file.BatchNum, &file.BatchSize, &file.Status, &file.InputRegions, &file.IncludeTravelers, &file.ExcludeRegions); err != nil {
+		if err := row.Scan(&file.BucketName, &file.Filename, &file.BatchID, &file.OutputRegion, &file.BatchNum, &file.BatchSize, &file.Status, &file.InputRegions, &file.IncludeTravelers, &file.ExcludeRegions, &file.OnlyNonTravelers); err != nil {
 			if errors.Is(err, pgx.ErrNoRows) {
 				return database.ErrNotFound
 			}
@@ -845,11 +846,11 @@ func addExportFile(ctx context.Context, tx pgx.Tx, ef *model.ExportFile) error {
 	tag, err := tx.Exec(ctx, `
 		INSERT INTO
 			ExportFile
-			(bucket_name, filename, batch_id, output_region, batch_num, batch_size, status, input_regions, include_travelers, exclude_regions)
+			(bucket_name, filename, batch_id, output_region, batch_num, batch_size, status, input_regions, include_travelers, exclude_regions, only_non_travelers)
 		VALUES
-			($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
+			($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
 		ON CONFLICT (filename) DO NOTHING
-		`, ef.BucketName, ef.Filename, ef.BatchID, ef.OutputRegion, ef.BatchNum, ef.BatchSize, ef.Status, ef.InputRegions, ef.IncludeTravelers, ef.ExcludeRegions)
+		`, ef.BucketName, ef.Filename, ef.BatchID, ef.OutputRegion, ef.BatchNum, ef.BatchSize, ef.Status, ef.InputRegions, ef.IncludeTravelers, ef.ExcludeRegions, ef.OnlyNonTravelers)
 	if err != nil {
 		return fmt.Errorf("inserting to ExportFile: %w", err)
 	}

internal/export/database/export_test.go

@@ -546,6 +546,7 @@ func TestTravelerKeys(t *testing.T) {
 // TestExcludeRegions ensures excluded regions are excluded.
 func TestExcludeRegions(t *testing.T) {
 	inclRegion, exclRegion := "US", "EU"
+	inclRegions, exclRegions, bothRegions := []string{inclRegion}, []string{exclRegion}, []string{inclRegion, exclRegion}
 	now := time.Now()
 	startTimestamp := now.Truncate(time.Hour).Add(-2 * time.Hour)
 	endTimestamp := startTimestamp.Add(time.Hour)
@@ -571,12 +572,17 @@ func TestExcludeRegions(t *testing.T) {
 	}
 
 	tests := []struct {
-		name  string
-		users []*publishmodel.Exposure
-		want  []string
+		name        string
+		users       []*publishmodel.Exposure
+		inclRegions []string
+		exclRegions []string
+		nonTraveler bool
+		want        []string
 	}{
-		{"1main,1ext", []*publishmodel.Exposure{mainUser, extUser}, []string{"aaa"}},
-		{"1main,2ext", []*publishmodel.Exposure{mainUser, extUser, extTravUser}, []string{"aaa"}},
+		{"1main,1ext", []*publishmodel.Exposure{mainUser, extUser}, inclRegions, exclRegions, false, []string{"aaa"}},
+		{"1main,2ext", []*publishmodel.Exposure{mainUser, extUser, extTravUser}, inclRegions, exclRegions, false, []string{"aaa"}},
+		{"1main,2ext,notravel", []*publishmodel.Exposure{mainUser, extUser, extTravUser}, inclRegions, nil, true, []string{"aaa"}},
+		{"1main,2ext,notravel,2regions", []*publishmodel.Exposure{mainUser, extUser, extTravUser}, bothRegions, nil, true, []string{"aaa", "ccc"}},
 	}
 
 	for _, test := range tests {
@@ -641,10 +647,11 @@ func TestExcludeRegions(t *testing.T) {
 			// Lookup the keys; they must be only the key created_at the startTimestamp
 			// (because start is inclusive, end is exclusive).
 			criteria := publishdb.IterateExposuresCriteria{
-				IncludeRegions: []string{inclRegion},
-				SinceTimestamp: leased.StartTimestamp,
-				UntilTimestamp: leased.EndTimestamp,
-				ExcludeRegions: []string{exclRegion},
+				IncludeRegions:   test.inclRegions,
+				SinceTimestamp:   leased.StartTimestamp,
+				UntilTimestamp:   leased.EndTimestamp,
+				ExcludeRegions:   test.exclRegions,
+				OnlyNonTravelers: test.nonTraveler,
 			}
 
 			var got []*publishmodel.Exposure
@@ -660,6 +667,7 @@ func TestExcludeRegions(t *testing.T) {
 			for _, exp := range got {
 				keys = append(keys, string(exp.ExposureKey))
 			}
+			sort.Strings(keys)
 			if !reflect.DeepEqual(test.want, keys) {
 				t.Fatalf("%v got = %v, want %v", test.name, keys, test.want)
 			}

internal/export/model/export_model.go

@@ -41,6 +41,7 @@ type ExportConfig struct {
 	InputRegions     []string
 	ExcludeRegions   []string
 	IncludeTravelers bool
+	OnlyNonTravelers bool
 	From             time.Time
 	Thru             time.Time
 	SignatureInfoIDs []int64
@@ -110,6 +111,7 @@ type ExportBatch struct {
 	OutputRegion     string
 	InputRegions     []string
 	IncludeTravelers bool
+	OnlyNonTravelers bool
 	ExcludeRegions   []string
 	Status           string
 	LeaseExpires     time.Time
@@ -129,6 +131,7 @@ type ExportFile struct {
 	OutputRegion     string
 	InputRegions     []string
 	IncludeTravelers bool
+	OnlyNonTravelers bool
 	ExcludeRegions   []string
 	BatchNum         int
 	BatchSize        int

internal/export/worker.go

@@ -120,6 +120,7 @@ func (s *Server) exportBatch(ctx context.Context, eb *model.ExportBatch, emitInd
 		UntilTimestamp:      eb.EndTimestamp,
 		IncludeRegions:      eb.EffectiveInputRegions(),
 		IncludeTravelers:    eb.IncludeTravelers, // Travelers are included from "any" region.
+		OnlyNonTravelers:    eb.OnlyNonTravelers,
 		ExcludeRegions:      eb.ExcludeRegions,
 		OnlyLocalProvenance: false, // include federated ids
 		OnlyRevisedKeys:     false,

internal/publish/database/exposure.go

@@ -73,6 +73,7 @@ func New(db *database.DB) *PublishDB {
 type IterateExposuresCriteria struct {
 	IncludeRegions   []string
 	IncludeTravelers bool // Include records in the IncludeRegions OR travalers
+	OnlyNonTravelers bool
 	OnlyTravelers    bool // Only includes records marked as travelers.
 	ExcludeRegions   []string
 	SinceTimestamp   time.Time
@@ -236,6 +237,11 @@ func generateExposureQuery(criteria IterateExposuresCriteria) (string, []interfa
 		q += fmt.Sprintf(" AND traveler = $%d", len(args))
 	}
 
+	if criteria.OnlyNonTravelers {
+		args = append(args, false)
+		q += fmt.Sprintf(" AND traveler = $%d", len(args))
+	}
+
 	if criteria.OnlyRevisedKeys {
 		q += " ORDER BY revised_at"
 	} else {

migrations/000046_AddOnlyNonTravelers.down.sql

@@ -0,0 +1,26 @@
+-- Copyright 2020 Google LLC
+--
+-- Licensed under the Apache License, Version 2.0 (the "License");
+-- you may not use this file except in compliance with the License.
+-- You may obtain a copy of the License at
+--
+--      http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing, software
+-- distributed under the License is distributed on an "AS IS" BASIS,
+-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+-- See the License for the specific language governing permissions and
+-- limitations under the License.
+
+BEGIN;
+
+ALTER TABLE exportconfig
+    DROP COLUMN only_non_travelers;
+
+ALTER TABLE exportbatch
+    DROP COLUMN only_non_travelers;
+
+ALTER TABLE exportfile
+    DROP COLUMN only_non_travelers;
+
+END;

migrations/000046_AddOnlyNonTravelers.up.sql

@@ -0,0 +1,26 @@
+-- Copyright 2020 Google LLC
+--
+-- Licensed under the Apache License, Version 2.0 (the "License");
+-- you may not use this file except in compliance with the License.
+-- You may obtain a copy of the License at
+--
+--      http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing, software
+-- distributed under the License is distributed on an "AS IS" BASIS,
+-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+-- See the License for the specific language governing permissions and
+-- limitations under the License.
+
+BEGIN;
+
+ALTER TABLE exportconfig
+    ADD COLUMN only_non_travelers BOOL DEFAULT FALSE;
+
+ALTER TABLE exportbatch
+    ADD COLUMN only_non_travelers BOOL DEFAULT FALSE;
+
+ALTER TABLE exportfile
+    ADD COLUMN only_non_travelers BOOL DEFAULT FALSE;
+
+END;

tools/admin-console/templates/export.html

@@ -58,6 +58,19 @@
         </small>
       </div>
 
+      <div class="form-group">
+        <label for="only-non-travelers">Only Non Travelers</label>
+        <select name="OnlyNonTravelers" id="only-non-travelers" class="form-control custom-select">
+          <option value="true" {{if .export.OnlyNonTravelers}}selected{{end}}>Yes</option>
+          <option value="false" {{if not .export.OnlyNonTravelers}}selected{{end}}>No</option>
+        </select>
+        <small class="form-text text-muted">
+		  Should only federeated-in non-traveler keys be included in this
+		  export. Note, setting this with 'Include Travelers' will result in an
+		  error.
+        </small>
+      </div>
+
       <div class="form-label-group">
         <textarea name="ExcludeRegions" id="exclude-regions" rows="3"
           placeholder="Exclude regions" class="form-control">{{.export.ExcludeRegionsOnePerLine}}</textarea>