View Javadoc
1   /*
2    *
3    * The DbUnit Database Testing Framework
4    * Copyright (C)2002-2026, DbUnit.org
5    *
6    * This library is free software; you can redistribute it and/or
7    * modify it under the terms of the GNU Lesser General Public
8    * License as published by the Free Software Foundation; either
9    * version 2.1 of the License, or (at your option) any later version.
10   *
11   * This library is distributed in the hope that it will be useful,
12   * but WITHOUT ANY WARRANTY; without even the implied warranty of
13   * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
14   * Lesser General Public License for more details.
15   *
16   * You should have received a copy of the GNU Lesser General Public
17   * License along with this library; if not, write to the Free Software
18   * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
19   *
20   */
21  package org.dbunit.junit.jupiter;
22  
23  import java.io.PrintWriter;
24  import java.lang.reflect.InvocationHandler;
25  import java.lang.reflect.InvocationTargetException;
26  import java.lang.reflect.Proxy;
27  import java.sql.Connection;
28  import java.sql.DriverManager;
29  import java.sql.SQLException;
30  import java.util.concurrent.atomic.AtomicInteger;
31  import java.util.logging.Logger;
32  
33  import javax.sql.DataSource;
34  
35  import org.dbunit.DatabaseProfile;
36  
37  /**
38   * A test {@link DataSource} that opens real {@link DriverManager} connections against a
39   * {@link DatabaseProfile} and accounts for every one: how many are open at once (peak), and
40   * whether they all get closed. Each connection it hands back is a proxy whose {@code close()}
41   * closes the real connection and frees its slot.
42   *
43   * <p>When constructed with a {@code maxConcurrent} cap, {@code getConnection()} throws rather
44   * than block once that many are already open, so a runtime that tries to hold more connections
45   * than expected fails fast instead of deadlocking the test.
46   *
47   * <p>Sequential use only - no synchronization beyond the atomic counters.
48   */
49  final class CountingDataSource implements DataSource
50  {
51      private final DatabaseProfile profile;
52      private final int maxConcurrent;
53      private final AtomicInteger inUse = new AtomicInteger();
54      private final AtomicInteger peak = new AtomicInteger();
55      private final AtomicInteger opened = new AtomicInteger();
56      private final AtomicInteger closed = new AtomicInteger();
57  
58      CountingDataSource(final DatabaseProfile profile)
59      {
60          this(profile, 0);
61      }
62  
63      CountingDataSource(final DatabaseProfile profile, final int maxConcurrent)
64      {
65          this.profile = profile;
66          this.maxConcurrent = maxConcurrent;
67      }
68  
69      int peakConcurrent()
70      {
71          return peak.get();
72      }
73  
74      int leaked()
75      {
76          return opened.get() - closed.get();
77      }
78  
79      @Override
80      public Connection getConnection() throws SQLException
81      {
82          final int now = inUse.incrementAndGet();
83          if (maxConcurrent > 0 && now > maxConcurrent)
84          {
85              inUse.decrementAndGet();
86              throw new SQLException("CountingDataSource capped at " + maxConcurrent
87                      + " concurrent connections; connection " + now + " was requested while "
88                      + (now - 1) + " were still open");
89          }
90          peak.accumulateAndGet(now, Math::max);
91          final Connection real;
92          try
93          {
94              real = DriverManager.getConnection(profile.getConnectionUrl(), profile.getUser(),
95                      profile.getPassword());
96          }
97          catch (final SQLException e)
98          {
99              inUse.decrementAndGet();
100             throw e;
101         }
102         opened.incrementAndGet();
103         return slotFreeingProxy(real);
104     }
105 
106     private Connection slotFreeingProxy(final Connection real)
107     {
108         final InvocationHandler handler = (proxy, method, args) ->
109         {
110             if ("close".equals(method.getName()) && (args == null || args.length == 0))
111             {
112                 if (!real.isClosed())
113                 {
114                     real.close();
115                     closed.incrementAndGet();
116                     inUse.decrementAndGet();
117                 }
118                 return null;
119             }
120             try
121             {
122                 return method.invoke(real, args);
123             }
124             catch (final InvocationTargetException e)
125             {
126                 throw e.getCause();
127             }
128         };
129         return (Connection) Proxy.newProxyInstance(getClass().getClassLoader(),
130                 new Class<?>[] {Connection.class}, handler);
131     }
132 
133     @Override
134     public Connection getConnection(final String username, final String password)
135             throws SQLException
136     {
137         return getConnection();
138     }
139 
140     @Override
141     public PrintWriter getLogWriter()
142     {
143         return null;
144     }
145 
146     @Override
147     public void setLogWriter(final PrintWriter out)
148     {
149     }
150 
151     @Override
152     public void setLoginTimeout(final int seconds)
153     {
154     }
155 
156     @Override
157     public int getLoginTimeout()
158     {
159         return 0;
160     }
161 
162     @Override
163     public Logger getParentLogger()
164     {
165         return Logger.getLogger("");
166     }
167 
168     @Override
169     public <T> T unwrap(final Class<T> iface)
170     {
171         throw new UnsupportedOperationException();
172     }
173 
174     @Override
175     public boolean isWrapperFor(final Class<?> iface)
176     {
177         return false;
178     }
179 }