Reuse temporary database in local postgres in fake environment (between specific test runs). (#9084)

isoos · web-flow · commit 7f9e0eb557c0 · 2025-11-20T21:43:11.000+01:00
diff --git a/app/lib/database/database.dart b/app/lib/database/database.dart
@@ -24,16 +24,19 @@ void registerPrimaryDatabase(PrimaryDatabase database) =>
     ss.register(#_primaryDatabase, database);
 
 /// The active primary database service.
-PrimaryDatabase? get primaryDatabase =>
+PrimaryDatabase? get primaryDatabase => _lookupPrimaryDatabase();
+
+PrimaryDatabase? _lookupPrimaryDatabase() =>
     ss.lookup(#_primaryDatabase) as PrimaryDatabase?;
 
 /// Access to the primary database connection and object mapping.
 class PrimaryDatabase {
   final Pool _pg;
   final DatabaseAdapter _adapter;
   final Database<PrimarySchema> db;
+  final Future<void> Function()? _closeFn;
 
-  PrimaryDatabase._(this._pg, this._adapter, this.db);
+  PrimaryDatabase._(this._pg, this._adapter, this.db, this._closeFn);
 
   /// Gets the connection string either from the environment variable or from
   /// the secret backend, connects to it and registers the primary database
@@ -43,54 +46,71 @@ class PrimaryDatabase {
       // Production is not configured for postgresql yet.
       return;
     }
-    var connectionString =
+    if (_lookupPrimaryDatabase() != null) {
+      // Already initialized, must be in a local test environment.
+      assert(activeConfiguration.isFakeOrTest);
+      return;
+    }
+    final connectionString =
         envConfig.pubPostgresUrl ??
         (await secretBackend.lookup(SecretKey.postgresConnectionString));
     if (connectionString == null && activeConfiguration.isStaging) {
       // Staging may not have the connection string set yet.
       return;
     }
+    final database = await createAndInit(url: connectionString);
+    registerPrimaryDatabase(database);
+    ss.registerScopeExitCallback(database.close);
+  }
+
+  /// Creates and initializes a [PrimaryDatabase] instance.
+  ///
+  /// When [url] is not provided, it will start a new local postgresql instance, or
+  /// if it detects an existing one, connects to it.
+  ///
+  /// When NOT running in the AppEngine environment (e.g. testing or local fake),
+  /// the initilization will create a new database, which will be dropped when the
+  /// [close] method is called.
+  static Future<PrimaryDatabase> createAndInit({String? url}) async {
     // The scope-specific custom database. We are creating a custom database for
     // each test run, in order to provide full isolation, however, this must not
     // be used in Appengine.
     String? customDb;
-    if (connectionString == null) {
-      (connectionString, customDb) = await _startOrUseLocalPostgresInDocker();
+    if (url == null) {
+      (url, customDb) = await _startOrUseLocalPostgresInDocker();
     }
     if (customDb == null && !envConfig.isRunningInAppengine) {
-      customDb = await _createCustomDatabase(connectionString);
+      customDb = await _createCustomDatabase(url);
     }
 
+    final originalUrl = url;
     if (customDb != null) {
       if (envConfig.isRunningInAppengine) {
         throw StateError('Should not use custom database inside AppEngine.');
       }
 
-      final originalUrl = connectionString;
-      connectionString = Uri.parse(
-        connectionString,
-      ).replace(path: customDb).toString();
-      ss.registerScopeExitCallback(() async {
-        await _dropCustomDatabase(originalUrl, customDb!);
-      });
+      url = Uri.parse(url).replace(path: customDb).toString();
     }
 
-    final database = await _fromConnectionString(connectionString);
-    registerPrimaryDatabase(database);
-    ss.registerScopeExitCallback(database.close);
-  }
+    Future<void> closeFn() async {
+      if (customDb != null) {
+        await _dropCustomDatabase(originalUrl, customDb);
+      }
+    }
 
-  static Future<PrimaryDatabase> _fromConnectionString(String value) async {
-    final pg = Pool.withUrl(value);
+    final pg = Pool.withUrl(url);
     final adapter = DatabaseAdapter.postgres(pg);
     final db = Database<PrimarySchema>(adapter, SqlDialect.postgres());
     await db.createTables();
-    return PrimaryDatabase._(pg, adapter, db);
+    return PrimaryDatabase._(pg, adapter, db, closeFn);
   }
 
   Future<void> close() async {
     await _adapter.close();
     await _pg.close();
+    if (_closeFn != null) {
+      await _closeFn();
+    }
   }
 
   @visibleForTesting
diff --git a/app/lib/service/services.dart b/app/lib/service/services.dart
@@ -145,6 +145,7 @@ Future<R> withFakeServices<R>({
   MemDatastore? datastore,
   MemStorage? storage,
   FakeCloudCompute? cloudCompute,
+  PrimaryDatabase? primaryDatabase,
 }) async {
   if (!envConfig.isRunningLocally) {
     throw StateError("Mustn't use fake services inside AppEngine.");
@@ -156,6 +157,9 @@ Future<R> withFakeServices<R>({
         register(#appengine.context, FakeClientContext());
         registerDbService(DatastoreDB(datastore!));
         registerStorageService(RetryEnforcerStorage(storage!));
+        if (primaryDatabase != null) {
+          registerPrimaryDatabase(primaryDatabase);
+        }
         IOServer? frontendServer;
         IOServer? searchServer;
         if (configuration == null) {
diff --git a/app/test/shared/test_services.dart b/app/test/shared/test_services.dart
@@ -11,6 +11,7 @@ import 'package:gcloud/db.dart';
 import 'package:gcloud/service_scope.dart';
 import 'package:meta/meta.dart';
 import 'package:pub_dev/account/models.dart';
+import 'package:pub_dev/database/database.dart';
 import 'package:pub_dev/fake/backend/fake_auth_provider.dart';
 import 'package:pub_dev/fake/backend/fake_download_counts.dart';
 import 'package:pub_dev/fake/backend/fake_email_sender.dart';
@@ -77,6 +78,28 @@ final class FakeAppengineEnv {
   final _storage = MemStorage();
   final _datastore = MemDatastore();
   final _cloudCompute = FakeCloudCompute();
+  final PrimaryDatabase _primaryDatabase;
+
+  FakeAppengineEnv._(this._primaryDatabase);
+
+  /// Initializes, provides and then disposes a fake environment, preserving
+  /// the backing databases, allowing the use of new runtimes and other dynamic
+  /// features.
+  static Future<T> withEnv<T>(
+    Future<T> Function(FakeAppengineEnv env) fn,
+  ) async {
+    final database = await PrimaryDatabase.createAndInit();
+    final env = FakeAppengineEnv._(database);
+    try {
+      return await fn(env);
+    } finally {
+      await env._dispose();
+    }
+  }
+
+  Future<void> _dispose() async {
+    await _primaryDatabase.close();
+  }
 
   /// Create a service scope with fake services and run [fn] in it.
   ///
@@ -102,6 +125,7 @@ final class FakeAppengineEnv {
               datastore: _datastore,
               storage: _storage,
               cloudCompute: _cloudCompute,
+              primaryDatabase: _primaryDatabase,
               fn: () async {
                 registerStaticFileCacheForTest(_staticFileCacheForTesting);
 
@@ -171,19 +195,19 @@ void testWithProfile(
   Iterable<Pattern>? expectedLogMessages,
   dynamic skip,
 }) {
-  final env = FakeAppengineEnv();
-
   scopedTest(
     name,
     () async {
       setupDebugEnvBasedLogging();
-      await env.run(
-        fn,
-        testProfile: testProfile ?? defaultTestProfile,
-        importSource: importSource,
-        processJobsWithFakeRunners: processJobsWithFakeRunners,
-        integrityProblem: integrityProblem,
-      );
+      await FakeAppengineEnv.withEnv((env) async {
+        await env.run(
+          fn,
+          testProfile: testProfile ?? defaultTestProfile,
+          importSource: importSource,
+          processJobsWithFakeRunners: processJobsWithFakeRunners,
+          integrityProblem: integrityProblem,
+        );
+      });
     },
     expectedLogMessages: expectedLogMessages,
     timeout: timeout,
diff --git a/app/test/task/fallback_test.dart b/app/test/task/fallback_test.dart
@@ -16,38 +16,45 @@ import '../shared/test_services.dart';
 void main() {
   group('task fallback test', () {
     test('analysis fallback', () async {
-      final env = FakeAppengineEnv();
-      await env.run(
-        testProfile: TestProfile(
-          generatedPackages: [
-            GeneratedTestPackage(
-              name: 'oxygen',
-              versions: [GeneratedTestVersion(version: '1.0.0')],
-            ),
-          ],
-          defaultUser: adminAtPubDevEmail,
-        ),
-        processJobsWithFakeRunners: true,
-        runtimeVersions: ['2023.08.24'],
-        () async {
+      await FakeAppengineEnv.withEnv((env) async {
+        await env.run(
+          testProfile: TestProfile(
+            generatedPackages: [
+              GeneratedTestPackage(
+                name: 'oxygen',
+                versions: [GeneratedTestVersion(version: '1.0.0')],
+              ),
+            ],
+            defaultUser: adminAtPubDevEmail,
+          ),
+          processJobsWithFakeRunners: true,
+          runtimeVersions: ['2023.08.24'],
+          () async {
+            final card = await scoreCardBackend.getScoreCardData(
+              'oxygen',
+              '1.0.0',
+            );
+            expect(card.runtimeVersion, '2023.08.24');
+          },
+        );
+
+        await env.run(runtimeVersions: ['2023.08.25', '2023.08.24'], () async {
+          // fallback into accepted runtime works
           final card = await scoreCardBackend.getScoreCardData(
             'oxygen',
             '1.0.0',
           );
           expect(card.runtimeVersion, '2023.08.24');
-        },
-      );
-
-      await env.run(runtimeVersions: ['2023.08.25', '2023.08.24'], () async {
-        // fallback into accepted runtime works
-        final card = await scoreCardBackend.getScoreCardData('oxygen', '1.0.0');
-        expect(card.runtimeVersion, '2023.08.24');
-      });
+        });
 
-      await env.run(runtimeVersions: ['2023.08.26', '2023.08.23'], () async {
-        // fallback into non-accepted runtime doesn't work
-        final card = await scoreCardBackend.getScoreCardData('oxygen', '1.0.0');
-        expect(card.runtimeVersion, '2023.08.26');
+        await env.run(runtimeVersions: ['2023.08.26', '2023.08.23'], () async {
+          // fallback into non-accepted runtime doesn't work
+          final card = await scoreCardBackend.getScoreCardData(
+            'oxygen',
+            '1.0.0',
+          );
+          expect(card.runtimeVersion, '2023.08.26');
+        });
       });
     });
   });
