Add ability to execute per-connection SQL.

Developers have been able to register custom collators using syntax
like "SELECT icu_load_collation()", but collators are registered per
database connection.

Since we don't expose any details APIs for interacting with connection
pools directly, developers can end up with flaky behavior as their
queries rotate through the pool of connections, as only a subset of
connections will have their collation registered.

This solve this, we add a new execPerConnectionSQL() method to
ensure that a given statement is executed on all current and future
database connections.

Bug: 152005629
Test: atest CtsDatabaseTestCases:android.database.sqlite.cts.SQLiteDatabaseTest
Change-Id: I459fb7b18660d2a04eec92d1e9cc410d769e361d

3 files changed, 81 insertions, 1 deletions

diff --git a/core/java/android/database/sqlite/SQLiteConnection.java b/core/java/android/database/sqlite/SQLiteConnection.java
index 796cfdce2c0d..bcb3934a5b08 100644
--- a/core/java/android/database/sqlite/SQLiteConnection.java
+++ b/core/java/android/database/sqlite/SQLiteConnection.java
@@ -28,6 +28,7 @@ import android.os.SystemClock;
 import android.os.Trace;
 import android.util.Log;
 import android.util.LruCache;
+import android.util.Pair;
 import android.util.Printer;
 
 import dalvik.system.BlockGuard;
@@ -230,6 +231,7 @@ public final class SQLiteConnection implements CancellationSignal.OnCancelListen
         setAutoCheckpointInterval();
         setLocaleFromConfiguration();
         setCustomFunctionsFromConfiguration();
+        executePerConnectionSqlFromConfiguration(0);
     }
 
     private void dispose(boolean finalized) {
@@ -468,6 +470,24 @@ public final class SQLiteConnection implements CancellationSignal.OnCancelListen
         }
     }
 
+    private void executePerConnectionSqlFromConfiguration(int startIndex) {
+        for (int i = startIndex; i < mConfiguration.perConnectionSql.size(); i++) {
+            final Pair<String, Object[]> statement = mConfiguration.perConnectionSql.get(i);
+            final int type = DatabaseUtils.getSqlStatementType(statement.first);
+            switch (type) {
+                case DatabaseUtils.STATEMENT_SELECT:
+                    executeForString(statement.first, statement.second, null);
+                    break;
+                case DatabaseUtils.STATEMENT_PRAGMA:
+                    execute(statement.first, statement.second, null);
+                    break;
+                default:
+                    throw new IllegalArgumentException(
+                            "Unsupported configuration statement: " + statement);
+            }
+        }
+    }
+
     private void checkDatabaseWiped() {
         if (!SQLiteGlobal.checkDbWipe()) {
             return;
@@ -513,6 +533,9 @@ public final class SQLiteConnection implements CancellationSignal.OnCancelListen
                 .equals(mConfiguration.customScalarFunctions);
         boolean customAggregateFunctionsChanged = !configuration.customAggregateFunctions
                 .equals(mConfiguration.customAggregateFunctions);
+        final int oldSize = mConfiguration.perConnectionSql.size();
+        final int newSize = configuration.perConnectionSql.size();
+        boolean perConnectionSqlChanged = newSize > oldSize;
 
         // Update configuration parameters.
         mConfiguration.updateParametersFrom(configuration);
@@ -532,6 +555,9 @@ public final class SQLiteConnection implements CancellationSignal.OnCancelListen
         if (customScalarFunctionsChanged || customAggregateFunctionsChanged) {
             setCustomFunctionsFromConfiguration();
         }
+        if (perConnectionSqlChanged) {
+            executePerConnectionSqlFromConfiguration(oldSize);
+        }
     }
 
     // Called by SQLiteConnectionPool only.
diff --git a/core/java/android/database/sqlite/SQLiteDatabase.java b/core/java/android/database/sqlite/SQLiteDatabase.java
index 458914efcbbd..24ac1527779e 100644
--- a/core/java/android/database/sqlite/SQLiteDatabase.java
+++ b/core/java/android/database/sqlite/SQLiteDatabase.java
@@ -1047,6 +1047,40 @@ public final class SQLiteDatabase extends SQLiteClosable {
     }
 
     /**
+     * Execute the given SQL statement on all connections to this database.
+     * <p>
+     * This statement will be immediately executed on all existing connections,
+     * and will be automatically executed on all future connections.
+     * <p>
+     * Some example usages are changes like {@code PRAGMA trusted_schema=OFF} or
+     * functions like {@code SELECT icu_load_collation()}. If you execute these
+     * statements using {@link #execSQL} then they will only apply to a single
+     * database connection; using this method will ensure that they are
+     * uniformly applied to all current and future connections.
+     *
+     * @param sql The SQL statement to be executed. Multiple statements
+     *            separated by semicolons are not supported.
+     * @param bindArgs The arguments that should be bound to the SQL statement.
+     */
+    public void execPerConnectionSQL(@NonNull String sql, @Nullable Object[] bindArgs)
+            throws SQLException {
+        Objects.requireNonNull(sql);
+
+        synchronized (mLock) {
+            throwIfNotOpenLocked();
+
+            final int index = mConfigurationLocked.perConnectionSql.size();
+            mConfigurationLocked.perConnectionSql.add(Pair.create(sql, bindArgs));
+            try {
+                mConnectionPoolLocked.reconfigure(mConfigurationLocked);
+            } catch (RuntimeException ex) {
+                mConfigurationLocked.perConnectionSql.remove(index);
+                throw ex;
+            }
+        }
+    }
+
+    /**
      * Gets the database version.
      *
      * @return the database version
@@ -1788,6 +1822,12 @@ public final class SQLiteDatabase extends SQLiteClosable {
      * using "PRAGMA journal_mode'<value>" statement if your app is using
      * {@link #enableWriteAheadLogging()}
      * </p>
+     * <p>
+     * Note that {@code PRAGMA} values which apply on a per-connection basis
+     * should <em>not</em> be configured using this method; you should instead
+     * use {@link #execPerConnectionSQL} to ensure that they are uniformly
+     * applied to all current and future connections.
+     * </p>
      *
      * @param sql the SQL statement to be executed. Multiple statements separated by semicolons are
      * not supported.
@@ -1834,6 +1874,12 @@ public final class SQLiteDatabase extends SQLiteClosable {
      * using "PRAGMA journal_mode'<value>" statement if your app is using
      * {@link #enableWriteAheadLogging()}
      * </p>
+     * <p>
+     * Note that {@code PRAGMA} values which apply on a per-connection basis
+     * should <em>not</em> be configured using this method; you should instead
+     * use {@link #execPerConnectionSQL} to ensure that they are uniformly
+     * applied to all current and future connections.
+     * </p>
      *
      * @param sql the SQL statement to be executed. Multiple statements separated by semicolons are
      * not supported.
diff --git a/core/java/android/database/sqlite/SQLiteDatabaseConfiguration.java b/core/java/android/database/sqlite/SQLiteDatabaseConfiguration.java
index b11942abe0c7..21c21c902fed 100644
--- a/core/java/android/database/sqlite/SQLiteDatabaseConfiguration.java
+++ b/core/java/android/database/sqlite/SQLiteDatabaseConfiguration.java
@@ -18,9 +18,10 @@ package android.database.sqlite;
 
 import android.compat.annotation.UnsupportedAppUsage;
 import android.util.ArrayMap;
+import android.util.Pair;
 
+import java.util.ArrayList;
 import java.util.Locale;
-import java.util.Map;
 import java.util.function.BinaryOperator;
 import java.util.function.UnaryOperator;
 import java.util.regex.Pattern;
@@ -102,6 +103,11 @@ public final class SQLiteDatabaseConfiguration {
             = new ArrayMap<>();
 
     /**
+     * The statements to execute to initialize each connection.
+     */
+    public final ArrayList<Pair<String, Object[]>> perConnectionSql = new ArrayList<>();
+
+    /**
      * The size in bytes of each lookaside slot
      *
      * <p>If negative, the default lookaside configuration will be used
@@ -194,6 +200,8 @@ public final class SQLiteDatabaseConfiguration {
         customScalarFunctions.putAll(other.customScalarFunctions);
         customAggregateFunctions.clear();
         customAggregateFunctions.putAll(other.customAggregateFunctions);
+        perConnectionSql.clear();
+        perConnectionSql.addAll(other.perConnectionSql);
         lookasideSlotSize = other.lookasideSlotSize;
         lookasideSlotCount = other.lookasideSlotCount;
         idleConnectionTimeoutMs = other.idleConnectionTimeoutMs;