sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlEnv.java - beam - Git at Google

 /*
  * 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.vendor.guava.v20_0.com.google.common.base.Preconditions.checkNotNull;

 import java.lang.reflect.Method;
 import java.sql.SQLException;
 import java.util.AbstractMap.SimpleEntry;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.ServiceLoader;
 import java.util.Set;
 import org.apache.beam.sdk.annotations.Experimental;
 import org.apache.beam.sdk.annotations.Internal;
 import org.apache.beam.sdk.extensions.sql.BeamSqlTable;
 import org.apache.beam.sdk.extensions.sql.BeamSqlUdf;
 import org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode;
 import org.apache.beam.sdk.extensions.sql.impl.udf.BeamBuiltinFunctionProvider;
 import org.apache.beam.sdk.extensions.sql.meta.provider.ReadOnlyTableProvider;
 import org.apache.beam.sdk.extensions.sql.meta.provider.TableProvider;
 import org.apache.beam.sdk.extensions.sql.meta.provider.UdfUdafProvider;
 import org.apache.beam.sdk.extensions.sql.meta.store.InMemoryMetaStore;
 import org.apache.beam.sdk.transforms.Combine.CombineFn;
 import org.apache.beam.sdk.transforms.SerializableFunction;
 import org.apache.beam.vendor.guava.v20_0.com.google.common.base.Strings;
 import org.apache.calcite.jdbc.CalcitePrepare;
 import org.apache.calcite.plan.RelOptUtil;
 import org.apache.calcite.schema.Function;
 import org.apache.calcite.sql.SqlExecutableStatement;

 /**
  * Contains the metadata of tables/UDF functions, and exposes APIs to
  * query/validate/optimize/translate SQL statements.
  */
 @Internal
 @Experimental
 public class BeamSqlEnv {
   JdbcConnection connection;
   QueryPlanner planner;

   private BeamSqlEnv(JdbcConnection connection, QueryPlanner planner) {
     this.connection = connection;
     this.planner = planner;
   }

   /** Creates a builder with the default schema backed by the table provider. */
   public static BeamSqlEnvBuilder builder(TableProvider tableProvider) {
     return new BeamSqlEnvBuilder(tableProvider);
   }

   public static BeamSqlEnv readOnly(String tableType, Map<String, BeamSqlTable> tables) {
     return withTableProvider(new ReadOnlyTableProvider(tableType, tables));
   }

   public static BeamSqlEnv withTableProvider(TableProvider tableProvider) {
     return builder(tableProvider).build();
   }

   public static BeamSqlEnv inMemory(TableProvider... tableProviders) {
     InMemoryMetaStore inMemoryMetaStore = new InMemoryMetaStore();
     for (TableProvider tableProvider : tableProviders) {
       inMemoryMetaStore.registerProvider(tableProvider);
     }

     return withTableProvider(inMemoryMetaStore);
   }

   public BeamRelNode parseQuery(String query) throws ParseException {
     return planner.convertToBeamRel(query);
   }

   public boolean isDdl(String sqlStatement) throws ParseException {
     return planner.parse(sqlStatement) instanceof SqlExecutableStatement;
   }

   public void executeDdl(String sqlStatement) throws ParseException {
     SqlExecutableStatement ddl = (SqlExecutableStatement) planner.parse(sqlStatement);
     ddl.execute(getContext());
   }

   public CalcitePrepare.Context getContext() {
     return connection.createPrepareContext();
   }

   public Map<String, String> getPipelineOptions() {
     return connection.getPipelineOptionsMap();
   }

   public String explain(String sqlString) throws ParseException {
     try {
       return RelOptUtil.toString(planner.convertToBeamRel(sqlString));
     } catch (Exception e) {
       throw new ParseException("Unable to parse statement", e);
     }
   }

   /** BeamSqlEnv's Builder. */
   public static class BeamSqlEnvBuilder {
     private static final String CALCITE_PLANNER =
         "org.apache.beam.sdk.extensions.sql.impl.CalciteQueryPlanner";
     private String queryPlannerClassName;
     private TableProvider defaultTableProvider;
     private String currentSchemaName;
     private Map<String, TableProvider> schemaMap;
     private Set<Map.Entry<String, Function>> functionSet;
     private boolean autoLoadBuiltinFunctions;
     private boolean autoLoadUdfs;

     private BeamSqlEnvBuilder(TableProvider tableProvider) {
       checkNotNull(tableProvider, "Table provider for the default schema must be sets.");

       defaultTableProvider = tableProvider;
       queryPlannerClassName = CALCITE_PLANNER;
       schemaMap = new HashMap<>();
       functionSet = new HashSet<>();
       autoLoadUdfs = false;
       autoLoadBuiltinFunctions = false;
     }

     /** Add a top-level schema backed by the table provider. */
     public BeamSqlEnvBuilder addSchema(String name, TableProvider tableProvider) {
       if (schemaMap.containsKey(name)) {
         throw new RuntimeException("Schema " + name + " is registered twice.");
       }

       schemaMap.put(name, tableProvider);
       return this;
     }

     /** Set the current (default) schema. */
     public BeamSqlEnvBuilder setCurrentSchema(String name) {
       currentSchemaName = name;
       return this;
     }

     /** Register a UDF function which can be used in SQL expression. */
     public BeamSqlEnvBuilder addUdf(String functionName, Class<?> clazz, String method) {
       functionSet.add(new SimpleEntry<>(functionName, UdfImpl.create(clazz, method)));
       return this;
     }

     /** Register a UDF function which can be used in SQL expression. */
     public BeamSqlEnvBuilder addUdf(String functionName, Class<? extends BeamSqlUdf> clazz) {
       return addUdf(functionName, clazz, BeamSqlUdf.UDF_METHOD);
     }

     /** Register a UDF function which can be used in SQL expression. */
     public BeamSqlEnvBuilder addUdf(String functionName, SerializableFunction sfn) {
       return addUdf(functionName, sfn.getClass(), "apply");
     }

     /**
      * Register a UDAF function which can be used in GROUP-BY expression.
      *
      * <p>See {@link CombineFn} on how to implement a UDAF.
      */
     public BeamSqlEnvBuilder addUdaf(String functionName, CombineFn combineFn) {
       functionSet.add(new SimpleEntry<>(functionName, new UdafImpl(combineFn)));
       return this;
     }

     /** Load UDF/UDAFs from {@link UdfUdafProvider}. */
     public BeamSqlEnvBuilder autoLoadUserDefinedFunctions() {
       autoLoadUdfs = true;
       return this;
     }

     /** Load Beam SQL built-in functions defined in {@link BeamBuiltinFunctionProvider}. */
     public BeamSqlEnvBuilder autoLoadBuiltinFunctions() {
       autoLoadBuiltinFunctions = true;
       return this;
     }

     public BeamSqlEnvBuilder setQueryPlannerClassName(String name) {
       queryPlannerClassName = name;
       return this;
     }

     /**
      * Build function to create an instance of BeamSqlEnv based on preset fields.
      *
      * @return BeamSqlEnv.
      */
     public BeamSqlEnv build() {

       JdbcConnection jdbcConnection = JdbcDriver.connect(defaultTableProvider);

       configureSchemas(jdbcConnection);

       loadBeamBuiltinFunctions();

       loadUdfs();

       addUdfsUdafs(jdbcConnection);

       QueryPlanner planner = instantiatePlanner(jdbcConnection);

       return new BeamSqlEnv(jdbcConnection, planner);
     }

     private void configureSchemas(JdbcConnection jdbcConnection) {
       // SetSchema adds the schema with the specified name
       // backed by the table provider.
       // Does not update the current default schema.
       schemaMap.forEach(jdbcConnection::setSchema);

       if (Strings.isNullOrEmpty(currentSchemaName)) {
         return;
       }

       try {
         jdbcConnection.setSchema(currentSchemaName);
       } catch (SQLException e) {
         throw new RuntimeException(e);
       }
     }

     private void loadBeamBuiltinFunctions() {
       if (!autoLoadBuiltinFunctions) {
         return;
       }

       for (BeamBuiltinFunctionProvider provider :
           ServiceLoader.load(BeamBuiltinFunctionProvider.class)) {
         loadBuiltinUdf(provider.getBuiltinMethods());
       }
     }

     private void loadBuiltinUdf(Map<String, List<Method>> methods) {
       for (Map.Entry<String, List<Method>> entry : methods.entrySet()) {
         for (Method method : entry.getValue()) {
           functionSet.add(new SimpleEntry<>(entry.getKey(), UdfImpl.create(method)));
         }
       }
     }

     private void loadUdfs() {
       if (!autoLoadUdfs) {
         return;
       }

       ServiceLoader.load(UdfUdafProvider.class)
           .forEach(
               ins -> {
                 ins.getBeamSqlUdfs().forEach(this::addUdf);
                 ins.getSerializableFunctionUdfs().forEach(this::addUdf);
                 ins.getUdafs().forEach(this::addUdaf);
               });
     }

     private void addUdfsUdafs(JdbcConnection connection) {
       for (Map.Entry<String, Function> functionEntry : functionSet) {
         connection.getCurrentSchemaPlus().add(functionEntry.getKey(), functionEntry.getValue());
       }
     }

     private QueryPlanner instantiatePlanner(JdbcConnection jdbcConnection) {

       if (queryPlannerClassName.equals(CALCITE_PLANNER)) {
         return new CalciteQueryPlanner(jdbcConnection);
       }

       try {
         return (QueryPlanner)
             Class.forName(queryPlannerClassName)
                 .getConstructor(JdbcConnection.class)
                 .newInstance(jdbcConnection);
       } catch (Exception e) {
         throw new RuntimeException(
             String.format("Cannot construct query planner %s", queryPlannerClassName), e);
       }
     }
   }
 }
	/*
	* 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.vendor.guava.v20_0.com.google.common.base.Preconditions.checkNotNull;

	import java.lang.reflect.Method;
	import java.sql.SQLException;
	import java.util.AbstractMap.SimpleEntry;
	import java.util.HashMap;
	import java.util.HashSet;
	import java.util.List;
	import java.util.Map;
	import java.util.ServiceLoader;
	import java.util.Set;
	import org.apache.beam.sdk.annotations.Experimental;
	import org.apache.beam.sdk.annotations.Internal;
	import org.apache.beam.sdk.extensions.sql.BeamSqlTable;
	import org.apache.beam.sdk.extensions.sql.BeamSqlUdf;
	import org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode;
	import org.apache.beam.sdk.extensions.sql.impl.udf.BeamBuiltinFunctionProvider;
	import org.apache.beam.sdk.extensions.sql.meta.provider.ReadOnlyTableProvider;
	import org.apache.beam.sdk.extensions.sql.meta.provider.TableProvider;
	import org.apache.beam.sdk.extensions.sql.meta.provider.UdfUdafProvider;
	import org.apache.beam.sdk.extensions.sql.meta.store.InMemoryMetaStore;
	import org.apache.beam.sdk.transforms.Combine.CombineFn;
	import org.apache.beam.sdk.transforms.SerializableFunction;
	import org.apache.beam.vendor.guava.v20_0.com.google.common.base.Strings;
	import org.apache.calcite.jdbc.CalcitePrepare;
	import org.apache.calcite.plan.RelOptUtil;
	import org.apache.calcite.schema.Function;
	import org.apache.calcite.sql.SqlExecutableStatement;

	/**
	* Contains the metadata of tables/UDF functions, and exposes APIs to
	* query/validate/optimize/translate SQL statements.
	*/
	@Internal
	@Experimental
	public class BeamSqlEnv {
	JdbcConnection connection;
	QueryPlanner planner;

	private BeamSqlEnv(JdbcConnection connection, QueryPlanner planner) {
	this.connection = connection;
	this.planner = planner;
	}

	/** Creates a builder with the default schema backed by the table provider. */
	public static BeamSqlEnvBuilder builder(TableProvider tableProvider) {
	return new BeamSqlEnvBuilder(tableProvider);
	}

	public static BeamSqlEnv readOnly(String tableType, Map<String, BeamSqlTable> tables) {
	return withTableProvider(new ReadOnlyTableProvider(tableType, tables));
	}

	public static BeamSqlEnv withTableProvider(TableProvider tableProvider) {
	return builder(tableProvider).build();
	}

	public static BeamSqlEnv inMemory(TableProvider... tableProviders) {
	InMemoryMetaStore inMemoryMetaStore = new InMemoryMetaStore();
	for (TableProvider tableProvider : tableProviders) {
	inMemoryMetaStore.registerProvider(tableProvider);
	}

	return withTableProvider(inMemoryMetaStore);
	}

	public BeamRelNode parseQuery(String query) throws ParseException {
	return planner.convertToBeamRel(query);
	}

	public boolean isDdl(String sqlStatement) throws ParseException {
	return planner.parse(sqlStatement) instanceof SqlExecutableStatement;
	}

	public void executeDdl(String sqlStatement) throws ParseException {
	SqlExecutableStatement ddl = (SqlExecutableStatement) planner.parse(sqlStatement);
	ddl.execute(getContext());
	}

	public CalcitePrepare.Context getContext() {
	return connection.createPrepareContext();
	}

	public Map<String, String> getPipelineOptions() {
	return connection.getPipelineOptionsMap();
	}

	public String explain(String sqlString) throws ParseException {
	try {
	return RelOptUtil.toString(planner.convertToBeamRel(sqlString));
	} catch (Exception e) {
	throw new ParseException("Unable to parse statement", e);
	}
	}

	/** BeamSqlEnv's Builder. */
	public static class BeamSqlEnvBuilder {
	private static final String CALCITE_PLANNER =
	"org.apache.beam.sdk.extensions.sql.impl.CalciteQueryPlanner";
	private String queryPlannerClassName;
	private TableProvider defaultTableProvider;
	private String currentSchemaName;
	private Map<String, TableProvider> schemaMap;
	private Set<Map.Entry<String, Function>> functionSet;
	private boolean autoLoadBuiltinFunctions;
	private boolean autoLoadUdfs;

	private BeamSqlEnvBuilder(TableProvider tableProvider) {
	checkNotNull(tableProvider, "Table provider for the default schema must be sets.");

	defaultTableProvider = tableProvider;
	queryPlannerClassName = CALCITE_PLANNER;
	schemaMap = new HashMap<>();
	functionSet = new HashSet<>();
	autoLoadUdfs = false;
	autoLoadBuiltinFunctions = false;
	}

	/** Add a top-level schema backed by the table provider. */
	public BeamSqlEnvBuilder addSchema(String name, TableProvider tableProvider) {
	if (schemaMap.containsKey(name)) {
	throw new RuntimeException("Schema " + name + " is registered twice.");
	}

	schemaMap.put(name, tableProvider);
	return this;
	}

	/** Set the current (default) schema. */
	public BeamSqlEnvBuilder setCurrentSchema(String name) {
	currentSchemaName = name;
	return this;
	}

	/** Register a UDF function which can be used in SQL expression. */
	public BeamSqlEnvBuilder addUdf(String functionName, Class<?> clazz, String method) {
	functionSet.add(new SimpleEntry<>(functionName, UdfImpl.create(clazz, method)));
	return this;
	}

	/** Register a UDF function which can be used in SQL expression. */
	public BeamSqlEnvBuilder addUdf(String functionName, Class<? extends BeamSqlUdf> clazz) {
	return addUdf(functionName, clazz, BeamSqlUdf.UDF_METHOD);
	}

	/** Register a UDF function which can be used in SQL expression. */
	public BeamSqlEnvBuilder addUdf(String functionName, SerializableFunction sfn) {
	return addUdf(functionName, sfn.getClass(), "apply");
	}

	/**
	* Register a UDAF function which can be used in GROUP-BY expression.
	*
	* <p>See {@link CombineFn} on how to implement a UDAF.
	*/
	public BeamSqlEnvBuilder addUdaf(String functionName, CombineFn combineFn) {
	functionSet.add(new SimpleEntry<>(functionName, new UdafImpl(combineFn)));
	return this;
	}

	/** Load UDF/UDAFs from {@link UdfUdafProvider}. */
	public BeamSqlEnvBuilder autoLoadUserDefinedFunctions() {
	autoLoadUdfs = true;
	return this;
	}

	/** Load Beam SQL built-in functions defined in {@link BeamBuiltinFunctionProvider}. */
	public BeamSqlEnvBuilder autoLoadBuiltinFunctions() {
	autoLoadBuiltinFunctions = true;
	return this;
	}

	public BeamSqlEnvBuilder setQueryPlannerClassName(String name) {
	queryPlannerClassName = name;
	return this;
	}

	/**
	* Build function to create an instance of BeamSqlEnv based on preset fields.
	*
	* @return BeamSqlEnv.
	*/
	public BeamSqlEnv build() {

	JdbcConnection jdbcConnection = JdbcDriver.connect(defaultTableProvider);

	configureSchemas(jdbcConnection);

	loadBeamBuiltinFunctions();

	loadUdfs();

	addUdfsUdafs(jdbcConnection);

	QueryPlanner planner = instantiatePlanner(jdbcConnection);

	return new BeamSqlEnv(jdbcConnection, planner);
	}

	private void configureSchemas(JdbcConnection jdbcConnection) {
	// SetSchema adds the schema with the specified name
	// backed by the table provider.
	// Does not update the current default schema.
	schemaMap.forEach(jdbcConnection::setSchema);

	if (Strings.isNullOrEmpty(currentSchemaName)) {
	return;
	}

	try {
	jdbcConnection.setSchema(currentSchemaName);
	} catch (SQLException e) {
	throw new RuntimeException(e);
	}
	}

	private void loadBeamBuiltinFunctions() {
	if (!autoLoadBuiltinFunctions) {
	return;
	}

	for (BeamBuiltinFunctionProvider provider :
	ServiceLoader.load(BeamBuiltinFunctionProvider.class)) {
	loadBuiltinUdf(provider.getBuiltinMethods());
	}
	}

	private void loadBuiltinUdf(Map<String, List<Method>> methods) {
	for (Map.Entry<String, List<Method>> entry : methods.entrySet()) {
	for (Method method : entry.getValue()) {
	functionSet.add(new SimpleEntry<>(entry.getKey(), UdfImpl.create(method)));
	}
	}
	}

	private void loadUdfs() {
	if (!autoLoadUdfs) {
	return;
	}

	ServiceLoader.load(UdfUdafProvider.class)
	.forEach(
	ins -> {
	ins.getBeamSqlUdfs().forEach(this::addUdf);
	ins.getSerializableFunctionUdfs().forEach(this::addUdf);
	ins.getUdafs().forEach(this::addUdaf);
	});
	}

	private void addUdfsUdafs(JdbcConnection connection) {
	for (Map.Entry<String, Function> functionEntry : functionSet) {
	connection.getCurrentSchemaPlus().add(functionEntry.getKey(), functionEntry.getValue());
	}
	}

	private QueryPlanner instantiatePlanner(JdbcConnection jdbcConnection) {

	if (queryPlannerClassName.equals(CALCITE_PLANNER)) {
	return new CalciteQueryPlanner(jdbcConnection);
	}

	try {
	return (QueryPlanner)
	Class.forName(queryPlannerClassName)
	.getConstructor(JdbcConnection.class)
	.newInstance(jdbcConnection);
	} catch (Exception e) {
	throw new RuntimeException(
	String.format("Cannot construct query planner %s", queryPlannerClassName), e);
	}
	}
	}
	}