Add custom scalar/aggregate functions to SQLite.

SQLite ships with a handful of basic functions, such as UPPER() as
a scalar function and MAX() as a aggregate function.  We now have
several use-cases for adding custom functions, where it's otherwise
prohibitively expensive to perform post-processing on the returned
Cursor, as that requires copying processed data to yet another
MatrixCursor before returning to apps.

This change adds the ability for developers to register custom
scalar and aggregate functions on databases that they've opened;
some contrived examples are scalar functions like REVERSE() for
reversing a string, or aggregate functions like STDDEV().

To give developers the most flexibility, we use the Java functional
interfaces for defining these operations, as developers may already
be familiar with the contracts of those methods.  This also opens
the door to quickly adapting existing code through utility methods
like BinaryOperator.minBy(Comparator).

Bug: 142564473
Test: atest CtsDatabaseTestCases:android.database.sqlite.cts.SQLiteDatabaseTest
Change-Id: I9fa0e60ec77bab676396729cc9cb8ba8aaf56224

3 files changed, 120 insertions, 42 deletions

diff --git a/core/java/android/database/sqlite/SQLiteConnection.java b/core/java/android/database/sqlite/SQLiteConnection.java
index f7222750b89b..796cfdce2c0d 100644
--- a/core/java/android/database/sqlite/SQLiteConnection.java
+++ b/core/java/android/database/sqlite/SQLiteConnection.java
@@ -39,6 +39,8 @@ import java.text.SimpleDateFormat;
 import java.util.ArrayList;
 import java.util.Date;
 import java.util.Map;
+import java.util.function.BinaryOperator;
+import java.util.function.UnaryOperator;
 
 /**
  * Represents a SQLite database connection.
@@ -123,8 +125,10 @@ public final class SQLiteConnection implements CancellationSignal.OnCancelListen
             boolean enableTrace, boolean enableProfile, int lookasideSlotSize,
             int lookasideSlotCount);
     private static native void nativeClose(long connectionPtr);
-    private static native void nativeRegisterCustomFunction(long connectionPtr,
-            SQLiteCustomFunction function);
+    private static native void nativeRegisterCustomScalarFunction(long connectionPtr,
+            String name, UnaryOperator<String> function);
+    private static native void nativeRegisterCustomAggregateFunction(long connectionPtr,
+            String name, BinaryOperator<String> function);
     private static native void nativeRegisterLocalizedCollators(long connectionPtr, String locale);
     private static native long nativePrepareStatement(long connectionPtr, String sql);
     private static native void nativeFinalizeStatement(long connectionPtr, long statementPtr);
@@ -225,13 +229,7 @@ public final class SQLiteConnection implements CancellationSignal.OnCancelListen
         setJournalSizeLimit();
         setAutoCheckpointInterval();
         setLocaleFromConfiguration();
-
-        // Register custom functions.
-        final int functionCount = mConfiguration.customFunctions.size();
-        for (int i = 0; i < functionCount; i++) {
-            SQLiteCustomFunction function = mConfiguration.customFunctions.get(i);
-            nativeRegisterCustomFunction(mConnectionPtr, function);
-        }
+        setCustomFunctionsFromConfiguration();
     }
 
     private void dispose(boolean finalized) {
@@ -457,6 +455,19 @@ public final class SQLiteConnection implements CancellationSignal.OnCancelListen
         }
     }
 
+    private void setCustomFunctionsFromConfiguration() {
+        for (int i = 0; i < mConfiguration.customScalarFunctions.size(); i++) {
+            nativeRegisterCustomScalarFunction(mConnectionPtr,
+                    mConfiguration.customScalarFunctions.keyAt(i),
+                    mConfiguration.customScalarFunctions.valueAt(i));
+        }
+        for (int i = 0; i < mConfiguration.customAggregateFunctions.size(); i++) {
+            nativeRegisterCustomAggregateFunction(mConnectionPtr,
+                    mConfiguration.customAggregateFunctions.keyAt(i),
+                    mConfiguration.customAggregateFunctions.valueAt(i));
+        }
+    }
+
     private void checkDatabaseWiped() {
         if (!SQLiteGlobal.checkDbWipe()) {
             return;
@@ -491,15 +502,6 @@ public final class SQLiteConnection implements CancellationSignal.OnCancelListen
     void reconfigure(SQLiteDatabaseConfiguration configuration) {
         mOnlyAllowReadOnlyOperations = false;
 
-        // Register custom functions.
-        final int functionCount = configuration.customFunctions.size();
-        for (int i = 0; i < functionCount; i++) {
-            SQLiteCustomFunction function = configuration.customFunctions.get(i);
-            if (!mConfiguration.customFunctions.contains(function)) {
-                nativeRegisterCustomFunction(mConnectionPtr, function);
-            }
-        }
-
         // Remember what changed.
         boolean foreignKeyModeChanged = configuration.foreignKeyConstraintsEnabled
                 != mConfiguration.foreignKeyConstraintsEnabled;
@@ -507,6 +509,10 @@ public final class SQLiteConnection implements CancellationSignal.OnCancelListen
                 & (SQLiteDatabase.ENABLE_WRITE_AHEAD_LOGGING
                 | SQLiteDatabase.ENABLE_LEGACY_COMPATIBILITY_WAL)) != 0;
         boolean localeChanged = !configuration.locale.equals(mConfiguration.locale);
+        boolean customScalarFunctionsChanged = !configuration.customScalarFunctions
+                .equals(mConfiguration.customScalarFunctions);
+        boolean customAggregateFunctionsChanged = !configuration.customAggregateFunctions
+                .equals(mConfiguration.customAggregateFunctions);
 
         // Update configuration parameters.
         mConfiguration.updateParametersFrom(configuration);
@@ -514,20 +520,18 @@ public final class SQLiteConnection implements CancellationSignal.OnCancelListen
         // Update prepared statement cache size.
         mPreparedStatementCache.resize(configuration.maxSqlCacheSize);
 
-        // Update foreign key mode.
         if (foreignKeyModeChanged) {
             setForeignKeyModeFromConfiguration();
         }
-
-        // Update WAL.
         if (walModeChanged) {
             setWalModeFromConfiguration();
         }
-
-        // Update locale.
         if (localeChanged) {
             setLocaleFromConfiguration();
         }
+        if (customScalarFunctionsChanged || customAggregateFunctionsChanged) {
+            setCustomFunctionsFromConfiguration();
+        }
     }
 
     // Called by SQLiteConnectionPool only.
diff --git a/core/java/android/database/sqlite/SQLiteDatabase.java b/core/java/android/database/sqlite/SQLiteDatabase.java
index 44c78aa783a7..458914efcbbd 100644
--- a/core/java/android/database/sqlite/SQLiteDatabase.java
+++ b/core/java/android/database/sqlite/SQLiteDatabase.java
@@ -62,6 +62,8 @@ import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
 import java.util.WeakHashMap;
+import java.util.function.BinaryOperator;
+import java.util.function.UnaryOperator;
 
 /**
  * Exposes methods to manage a SQLite database.
@@ -958,26 +960,87 @@ public final class SQLiteDatabase extends SQLiteClosable {
     }
 
     /**
-     * Registers a CustomFunction callback as a function that can be called from
-     * SQLite database triggers.
-     *
-     * @param name the name of the sqlite3 function
-     * @param numArgs the number of arguments for the function
-     * @param function callback to call when the function is executed
-     * @hide
-     */
-    public void addCustomFunction(String name, int numArgs, CustomFunction function) {
-        // Create wrapper (also validates arguments).
-        SQLiteCustomFunction wrapper = new SQLiteCustomFunction(name, numArgs, function);
+     * Register a custom scalar function that can be called from SQL
+     * expressions.
+     * <p>
+     * For example, registering a custom scalar function named {@code REVERSE}
+     * could be used in a query like
+     * {@code SELECT REVERSE(name) FROM employees}.
+     * <p>
+     * When attempting to register multiple functions with the same function
+     * name, SQLite will replace any previously defined functions with the
+     * latest definition, regardless of what function type they are. SQLite does
+     * not support unregistering functions.
+     *
+     * @param functionName Case-insensitive name to register this function
+     *            under, limited to 255 UTF-8 bytes in length.
+     * @param scalarFunction Functional interface that will be invoked when the
+     *            function name is used by a SQL statement. The argument values
+     *            from the SQL statement are passed to the functional interface,
+     *            and the return values from the functional interface are
+     *            returned back into the SQL statement.
+     * @throws SQLiteException if the custom function could not be registered.
+     * @see #setCustomAggregateFunction(String, BinaryOperator)
+     */
+    public void setCustomScalarFunction(@NonNull String functionName,
+            @NonNull UnaryOperator<String> scalarFunction) throws SQLiteException {
+        Objects.requireNonNull(functionName);
+        Objects.requireNonNull(scalarFunction);
+
+        synchronized (mLock) {
+            throwIfNotOpenLocked();
+
+            mConfigurationLocked.customScalarFunctions.put(functionName, scalarFunction);
+            try {
+                mConnectionPoolLocked.reconfigure(mConfigurationLocked);
+            } catch (RuntimeException ex) {
+                mConfigurationLocked.customScalarFunctions.remove(functionName);
+                throw ex;
+            }
+        }
+    }
+
+    /**
+     * Register a custom aggregate function that can be called from SQL
+     * expressions.
+     * <p>
+     * For example, registering a custom aggregation function named
+     * {@code LONGEST} could be used in a query like
+     * {@code SELECT LONGEST(name) FROM employees}.
+     * <p>
+     * The implementation of this method follows the reduction flow outlined in
+     * {@link java.util.stream.Stream#reduce(BinaryOperator)}, and the custom
+     * aggregation function is expected to be an associative accumulation
+     * function, as defined by that class.
+     * <p>
+     * When attempting to register multiple functions with the same function
+     * name, SQLite will replace any previously defined functions with the
+     * latest definition, regardless of what function type they are. SQLite does
+     * not support unregistering functions.
+     *
+     * @param functionName Case-insensitive name to register this function
+     *            under, limited to 255 UTF-8 bytes in length.
+     * @param aggregateFunction Functional interface that will be invoked when
+     *            the function name is used by a SQL statement. The argument
+     *            values from the SQL statement are passed to the functional
+     *            interface, and the return values from the functional interface
+     *            are returned back into the SQL statement.
+     * @throws SQLiteException if the custom function could not be registered.
+     * @see #setCustomScalarFunction(String, UnaryOperator)
+     */
+    public void setCustomAggregateFunction(@NonNull String functionName,
+            @NonNull BinaryOperator<String> aggregateFunction) throws SQLiteException {
+        Objects.requireNonNull(functionName);
+        Objects.requireNonNull(aggregateFunction);
 
         synchronized (mLock) {
             throwIfNotOpenLocked();
 
-            mConfigurationLocked.customFunctions.add(wrapper);
+            mConfigurationLocked.customAggregateFunctions.put(functionName, aggregateFunction);
             try {
                 mConnectionPoolLocked.reconfigure(mConfigurationLocked);
             } catch (RuntimeException ex) {
-                mConfigurationLocked.customFunctions.remove(wrapper);
+                mConfigurationLocked.customAggregateFunctions.remove(functionName);
                 throw ex;
             }
         }
diff --git a/core/java/android/database/sqlite/SQLiteDatabaseConfiguration.java b/core/java/android/database/sqlite/SQLiteDatabaseConfiguration.java
index 6a52b72a9e1c..b11942abe0c7 100644
--- a/core/java/android/database/sqlite/SQLiteDatabaseConfiguration.java
+++ b/core/java/android/database/sqlite/SQLiteDatabaseConfiguration.java
@@ -17,9 +17,12 @@
 package android.database.sqlite;
 
 import android.compat.annotation.UnsupportedAppUsage;
+import android.util.ArrayMap;
 
-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;
 
 /**
@@ -87,10 +90,16 @@ public final class SQLiteDatabaseConfiguration {
     public boolean foreignKeyConstraintsEnabled;
 
     /**
-     * The custom functions to register.
+     * The custom scalar functions to register.
      */
-    public final ArrayList<SQLiteCustomFunction> customFunctions =
-            new ArrayList<SQLiteCustomFunction>();
+    public final ArrayMap<String, UnaryOperator<String>> customScalarFunctions
+            = new ArrayMap<>();
+
+    /**
+     * The custom aggregate functions to register.
+     */
+    public final ArrayMap<String, BinaryOperator<String>> customAggregateFunctions
+            = new ArrayMap<>();
 
     /**
      * The size in bytes of each lookaside slot
@@ -181,8 +190,10 @@ public final class SQLiteDatabaseConfiguration {
         maxSqlCacheSize = other.maxSqlCacheSize;
         locale = other.locale;
         foreignKeyConstraintsEnabled = other.foreignKeyConstraintsEnabled;
-        customFunctions.clear();
-        customFunctions.addAll(other.customFunctions);
+        customScalarFunctions.clear();
+        customScalarFunctions.putAll(other.customScalarFunctions);
+        customAggregateFunctions.clear();
+        customAggregateFunctions.putAll(other.customAggregateFunctions);
         lookasideSlotSize = other.lookasideSlotSize;
         lookasideSlotCount = other.lookasideSlotCount;
         idleConnectionTimeoutMs = other.idleConnectionTimeoutMs;