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
diff --git a/api/current.txt b/api/current.txt
index ff74ce8..d706754 100644
--- a/api/current.txt
+++ b/api/current.txt
@@ -13391,6 +13391,7 @@
     method public void disableWriteAheadLogging();
     method public boolean enableWriteAheadLogging();
     method public void endTransaction();
+    method public void execPerConnectionSQL(@NonNull String, @Nullable Object[]) throws android.database.SQLException;
     method public void execSQL(String) throws android.database.SQLException;
     method public void execSQL(String, Object[]) throws android.database.SQLException;
     method public static String findEditTable(String);
diff --git a/core/java/android/database/sqlite/SQLiteConnection.java b/core/java/android/database/sqlite/SQLiteConnection.java
index 796cfdc..bcb3934 100644
--- a/core/java/android/database/sqlite/SQLiteConnection.java
+++ b/core/java/android/database/sqlite/SQLiteConnection.java
@@ -28,6 +28,7 @@
 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 @@
         setAutoCheckpointInterval();
         setLocaleFromConfiguration();
         setCustomFunctionsFromConfiguration();
+        executePerConnectionSqlFromConfiguration(0);
     }
 
     private void dispose(boolean finalized) {
@@ -468,6 +470,24 @@
         }
     }
 
+    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 @@
                 .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 @@
         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 458914e..24ac152 100644
--- a/core/java/android/database/sqlite/SQLiteDatabase.java
+++ b/core/java/android/database/sqlite/SQLiteDatabase.java
@@ -1047,6 +1047,40 @@
     }
 
     /**
+     * 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 @@
      * 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 @@
      * 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 b11942a..21c21c9 100644
--- a/core/java/android/database/sqlite/SQLiteDatabaseConfiguration.java
+++ b/core/java/android/database/sqlite/SQLiteDatabaseConfiguration.java
@@ -18,9 +18,10 @@
 
 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 @@
             = 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 @@
         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;