Refactoring Android datastore interface (#26982)

For #26219 Refactoring the interface between Android service and Android datastore to use 1 common datastore interface: `fleet.AndroidDatastore` These changes are based on feedback from the recent Backend Sync. ```mermaid --- title: Partial class diagram --- classDiagram direction LR class `android.Service` <<interface>> `android.Service` class `android/service.Service` `android/service.Service` ..|> `android.Service`: implements class `fleet.AndroidDatastore` <<interface>> `fleet.AndroidDatastore` class `fleet.Datastore` <<interface>> `fleet.Datastore` class `android.Datastore` <<interface>> `android.Datastore` `android/service.Service` *-- `fleet.AndroidDatastore`: USES (THIS IS THE KEY CHANGE) `fleet.Datastore` *-- `fleet.AndroidDatastore`: contains `mysql.Datastore` ..|> `fleet.Datastore`: implements `fleet.AndroidDatastore` *-- `android.Datastore`: contains `mysql.Datastore` *-- `android.Datastore`: contains `android/mysql.Datastore` ..|> `android.Datastore`: implements ```
2025-03-13 14:28:52 -05:00
parent 996aa769a7
commit a86253d2bf
11 changed files with 240 additions and 131 deletions
@@ -306,14 +306,6 @@ type OSVersionFunc func(ctx context.Context, osVersionID uint, teamFilter *fleet

 type UpdateOSVersionsFunc func(ctx context.Context) error

-type GetAndroidDSFunc func() android.Datastore
-
-type NewAndroidHostFunc func(ctx context.Context, host *fleet.AndroidHost) (*fleet.AndroidHost, error)
-
-type UpdateAndroidHostFunc func(ctx context.Context, host *fleet.AndroidHost, fromEnroll bool) error
-
-type AndroidHostLiteFunc func(ctx context.Context, enterpriseSpecificID string) (*fleet.AndroidHost, error)
-
 type CountHostsInTargetsFunc func(ctx context.Context, filter fleet.TeamFilter, targets fleet.HostTargets, now time.Time) (fleet.TargetMetrics, error)

 type HostIDsInTargetsFunc func(ctx context.Context, filter fleet.TeamFilter, targets fleet.HostTargets) ([]uint, error)
@@ -1256,10 +1248,34 @@ type ExpandEmbeddedSecretsFunc func(ctx context.Context, document string) (strin

 type ExpandEmbeddedSecretsAndUpdatedAtFunc func(ctx context.Context, document string) (string, *time.Time, error)

-type SetAndroidEnabledAndConfiguredFunc func(ctx context.Context, configured bool) error
+type CreateEnterpriseFunc func(ctx context.Context, userID uint) (uint, error)
+
+type GetEnterpriseByIDFunc func(ctx context.Context, ID uint) (*android.EnterpriseDetails, error)
+
+type GetEnterpriseBySignupTokenFunc func(ctx context.Context, signupToken string) (*android.EnterpriseDetails, error)
+
+type GetEnterpriseFunc func(ctx context.Context) (*android.Enterprise, error)
+
+type UpdateEnterpriseFunc func(ctx context.Context, enterprise *android.EnterpriseDetails) error
+
+type DeleteAllEnterprisesFunc func(ctx context.Context) error
+
+type DeleteOtherEnterprisesFunc func(ctx context.Context, ID uint) error
+
+type CreateDeviceTxFunc func(ctx context.Context, tx sqlx.ExtContext, device *android.Device) (*android.Device, error)
+
+type UpdateDeviceTxFunc func(ctx context.Context, tx sqlx.ExtContext, device *android.Device) error
+
+type AndroidHostLiteFunc func(ctx context.Context, enterpriseSpecificID string) (*fleet.AndroidHost, error)

 type BulkSetAndroidHostsUnenrolledFunc func(ctx context.Context) error

+type NewAndroidHostFunc func(ctx context.Context, host *fleet.AndroidHost) (*fleet.AndroidHost, error)
+
+type SetAndroidEnabledAndConfiguredFunc func(ctx context.Context, configured bool) error
+
+type UpdateAndroidHostFunc func(ctx context.Context, host *fleet.AndroidHost, fromEnroll bool) error
+
 type DataStore struct {
 	HealthCheckFunc        HealthCheckFunc
 	HealthCheckFuncInvoked bool
@@ -1687,18 +1703,6 @@ type DataStore struct {
 	UpdateOSVersionsFunc        UpdateOSVersionsFunc
 	UpdateOSVersionsFuncInvoked bool

-	GetAndroidDSFunc        GetAndroidDSFunc
-	GetAndroidDSFuncInvoked bool
-
-	NewAndroidHostFunc        NewAndroidHostFunc
-	NewAndroidHostFuncInvoked bool
-
-	UpdateAndroidHostFunc        UpdateAndroidHostFunc
-	UpdateAndroidHostFuncInvoked bool
-
-	AndroidHostLiteFunc        AndroidHostLiteFunc
-	AndroidHostLiteFuncInvoked bool
-
 	CountHostsInTargetsFunc        CountHostsInTargetsFunc
 	CountHostsInTargetsFuncInvoked bool

@@ -3112,12 +3116,48 @@ type DataStore struct {
 	ExpandEmbeddedSecretsAndUpdatedAtFunc        ExpandEmbeddedSecretsAndUpdatedAtFunc
 	ExpandEmbeddedSecretsAndUpdatedAtFuncInvoked bool

-	SetAndroidEnabledAndConfiguredFunc        SetAndroidEnabledAndConfiguredFunc
-	SetAndroidEnabledAndConfiguredFuncInvoked bool
+	CreateEnterpriseFunc        CreateEnterpriseFunc
+	CreateEnterpriseFuncInvoked bool
+
+	GetEnterpriseByIDFunc        GetEnterpriseByIDFunc
+	GetEnterpriseByIDFuncInvoked bool
+
+	GetEnterpriseBySignupTokenFunc        GetEnterpriseBySignupTokenFunc
+	GetEnterpriseBySignupTokenFuncInvoked bool
+
+	GetEnterpriseFunc        GetEnterpriseFunc
+	GetEnterpriseFuncInvoked bool
+
+	UpdateEnterpriseFunc        UpdateEnterpriseFunc
+	UpdateEnterpriseFuncInvoked bool
+
+	DeleteAllEnterprisesFunc        DeleteAllEnterprisesFunc
+	DeleteAllEnterprisesFuncInvoked bool
+
+	DeleteOtherEnterprisesFunc        DeleteOtherEnterprisesFunc
+	DeleteOtherEnterprisesFuncInvoked bool
+
+	CreateDeviceTxFunc        CreateDeviceTxFunc
+	CreateDeviceTxFuncInvoked bool
+
+	UpdateDeviceTxFunc        UpdateDeviceTxFunc
+	UpdateDeviceTxFuncInvoked bool
+
+	AndroidHostLiteFunc        AndroidHostLiteFunc
+	AndroidHostLiteFuncInvoked bool

 	BulkSetAndroidHostsUnenrolledFunc        BulkSetAndroidHostsUnenrolledFunc
 	BulkSetAndroidHostsUnenrolledFuncInvoked bool

+	NewAndroidHostFunc        NewAndroidHostFunc
+	NewAndroidHostFuncInvoked bool
+
+	SetAndroidEnabledAndConfiguredFunc        SetAndroidEnabledAndConfiguredFunc
+	SetAndroidEnabledAndConfiguredFuncInvoked bool
+
+	UpdateAndroidHostFunc        UpdateAndroidHostFunc
+	UpdateAndroidHostFuncInvoked bool
+
 	mu sync.Mutex
 }

@@ -4115,34 +4155,6 @@ func (s *DataStore) UpdateOSVersions(ctx context.Context) error {
 	return s.UpdateOSVersionsFunc(ctx)
 }

-func (s *DataStore) GetAndroidDS() android.Datastore {
-	s.mu.Lock()
-	s.GetAndroidDSFuncInvoked = true
-	s.mu.Unlock()
-	return s.GetAndroidDSFunc()
-}
-
-func (s *DataStore) NewAndroidHost(ctx context.Context, host *fleet.AndroidHost) (*fleet.AndroidHost, error) {
-	s.mu.Lock()
-	s.NewAndroidHostFuncInvoked = true
-	s.mu.Unlock()
-	return s.NewAndroidHostFunc(ctx, host)
-}
-
-func (s *DataStore) UpdateAndroidHost(ctx context.Context, host *fleet.AndroidHost, fromEnroll bool) error {
-	s.mu.Lock()
-	s.UpdateAndroidHostFuncInvoked = true
-	s.mu.Unlock()
-	return s.UpdateAndroidHostFunc(ctx, host, fromEnroll)
-}
-
-func (s *DataStore) AndroidHostLite(ctx context.Context, enterpriseSpecificID string) (*fleet.AndroidHost, error) {
-	s.mu.Lock()
-	s.AndroidHostLiteFuncInvoked = true
-	s.mu.Unlock()
-	return s.AndroidHostLiteFunc(ctx, enterpriseSpecificID)
-}
-
 func (s *DataStore) CountHostsInTargets(ctx context.Context, filter fleet.TeamFilter, targets fleet.HostTargets, now time.Time) (fleet.TargetMetrics, error) {
 	s.mu.Lock()
 	s.CountHostsInTargetsFuncInvoked = true
@@ -7440,11 +7452,74 @@ func (s *DataStore) ExpandEmbeddedSecretsAndUpdatedAt(ctx context.Context, docum
 	return s.ExpandEmbeddedSecretsAndUpdatedAtFunc(ctx, document)
 }

-func (s *DataStore) SetAndroidEnabledAndConfigured(ctx context.Context, configured bool) error {
+func (s *DataStore) CreateEnterprise(ctx context.Context, userID uint) (uint, error) {
 	s.mu.Lock()
-	s.SetAndroidEnabledAndConfiguredFuncInvoked = true
+	s.CreateEnterpriseFuncInvoked = true
 	s.mu.Unlock()
-	return s.SetAndroidEnabledAndConfiguredFunc(ctx, configured)
+	return s.CreateEnterpriseFunc(ctx, userID)
+}
+
+func (s *DataStore) GetEnterpriseByID(ctx context.Context, ID uint) (*android.EnterpriseDetails, error) {
+	s.mu.Lock()
+	s.GetEnterpriseByIDFuncInvoked = true
+	s.mu.Unlock()
+	return s.GetEnterpriseByIDFunc(ctx, ID)
+}
+
+func (s *DataStore) GetEnterpriseBySignupToken(ctx context.Context, signupToken string) (*android.EnterpriseDetails, error) {
+	s.mu.Lock()
+	s.GetEnterpriseBySignupTokenFuncInvoked = true
+	s.mu.Unlock()
+	return s.GetEnterpriseBySignupTokenFunc(ctx, signupToken)
+}
+
+func (s *DataStore) GetEnterprise(ctx context.Context) (*android.Enterprise, error) {
+	s.mu.Lock()
+	s.GetEnterpriseFuncInvoked = true
+	s.mu.Unlock()
+	return s.GetEnterpriseFunc(ctx)
+}
+
+func (s *DataStore) UpdateEnterprise(ctx context.Context, enterprise *android.EnterpriseDetails) error {
+	s.mu.Lock()
+	s.UpdateEnterpriseFuncInvoked = true
+	s.mu.Unlock()
+	return s.UpdateEnterpriseFunc(ctx, enterprise)
+}
+
+func (s *DataStore) DeleteAllEnterprises(ctx context.Context) error {
+	s.mu.Lock()
+	s.DeleteAllEnterprisesFuncInvoked = true
+	s.mu.Unlock()
+	return s.DeleteAllEnterprisesFunc(ctx)
+}
+
+func (s *DataStore) DeleteOtherEnterprises(ctx context.Context, ID uint) error {
+	s.mu.Lock()
+	s.DeleteOtherEnterprisesFuncInvoked = true
+	s.mu.Unlock()
+	return s.DeleteOtherEnterprisesFunc(ctx, ID)
+}
+
+func (s *DataStore) CreateDeviceTx(ctx context.Context, tx sqlx.ExtContext, device *android.Device) (*android.Device, error) {
+	s.mu.Lock()
+	s.CreateDeviceTxFuncInvoked = true
+	s.mu.Unlock()
+	return s.CreateDeviceTxFunc(ctx, tx, device)
+}
+
+func (s *DataStore) UpdateDeviceTx(ctx context.Context, tx sqlx.ExtContext, device *android.Device) error {
+	s.mu.Lock()
+	s.UpdateDeviceTxFuncInvoked = true
+	s.mu.Unlock()
+	return s.UpdateDeviceTxFunc(ctx, tx, device)
+}
+
+func (s *DataStore) AndroidHostLite(ctx context.Context, enterpriseSpecificID string) (*fleet.AndroidHost, error) {
+	s.mu.Lock()
+	s.AndroidHostLiteFuncInvoked = true
+	s.mu.Unlock()
+	return s.AndroidHostLiteFunc(ctx, enterpriseSpecificID)
 }

 func (s *DataStore) BulkSetAndroidHostsUnenrolled(ctx context.Context) error {
@@ -7453,3 +7528,24 @@ func (s *DataStore) BulkSetAndroidHostsUnenrolled(ctx context.Context) error {
 	s.mu.Unlock()
 	return s.BulkSetAndroidHostsUnenrolledFunc(ctx)
 }
+
+func (s *DataStore) NewAndroidHost(ctx context.Context, host *fleet.AndroidHost) (*fleet.AndroidHost, error) {
+	s.mu.Lock()
+	s.NewAndroidHostFuncInvoked = true
+	s.mu.Unlock()
+	return s.NewAndroidHostFunc(ctx, host)
+}
+
+func (s *DataStore) SetAndroidEnabledAndConfigured(ctx context.Context, configured bool) error {
+	s.mu.Lock()
+	s.SetAndroidEnabledAndConfiguredFuncInvoked = true
+	s.mu.Unlock()
+	return s.SetAndroidEnabledAndConfiguredFunc(ctx, configured)
+}
+
+func (s *DataStore) UpdateAndroidHost(ctx context.Context, host *fleet.AndroidHost, fromEnroll bool) error {
+	s.mu.Lock()
+	s.UpdateAndroidHostFuncInvoked = true
+	s.mu.Unlock()
+	return s.UpdateAndroidHostFunc(ctx, host, fromEnroll)
+}
@@ -377,6 +377,17 @@ func main() {
 		fatal(err)
 	}

+	// Remove duplicates from fns
+	uniqueFns := make(map[string]Func, len(fns))
+	dedupedFns := make([]Func, 0, len(fns))
+	for _, fn := range fns {
+		if _, exists := uniqueFns[fn.Name]; !exists {
+			uniqueFns[fn.Name] = fn
+			dedupedFns = append(dedupedFns, fn)
+		}
+	}
+	fns = dedupedFns
+
 	src := genStubs(recv, fns)
 	recName := strings.SplitN(recv, " ", 2)
 	name := strings.TrimPrefix(recName[1], "*")