From 8f7e0c4c2fdcd36b1a0e9e3f0131442f0891dae6 Mon Sep 17 00:00:00 2001
From: Gavin King <gavin@hibernate.org>
Date: Mon, 13 Sep 2021 13:39:04 +0200
Subject: [PATCH] HHH-6054 introduce notion of a "root" tenant

---
 .../spi/CurrentTenantIdentifierResolver.java  |  9 ++++++++
 .../org/hibernate/internal/SessionImpl.java   | 11 ++++++---
 .../hibernate/tuple/TenantIdGeneration.java   | 23 ++++++++++++++-----
 3 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/hibernate-core/src/main/java/org/hibernate/context/spi/CurrentTenantIdentifierResolver.java b/hibernate-core/src/main/java/org/hibernate/context/spi/CurrentTenantIdentifierResolver.java
index dc8723ed29..f8d9d21304 100644
--- a/hibernate-core/src/main/java/org/hibernate/context/spi/CurrentTenantIdentifierResolver.java
+++ b/hibernate-core/src/main/java/org/hibernate/context/spi/CurrentTenantIdentifierResolver.java
@@ -31,4 +31,13 @@ public interface CurrentTenantIdentifierResolver {
 	 * @see org.hibernate.context.TenantIdentifierMismatchException
 	 */
 	public boolean validateExistingCurrentSessions();
+
+	/**
+	 * Does the given tenant id represent a "root" tenant with access to all partitions.
+	 * @param tenantId a tenant id produced by {@link #resolveCurrentTenantIdentifier()}
+	 * @return true is this is root tenant
+	 */
+	public default boolean isRoot(String tenantId) {
+		return false;
+	}
 }
diff --git a/hibernate-core/src/main/java/org/hibernate/internal/SessionImpl.java b/hibernate-core/src/main/java/org/hibernate/internal/SessionImpl.java
index 1ced767c0f..8466e3f574 100644
--- a/hibernate-core/src/main/java/org/hibernate/internal/SessionImpl.java
+++ b/hibernate-core/src/main/java/org/hibernate/internal/SessionImpl.java
@@ -51,6 +51,7 @@ import org.hibernate.UnknownProfileException;
 import org.hibernate.UnresolvableObjectException;
 import org.hibernate.cfg.AvailableSettings;
 import org.hibernate.collection.spi.PersistentCollection;
+import org.hibernate.context.spi.CurrentTenantIdentifierResolver;
 import org.hibernate.engine.internal.StatefulPersistenceContext;
 import org.hibernate.engine.jdbc.LobCreator;
 import org.hibernate.engine.jdbc.NonContextualLobCreator;
@@ -249,9 +250,13 @@ public class SessionImpl
 				throw new HibernateException( "SessionFactory configured for multi-tenancy, but no tenant identifier specified" );
 			}
 			else {
-				getLoadQueryInfluencers()
-						.enableFilter( TenantIdBinder.FILTER_NAME )
-						.setParameter( TenantIdBinder.PARAMETER_NAME, tenantIdentifier );
+				CurrentTenantIdentifierResolver resolver = factory.getCurrentTenantIdentifierResolver();
+				if ( resolver==null || !resolver.isRoot(tenantIdentifier) ) {
+					// turn on the filter, unless this is the "root" tenant with access to all partitions
+					getLoadQueryInfluencers()
+							.enableFilter( TenantIdBinder.FILTER_NAME )
+							.setParameter( TenantIdBinder.PARAMETER_NAME, tenantIdentifier );
+				}
 			}
 		}
 
diff --git a/hibernate-core/src/main/java/org/hibernate/tuple/TenantIdGeneration.java b/hibernate-core/src/main/java/org/hibernate/tuple/TenantIdGeneration.java
index 30cddfd2c4..f1815f940b 100644
--- a/hibernate-core/src/main/java/org/hibernate/tuple/TenantIdGeneration.java
+++ b/hibernate-core/src/main/java/org/hibernate/tuple/TenantIdGeneration.java
@@ -9,6 +9,8 @@ package org.hibernate.tuple;
 import org.hibernate.PropertyValueException;
 import org.hibernate.Session;
 import org.hibernate.annotations.TenantId;
+import org.hibernate.context.spi.CurrentTenantIdentifierResolver;
+import org.hibernate.engine.spi.SessionFactoryImplementor;
 
 /**
  * Value generation implementation for {@link TenantId}.
@@ -44,12 +46,21 @@ public class TenantIdGeneration implements AnnotationValueGeneration<TenantId>,
 	@Override
 	public Object generateValue(Session session, Object owner, Object currentValue) {
 		String identifier = session.getTenantIdentifier();
-		if ( currentValue != null && !currentValue.equals(identifier) ) {
-			throw new PropertyValueException(
-					"assigned tenant id differs from current tenant id: "
-							+ currentValue + "!=" + identifier,
-					entityName, propertyName
-			);
+		if ( currentValue != null ) {
+			CurrentTenantIdentifierResolver resolver =
+					((SessionFactoryImplementor) session.getSessionFactory())
+							.getCurrentTenantIdentifierResolver();
+			if ( resolver!=null && resolver.isRoot( session.getTenantIdentifier() ) ) {
+				// the "root" tenant is allowed to set the tenant id explicitly
+				return currentValue;
+			}
+			if ( !currentValue.equals(identifier) ) {
+				throw new PropertyValueException(
+						"assigned tenant id differs from current tenant id: "
+								+ currentValue + "!=" + identifier,
+						entityName, propertyName
+				);
+			}
 		}
 		return identifier;
 	}