Skip to content

Commit a2dec7d

Browse files
marko-bekhtamarko-bekhta
committed
HHH-19849 Add an SPI that allows attaching session-scoped "extension storage" to the session/statelesssession implementors
1 parent 2b39aae commit a2dec7d

File tree

7 files changed

+247
-0
lines changed

7 files changed

+247
-0
lines changed

‎hibernate-core/src/main/java/org/hibernate/engine/spi/ExtensionStorage.java‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
* Copyright Red Hat Inc. and Hibernate Authors
4+
*/
5+
package org.hibernate.engine.spi;
6+
7+
import org.hibernate.Incubating;
8+
9+
import java.util.function.Supplier;
10+
11+
/**
12+
* Marker interface for extensions to register themselves within a session instance.
13+
* Extension storage implementations <b>must</b> have a public, no-argument constructor,
14+
* or otherwise be retrieved/registered with a session via {@link SharedSessionContractImplementor#getExtensionStorage(Class, Supplier)}
15+
*
16+
* @see SharedSessionContractImplementor#getExtensionStorage(Class)
17+
* @see SharedSessionContractImplementor#getExtensionStorage(Class, Supplier)
18+
*/
19+
@Incubating
20+
public interface ExtensionStorage {
21+
}

‎hibernate-core/src/main/java/org/hibernate/engine/spi/SessionDelegatorBaseImpl.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,7 @@
7575
import java.util.Set;
7676
import java.util.TimeZone;
7777
import java.util.UUID;
78+
import java.util.function.Supplier;
7879

7980
/**
8081
* A wrapper class that delegates all method invocations to a delegate instance of
@@ -517,6 +518,16 @@ public RootGraphImplementor<?> getEntityGraph(String graphName) {
517518
return delegate.getEntityGraph( graphName );
518519
}
519520

521+
@Override
522+
public <T extends ExtensionStorage> T getExtensionStorage(Class<T> extension) {
523+
return delegate.getExtensionStorage( extension );
524+
}
525+
526+
@Override
527+
public <T extends ExtensionStorage> T getExtensionStorage(Class<T> extension, Supplier<T> createIfMissing) {
528+
return delegate.getExtensionStorage( extension, createIfMissing );
529+
}
530+
520531
@Override
521532
public <T> QueryImplementor<T> createQuery(CriteriaSelect<T> selectQuery) {
522533
return delegate.createQuery( selectQuery );

‎hibernate-core/src/main/java/org/hibernate/engine/spi/SharedSessionContractImplementor.java‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66

77
import java.util.Set;
88
import java.util.UUID;
9+
import java.util.function.Supplier;
10+
911
import jakarta.persistence.TransactionRequiredException;
1012
import org.checkerframework.checker.nullness.qual.Nullable;
1113

@@ -621,4 +623,25 @@ default boolean isStatelessSession() {
621623

622624
@Override
623625
RootGraphImplementor<?> getEntityGraph(String graphName);
626+
627+
/**
628+
* Allows accessing session scoped extension storages of the particular session instance.
629+
*
630+
* @param extension The extension storage attached to the current session.
631+
* @param <T> The type of the extension storage.
632+
*/
633+
@Incubating
634+
<T extends ExtensionStorage> T getExtensionStorage(Class<T> extension);
635+
636+
/**
637+
* Allows accessing session scoped extension storages of the particular session instance.
638+
*
639+
* @param extension The extension storage attached to the current session.
640+
* @param createIfMissing Creates a storage extension using the supplier,
641+
* if the current session does not yet have the particular storage type attached to this session.
642+
* @param <T> The type of the extension storage.
643+
*/
644+
@Incubating
645+
<T extends ExtensionStorage> T getExtensionStorage(Class<T> extension, Supplier<T> createIfMissing);
646+
624647
}

‎hibernate-core/src/main/java/org/hibernate/engine/spi/SharedSessionDelegatorBaseImpl.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@
5252
import java.util.Set;
5353
import java.util.TimeZone;
5454
import java.util.UUID;
55+
import java.util.function.Supplier;
5556

5657
/**
5758
* A wrapper class that delegates all method invocations to a delegate instance of
@@ -663,6 +664,16 @@ public RootGraphImplementor<?> getEntityGraph(String graphName) {
663664
return delegate.getEntityGraph( graphName );
664665
}
665666

667+
@Override
668+
public <T extends ExtensionStorage> T getExtensionStorage(Class<T> extension) {
669+
return delegate.getExtensionStorage( extension );
670+
}
671+
672+
@Override
673+
public <T extends ExtensionStorage> T getExtensionStorage(Class<T> extension, Supplier<T> createIfMissing) {
674+
return delegate.getExtensionStorage( extension, createIfMissing );
675+
}
676+
666677
@Override
667678
public <T> List<EntityGraph<? super T>> getEntityGraphs(Class<T> entityClass) {
668679
return delegate.getEntityGraphs( entityClass );

‎hibernate-core/src/main/java/org/hibernate/internal/AbstractSharedSessionContract.java‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
import org.hibernate.engine.jdbc.spi.JdbcServices;
4242
import org.hibernate.engine.spi.EntityKey;
4343
import org.hibernate.engine.spi.ExceptionConverter;
44+
import org.hibernate.engine.spi.ExtensionStorage;
4445
import org.hibernate.engine.spi.LoadQueryInfluencers;
4546
import org.hibernate.engine.spi.SessionEventListenerManager;
4647
import org.hibernate.engine.spi.SessionFactoryImplementor;
@@ -109,14 +110,18 @@
109110
import java.io.ObjectInputStream;
110111
import java.io.ObjectOutputStream;
111112
import java.io.Serial;
113+
import java.lang.reflect.InvocationTargetException;
112114
import java.sql.Connection;
113115
import java.sql.SQLException;
116+
import java.util.HashMap;
114117
import java.util.List;
115118
import java.util.Locale;
119+
import java.util.Map;
116120
import java.util.Objects;
117121
import java.util.TimeZone;
118122
import java.util.UUID;
119123
import java.util.function.Function;
124+
import java.util.function.Supplier;
120125

121126
import static java.lang.Boolean.TRUE;
122127
import static org.hibernate.boot.model.naming.Identifier.toIdentifier;
@@ -186,6 +191,8 @@ public abstract class AbstractSharedSessionContract implements SharedSessionCont
186191
private transient ExceptionConverter exceptionConverter;
187192
private transient SessionAssociationMarkers sessionAssociationMarkers;
188193

194+
private transient Map<Class<?>, Object> extensionStorages;
195+
189196
public AbstractSharedSessionContract(SessionFactoryImpl factory, SessionCreationOptions options) {
190197
this.factory = factory;
191198

@@ -1704,6 +1711,39 @@ public SessionAssociationMarkers getSessionAssociationMarkers() {
17041711
return sessionAssociationMarkers;
17051712
}
17061713

1714+
@Override
1715+
public <T extends ExtensionStorage> T getExtensionStorage(Class<T> extension) {
1716+
if ( extensionStorages == null ) {
1717+
extensionStorages = new HashMap<>();
1718+
}
1719+
Object storage = extensionStorages.get( extension );;
1720+
if ( storage == null ) {
1721+
try {
1722+
storage = extension.getDeclaredConstructor().newInstance();
1723+
extensionStorages.put( extension, storage );
1724+
}
1725+
catch (InvocationTargetException | InstantiationException | IllegalAccessException | NoSuchMethodException e) {
1726+
throw SESSION_LOGGER.unableToCreateSessionExtensionStorage( extension, e.getMessage(), e );
1727+
}
1728+
}
1729+
1730+
return extension.cast( storage );
1731+
}
1732+
1733+
@Override
1734+
public <T extends ExtensionStorage> T getExtensionStorage(Class<T> extension, Supplier<T> createIfMissing) {
1735+
if ( extensionStorages == null ) {
1736+
extensionStorages = new HashMap<>();
1737+
}
1738+
Object storage = extensionStorages.get( extension );
1739+
if ( storage == null ) {
1740+
storage = createIfMissing.get();
1741+
extensionStorages.put( extension, storage );
1742+
}
1743+
1744+
return extension.cast( storage );
1745+
}
1746+
17071747
@Serial
17081748
private void writeObject(ObjectOutputStream oos) throws IOException {
17091749
SESSION_LOGGER.serializingSession( getSessionIdentifier() );

‎hibernate-core/src/main/java/org/hibernate/internal/SessionLogging.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
*/
55
package org.hibernate.internal;
66

7+
import org.hibernate.HibernateException;
78
import org.hibernate.Internal;
89
import org.hibernate.internal.log.SubSystemLogging;
910

@@ -182,4 +183,7 @@ public interface SessionLogging extends BasicLogger {
182183
@LogMessage(level = TRACE)
183184
@Message("Collection fetched")
184185
void collectionFetched();
186+
187+
@Message("Failed to create a storage of type %s: %s")
188+
HibernateException unableToCreateSessionExtensionStorage(Class<?> extension, String message, @Cause Exception reason);
185189
}

‎hibernate-core/src/test/java/org/hibernate/orm/test/engine/spi/SessionExtensionTest.java‎

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
* Copyright Red Hat Inc. and Hibernate Authors
4+
*/
5+
package org.hibernate.orm.test.engine.spi;
6+
7+
import jakarta.persistence.Id;
8+
import org.hibernate.HibernateException;
9+
import org.hibernate.engine.spi.ExtensionStorage;
10+
import org.hibernate.testing.orm.junit.DomainModel;
11+
import org.hibernate.testing.orm.junit.SessionFactory;
12+
import org.hibernate.testing.orm.junit.SessionFactoryScope;
13+
import org.junit.jupiter.api.Test;
14+
15+
import java.util.HashMap;
16+
import java.util.Map;
17+
18+
import static org.assertj.core.api.Assertions.assertThat;
19+
import static org.assertj.core.api.Assertions.assertThatThrownBy;
20+
21+
@DomainModel(annotatedClasses = {
22+
SessionExtensionTest.UselessEntity.class,
23+
})
24+
@SessionFactory
25+
public class SessionExtensionTest {
26+
27+
@Test
28+
public void smoke(SessionFactoryScope scope) {
29+
scope.inSession( sessionImplementor -> {
30+
sessionImplementor.getExtensionStorage( MyExtensionStorage.class )
31+
.add( new Extension( 1 ) );
32+
33+
assertThat( sessionImplementor.getExtensionStorage( MyExtensionStorage.class ).get( 1 ) )
34+
.isNotNull()
35+
.isEqualTo( new Extension( 1 ) );
36+
} );
37+
38+
scope.inStatelessSession( sessionImplementor -> {
39+
sessionImplementor.getExtensionStorage( MyExtensionStorage.class )
40+
.add( new Extension( 1 ) );
41+
42+
assertThat( sessionImplementor.getExtensionStorage( MyExtensionStorage.class ).get( 1 ) )
43+
.isNotNull()
44+
.isEqualTo( new Extension( 1 ) );
45+
} );
46+
}
47+
48+
@Test
49+
public void failing(SessionFactoryScope scope) {
50+
scope.inSession( sessionImplementor -> {
51+
assertThatThrownBy(
52+
() -> sessionImplementor.getExtensionStorage( MyFailingExtensionStorage.class ) )
53+
.isInstanceOf( HibernateException.class )
54+
.cause().hasCauseInstanceOf( UnsupportedOperationException.class );
55+
} );
56+
57+
scope.inStatelessSession( sessionImplementor -> {
58+
assertThatThrownBy(
59+
() -> sessionImplementor.getExtensionStorage( MyFailingExtensionStorage.class ) )
60+
.isInstanceOf( HibernateException.class )
61+
.cause().hasCauseInstanceOf( UnsupportedOperationException.class );
62+
} );
63+
}
64+
65+
@Test
66+
public void supplier(SessionFactoryScope scope) {
67+
scope.inSession( sessionImplementor -> {
68+
sessionImplementor.getExtensionStorage( MyFailingExtensionStorage.class,
69+
() -> new MyFailingExtensionStorage( new HashMap<>() ) )
70+
.add( new Extension( 1 ) );
71+
72+
assertThat( sessionImplementor.getExtensionStorage( MyFailingExtensionStorage.class ).get( 1 ) )
73+
.isNotNull()
74+
.isEqualTo( new Extension( 1 ) );
75+
76+
assertThat( sessionImplementor.getExtensionStorage( MyFailingExtensionStorage.class,
77+
() -> new MyFailingExtensionStorage( new HashMap<>() ) ).get( 1 ) )
78+
.isNotNull()
79+
.isEqualTo( new Extension( 1 ) );
80+
} );
81+
82+
scope.inStatelessSession( sessionImplementor -> {
83+
sessionImplementor.getExtensionStorage( MyFailingExtensionStorage.class,
84+
() -> new MyFailingExtensionStorage( new HashMap<>() ) )
85+
.add( new Extension( 1 ) );
86+
87+
assertThat( sessionImplementor.getExtensionStorage( MyFailingExtensionStorage.class ).get( 1 ) )
88+
.isNotNull()
89+
.isEqualTo( new Extension( 1 ) );
90+
91+
assertThat( sessionImplementor.getExtensionStorage( MyFailingExtensionStorage.class,
92+
() -> new MyFailingExtensionStorage( new HashMap<>() ) ).get( 1 ) )
93+
.isNotNull()
94+
.isEqualTo( new Extension( 1 ) );
95+
} );
96+
}
97+
98+
public static class MyExtensionStorage implements ExtensionStorage {
99+
Map<Integer, Extension> extensions = new HashMap<>();
100+
101+
public void add(Extension extension) {
102+
extensions.put( extension.number, extension );
103+
}
104+
105+
public Extension get(int number) {
106+
return extensions.get( number );
107+
}
108+
}
109+
110+
public static class MyFailingExtensionStorage implements ExtensionStorage {
111+
Map<Integer, Extension> extensions = new HashMap<>();
112+
113+
public MyFailingExtensionStorage() {
114+
throw new UnsupportedOperationException();
115+
}
116+
117+
MyFailingExtensionStorage(Map<Integer, Extension> extensions) {
118+
this.extensions = extensions;
119+
}
120+
121+
public void add(Extension extension) {
122+
extensions.put( extension.number, extension );
123+
}
124+
125+
public Extension get(int number) {
126+
return extensions.get( number );
127+
}
128+
}
129+
130+
public record Extension(int number) {
131+
}
132+
133+
static class UselessEntity {
134+
@Id
135+
Long id;
136+
}
137+
}

0 commit comments

Comments
 (0)