nmode's Git Repositories - signal-cli/blob - lib/src/main/java/org/asamk/signal/manager/storage/keyValue/KeyValueStore.java

   1 package org.asamk.signal.manager.storage.keyValue;
   2
   3 import org.asamk.signal.manager.storage.Database;
   4 import org.asamk.signal.manager.storage.Utils;
   5 import org.slf4j.Logger;
   6 import org.slf4j.LoggerFactory;
   7
   8 import java.sql.Connection;
   9 import java.sql.PreparedStatement;
  10 import java.sql.ResultSet;
  11 import java.sql.SQLException;
  12 import java.sql.Types;
  13
  14 public class KeyValueStore {
  15
  16     private static final String TABLE_KEY_VALUE = "key_value";
  17     private static final Logger logger = LoggerFactory.getLogger(KeyValueStore.class);
  18
  19     private final Database database;
  20
  21     public static void createSql(Connection connection) throws SQLException {
  22         // When modifying the CREATE statement here, also add a migration in AccountDatabase.java
  23         try (final var statement = connection.createStatement()) {
  24             statement.executeUpdate("""
  25                                     CREATE TABLE key_value (
  26                                       _id INTEGER PRIMARY KEY,
  27                                       key TEXT UNIQUE NOT NULL,
  28                                       value ANY
  29                                     ) STRICT;
  30                                     """);
  31         }
  32     }
  33
  34     public KeyValueStore(final Database database) {
  35         this.database = database;
  36     }
  37
  38     public <T> T getEntry(KeyValueEntry<T> key) {
  39         try (final var connection = database.getConnection()) {
  40             return getEntry(connection, key);
  41         } catch (SQLException e) {
  42             throw new RuntimeException("Failed read from pre_key store", e);
  43         }
  44     }
  45
  46     public <T> void storeEntry(KeyValueEntry<T> key, T value) {
  47         try (final var connection = database.getConnection()) {
  48             storeEntry(connection, key, value);
  49         } catch (SQLException e) {
  50             throw new RuntimeException("Failed update key_value store", e);
  51         }
  52     }
  53
  54     private <T> T getEntry(final Connection connection, final KeyValueEntry<T> key) throws SQLException {
  55         final var sql = (
  56                 """
  57                 SELECT key, value
  58                 FROM %s p
  59                 WHERE p.key = ?
  60                 """
  61         ).formatted(TABLE_KEY_VALUE);
  62         try (final var statement = connection.prepareStatement(sql)) {
  63             statement.setString(1, key.key());
  64
  65             final var result = Utils.executeQueryForOptional(statement,
  66                     resultSet -> readValueFromResultSet(key, resultSet)).orElse(null);
  67
  68             if (result == null) {
  69                 return key.defaultValue();
  70             }
  71             return result;
  72         }
  73     }
  74
  75     private <T> void storeEntry(
  76             final Connection connection, final KeyValueEntry<T> key, final T value
  77     ) throws SQLException {
  78         final var sql = (
  79                 """
  80                 INSERT INTO %s (key, value)
  81                 VALUES (?1, ?2)
  82                 ON CONFLICT (key) DO UPDATE SET value=excluded.value
  83                 """
  84         ).formatted(TABLE_KEY_VALUE);
  85         try (final var statement = connection.prepareStatement(sql)) {
  86             statement.setString(1, key.key());
  87             setParameterValue(statement, 2, key.clazz(), value);
  88             statement.executeUpdate();
  89         }
  90     }
  91
  92     @SuppressWarnings("unchecked")
  93     private static <T> T readValueFromResultSet(
  94             final KeyValueEntry<T> key, final ResultSet resultSet
  95     ) throws SQLException {
  96         Object value;
  97         final var clazz = key.clazz();
  98         if (clazz == int.class || clazz == Integer.class) {
  99             value = resultSet.getInt("value");
 100         } else if (clazz == long.class || clazz == Long.class) {
 101             value = resultSet.getLong("value");
 102         } else if (clazz == boolean.class || clazz == Boolean.class) {
 103             value = resultSet.getBoolean("value");
 104         } else if (clazz == byte[].class || clazz == Byte[].class) {
 105             value = resultSet.getBytes("value");
 106         } else if (clazz == String.class) {
 107             value = resultSet.getString("value");
 108         } else if (Enum.class.isAssignableFrom(clazz)) {
 109             final var name = resultSet.getString("value");
 110             if (name == null) {
 111                 value = null;
 112             } else {
 113                 try {
 114                     value = Enum.valueOf((Class<Enum>) key.clazz(), name);
 115                 } catch (IllegalArgumentException e) {
 116                     logger.debug("Read invalid enum value from store, ignoring: {} for {}", name, key.clazz());
 117                     value = null;
 118                 }
 119             }
 120         } else {
 121             throw new AssertionError("Invalid key type " + clazz.getSimpleName());
 122         }
 123         if (resultSet.wasNull()) {
 124             return null;
 125         }
 126         return (T) value;
 127     }
 128
 129     private static <T> void setParameterValue(
 130             final PreparedStatement statement, final int parameterIndex, final Class<T> clazz, final T value
 131     ) throws SQLException {
 132         if (clazz == int.class || clazz == Integer.class) {
 133             if (value == null) {
 134                 statement.setNull(parameterIndex, Types.INTEGER);
 135             } else {
 136                 statement.setInt(parameterIndex, (int) value);
 137             }
 138         } else if (clazz == long.class || clazz == Long.class) {
 139             if (value == null) {
 140                 statement.setNull(parameterIndex, Types.INTEGER);
 141             } else {
 142                 statement.setLong(parameterIndex, (long) value);
 143             }
 144         } else if (clazz == boolean.class || clazz == Boolean.class) {
 145             if (value == null) {
 146                 statement.setNull(parameterIndex, Types.BOOLEAN);
 147             } else {
 148                 statement.setBoolean(parameterIndex, (boolean) value);
 149             }
 150         } else if (clazz == byte[].class || clazz == Byte[].class) {
 151             if (value == null) {
 152                 statement.setNull(parameterIndex, Types.BLOB);
 153             } else {
 154                 statement.setBytes(parameterIndex, (byte[]) value);
 155             }
 156         } else if (clazz == String.class) {
 157             if (value == null) {
 158                 statement.setNull(parameterIndex, Types.VARCHAR);
 159             } else {
 160                 statement.setString(parameterIndex, (String) value);
 161             }
 162         } else if (Enum.class.isAssignableFrom(clazz)) {
 163             if (value == null) {
 164                 statement.setNull(parameterIndex, Types.VARCHAR);
 165             } else {
 166                 statement.setString(parameterIndex, ((Enum<?>) value).name());
 167             }
 168         } else {
 169             throw new AssertionError("Invalid key type " + clazz.getSimpleName());
 170         }
 171     }
 172 }