is: Refactor model and query context handling in bunstore

pkg/identityserver/bunstore/account_store.go

@@ -77,10 +77,7 @@ func (s *baseStore) getAccountModel(
 	accountType, uid string,
 ) (*Account, error) {
 	model := &Account{}
-	selectQuery := s.DB.NewSelect().
-		Model(model).
-		Apply(selectWithSoftDeletedFromContext(ctx)).
-		Apply(selectWithContext(ctx)).
+	selectQuery := s.newSelectModel(ctx, model).
 		Where("?TableAlias.account_type = ?", accountType).
 		Where("?TableAlias.uid = ?", uid)
 

pkg/identityserver/bunstore/api_key_store.go

@@ -121,9 +121,7 @@ func (s *apiKeyStore) listAPIKeysBy(
 	by func(*bun.SelectQuery) *bun.SelectQuery,
 ) ([]*ttnpb.APIKey, error) {
 	models := []*APIKey{}
-	selectQuery := s.DB.NewSelect().
-		Model(&models).
-		Apply(by)
+	selectQuery := newSelectModels(ctx, s.DB, &models).Apply(by)
 
 	// Count the total number of results.
 	count, err := selectQuery.Count(ctx)
@@ -199,9 +197,7 @@ func (s *apiKeyStore) getAPIKeyModelBy(
 	by func(*bun.SelectQuery) *bun.SelectQuery,
 ) (*APIKey, error) {
 	model := &APIKey{}
-	selectQuery := s.DB.NewSelect().
-		Model(model).
-		Apply(by)
+	selectQuery := s.newSelectModel(ctx, model).Apply(by)
 
 	if err := selectQuery.Scan(ctx); err != nil {
 		return nil, wrapDriverError(err)

pkg/identityserver/bunstore/application_store.go

@@ -246,10 +246,7 @@ func (*applicationStore) selectWithFields(q *bun.SelectQuery, fieldMask store.Fi
 }
 
 func (s *applicationStore) CountApplications(ctx context.Context) (uint64, error) {
-	selectQuery := s.DB.NewSelect().
-		Model(&Application{}).
-		Apply(selectWithSoftDeletedFromContext(ctx)).
-		Apply(selectWithContext(ctx))
+	selectQuery := s.newSelectModel(ctx, &Application{})
 
 	// Count the total number of results.
 	count, err := selectQuery.Count(ctx)
@@ -266,10 +263,7 @@ func (s *applicationStore) listApplicationsBy(
 	fieldMask store.FieldMask,
 ) ([]*ttnpb.Application, error) {
 	models := []*Application{}
-	selectQuery := s.DB.NewSelect().
-		Model(&models).
-		Apply(selectWithSoftDeletedFromContext(ctx)).
-		Apply(by)
+	selectQuery := newSelectModels(ctx, s.DB, &models).Apply(by)
 
 	// Count the total number of results.
 	count, err := selectQuery.Count(ctx)
@@ -312,10 +306,9 @@ func (s *applicationStore) listApplicationsBy(
 }
 
 func (*applicationStore) selectWithID(
-	ctx context.Context, ids ...string,
+	_ context.Context, ids ...string,
 ) func(*bun.SelectQuery) *bun.SelectQuery {
 	return func(q *bun.SelectQuery) *bun.SelectQuery {
-		q = q.Apply(selectWithContext(ctx))
 		switch len(ids) {
 		case 0:
 			return q
@@ -344,10 +337,7 @@ func (s *applicationStore) getApplicationModelBy(
 	fieldMask store.FieldMask,
 ) (*Application, error) {
 	model := &Application{}
-	selectQuery := s.DB.NewSelect().
-		Model(model).
-		Apply(selectWithSoftDeletedFromContext(ctx)).
-		Apply(by)
+	selectQuery := s.newSelectModel(ctx, model).Apply(by)
 
 	selectQuery, err := s.selectWithFields(selectQuery, fieldMask)
 	if err != nil {

pkg/identityserver/bunstore/attribute_store.go

@@ -36,6 +36,8 @@ type Attribute struct {
 	Value string `bun:"value,notnull"`
 }
 
+func (Attribute) _isModel() {} // It doesn't embed Model, but it's still a model.
+
 // AttributeSlice is a slice of Attributes.
 type AttributeSlice []*Attribute
 

pkg/identityserver/bunstore/client_store.go

@@ -274,10 +274,7 @@ func (*clientStore) selectWithFields(q *bun.SelectQuery, fieldMask store.FieldMa
 }
 
 func (s *clientStore) CountClients(ctx context.Context) (uint64, error) {
-	selectQuery := s.DB.NewSelect().
-		Model(&Client{}).
-		Apply(selectWithSoftDeletedFromContext(ctx)).
-		Apply(selectWithContext(ctx))
+	selectQuery := s.newSelectModel(ctx, &Client{})
 
 	// Count the total number of results.
 	count, err := selectQuery.Count(ctx)
@@ -294,10 +291,7 @@ func (s *clientStore) listClientsBy(
 	fieldMask store.FieldMask,
 ) ([]*ttnpb.Client, error) {
 	models := []*Client{}
-	selectQuery := s.DB.NewSelect().
-		Model(&models).
-		Apply(selectWithSoftDeletedFromContext(ctx)).
-		Apply(by)
+	selectQuery := newSelectModels(ctx, s.DB, &models).Apply(by)
 
 	// Count the total number of results.
 	count, err := selectQuery.Count(ctx)
@@ -340,10 +334,9 @@ func (s *clientStore) listClientsBy(
 }
 
 func (*clientStore) selectWithID(
-	ctx context.Context, ids ...string,
+	_ context.Context, ids ...string,
 ) func(*bun.SelectQuery) *bun.SelectQuery {
 	return func(q *bun.SelectQuery) *bun.SelectQuery {
-		q = q.Apply(selectWithContext(ctx))
 		switch len(ids) {
 		case 0:
 			return q
@@ -372,10 +365,7 @@ func (s *clientStore) getClientModelBy(
 	fieldMask store.FieldMask,
 ) (*Client, error) {
 	model := &Client{}
-	selectQuery := s.DB.NewSelect().
-		Model(model).
-		Apply(selectWithSoftDeletedFromContext(ctx)).
-		Apply(by)
+	selectQuery := s.newSelectModel(ctx, model).Apply(by)
 
 	selectQuery, err := s.selectWithFields(selectQuery, fieldMask)
 	if err != nil {

pkg/identityserver/bunstore/contact_info_store.go

@@ -48,6 +48,8 @@ type ContactInfo struct {
 	ValidatedAt *time.Time `bun:"validated_at"`
 }
 
+func (ContactInfo) _isModel() {} // It doesn't embed Model, but it's still a model.
+
 func contactInfoFromPB(pb *ttnpb.ContactInfo, entityType, entityID string) *ContactInfo {
 	return &ContactInfo{
 		EntityType:    entityType,
@@ -223,8 +225,7 @@ func (s *contactInfoStore) getContactInfoModelsBy(
 	by func(*bun.SelectQuery) *bun.SelectQuery,
 ) ([]*ContactInfo, error) {
 	models := []*ContactInfo{}
-	selectQuery := s.DB.NewSelect().
-		Model(&models).
+	selectQuery := newSelectModels(ctx, s.DB, &models).
 		Apply(by)
 
 	err := selectQuery.Scan(ctx)
@@ -332,8 +333,7 @@ func (s *contactInfoStore) CreateValidation(
 		return nil, err
 	}
 
-	n, err := s.DB.NewSelect().
-		Model(&ContactInfoValidation{}).
+	n, err := s.newSelectModel(ctx, &ContactInfoValidation{}).
 		Where("entity_type = ? AND entity_id = ?", entityType, entityUUID).
 		Where("contact_method = ? AND LOWER(value) = LOWER(?)", contactMethod, value).
 		Where("expires_at IS NULL OR expires_at > NOW()").
@@ -381,8 +381,7 @@ func (s *contactInfoStore) Validate(ctx context.Context, validation *ttnpb.Conta
 
 	model := &ContactInfoValidation{}
 
-	err := s.DB.NewSelect().
-		Model(model).
+	err := s.newSelectModel(ctx, model).
 		Where("reference = ? AND token = ?", validation.Id, validation.Token).
 		Scan(ctx)
 	if err != nil {

pkg/identityserver/bunstore/end_device_store.go

@@ -328,10 +328,7 @@ func (s *endDeviceStore) listEndDevicesBy(
 	fieldMask store.FieldMask,
 ) ([]*ttnpb.EndDevice, error) {
 	models := []*EndDevice{}
-	selectQuery := s.DB.NewSelect().
-		Model(&models).
-		Apply(selectWithSoftDeletedFromContext(ctx)).
-		Apply(by)
+	selectQuery := newSelectModels(ctx, s.DB, &models).Apply(by)
 
 	// Count the total number of results.
 	count, err := selectQuery.Count(ctx)
@@ -374,10 +371,9 @@ func (s *endDeviceStore) listEndDevicesBy(
 }
 
 func (*endDeviceStore) selectWithID(
-	ctx context.Context, applicationID string, deviceIDs ...string,
+	_ context.Context, applicationID string, deviceIDs ...string,
 ) func(*bun.SelectQuery) *bun.SelectQuery {
 	return func(q *bun.SelectQuery) *bun.SelectQuery {
-		q = q.Apply(selectWithContext(ctx))
 		q = q.Where("?TableAlias.application_id = ?", applicationID)
 		switch len(deviceIDs) {
 		case 0:
@@ -417,11 +413,7 @@ func (s *endDeviceStore) CountEndDevices(ctx context.Context, ids *ttnpb.Applica
 		by = s.selectWithID(ctx, ids.GetApplicationId())
 	}
 
-	models := []*EndDevice{}
-	selectQuery := s.DB.NewSelect().
-		Model(&models).
-		Apply(selectWithSoftDeletedFromContext(ctx)).
-		Apply(by)
+	selectQuery := s.newSelectModel(ctx, &EndDevice{}).Apply(by)
 
 	// Count the total number of results.
 	count, err := selectQuery.Count(ctx)
@@ -485,10 +477,7 @@ func (s *endDeviceStore) getEndDeviceModelBy(
 	fieldMask store.FieldMask,
 ) (*EndDevice, error) {
 	model := &EndDevice{}
-	selectQuery := s.DB.NewSelect().
-		Model(model).
-		Apply(selectWithSoftDeletedFromContext(ctx)).
-		Apply(by)
+	selectQuery := s.newSelectModel(ctx, model).Apply(by)
 
 	selectQuery, err := s.selectWithFields(selectQuery, fieldMask)
 	if err != nil {

pkg/identityserver/bunstore/entity_search.go

@@ -71,7 +71,7 @@ func (s *entitySearch) queryStringQuery(
 }
 
 func (s *entitySearch) selectWithMetaFields(
-	_ context.Context, entityType string, req metaFields,
+	ctx context.Context, entityType string, req metaFields,
 ) func(*bun.SelectQuery) *bun.SelectQuery {
 	return func(q *bun.SelectQuery) *bun.SelectQuery {
 		if v := req.GetIdContains(); v != "" {
@@ -91,8 +91,7 @@ func (s *entitySearch) selectWithMetaFields(
 			q = q.Where(s.ftsQuery("description"), v)
 		}
 		if kv := req.GetAttributesContain(); len(kv) > 0 {
-			attrQuery := s.DB.NewSelect().
-				Model(&Attribute{}).
+			attrQuery := s.newSelectModel(ctx, &Attribute{}).
 				Column("entity_id")
 			switch entityType {
 			case "end_device":

pkg/identityserver/bunstore/entity_store.go

@@ -48,28 +48,23 @@ func newEntityStore(baseStore *baseStore) *entityStore {
 }
 
 func (s *entityStore) query(ctx context.Context, entityType string, entityIDs ...string) *bun.SelectQuery {
-	query := s.DB.NewSelect()
+	var query *bun.SelectQuery
 
 	switch entityType {
 	case "application":
-		query = query.
-			Model(&Application{}).
+		query = s.newSelectModel(ctx, &Application{}).
 			Column("id").
 			Apply(s.applicationStore.selectWithID(ctx, entityIDs...))
 	case "client":
-		query = query.
-			Model(&Client{}).
+		query = s.newSelectModel(ctx, &Client{}).
 			Column("id").
 			Apply(s.clientStore.selectWithID(ctx, entityIDs...))
 	case "gateway":
-		query = query.
-			Model(&Gateway{}).
+		query = s.newSelectModel(ctx, &Gateway{}).
 			Column("id").
 			Apply(s.gatewayStore.selectWithID(ctx, entityIDs...))
 	case "organization", "user":
-		query = query.
-			Model(&Account{}).
-			Apply(selectWithContext(ctx)).
+		query = s.newSelectModel(ctx, &Account{}).
 			Column("account_id").
 			Where("?TableAlias.account_type = ?", entityType)
 		switch len(entityIDs) {
@@ -81,8 +76,6 @@ func (s *entityStore) query(ctx context.Context, entityType string, entityIDs ..
 		}
 	}
 
-	query = query.Apply(selectWithSoftDeletedFromContext(ctx))
-
 	return query
 }
 
@@ -141,36 +134,28 @@ func (s *entityStore) getEntityUUIDs(ctx context.Context, entityType string, ent
 }
 
 func (s *entityStore) getEntityID(ctx context.Context, entityType, entityUUID string) (string, error) {
-	query := s.DB.NewSelect()
+	var query *bun.SelectQuery
 
 	switch entityType {
 	case "application":
-		query = query.
-			Model(&Application{}).
+		query = s.newSelectModel(ctx, &Application{}).
 			Column("application_id").
 			Where("id = ?", entityUUID)
 	case "client":
-		query = query.
-			Model(&Client{}).
+		query = s.newSelectModel(ctx, &Client{}).
 			Column("client_id").
 			Where("id = ?", entityUUID)
 	case "gateway":
-		query = query.
-			Model(&Gateway{}).
+		query = s.newSelectModel(ctx, &Gateway{}).
 			Column("gateway_id").
 			Where("id = ?", entityUUID)
 	case "organization", "user":
-		query = query.
-			Model(&Account{}).
+		query = s.newSelectModel(ctx, &Account{}).
 			Column("uid").
 			Where("account_id = ?", entityUUID).
 			Where("?TableAlias.account_type = ?", entityType)
 	}
 
-	query = query.
-		Apply(selectWithContext(ctx)).
-		Apply(selectWithSoftDeletedFromContext(ctx))
-
 	var friendlyID string
 	err := query.Scan(ctx, &friendlyID)
 	if err != nil {

pkg/identityserver/bunstore/eui_store.go

@@ -67,8 +67,7 @@ func (s *euiStore) CreateEUIBlock(
 	defer span.End()
 
 	model := &EUIBlock{}
-	selectQuery := s.DB.NewSelect().
-		Model(model).
+	selectQuery := s.newSelectModel(ctx, model).
 		Where("type = ?", euiType).
 		Where("LOWER(start_eui) = LOWER(?)", prefix.EUI64.String())
 
@@ -144,8 +143,7 @@ func (s *euiStore) IssueDevEUIForApplication(
 			)
 		}
 
-		selectQuery := s.DB.NewSelect().
-			Model(euiBlockModel).
+		selectQuery := s.newSelectModel(ctx, euiBlockModel).
 			For("UPDATE").
 			Where("type = ?", "dev_eui").
 			Where("current_counter <= end_counter").