001/*
002 * Copyright 2002-2020 the original author or authors.
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 *      https://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016
017package org.springframework.jdbc.datasource.lookup;
018
019import java.sql.Connection;
020import java.sql.SQLException;
021import java.util.Collections;
022import java.util.HashMap;
023import java.util.Map;
024
025import javax.sql.DataSource;
026
027import org.springframework.beans.factory.InitializingBean;
028import org.springframework.jdbc.datasource.AbstractDataSource;
029import org.springframework.lang.Nullable;
030import org.springframework.util.Assert;
031
032/**
033 * Abstract {@link javax.sql.DataSource} implementation that routes {@link #getConnection()}
034 * calls to one of various target DataSources based on a lookup key. The latter is usually
035 * (but not necessarily) determined through some thread-bound transaction context.
036 *
037 * @author Juergen Hoeller
038 * @since 2.0.1
039 * @see #setTargetDataSources
040 * @see #setDefaultTargetDataSource
041 * @see #determineCurrentLookupKey()
042 */
043public abstract class AbstractRoutingDataSource extends AbstractDataSource implements InitializingBean {
044
045        @Nullable
046        private Map<Object, Object> targetDataSources;
047
048        @Nullable
049        private Object defaultTargetDataSource;
050
051        private boolean lenientFallback = true;
052
053        private DataSourceLookup dataSourceLookup = new JndiDataSourceLookup();
054
055        @Nullable
056        private Map<Object, DataSource> resolvedDataSources;
057
058        @Nullable
059        private DataSource resolvedDefaultDataSource;
060
061
062        /**
063         * Specify the map of target DataSources, with the lookup key as key.
064         * The mapped value can either be a corresponding {@link javax.sql.DataSource}
065         * instance or a data source name String (to be resolved via a
066         * {@link #setDataSourceLookup DataSourceLookup}).
067         * <p>The key can be of arbitrary type; this class implements the
068         * generic lookup process only. The concrete key representation will
069         * be handled by {@link #resolveSpecifiedLookupKey(Object)} and
070         * {@link #determineCurrentLookupKey()}.
071         */
072        public void setTargetDataSources(Map<Object, Object> targetDataSources) {
073                this.targetDataSources = targetDataSources;
074        }
075
076        /**
077         * Specify the default target DataSource, if any.
078         * <p>The mapped value can either be a corresponding {@link javax.sql.DataSource}
079         * instance or a data source name String (to be resolved via a
080         * {@link #setDataSourceLookup DataSourceLookup}).
081         * <p>This DataSource will be used as target if none of the keyed
082         * {@link #setTargetDataSources targetDataSources} match the
083         * {@link #determineCurrentLookupKey()} current lookup key.
084         */
085        public void setDefaultTargetDataSource(Object defaultTargetDataSource) {
086                this.defaultTargetDataSource = defaultTargetDataSource;
087        }
088
089        /**
090         * Specify whether to apply a lenient fallback to the default DataSource
091         * if no specific DataSource could be found for the current lookup key.
092         * <p>Default is "true", accepting lookup keys without a corresponding entry
093         * in the target DataSource map - simply falling back to the default DataSource
094         * in that case.
095         * <p>Switch this flag to "false" if you would prefer the fallback to only apply
096         * if the lookup key was {@code null}. Lookup keys without a DataSource
097         * entry will then lead to an IllegalStateException.
098         * @see #setTargetDataSources
099         * @see #setDefaultTargetDataSource
100         * @see #determineCurrentLookupKey()
101         */
102        public void setLenientFallback(boolean lenientFallback) {
103                this.lenientFallback = lenientFallback;
104        }
105
106        /**
107         * Set the DataSourceLookup implementation to use for resolving data source
108         * name Strings in the {@link #setTargetDataSources targetDataSources} map.
109         * <p>Default is a {@link JndiDataSourceLookup}, allowing the JNDI names
110         * of application server DataSources to be specified directly.
111         */
112        public void setDataSourceLookup(@Nullable DataSourceLookup dataSourceLookup) {
113                this.dataSourceLookup = (dataSourceLookup != null ? dataSourceLookup : new JndiDataSourceLookup());
114        }
115
116
117        @Override
118        public void afterPropertiesSet() {
119                if (this.targetDataSources == null) {
120                        throw new IllegalArgumentException("Property 'targetDataSources' is required");
121                }
122                this.resolvedDataSources = new HashMap<>(this.targetDataSources.size());
123                this.targetDataSources.forEach((key, value) -> {
124                        Object lookupKey = resolveSpecifiedLookupKey(key);
125                        DataSource dataSource = resolveSpecifiedDataSource(value);
126                        this.resolvedDataSources.put(lookupKey, dataSource);
127                });
128                if (this.defaultTargetDataSource != null) {
129                        this.resolvedDefaultDataSource = resolveSpecifiedDataSource(this.defaultTargetDataSource);
130                }
131        }
132
133        /**
134         * Resolve the given lookup key object, as specified in the
135         * {@link #setTargetDataSources targetDataSources} map, into
136         * the actual lookup key to be used for matching with the
137         * {@link #determineCurrentLookupKey() current lookup key}.
138         * <p>The default implementation simply returns the given key as-is.
139         * @param lookupKey the lookup key object as specified by the user
140         * @return the lookup key as needed for matching
141         */
142        protected Object resolveSpecifiedLookupKey(Object lookupKey) {
143                return lookupKey;
144        }
145
146        /**
147         * Resolve the specified data source object into a DataSource instance.
148         * <p>The default implementation handles DataSource instances and data source
149         * names (to be resolved via a {@link #setDataSourceLookup DataSourceLookup}).
150         * @param dataSource the data source value object as specified in the
151         * {@link #setTargetDataSources targetDataSources} map
152         * @return the resolved DataSource (never {@code null})
153         * @throws IllegalArgumentException in case of an unsupported value type
154         */
155        protected DataSource resolveSpecifiedDataSource(Object dataSource) throws IllegalArgumentException {
156                if (dataSource instanceof DataSource) {
157                        return (DataSource) dataSource;
158                }
159                else if (dataSource instanceof String) {
160                        return this.dataSourceLookup.getDataSource((String) dataSource);
161                }
162                else {
163                        throw new IllegalArgumentException(
164                                        "Illegal data source value - only [javax.sql.DataSource] and String supported: " + dataSource);
165                }
166        }
167
168        /**
169         * Return the resolved target DataSources that this router manages.
170         * @return an unmodifiable map of resolved lookup keys and DataSources
171         * @throws IllegalStateException if the target DataSources are not resolved yet
172         * @since 5.2.9
173         * @see #setTargetDataSources
174         */
175        public Map<Object, DataSource> getResolvedDataSources() {
176                Assert.state(this.resolvedDataSources != null, "DataSources not resolved yet - call afterPropertiesSet");
177                return Collections.unmodifiableMap(this.resolvedDataSources);
178        }
179
180        /**
181         * Return the resolved default target DataSource, if any.
182         * @return the default DataSource, or {@code null} if none or not resolved yet
183         * @since 5.2.9
184         * @see #setDefaultTargetDataSource
185         */
186        @Nullable
187        public DataSource getResolvedDefaultDataSource() {
188                return this.resolvedDefaultDataSource;
189        }
190
191
192        @Override
193        public Connection getConnection() throws SQLException {
194                return determineTargetDataSource().getConnection();
195        }
196
197        @Override
198        public Connection getConnection(String username, String password) throws SQLException {
199                return determineTargetDataSource().getConnection(username, password);
200        }
201
202        @Override
203        @SuppressWarnings("unchecked")
204        public <T> T unwrap(Class<T> iface) throws SQLException {
205                if (iface.isInstance(this)) {
206                        return (T) this;
207                }
208                return determineTargetDataSource().unwrap(iface);
209        }
210
211        @Override
212        public boolean isWrapperFor(Class<?> iface) throws SQLException {
213                return (iface.isInstance(this) || determineTargetDataSource().isWrapperFor(iface));
214        }
215
216        /**
217         * Retrieve the current target DataSource. Determines the
218         * {@link #determineCurrentLookupKey() current lookup key}, performs
219         * a lookup in the {@link #setTargetDataSources targetDataSources} map,
220         * falls back to the specified
221         * {@link #setDefaultTargetDataSource default target DataSource} if necessary.
222         * @see #determineCurrentLookupKey()
223         */
224        protected DataSource determineTargetDataSource() {
225                Assert.notNull(this.resolvedDataSources, "DataSource router not initialized");
226                Object lookupKey = determineCurrentLookupKey();
227                DataSource dataSource = this.resolvedDataSources.get(lookupKey);
228                if (dataSource == null && (this.lenientFallback || lookupKey == null)) {
229                        dataSource = this.resolvedDefaultDataSource;
230                }
231                if (dataSource == null) {
232                        throw new IllegalStateException("Cannot determine target DataSource for lookup key [" + lookupKey + "]");
233                }
234                return dataSource;
235        }
236
237        /**
238         * Determine the current lookup key. This will typically be
239         * implemented to check a thread-bound transaction context.
240         * <p>Allows for arbitrary keys. The returned key needs
241         * to match the stored lookup key type, as resolved by the
242         * {@link #resolveSpecifiedLookupKey} method.
243         */
244        @Nullable
245        protected abstract Object determineCurrentLookupKey();
246
247}