Added SQL

Author: cgivre

exec/java-exec/src/main/java/org/apache/drill/exec/ops/QueryContext.java

@@ -35,7 +35,7 @@
 import org.apache.drill.exec.physical.resultSet.ResultSetLoader;
 import org.apache.drill.exec.physical.resultSet.impl.ResultSetLoaderImpl;
 import org.apache.drill.exec.physical.resultSet.impl.ResultSetLoaderImpl.ResultSetOptions;
-import org.apache.drill.exec.planner.common.DaffodilSchemaRegistry;
+import org.apache.drill.exec.schema.daffodil.RemoteDaffodilSchemaRegistry;
 import org.apache.drill.exec.planner.physical.PlannerSettings;
 import org.apache.drill.exec.planner.sql.DrillOperatorTable;
 import org.apache.drill.exec.proto.BitControl.QueryContextInformation;
@@ -318,8 +318,8 @@ public RemoteFunctionRegistry getRemoteFunctionRegistry() {
     return drillbitContext.getRemoteFunctionRegistry();
   }
 
-  public DaffodilSchemaRegistry getDaffodilSchemaRegistry() {
-    return drillbitContext.getDaffodilSchemaProvider();
+  public RemoteDaffodilSchemaRegistry getDaffodilSchemaRegistry() {
+    return drillbitContext.getRemoteDaffodilSchemaRegistry();
   }
 
   @Override

exec/java-exec/src/main/java/org/apache/drill/exec/planner/sql/handlers/CreateDaffodilSchemaHandler.java

@@ -0,0 +1,152 @@
+/*
+ * 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.drill.exec.planner.sql.handlers;
+
+import com.google.common.collect.Lists;
+import org.apache.calcite.sql.SqlCharStringLiteral;
+import org.apache.calcite.sql.SqlNode;
+import org.apache.drill.common.exceptions.DrillRuntimeException;
+import org.apache.drill.exec.exception.VersionMismatchException;
+import org.apache.drill.exec.physical.PhysicalPlan;
+import org.apache.drill.exec.planner.sql.DirectPlan;
+import org.apache.drill.exec.planner.sql.parser.SqlDropDaffodilSchema;
+import org.apache.drill.exec.proto.UserBitShared.Jar;
+import org.apache.drill.exec.proto.UserBitShared.Registry;
+import org.apache.drill.exec.schema.daffodil.RemoteDaffodilSchemaRegistry;
+import org.apache.drill.exec.store.sys.store.DataChangeVersion;
+import org.apache.drill.exec.work.foreman.ForemanSetupException;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.util.List;
+
+public class DropDaffodilSchemaHandler extends DefaultSqlHandler {
+  private static Logger logger = LoggerFactory.getLogger(DropDaffodilSchemaHandler.class);
+
+  public DropDaffodilSchemaHandler(SqlHandlerConfig config) {
+    super(config);
+  }
+
+  /**
+   * Unregisters Daffodil schema JARs dynamically. Process consists of several steps:
+   * <ol>
+   * <li>Registering jar in jar registry to ensure that the jar is not being unregistered elsewhere.</li>
+   * <li>Starts remote unregistration process, gets list of all jars and excludes jar to be deleted.</li>
+   * <li>Removes jar from registry area.</li>
+   * </ol>
+   *
+   * Only jars registered dynamically can be unregistered.
+   *
+   * Limitation: before jar unregistration make sure no one is using schemas from this jar.
+   * There is no guarantee that running queries will finish successfully or give correct result.
+   *
+   * @return - Single row indicating successful unregistration, raise exception otherwise
+   */
+  @Override
+  public PhysicalPlan getPlan(SqlNode sqlNode) throws ForemanSetupException, IOException {
+    SqlDropDaffodilSchema node = unwrap(sqlNode, SqlDropDaffodilSchema.class);
+    String jarName = ((SqlCharStringLiteral) node.getJar()).toValue();
+    RemoteDaffodilSchemaRegistry remoteSchemaRegistry = context.getDaffodilSchemaRegistry();
+
+    boolean inProgress = false;
+    try {
+      final String action = remoteSchemaRegistry.addToJars(jarName, RemoteDaffodilSchemaRegistry.Action.UNREGISTRATION);
+      if (!(inProgress = action == null)) {
+        return DirectPlan.createDirectPlan(context, false, String.format("Jar with %s name is used. Action: %s", jarName, action));
+      }
+
+      Jar deletedJar = unregister(jarName, remoteSchemaRegistry);
+      if (deletedJar == null) {
+        return DirectPlan.createDirectPlan(context, false, String.format("Jar %s is not registered in remote registry", jarName));
+      }
+
+      removeJarFromArea(jarName, remoteSchemaRegistry.getFs(), remoteSchemaRegistry.getRegistryArea());
+
+      return DirectPlan.createDirectPlan(context, true,
+          String.format("Daffodil schema jar %s has been unregistered successfully.", jarName));
+
+    } catch (Exception e) {
+      logger.error("Error during Daffodil schema unregistration", e);
+      return DirectPlan.createDirectPlan(context, false, e.getMessage());
+    } finally {
+      if (inProgress) {
+        remoteSchemaRegistry.removeFromJars(jarName);
+      }
+    }
+  }
+
+  /**
+   * Gets remote schema registry with version.
+   * Version is used to ensure that we update the same registry we removed jars from.
+   * Looks for a jar to be deleted, if found,
+   * attempts to update remote registry with list of jars, that excludes jar to be deleted.
+   * If during update {@link VersionMismatchException} was detected,
+   * attempts to repeat unregistration process till retry attempts exceeds the limit.
+   * If retry attempts number hits 0, throws exception that failed to update remote schema registry.
+   *
+   * @param jarName jar name
+   * @param remoteSchemaRegistry remote schema registry
+   * @return jar that was unregistered, null otherwise
+   */
+  private Jar unregister(String jarName, RemoteDaffodilSchemaRegistry remoteSchemaRegistry) {
+    int retryAttempts = remoteSchemaRegistry.getRetryAttempts();
+    while (retryAttempts >= 0) {
+      DataChangeVersion version = new DataChangeVersion();
+      Registry registry = remoteSchemaRegistry.getRegistry(version);
+      Jar jarToBeDeleted = null;
+      List<Jar> jars = Lists.newArrayList();
+      for (Jar j : registry.getJarList()) {
+        if (j.getName().equals(jarName)) {
+          jarToBeDeleted = j;
+        } else {
+          jars.add(j);
+        }
+      }
+      if (jarToBeDeleted == null) {
+        return null;
+      }
+      Registry updatedRegistry = Registry.newBuilder().addAllJar(jars).build();
+      try {
+        remoteSchemaRegistry.updateRegistry(updatedRegistry, version);
+        return jarToBeDeleted;
+      } catch (VersionMismatchException ex) {
+        logger.debug("Failed to update schema registry during unregistration, version mismatch was detected.", ex);
+        retryAttempts--;
+      }
+    }
+    throw new DrillRuntimeException("Failed to update remote schema registry. Exceeded retry attempts limit.");
+  }
+
+  /**
+   * Removes jar from indicated area, in case of error log it and proceeds.
+   *
+   * @param jarName jar name
+   * @param fs file system
+   * @param area path to area
+   */
+  private void removeJarFromArea(String jarName, FileSystem fs, Path area) {
+    try {
+      fs.delete(new Path(area, jarName), false);
+    } catch (IOException e) {
+      logger.error("Error removing jar {} from area {}", jarName, area.toUri().getPath());
+    }
+  }
+}

exec/java-exec/src/main/java/org/apache/drill/exec/planner/sql/handlers/DropDaffodilSchemaHandler.java

@@ -27,7 +27,7 @@
 import org.apache.calcite.sql.SqlWriter;
 import org.apache.calcite.sql.parser.SqlParserPos;
 import org.apache.drill.exec.planner.sql.handlers.AbstractSqlHandler;
-import org.apache.drill.exec.planner.sql.handlers.DropFunctionHandler;
+import org.apache.drill.exec.planner.sql.handlers.DropDaffodilSchemaHandler;
 import org.apache.drill.exec.planner.sql.handlers.SqlHandlerConfig;
 
 import java.util.List;
@@ -72,7 +72,7 @@ public void unparse(SqlWriter writer, int leftPrec, int rightPrec) {
 
   @Override
   public AbstractSqlHandler getSqlHandler(SqlHandlerConfig config) {
-    return new DropFunctionHandler(config);
+    return new DropDaffodilSchemaHandler(config);
   }
 
   public SqlNode getJar() { return jar; }

exec/java-exec/src/main/java/org/apache/drill/exec/planner/sql/parser/SqlDropDaffodilSchema.java

@@ -17,60 +17,40 @@
  */
 package org.apache.drill.exec.schema.daffodil;
 
-import org.apache.commons.lang3.StringUtils;
 import org.apache.drill.common.AutoCloseables;
 import org.apache.drill.common.config.DrillConfig;
 import org.apache.drill.common.scanner.ClassPathScanner;
 import org.apache.drill.common.scanner.persistence.ScanResult;
-import org.apache.drill.exec.oauth.PersistentTokenRegistry;
-import org.apache.drill.exec.oauth.TokenRegistry;
+import org.apache.drill.exec.coord.ClusterCoordinator;
 import org.apache.drill.exec.server.DrillbitContext;
+import org.apache.drill.exec.store.sys.PersistentStoreProvider;
 
 /**
- * Class for managing daffodil schemata.  Schemata will be obtained via INSTALL/CREATE SCHEMA queries.
+ * Class for managing Daffodil schemata.  Schemata will be obtained via CREATE DAFFODIL SCHEMA queries.
  */
 public class DaffodilSchemaProvider implements AutoCloseable {
-  private static final String STORAGE_REGISTRY_PATH = "daffodil_schema";
 
-  private final DrillbitContext context;
-
-  private PersistentTokenRegistry daffodilSchemaRegistry;
+  private RemoteDaffodilSchemaRegistry remoteDaffodilSchemaRegistry;
 
   public DaffodilSchemaProvider(DrillbitContext context) {
-    this(context.getConfig(), ClassPathScanner.fromPrescan(context.getConfig()));
+    this(context.getConfig(), context.getStoreProvider(), context.getClusterCoordinator());
   }
 
   public DaffodilSchemaProvider(DrillConfig config, ScanResult classpathScan) {
-    this(config, classpathScan, null);
+    // This constructor is incomplete - needs StoreProvider and ClusterCoordinator
   }
 
-  public DaffodilSchemaProvider(DrillConfig config, ScanResult classpathScan, String username) {
-
+  public DaffodilSchemaProvider(DrillConfig config, PersistentStoreProvider storeProvider, ClusterCoordinator coordinator) {
+    this.remoteDaffodilSchemaRegistry = new RemoteDaffodilSchemaRegistry();
+    this.remoteDaffodilSchemaRegistry.init(config, storeProvider, coordinator);
   }
 
-  public TokenRegistry getDaffodilSchemaRegistry(String username) {
-    if (daffodilSchemaRegistry == null) {
-      initRemoteRegistries(username);
-    }
-    return daffodilSchemaRegistry;
-  }
-
-  private synchronized void initRemoteRegistries(String username) {
-    // Add the username to the path if present
-    String finalpath;
-    if (StringUtils.isNotEmpty(username)) {
-      finalpath = STORAGE_REGISTRY_PATH + "/" + username;
-    } else {
-      finalpath = STORAGE_REGISTRY_PATH;
-    }
-
-    if (daffodilSchemaRegistry == null) {
-      daffodilSchemaRegistry = new PersistentTokenRegistry(context, finalpath);
-    }
+  public RemoteDaffodilSchemaRegistry getRemoteDaffodilSchemaRegistry() {
+    return remoteDaffodilSchemaRegistry;
   }
 
   @Override
   public void close() throws Exception {
-    AutoCloseables.closeSilently(daffodilSchemaRegistry);
+    AutoCloseables.closeSilently(remoteDaffodilSchemaRegistry);
   }
 }