Allow migration of entity between shards

misk-hibernate/src/main/kotlin/misk/hibernate/CowritesDetector.kt

@@ -0,0 +1,15 @@
+package misk.hibernate
+
+interface CowritesDetector {
+
+  fun pushSuppressChecks()
+  fun popSuppressChecks()
+
+  companion object {
+    val NONE: CowritesDetector = object : CowritesDetector {
+      override fun pushSuppressChecks() {}
+
+      override fun popSuppressChecks() {}
+    }
+  }
+}

misk-hibernate/src/main/kotlin/misk/hibernate/DbChild.kt

@@ -9,4 +9,12 @@ package misk.hibernate
 interface DbChild<R : DbRoot<R>, T : DbChild<R, T>> : DbSharded<R, T> {
   override val id: Id<T>
     get() = gid.id
+
+  /**
+   * Create a new copy of this entity in the target entity group (possibly different shard).
+   */
+  fun migrate(shardMigrator: ShardMigrator, session: Session, newParentId: Id<R>): T {
+    throw UnsupportedOperationException(
+        "${this::class.qualifiedName} can't be migrated yet. You need to implement this method.")
+  }
 }

misk-hibernate/src/main/kotlin/misk/hibernate/DbRoot.kt

@@ -11,4 +11,9 @@ interface DbRoot<T : DbRoot<T>> : DbSharded<T, T> {
 
   override val rootId: Id<T>
     get() = id
+
+  /**
+   * Increase the version to trigger an "empty" update which forces Vitess to start a transaction.
+   */
+  abstract fun incrementEditCount()
 }

misk-hibernate/src/main/kotlin/misk/hibernate/EntityGroupTransactionOrder.kt

@@ -0,0 +1,44 @@
+package misk.hibernate
+
+/**
+ * Vitess multi shard transactions are guaranteed to be committed in the order that updates are
+ * issued. So if you have entity groups A and B, you issue an update against A followed by an
+ * update to B, then the transaction against the shard of A will be fully committed before the
+ * transaction against the shard of B. You can use this fact to make cross shard transactions safe
+ * by committing clean up tasks to the earlier shards.
+ */
+interface EntityGroupTransactionOrder {
+  /**
+   * Issues a "spurious"/empty update on the entity group root to enforce a transaction to be
+   * started on that entity group.
+   */
+  fun <P: DbRoot<P>> start(session: Session, entityGroup: P)
+
+  /**
+   * Checks that the transaction of the "leader" entity group will commit before the transaction of
+   * the "follower" entity group.
+   */
+  fun <P: DbRoot<P>> assertOrder(session: Session, leader: Id<P>, follower: Id<P>)
+
+  /**
+   * Suppress checks.
+   *
+   * Should be used like this to suppress checks of code that we deliberately do not need to be
+   * safe:
+   * <pre>
+   * transactionOrder.pushSuppressChecks();
+   * try {
+   * // unsafe code
+   * } finally {
+   * transactionOrder.popSuppressChecks();
+   * }
+  </pre> *
+   */
+  fun pushSuppressChecks()
+
+  /**
+   * Enable checks again.
+   */
+  fun popSuppressChecks()
+}
+

misk-hibernate/src/main/kotlin/misk/hibernate/EntityGroups.kt

@@ -0,0 +1,11 @@
+package misk.hibernate
+
+class EntityGroups {
+  fun entityGroupId(entity: Any): Id<*>? {
+    return when (entity) {
+      is DbRoot<*> -> entity.id
+      is DbChild<*, *> -> entity.rootId
+      else -> null
+    }
+  }
+}

misk-hibernate/src/main/kotlin/misk/hibernate/PersistenceMetadata.kt

@@ -0,0 +1,20 @@
+package misk.hibernate
+
+import com.google.common.collect.ImmutableSet
+import com.google.inject.ImplementedBy
+
+/** Persistence related metadata  */
+@ImplementedBy(PersistenceMetadataImpl::class)
+interface PersistenceMetadata {
+  /** Gets the table name for the given class  */
+  fun <T: DbEntity<T>> getTableName(entityType: T): String
+
+  /** Returns all of the columns of `entityType`.  */
+  fun <T: DbEntity<T>> getColumnNames(entityType: T): ImmutableSet<String>
+
+  /**
+   * Gets the column names for the given class and property name. Multiple names
+   * can be returned since single properties can be mapped to multiple columns.
+   */
+  fun <T: DbEntity<T>> getColumnNames(entityType: T, propertyName: String): Array<String>
+}

misk-hibernate/src/main/kotlin/misk/hibernate/PersistenceMetadataImpl.kt

@@ -0,0 +1,59 @@
+package misk.hibernate
+
+import com.google.common.collect.ImmutableSet
+import javax.inject.Inject
+import org.hibernate.SessionFactory
+import org.hibernate.persister.entity.AbstractEntityPersister
+
+import com.google.common.base.Preconditions.checkState
+
+internal class PersistenceMetadataImpl @Inject constructor(
+  private val sessionFactory: SessionFactory
+) : PersistenceMetadata {
+
+  override fun <T: DbEntity<T>> getTableName(entityType: T): String {
+    return hibernateMetadataForClass(entityType).tableName
+  }
+
+  override fun <T: DbEntity<T>> getColumnNames(
+    entityType: T
+  ): ImmutableSet<String> {
+    val result = ImmutableSet.builder<String>()
+    val classMetadata = hibernateMetadataForClass(entityType)
+
+    for (column in classMetadata.identifierColumnNames) {
+      result.add(column)
+    }
+
+    val propertyNames = classMetadata.propertyNames
+
+    for (i in propertyNames.indices) {
+      result.add(*classMetadata.getPropertyColumnNames(i))
+    }
+
+    return result.build()
+  }
+
+  override fun <T: DbEntity<T>> getColumnNames(
+    entityType: T,
+    propertyName: String
+  ): Array<String> {
+    return hibernateMetadataForClass(entityType).getPropertyColumnNames(propertyName)
+  }
+
+  private fun <T: DbEntity<T>> hibernateMetadataForClass(
+    entityType: T
+  ): AbstractEntityPersister {
+    val hibernateMetadata = sessionFactory.getClassMetadata(entityType::class.qualifiedName)
+
+    checkState(hibernateMetadata != null,
+        "${entityType::class.qualifiedName} does not map to a known entity type"
+    )
+
+    checkState(hibernateMetadata is AbstractEntityPersister,
+        "${entityType::class.qualifiedName} does not map to a persistent class"
+    )
+
+    return hibernateMetadata as AbstractEntityPersister
+  }
+}

misk-hibernate/src/main/kotlin/misk/hibernate/RealTransacter.kt

@@ -208,6 +208,38 @@ internal class RealTransacter private constructor(
       } as Id<T>
     }
 
+    override fun <T : DbEntity<T>> delete(entity: T) {
+      if (readOnly) {
+        throw IllegalStateException("Deleting isn't permitted in a read only session.")
+      } else {
+        when (entity) {
+          is DbChild<*, *> -> session.delete(entity)
+          is DbRoot<*> -> session.delete(entity)
+          is DbUnsharded<*> -> session.delete(entity)
+          else -> throw IllegalArgumentException(
+              "You need to sub-class one of [DbChild, DbRoot, DbUnsharded]")
+        }
+      }
+    }
+
+    override fun <T : DbEntity<T>> update(entity: T) {
+      if (readOnly) {
+        throw IllegalStateException("Updating isn't permitted in a read only session.")
+      } else {
+        when (entity) {
+          is DbChild<*, *> -> session.update(entity)
+          is DbRoot<*> -> session.update(entity)
+          is DbUnsharded<*> -> session.update(entity)
+          else -> throw IllegalArgumentException(
+              "You need to sub-class one of [DbChild, DbRoot, DbUnsharded]")
+        }
+      }
+    }
+
+    override fun flush() {
+      session.flush()
+    }
+
     override fun <T : DbEntity<T>> load(id: Id<T>, type: KClass<T>): T {
       return session.get(type.java, id)
     }

misk-hibernate/src/main/kotlin/misk/hibernate/Session.kt

@@ -9,6 +9,9 @@ interface Session {
    * @throws IllegalStateException when save is called on a read only session.
    */
   fun <T : DbEntity<T>> save(entity: T): Id<T>
+  fun <T: DbEntity<T>> delete(entity: T)
+  fun <T: DbEntity<T>> update(entity: T)
+  fun flush()
 
   fun <T : DbEntity<T>> load(id: Id<T>, type: KClass<T>): T
   fun <R : DbRoot<R>, T : DbSharded<R, T>> loadSharded(gid: Gid<R, T>, type: KClass<T>): T

misk-hibernate/src/main/kotlin/misk/hibernate/ShardMigrator.kt

@@ -0,0 +1,40 @@
+package misk.hibernate
+
+import javax.inject.Inject
+
+/**
+ * This class changes the parent id of an entity group child (the customer id, for example).
+ *
+ * This may result in the entity group child migrating between shards as the child will always
+ * end up on the same shard as the parent.
+ */
+class ShardMigrator {
+  @Inject private lateinit var transacter: Transacter
+  @Inject private lateinit var transactionOrderListener: EntityGroupTransactionOrder
+
+  fun <P : DbRoot<P>, C : DbChild<P, C>> migrate(
+    child: C,
+    oldParentId: Id<P>,
+    newParentId: Id<P>
+  ): C {
+    return transacter.transaction {
+      transactionOrderListener!!.assertOrder(it, newParentId, oldParentId)
+      updateParent(it, child, newParentId)
+    }
+  }
+
+  private fun <P : DbRoot<P>, C : DbChild<P, C>> updateParent(
+    session: Session,
+    child: C,
+    newParentId: Id<P>
+  ): C {
+    return transacter.transaction {
+      val newChild = child.migrate(this, it, newParentId)
+      session.delete(child)
+      session.flush()
+      session.save(newChild)
+      newChild
+    }
+  }
+}
+