sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/JdbcFactory.java

/*
 * 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.beam.sdk.extensions.sql.impl;

import static org.apache.beam.sdk.extensions.sql.impl.JdbcDriver.TOP_LEVEL_BEAM_SCHEMA;
import static org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.avatica.BuiltInConnectionProperty.TIME_ZONE;
import static org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.config.CalciteConnectionProperty.LEX;
import static org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.config.CalciteConnectionProperty.PARSER_FACTORY;
import static org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.config.CalciteConnectionProperty.SCHEMA;
import static org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.config.CalciteConnectionProperty.SCHEMA_FACTORY;
import static org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.config.CalciteConnectionProperty.TYPE_SYSTEM;

import java.util.Properties;
import org.apache.beam.sdk.extensions.sql.impl.parser.BeamSqlParser;
import org.apache.beam.sdk.extensions.sql.impl.planner.BeamRelDataTypeSystem;
import org.apache.beam.sdk.extensions.sql.meta.provider.TableProvider;
import org.apache.beam.sdk.util.ReleaseInfo;
import org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.adapter.java.JavaTypeFactory;
import org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.avatica.AvaticaConnection;
import org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.avatica.AvaticaFactory;
import org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.avatica.ConnectionProperty;
import org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.avatica.UnregisteredDriver;
import org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.config.Lex;
import org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.CalciteFactory;
import org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.CalciteSchema;
import org.checkerframework.checker.nullness.qual.Nullable;

/**
 * Implements {@link CalciteFactory} that is used by Clacite JDBC driver to instantiate different
 * JDBC objects, like connections, result sets, etc.
 *
 * <p>The purpose of this class is to intercept the connection creation and force a cache-less root
 * schema ({@link
 * org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.SimpleCalciteSchema}). Otherwise
 * Calcite uses {@link
 * org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.CachingCalciteSchema} that eagerly
 * caches table information. This behavior does not work well for dynamic table providers.
 */
class JdbcFactory extends CalciteFactoryWrapper {

  JdbcFactory(CalciteFactory factory) {
    super(factory);
  }

  static JdbcFactory wrap(CalciteFactory calciteFactory) {
    return new JdbcFactory(calciteFactory);
  }

  @Override
  public AvaticaConnection newConnection(
      UnregisteredDriver driver,
      AvaticaFactory avaticaFactory,
      String url,
      Properties info,
      @Nullable CalciteSchema rootSchema,
      @Nullable JavaTypeFactory typeFactory) {

    Properties connectionProps = ensureDefaultProperties(info);
    CalciteSchema actualRootSchema = rootSchema;
    if (rootSchema == null) {
      actualRootSchema = CalciteSchema.createRootSchema(true, false, "");
    }

    return super.newConnection(
        driver, avaticaFactory, url, connectionProps, actualRootSchema, typeFactory);
  }

  /**
   * Make sure required default properties are set.
   *
   * <p>Among other things sets up the parser class name, rel data type system and default schema
   * factory.
   *
   * <p>The specified Beam schema factory will be used by Calcite to create the initial top level
   * Beam schema. It can be later overridden by setting the schema via {@link
   * JdbcConnection#setSchema(String, TableProvider)}.
   */
  private Properties ensureDefaultProperties(Properties originalInfo) {
    Properties info = new Properties();
    info.putAll(originalInfo);

    setIfNull(info, TIME_ZONE, "UTC");
    setIfNull(info, LEX, Lex.JAVA.name());
    setIfNull(info, PARSER_FACTORY, BeamSqlParser.class.getName() + "#FACTORY");
    setIfNull(info, TYPE_SYSTEM, BeamRelDataTypeSystem.class.getName());
    setIfNull(info, SCHEMA, TOP_LEVEL_BEAM_SCHEMA);
    setIfNull(info, SCHEMA_FACTORY, BeamCalciteSchemaFactory.AllProviders.class.getName());
    setIfNull(info, "beam.userAgent", "BeamSQL/" + ReleaseInfo.getReleaseInfo().getVersion());

    return info;
  }

  private static void setIfNull(Properties info, ConnectionProperty key, String value) {
    setIfNull(info, key.camelName(), value);
  }

  private static void setIfNull(Properties info, String key, String value) {
    // A null value indicates the default. We want to override defaults only.
    if (info.getProperty(key) == null) {
      info.setProperty(key, value);
    }
  }
}