21
21
import static org .junit .Assert .fail ;
22
22
23
23
import com .datastax .oss .driver .api .core .CqlSession ;
24
+ import com .datastax .oss .driver .api .core .config .DriverConfigLoader ;
25
+ import com .datastax .oss .driver .api .core .context .DriverContext ;
24
26
import com .datastax .oss .driver .api .core .cql .PrepareRequest ;
25
27
import com .datastax .oss .driver .api .core .cql .PreparedStatement ;
28
+ import com .datastax .oss .driver .api .core .session .ProgrammaticArguments ;
29
+ import com .datastax .oss .driver .api .core .session .SessionBuilder ;
26
30
import com .datastax .oss .driver .api .testinfra .ccm .CustomCcmRule ;
27
31
import com .datastax .oss .driver .api .testinfra .session .SessionRule ;
28
32
import com .datastax .oss .driver .api .testinfra .session .SessionUtils ;
29
33
import com .datastax .oss .driver .categories .IsolatedTests ;
30
34
import com .datastax .oss .driver .internal .core .context .DefaultDriverContext ;
31
35
import com .datastax .oss .driver .internal .core .cql .CqlPrepareAsyncProcessor ;
36
+ import com .datastax .oss .driver .internal .core .cql .CqlPrepareSyncProcessor ;
37
+ import com .datastax .oss .driver .internal .core .session .BuiltInRequestProcessors ;
38
+ import com .datastax .oss .driver .internal .core .session .RequestProcessor ;
39
+ import com .datastax .oss .driver .internal .core .session .RequestProcessorRegistry ;
32
40
import com .datastax .oss .driver .shaded .guava .common .base .Predicates ;
33
41
import com .datastax .oss .driver .shaded .guava .common .cache .Cache ;
42
+ import com .datastax .oss .driver .shaded .guava .common .cache .CacheBuilder ;
34
43
import com .datastax .oss .driver .shaded .guava .common .collect .Iterables ;
44
+ import edu .umd .cs .findbugs .annotations .NonNull ;
45
+ import java .util .List ;
46
+ import java .util .Optional ;
35
47
import java .util .concurrent .CompletableFuture ;
36
48
import org .junit .After ;
49
+ import org .junit .AfterClass ;
37
50
import org .junit .Before ;
51
+ import org .junit .BeforeClass ;
38
52
import org .junit .Rule ;
39
53
import org .junit .Test ;
40
54
import org .junit .experimental .categories .Category ;
@@ -50,6 +64,69 @@ public class PreparedStatementCancellationIT {
50
64
51
65
@ Rule public TestRule chain = RuleChain .outerRule (ccmRule ).around (sessionRule );
52
66
67
+ private static class TestCqlPrepareAsyncProcessor extends CqlPrepareAsyncProcessor {
68
+
69
+ public TestCqlPrepareAsyncProcessor (@ NonNull Optional <DefaultDriverContext > context ) {
70
+ // Default CqlPrepareAsyncProcessor uses weak values here as well. We avoid doing so
71
+ // to prevent cache entries from unexpectedly disappearing mid-test.
72
+ super (context , CacheBuilder .newBuilder ());
73
+ }
74
+ }
75
+
76
+ private static class TestDefaultDriverContext extends DefaultDriverContext {
77
+ public TestDefaultDriverContext (
78
+ DriverConfigLoader configLoader , ProgrammaticArguments programmaticArguments ) {
79
+ super (configLoader , programmaticArguments );
80
+ }
81
+
82
+ @ Override
83
+ protected RequestProcessorRegistry buildRequestProcessorRegistry () {
84
+ // Re-create the processor cache to insert the TestCqlPrepareAsyncProcessor with it's strong
85
+ // prepared statement cache, see JAVA-3062
86
+ List <RequestProcessor <?, ?>> processors =
87
+ BuiltInRequestProcessors .createDefaultProcessors (this );
88
+ processors .removeIf ((processor ) -> processor instanceof CqlPrepareAsyncProcessor );
89
+ processors .removeIf ((processor ) -> processor instanceof CqlPrepareSyncProcessor );
90
+ CqlPrepareAsyncProcessor asyncProcessor =
91
+ new PreparedStatementCancellationIT .TestCqlPrepareAsyncProcessor (Optional .of (this ));
92
+ processors .add (2 , asyncProcessor );
93
+ processors .add (3 , new CqlPrepareSyncProcessor (asyncProcessor ));
94
+ return new RequestProcessorRegistry (
95
+ getSessionName (), processors .toArray (new RequestProcessor [0 ]));
96
+ }
97
+ }
98
+
99
+ private static class TestSessionBuilder extends SessionBuilder {
100
+
101
+ @ Override
102
+ protected Object wrap (@ NonNull CqlSession defaultSession ) {
103
+ return defaultSession ;
104
+ }
105
+
106
+ @ Override
107
+ protected DriverContext buildContext (
108
+ DriverConfigLoader configLoader , ProgrammaticArguments programmaticArguments ) {
109
+ return new PreparedStatementCancellationIT .TestDefaultDriverContext (
110
+ configLoader , programmaticArguments );
111
+ }
112
+ }
113
+
114
+ @ BeforeClass
115
+ public static void setupBeforeClass () {
116
+ System .setProperty (
117
+ SessionUtils .SESSION_BUILDER_CLASS_PROPERTY ,
118
+ PreparedStatementCancellationIT .class .getName ());
119
+ }
120
+
121
+ @ AfterClass
122
+ public static void teardownAfterClass () {
123
+ System .clearProperty (SessionUtils .SESSION_BUILDER_CLASS_PROPERTY );
124
+ }
125
+
126
+ public static SessionBuilder builder () {
127
+ return new PreparedStatementCancellationIT .TestSessionBuilder ();
128
+ }
129
+
53
130
@ Before
54
131
public void setup () {
55
132
0 commit comments