com.yugabyte.data.jdbc.repository.config.YugabyteDialectResolver Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of spring-data-yugabytedb-ysql Show documentation
Spring Data support for YugabyteDB YSQL.
The newest version!
/*
 * Copyright (c) Yugabyte, Inc.
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except 
 * in compliance with the License.  You may obtain a copy of the License at
 * 
 * http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software distributed under the License 
 * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 
 * or implied.  See the License for the specific language governing permissions and limitations 
 * under the License.
*/
package com.yugabyte.data.jdbc.repository.config;

import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.Optional;
import java.util.regex.Pattern;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.postgresql.util.PSQLException;
import org.springframework.data.jdbc.repository.config.DialectResolver.DefaultDialectProvider;
import org.springframework.data.relational.core.dialect.Dialect;
import org.springframework.jdbc.core.ConnectionCallback;
import org.springframework.jdbc.core.JdbcOperations;
import org.springframework.lang.Nullable;
import org.springframework.util.StringUtils;

import com.yugabyte.data.relational.core.dialect.YugabyteDialect;

/**
 * An SQL dialect for YugabyteDB.
 *
 * @author Nikhil Chandrappa
 * @since 2.3
 */
public class YugabyteDialectResolver {

	static public class YugabyteDialectProvider extends DefaultDialectProvider {

		private static final Log LOG = LogFactory.getLog(YugabyteDialectProvider.class);
		private static final String YUGABYTE_SERVERS_QUERY = "select * from yb_servers()";
		private static final String YB_SERVERS_FUNCTION_DOESNT_EXIST = "function yb_servers() does not exist";

		@Override
		public Optional getDialect(JdbcOperations operations) {

			Optional yugabyteDialect = Optional.ofNullable(
					operations.execute((ConnectionCallback) YugabyteDialectProvider::getDialect));

			if (yugabyteDialect.isPresent()) {
				return yugabyteDialect;
			}

			return super.getDialect(operations);
		}

		@Nullable
		private static Dialect getDialect(Connection connection) throws SQLException {
			
			Statement ybStatement = connection.createStatement();
			Dialect dialect = null;
			
			LOG.debug("Executing query against YB_SERVERS() system function.");
			try {

				ResultSet rs = ybStatement.executeQuery(YUGABYTE_SERVERS_QUERY);
				if (!rs.next()) {
					LOG.debug(
							"Query for YB_SERVERS() system function returned null. Falling back on DefaultDialectProvider.");
					return null;
				}

				// Determine client application is connecting to a YugabyteDB cluster.
				String hostName = rs.getString("host");
				if (StringUtils.hasText(hostName)) {
					dialect = YugabyteDialect.INSTANCE;
				}
			} catch (PSQLException psqlException) {

				Boolean functionNotExist = Pattern
						.compile(Pattern.quote(YB_SERVERS_FUNCTION_DOESNT_EXIST), Pattern.CASE_INSENSITIVE)
						.matcher(psqlException.getMessage()).find();

				if (functionNotExist) {
					LOG.debug("YB_SERVERS() system function doesn't exist. Falling back on DefaultDialectProvider.");
					return null;
				}
			}

			LOG.debug("Using YugabyteDB YSQL Dialect.");
			return dialect;
		}
	}

}