Rewrite StatesWithStatus to do a single db query

database/database.go

@@ -183,23 +183,6 @@ func (db DatastoreBacked) AllDomainStates() (states []DomainState, err error) {
 
 // StatesWithStatus returns the states of domains with the given status in the database.
 func (db DatastoreBacked) StatesWithStatus(status PreloadStatus) (domains []DomainState, err error) {
-	// Calling statesForQuery results in timeout. Instead, call domainsForQuery, which returns keys
-	// only, with filter on IncludeSubDomains, and reconstruct the DomainState explicitly.
-	domainNames, err := db.domainsForQuery(
-		datastore.NewQuery("DomainState").Filter("Status =", string(status)).Filter("IncludeSubDomains =", false))
-	if err != nil {
-		return domains, err
-	}
-	for _, domain := range domainNames {
-		domains = append(domains, DomainState{Name: domain, Status: status, IncludeSubDomains: false})
-	}
-	domainNames, err = db.domainsForQuery(
-		datastore.NewQuery("DomainState").Filter("Status =", string(status)).Filter("IncludeSubDomains =", true))
-	if err != nil {
-		return domains, err
-	}
-	for _, domain := range domainNames {
-		domains = append(domains, DomainState{Name: domain, Status: status, IncludeSubDomains: true})
-	}
-	return domains, err
+	return db.statesForQuery(
+		datastore.NewQuery("DomainState").Filter("Status =", string(status)))
 }

database/database_test.go

@@ -224,8 +224,13 @@ func TestStatesWithStatus(t *testing.T) {
 			t.Errorf("%s", err)
 		}
 		sort.Slice(domainStates, func(i, j int) bool { return domainStates[i].Name < domainStates[j].Name })
-		if !reflect.DeepEqual(domainStates, tt.domains) {
-			t.Errorf("not the list of expected domains for status %s: %#v", tt.status, domainStates)
+		if len(domainStates) != len(tt.domains) {
+			t.Errorf("Incorrect count of states for status %s", tt.status)
+		}
+		for i, domainState := range domainStates {
+			if !domainState.Equal(tt.domains[i]) {
+				t.Errorf("unexpected domain at position %d for status %s: %#v", i, tt.status, domainState)
+			}
 		}
 	}
 }

database/domainstate.go

@@ -57,6 +57,17 @@ func (s DomainState) MatchesWanted(wanted DomainState) bool {
 	return true
 }
 
+// Equal checks if the fields of `s` are equal to the fields of `s2`,
+// using == for all fields except for SubmissionDate, where Time.Equal is
+// used instead. This is a more strict check than MatchesWanted and is
+// intended for testing purposes.
+func (s DomainState) Equal(s2 DomainState) bool {
+	return s.Name == s2.Name && s.Status == s2.Status &&
+		s.Message == s2.Message &&
+		s.SubmissionDate.Equal(s2.SubmissionDate) &&
+		s.IncludeSubDomains == s2.IncludeSubDomains
+}
+
 // ToEntry converts a DomainState to a preloadlist.Entry.
 //
 // Only the name, preload status and include subdomains boolean is preserved during the conversion.