org.apache.pinot.client.PinotConnection Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of datacap-shaded-pinot Show documentation
DataCap - Shaded for Apache Pinot
There is a newer version: 2024.03.7
/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.pinot.client;

import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.List;
import org.apache.pinot.client.base.AbstractBaseConnection;
import org.apache.pinot.client.controller.PinotControllerTransport;
import org.apache.pinot.client.controller.response.ControllerTenantBrokerResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;


public class PinotConnection extends AbstractBaseConnection {

  private static final Logger LOGGER = LoggerFactory.getLogger(Connection.class);
  private org.apache.pinot.client.Connection _session;
  private boolean _closed;
  private String _controllerURL;
  private PinotControllerTransport _controllerTransport;

  PinotConnection(String controllerURL, PinotClientTransport transport, String tenant) {
    this(controllerURL, transport, tenant, null);
  }

  PinotConnection(String controllerURL, PinotClientTransport transport, String tenant,
      PinotControllerTransport controllerTransport) {
    _closed = false;
    _controllerURL = controllerURL;
    if (controllerTransport == null) {
      _controllerTransport = new PinotControllerTransport();
    } else {
      _controllerTransport = controllerTransport;
    }
    List brokers = getBrokerList(controllerURL, tenant);
    _session = new org.apache.pinot.client.Connection(brokers, transport);
  }

  public org.apache.pinot.client.Connection getSession() {
    return _session;
  }

  private List getBrokerList(String controllerURL, String tenant) {
    ControllerTenantBrokerResponse controllerTenantBrokerResponse =
        _controllerTransport.getBrokersFromController(controllerURL, tenant);
    return controllerTenantBrokerResponse.getBrokers();
  }

  @Override
  protected void validateState()
      throws SQLException {
    if (isClosed()) {
      throw new SQLException("Connection is already closed!");
    }
  }

  @Override
  public void close()
      throws SQLException {
    if (!isClosed()) {
      _session.close();
      _controllerTransport.close();
    }
    _controllerTransport = null;
    _session = null;
    _closed = true;
  }

  @Override
  public Statement createStatement()
      throws SQLException {
    validateState();
    return new PinotStatement(this);
  }

  @Override
  public PreparedStatement prepareStatement(String sql)
      throws SQLException {
    validateState();
    return new PinotPreparedStatement(this, sql);
  }

  @Override
  public boolean isClosed()
      throws SQLException {
    return _closed;
  }

  @Override
  public DatabaseMetaData getMetaData()
      throws SQLException {
    return new PinotConnectionMetaData(this, _controllerURL, _controllerTransport);
  }
}