exec/jdbc-all/src/test/java/org/apache/drill/jdbc/ITTestShadedJar.java - drill - 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.drill.jdbc;

 import org.apache.drill.test.BaseTest;
 import org.junit.ClassRule;
 import org.junit.Test;
 import org.junit.rules.TestWatcher;
 import org.junit.runner.Description;

 import java.io.BufferedReader;
 import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.lang.reflect.Field;
 import java.lang.reflect.Method;
 import java.net.MalformedURLException;
 import java.net.URL;
 import java.net.URLClassLoader;
 import java.sql.Connection;
 import java.sql.DatabaseMetaData;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.Vector;
 import java.util.concurrent.Semaphore;
 import java.util.stream.Collectors;

 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;

 // Note: this test cannot be run under surefire (or in IDEs). Execute
 // `mvn integration-test` in this module's directory to run it locally. See the
 // maven-failsafe-plugin config in this module's pom.xml for more information.
 public class ITTestShadedJar extends BaseTest {
   private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(ITTestShadedJar.class);

   private static DrillbitClassLoader drillbitLoader;
   private static ClassLoader rootClassLoader;
   private static int userPort;

   @ClassRule
   public static final TestWatcher testWatcher = new TestWatcher() {

     @Override
     protected void starting(Description description) {
       super.starting(description);

       try {
         drillbitLoader = new DrillbitClassLoader();
         drillbitLoader.loadClass("org.apache.commons.io.FileUtils");
         rootClassLoader = Thread.currentThread().getContextClassLoader();

         Class<?> clazz = drillbitLoader.loadClass("org.apache.drill.test.BaseTestQuery");
         Class<?> watcherClazz = drillbitLoader.loadClass("org.apache.drill.test.BaseDirTestWatcher");

         // Setup tmp dirs
         runMethod("starting", description);

         // Configure tmp dir
         Object watcher = clazz.getField("dirTestWatcher").get(null);
         Method method = watcherClazz.getDeclaredMethod("getTmpDir");
         File tmpDir = (File) method.invoke(watcher);
         System.setProperty("DRILL_CONF_DIR", tmpDir.getAbsolutePath());

         // Start the drill cluster
         try {
           runWithLoader("DrillbitStartThread", drillbitLoader);
         } catch (Exception e) {
           printClassesLoaded("root", rootClassLoader);
           throw e;
         }

         DrillbitStartThread.SEM.acquire();

         // After starting the drill cluster get the client port
         userPort = (Integer) clazz.getMethod("getUserPort").invoke(null);
       } catch (Exception e) {
         throw new RuntimeException(e);
       }
     }

     @Override
     protected void finished(Description description) {
       super.finished(description);
       done();
       runMethod("finished", description);
     }

     @Override
     protected void failed(Throwable e, Description description) {
       super.failed(e, description);
       done();
       runMethod("failed", description);
     }

     private void done() {
       try {
         runWithLoader("DrillbitStopThread", drillbitLoader);
         DrillbitStopThread.SEM.acquire();
       } catch (Exception e) {
         throw new RuntimeException(e);
       }
     }

     private void runMethod(String name, Description description) {
       try {
         Class<?> clazz = drillbitLoader.loadClass("org.apache.drill.test.BaseTestQuery");
         Class<?> watcherClazz = drillbitLoader.loadClass("org.junit.rules.TestWatcher");
         Object watcher = clazz.getField("dirTestWatcher").get(null);
         Method method = watcherClazz.getDeclaredMethod(name, Description.class);
         method.setAccessible(true);
         method.invoke(watcher, description);
       } catch (Exception e) {
         throw new RuntimeException(e);
       }
     }
   };

   private static URL getJdbcUrl() throws MalformedURLException {
     return new URL(
         String.format("%s../../target/drill-jdbc-all-%s.jar",
             ClassLoader.getSystemClassLoader().getResource("").toString(),
             System.getProperty("project.version")
             ));

   }

   @Test
   public void testDatabaseVersion() throws Exception {
     final URLClassLoader loader = URLClassLoader.newInstance(new URL[] {getJdbcUrl()});

     Class<?> clazz = loader.loadClass("org.apache.drill.jdbc.Driver");
     try {
       Driver driver = (Driver) clazz.newInstance();
       try (Connection c = driver.connect("jdbc:drill:drillbit=localhost:" + userPort, null)) {
         DatabaseMetaData metadata = c.getMetaData();
         assertEquals("Apache Drill JDBC Driver", metadata.getDriverName());
         assertEquals("Apache Drill Server", metadata.getDatabaseProductName());
       }
     } catch (Exception ex) {
       throw ex;
     }

   }

   @Test
   public void executeJdbcAllQuery() throws Exception {
     final URLClassLoader loader = URLClassLoader.newInstance(new URL[] {getJdbcUrl()});

     Class<?> clazz = loader.loadClass("org.apache.drill.jdbc.Driver");
     try {
       Driver driver = (Driver) clazz.newInstance();
       try (Connection c = driver.connect("jdbc:drill:drillbit=localhost:" + userPort, null)) {
         printQuery(c, "select * from cp.`types.json`");
       }
     } catch (Exception ex) {
       throw ex;
     }

   }

   @Test
   public void serviceFileContainsCorrectDriver() throws IOException {
     URLClassLoader loader = URLClassLoader.newInstance(new URL[]{getJdbcUrl()});
     try (InputStream resourceStream = loader.getResourceAsStream("META-INF/services/java.sql.Driver")) {
       assertNotNull("java.sql.Driver is not present in the jdbc jar", resourceStream);
       try (BufferedReader reader = new BufferedReader(new InputStreamReader(resourceStream))) {
         String driverClass = reader.lines().collect(Collectors.joining(System.lineSeparator()));
         assertEquals("org.apache.drill.jdbc.Driver", driverClass);
       }
     }
   }

   private static void printQuery(Connection c, String query) throws SQLException {
     final StringBuilder sb = new StringBuilder();

     try (Statement s = c.createStatement();
          ResultSet result = s.executeQuery(query)) {
       while (result.next()) {
         final int columnCount = result.getMetaData().getColumnCount();
         for(int i = 1; i < columnCount+1; i++){
           sb.append(result.getObject(i));
           sb.append('\t');
         }

         sb.append(result.getObject(1)).append('\n');
       }
     }

     logger.info(sb.toString());
   }

   private static int getClassesLoadedCount(ClassLoader classLoader) {
     try {
       Field f = ClassLoader.class.getDeclaredField("classes");
       f.setAccessible(true);
       Vector<Class<?>> classes = (Vector<Class<?>>) f.get(classLoader);
       return classes.size();
     } catch (Exception e) {
       logger.error("Exception ", e);
       return -1;
     }
   }

   private static void printClassesLoaded(String prefix, ClassLoader classLoader) {
     final StringBuilder sb = new StringBuilder();

     try {
       Field f = ClassLoader.class.getDeclaredField("classes");
       f.setAccessible(true);
       Vector<Class<?>> classes = (Vector<Class<?>>) f.get(classLoader);
       for (Class<?> c : classes) {
         sb.append(prefix).append(": ").append(c.getName()).append('\n');
       }
     } catch (Exception e) {
       logger.error("Exception ", e);
     }

     logger.info(sb.toString());
   }

   private static void runWithLoader(String name, ClassLoader loader) throws Exception {
     Class<?> clazz = loader.loadClass(ITTestShadedJar.class.getName() + "$" + name);
     Object instance = clazz.getDeclaredConstructors()[0].newInstance(loader);
     clazz.getMethod("go").invoke(instance);
   }

   public abstract static class AbstractLoaderThread extends Thread {
     private Exception ex;
     protected final ClassLoader loader;

     public AbstractLoaderThread(ClassLoader loader) {
       this.setContextClassLoader(loader);
       this.loader = loader;
     }

     @Override
     public final void run() {
       try {
         internalRun();
       } catch (Exception e) {
         this.ex = e;
       }
     }

     protected abstract void internalRun() throws Exception;

     public void go() throws Exception {
       start();
       join();
       if (ex != null) {
         throw ex;
       }
     }
   }

   public static class DrillbitStartThread extends AbstractLoaderThread {
     public static final Semaphore SEM = new Semaphore(0);

     public DrillbitStartThread(ClassLoader loader) {
       super(loader);
     }

     @Override
     protected void internalRun() throws Exception {
       Class<?> clazz = loader.loadClass("org.apache.drill.test.BaseTestQuery");
       clazz.getMethod("setupDefaultTestCluster").invoke(null);

       // loader.loadClass("org.apache.drill.exec.exception.SchemaChangeException");

       // execute a single query to make sure the drillbit is fully up
       clazz.getMethod("testNoResult", String.class, Object[].class)
           .invoke(null, "select * from (VALUES 1)", new Object[] {});

       SEM.release();
     }

   }

   public static class DrillbitStopThread extends AbstractLoaderThread {
     public static final Semaphore SEM = new Semaphore(0);

     public DrillbitStopThread(ClassLoader loader) {
       super(loader);
     }

     @Override
     protected void internalRun() throws Exception {
       Class<?> clazz = loader.loadClass("org.apache.drill.test.BaseTestQuery");
       clazz.getMethod("closeClient").invoke(null);
       SEM.release();
     }
   }

 }
	/*
	* 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.jdbc;

	import org.apache.drill.test.BaseTest;
	import org.junit.ClassRule;
	import org.junit.Test;
	import org.junit.rules.TestWatcher;
	import org.junit.runner.Description;

	import java.io.BufferedReader;
	import java.io.File;
	import java.io.IOException;
	import java.io.InputStream;
	import java.io.InputStreamReader;
	import java.lang.reflect.Field;
	import java.lang.reflect.Method;
	import java.net.MalformedURLException;
	import java.net.URL;
	import java.net.URLClassLoader;
	import java.sql.Connection;
	import java.sql.DatabaseMetaData;
	import java.sql.ResultSet;
	import java.sql.SQLException;
	import java.sql.Statement;
	import java.util.Vector;
	import java.util.concurrent.Semaphore;
	import java.util.stream.Collectors;

	import static org.junit.Assert.assertEquals;
	import static org.junit.Assert.assertNotNull;

	// Note: this test cannot be run under surefire (or in IDEs). Execute
	// `mvn integration-test` in this module's directory to run it locally. See the
	// maven-failsafe-plugin config in this module's pom.xml for more information.
	public class ITTestShadedJar extends BaseTest {
	private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(ITTestShadedJar.class);

	private static DrillbitClassLoader drillbitLoader;
	private static ClassLoader rootClassLoader;
	private static int userPort;

	@ClassRule
	public static final TestWatcher testWatcher = new TestWatcher() {

	@Override
	protected void starting(Description description) {
	super.starting(description);

	try {
	drillbitLoader = new DrillbitClassLoader();
	drillbitLoader.loadClass("org.apache.commons.io.FileUtils");
	rootClassLoader = Thread.currentThread().getContextClassLoader();

	Class<?> clazz = drillbitLoader.loadClass("org.apache.drill.test.BaseTestQuery");
	Class<?> watcherClazz = drillbitLoader.loadClass("org.apache.drill.test.BaseDirTestWatcher");

	// Setup tmp dirs
	runMethod("starting", description);

	// Configure tmp dir
	Object watcher = clazz.getField("dirTestWatcher").get(null);
	Method method = watcherClazz.getDeclaredMethod("getTmpDir");
	File tmpDir = (File) method.invoke(watcher);
	System.setProperty("DRILL_CONF_DIR", tmpDir.getAbsolutePath());

	// Start the drill cluster
	try {
	runWithLoader("DrillbitStartThread", drillbitLoader);
	} catch (Exception e) {
	printClassesLoaded("root", rootClassLoader);
	throw e;
	}

	DrillbitStartThread.SEM.acquire();

	// After starting the drill cluster get the client port
	userPort = (Integer) clazz.getMethod("getUserPort").invoke(null);
	} catch (Exception e) {
	throw new RuntimeException(e);
	}
	}

	@Override
	protected void finished(Description description) {
	super.finished(description);
	done();
	runMethod("finished", description);
	}

	@Override
	protected void failed(Throwable e, Description description) {
	super.failed(e, description);
	done();
	runMethod("failed", description);
	}

	private void done() {
	try {
	runWithLoader("DrillbitStopThread", drillbitLoader);
	DrillbitStopThread.SEM.acquire();
	} catch (Exception e) {
	throw new RuntimeException(e);
	}
	}

	private void runMethod(String name, Description description) {
	try {
	Class<?> clazz = drillbitLoader.loadClass("org.apache.drill.test.BaseTestQuery");
	Class<?> watcherClazz = drillbitLoader.loadClass("org.junit.rules.TestWatcher");
	Object watcher = clazz.getField("dirTestWatcher").get(null);
	Method method = watcherClazz.getDeclaredMethod(name, Description.class);
	method.setAccessible(true);
	method.invoke(watcher, description);
	} catch (Exception e) {
	throw new RuntimeException(e);
	}
	}
	};

	private static URL getJdbcUrl() throws MalformedURLException {
	return new URL(
	String.format("%s../../target/drill-jdbc-all-%s.jar",
	ClassLoader.getSystemClassLoader().getResource("").toString(),
	System.getProperty("project.version")
	));

	}

	@Test
	public void testDatabaseVersion() throws Exception {
	final URLClassLoader loader = URLClassLoader.newInstance(new URL[] {getJdbcUrl()});

	Class<?> clazz = loader.loadClass("org.apache.drill.jdbc.Driver");
	try {
	Driver driver = (Driver) clazz.newInstance();
	try (Connection c = driver.connect("jdbc:drill:drillbit=localhost:" + userPort, null)) {
	DatabaseMetaData metadata = c.getMetaData();
	assertEquals("Apache Drill JDBC Driver", metadata.getDriverName());
	assertEquals("Apache Drill Server", metadata.getDatabaseProductName());
	}
	} catch (Exception ex) {
	throw ex;
	}

	}

	@Test
	public void executeJdbcAllQuery() throws Exception {
	final URLClassLoader loader = URLClassLoader.newInstance(new URL[] {getJdbcUrl()});

	Class<?> clazz = loader.loadClass("org.apache.drill.jdbc.Driver");
	try {
	Driver driver = (Driver) clazz.newInstance();
	try (Connection c = driver.connect("jdbc:drill:drillbit=localhost:" + userPort, null)) {
	printQuery(c, "select * from cp.`types.json`");
	}
	} catch (Exception ex) {
	throw ex;
	}

	}

	@Test
	public void serviceFileContainsCorrectDriver() throws IOException {
	URLClassLoader loader = URLClassLoader.newInstance(new URL[]{getJdbcUrl()});
	try (InputStream resourceStream = loader.getResourceAsStream("META-INF/services/java.sql.Driver")) {
	assertNotNull("java.sql.Driver is not present in the jdbc jar", resourceStream);
	try (BufferedReader reader = new BufferedReader(new InputStreamReader(resourceStream))) {
	String driverClass = reader.lines().collect(Collectors.joining(System.lineSeparator()));
	assertEquals("org.apache.drill.jdbc.Driver", driverClass);
	}
	}
	}

	private static void printQuery(Connection c, String query) throws SQLException {
	final StringBuilder sb = new StringBuilder();

	try (Statement s = c.createStatement();
	ResultSet result = s.executeQuery(query)) {
	while (result.next()) {
	final int columnCount = result.getMetaData().getColumnCount();
	for(int i = 1; i < columnCount+1; i++){
	sb.append(result.getObject(i));
	sb.append('\t');
	}

	sb.append(result.getObject(1)).append('\n');
	}
	}

	logger.info(sb.toString());
	}

	private static int getClassesLoadedCount(ClassLoader classLoader) {
	try {
	Field f = ClassLoader.class.getDeclaredField("classes");
	f.setAccessible(true);
	Vector<Class<?>> classes = (Vector<Class<?>>) f.get(classLoader);
	return classes.size();
	} catch (Exception e) {
	logger.error("Exception ", e);
	return -1;
	}
	}

	private static void printClassesLoaded(String prefix, ClassLoader classLoader) {
	final StringBuilder sb = new StringBuilder();

	try {
	Field f = ClassLoader.class.getDeclaredField("classes");
	f.setAccessible(true);
	Vector<Class<?>> classes = (Vector<Class<?>>) f.get(classLoader);
	for (Class<?> c : classes) {
	sb.append(prefix).append(": ").append(c.getName()).append('\n');
	}
	} catch (Exception e) {
	logger.error("Exception ", e);
	}

	logger.info(sb.toString());
	}

	private static void runWithLoader(String name, ClassLoader loader) throws Exception {
	Class<?> clazz = loader.loadClass(ITTestShadedJar.class.getName() + "$" + name);
	Object instance = clazz.getDeclaredConstructors()[0].newInstance(loader);
	clazz.getMethod("go").invoke(instance);
	}

	public abstract static class AbstractLoaderThread extends Thread {
	private Exception ex;
	protected final ClassLoader loader;

	public AbstractLoaderThread(ClassLoader loader) {
	this.setContextClassLoader(loader);
	this.loader = loader;
	}

	@Override
	public final void run() {
	try {
	internalRun();
	} catch (Exception e) {
	this.ex = e;
	}
	}

	protected abstract void internalRun() throws Exception;

	public void go() throws Exception {
	start();
	join();
	if (ex != null) {
	throw ex;
	}
	}
	}

	public static class DrillbitStartThread extends AbstractLoaderThread {
	public static final Semaphore SEM = new Semaphore(0);

	public DrillbitStartThread(ClassLoader loader) {
	super(loader);
	}

	@Override
	protected void internalRun() throws Exception {
	Class<?> clazz = loader.loadClass("org.apache.drill.test.BaseTestQuery");
	clazz.getMethod("setupDefaultTestCluster").invoke(null);

	// loader.loadClass("org.apache.drill.exec.exception.SchemaChangeException");

	// execute a single query to make sure the drillbit is fully up
	clazz.getMethod("testNoResult", String.class, Object[].class)
	.invoke(null, "select * from (VALUES 1)", new Object[] {});

	SEM.release();
	}

	}

	public static class DrillbitStopThread extends AbstractLoaderThread {
	public static final Semaphore SEM = new Semaphore(0);

	public DrillbitStopThread(ClassLoader loader) {
	super(loader);
	}

	@Override
	protected void internalRun() throws Exception {
	Class<?> clazz = loader.loadClass("org.apache.drill.test.BaseTestQuery");
	clazz.getMethod("closeClient").invoke(null);
	SEM.release();
	}
	}

	}