src/test/com/cloudera/sqoop/TestMerge.java - sqoop - 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 com.cloudera.sqoop;

 import java.io.BufferedReader;
 import java.io.IOException;
 import java.io.InputStreamReader;
 import java.sql.Connection;
 import java.sql.Timestamp;
 import java.sql.PreparedStatement;
 import java.sql.SQLException;

 import java.util.List;

 import com.cloudera.sqoop.testutil.CommonArgs;
 import com.cloudera.sqoop.testutil.HsqldbTestServer;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.FileStatus;
 import org.apache.hadoop.fs.FileSystem;
 import org.apache.hadoop.fs.Path;

 import com.cloudera.sqoop.SqoopOptions.IncrementalMode;
 import com.cloudera.sqoop.manager.ConnManager;
 import com.cloudera.sqoop.testutil.BaseSqoopTestCase;
 import com.cloudera.sqoop.tool.CodeGenTool;
 import com.cloudera.sqoop.tool.ImportTool;
 import com.cloudera.sqoop.tool.MergeTool;
 import com.cloudera.sqoop.util.ClassLoaderStack;

 /**
  * Test that the merge tool works.
  */
 public class TestMerge extends BaseSqoopTestCase {

   private static final Log LOG =
       LogFactory.getLog(TestMerge.class.getName());

   protected ConnManager manager;
   protected Connection conn;

   public static final String SOURCE_DB_URL = "jdbc:hsqldb:mem:merge";

   @Override
   public void setUp() {
     super.setUp();
     manager = getManager();
     try {
       conn = manager.getConnection();
     } catch (SQLException e) {
       e.printStackTrace();
       throw new RuntimeException(e);
     }
   }

   public static final String TABLE_NAME = "MergeTable";

   public Configuration newConf() {
     Configuration conf = new Configuration();
     if (!BaseSqoopTestCase.isOnPhysicalCluster()) {
       conf.set(CommonArgs.FS_DEFAULT_NAME, CommonArgs.LOCAL_FS);
     }
     conf.set("mapred.job.tracker", "local");
     return conf;
   }

   /**
    * Create a SqoopOptions to connect to the manager.
    */
   public SqoopOptions getSqoopOptions(Configuration conf) {
     SqoopOptions options = new SqoopOptions(conf);
     options.setConnectString(HsqldbTestServer.getDbUrl());

     return options;
   }

   protected void createTable() throws SQLException {
     PreparedStatement s = conn.prepareStatement("DROP TABLE \"" + TABLE_NAME + "\" IF EXISTS");
     try {
       s.executeUpdate();
     } finally {
       s.close();
     }

     s = conn.prepareStatement("CREATE TABLE \"" + TABLE_NAME + "\" (id INT NOT NULL PRIMARY KEY, val INT, lastmod TIMESTAMP)");
     try {
       s.executeUpdate();
     } finally {
       s.close();
     }

     s = conn.prepareStatement("INSERT INTO \"" + TABLE_NAME + "\" VALUES (0, 0, NOW())");
     try {
       s.executeUpdate();
     } finally {
       s.close();
     }

     s = conn.prepareStatement("INSERT INTO \"" + TABLE_NAME + "\" VALUES (1, 42, NOW())");
     try {
       s.executeUpdate();
     } finally {
       s.close();
     }

     conn.commit();
   }

   public void testMerge() throws Exception {
     createTable();

     // Create a jar to use for the merging process; we'll load it
     // into the current thread CL for when this runs. This needs
     // to contain a different class name than used for the imports
     // due to classloaderstack issues in the same JVM.
     final String MERGE_CLASS_NAME = "ClassForMerging";
     SqoopOptions options = getSqoopOptions(newConf());
     options.setTableName(TABLE_NAME);
     options.setClassName(MERGE_CLASS_NAME);

     CodeGenTool codeGen = new CodeGenTool();
     Sqoop codeGenerator = new Sqoop(codeGen, options.getConf(), options);
     int ret = Sqoop.runSqoop(codeGenerator, new String[0]);
     if (0 != ret) {
       fail("Nonzero exit from codegen: " + ret);
     }

     List<String> jars = codeGen.getGeneratedJarFiles();
     String jarFileName = jars.get(0);

     // Now do the imports.

     Path warehouse = new Path(BaseSqoopTestCase.LOCAL_WAREHOUSE_DIR);

     options = getSqoopOptions(newConf());
     options.setTableName(TABLE_NAME);
     options.setNumMappers(1);

     // Do an import of this data into the "old" dataset.
     options.setTargetDir(new Path(warehouse, "merge-old").toString());
     options.setIncrementalMode(IncrementalMode.DateLastModified);
     options.setIncrementalTestColumn("LASTMOD");

     ImportTool importTool = new ImportTool();
     Sqoop importer = new Sqoop(importTool, options.getConf(), options);
     ret = Sqoop.runSqoop(importer, new String[0]);
     if (0 != ret) {
       fail("Initial import failed with exit code " + ret);
     }

     // Check that we got records that meet our expected values.
     assertRecordStartsWith("0,0,", "merge-old");
     assertRecordStartsWith("1,42,", "merge-old");

     long prevImportEnd = System.currentTimeMillis();

     Thread.sleep(25);

     // Modify the data in the warehouse.
     PreparedStatement s = conn.prepareStatement("UPDATE \"" + TABLE_NAME + "\" SET val=43, lastmod=NOW() WHERE id=1");
     try {
       s.executeUpdate();
       conn.commit();
     } finally {
       s.close();
     }

     s = conn.prepareStatement("INSERT INTO \"" + TABLE_NAME + "\" VALUES (3,313,NOW())");
     try {
       s.executeUpdate();
       conn.commit();
     } finally {
       s.close();
     }

     Thread.sleep(25);

     // Do another import, into the "new" dir.
     options = getSqoopOptions(newConf());
     options.setTableName(TABLE_NAME);
     options.setNumMappers(1);
     options.setTargetDir(new Path(warehouse, "merge-new").toString());
     options.setIncrementalMode(IncrementalMode.DateLastModified);
     options.setIncrementalTestColumn("LASTMOD");
     options.setIncrementalLastValue(new Timestamp(prevImportEnd).toString());

     importTool = new ImportTool();
     importer = new Sqoop(importTool, options.getConf(), options);
     ret = Sqoop.runSqoop(importer, new String[0]);
     if (0 != ret) {
       fail("Second import failed with exit code " + ret);
     }

     assertRecordStartsWith("1,43,", "merge-new");
     assertRecordStartsWith("3,313,", "merge-new");

     // Now merge the results!
     ClassLoaderStack.addJarFile(jarFileName, MERGE_CLASS_NAME);

     options = getSqoopOptions(newConf());
     options.setMergeOldPath(new Path(warehouse, "merge-old").toString());
     options.setMergeNewPath(new Path(warehouse, "merge-new").toString());
     options.setMergeKeyCol("ID");
     options.setTargetDir(new Path(warehouse, "merge-final").toString());
     options.setClassName(MERGE_CLASS_NAME);

     MergeTool mergeTool = new MergeTool();
     Sqoop merger = new Sqoop(mergeTool, options.getConf(), options);
     ret = Sqoop.runSqoop(merger, new String[0]);
     if (0 != ret) {
       fail("Merge failed with exit code " + ret);
     }

     assertRecordStartsWith("0,0,", "merge-final");
     assertRecordStartsWith("1,43,", "merge-final");
     assertRecordStartsWith("3,313,", "merge-final");
   }

   /**
    * @return true if the file specified by path 'p' contains a line
    * that starts with 'prefix'
    */
   protected boolean checkFileForLine(FileSystem fs, Path p, String prefix)
       throws IOException {
     BufferedReader r = new BufferedReader(new InputStreamReader(fs.open(p)));
     try {
       while (true) {
         String in = r.readLine();
         if (null == in) {
           break; // done with the file.
         }

         if (in.startsWith(prefix)) {
           return true;
         }
       }
     } finally {
       r.close();
     }

     return false;
   }

   /**
    * Return true if there's a file in 'dirName' with a line that starts with
    * 'prefix'.
    */
   protected boolean recordStartsWith(String prefix, String dirName)
       throws Exception {
     Path warehousePath = new Path(LOCAL_WAREHOUSE_DIR);
     Path targetPath = new Path(warehousePath, dirName);

     FileSystem fs = FileSystem.getLocal(new Configuration());
     FileStatus [] files = fs.listStatus(targetPath);

     if (null == files || files.length == 0) {
       fail("Got no import files!");
     }

     for (FileStatus stat : files) {
       Path p = stat.getPath();
       if (p.getName().startsWith("part-")) {
         if (checkFileForLine(fs, p, prefix)) {
           // We found the line. Nothing further to do.
           return true;
         }
       }
     }

     return false;
   }

   protected void assertRecordStartsWith(String prefix, String dirName)
       throws Exception {
     if (!recordStartsWith(prefix, dirName)) {
       fail("No record found that starts with " + prefix + " in " + dirName);
     }
   }
 }
	/**
	* 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 com.cloudera.sqoop;

	import java.io.BufferedReader;
	import java.io.IOException;
	import java.io.InputStreamReader;
	import java.sql.Connection;
	import java.sql.Timestamp;
	import java.sql.PreparedStatement;
	import java.sql.SQLException;

	import java.util.List;

	import com.cloudera.sqoop.testutil.CommonArgs;
	import com.cloudera.sqoop.testutil.HsqldbTestServer;
	import org.apache.commons.logging.Log;
	import org.apache.commons.logging.LogFactory;
	import org.apache.hadoop.conf.Configuration;
	import org.apache.hadoop.fs.FileStatus;
	import org.apache.hadoop.fs.FileSystem;
	import org.apache.hadoop.fs.Path;

	import com.cloudera.sqoop.SqoopOptions.IncrementalMode;
	import com.cloudera.sqoop.manager.ConnManager;
	import com.cloudera.sqoop.testutil.BaseSqoopTestCase;
	import com.cloudera.sqoop.tool.CodeGenTool;
	import com.cloudera.sqoop.tool.ImportTool;
	import com.cloudera.sqoop.tool.MergeTool;
	import com.cloudera.sqoop.util.ClassLoaderStack;

	/**
	* Test that the merge tool works.
	*/
	public class TestMerge extends BaseSqoopTestCase {

	private static final Log LOG =
	LogFactory.getLog(TestMerge.class.getName());

	protected ConnManager manager;
	protected Connection conn;

	public static final String SOURCE_DB_URL = "jdbc:hsqldb:mem:merge";

	@Override
	public void setUp() {
	super.setUp();
	manager = getManager();
	try {
	conn = manager.getConnection();
	} catch (SQLException e) {
	e.printStackTrace();
	throw new RuntimeException(e);
	}
	}

	public static final String TABLE_NAME = "MergeTable";

	public Configuration newConf() {
	Configuration conf = new Configuration();
	if (!BaseSqoopTestCase.isOnPhysicalCluster()) {
	conf.set(CommonArgs.FS_DEFAULT_NAME, CommonArgs.LOCAL_FS);
	}
	conf.set("mapred.job.tracker", "local");
	return conf;
	}

	/**
	* Create a SqoopOptions to connect to the manager.
	*/
	public SqoopOptions getSqoopOptions(Configuration conf) {
	SqoopOptions options = new SqoopOptions(conf);
	options.setConnectString(HsqldbTestServer.getDbUrl());

	return options;
	}

	protected void createTable() throws SQLException {
	PreparedStatement s = conn.prepareStatement("DROP TABLE \"" + TABLE_NAME + "\" IF EXISTS");
	try {
	s.executeUpdate();
	} finally {
	s.close();
	}

	s = conn.prepareStatement("CREATE TABLE \"" + TABLE_NAME + "\" (id INT NOT NULL PRIMARY KEY, val INT, lastmod TIMESTAMP)");
	try {
	s.executeUpdate();
	} finally {
	s.close();
	}

	s = conn.prepareStatement("INSERT INTO \"" + TABLE_NAME + "\" VALUES (0, 0, NOW())");
	try {
	s.executeUpdate();
	} finally {
	s.close();
	}

	s = conn.prepareStatement("INSERT INTO \"" + TABLE_NAME + "\" VALUES (1, 42, NOW())");
	try {
	s.executeUpdate();
	} finally {
	s.close();
	}

	conn.commit();
	}

	public void testMerge() throws Exception {
	createTable();

	// Create a jar to use for the merging process; we'll load it
	// into the current thread CL for when this runs. This needs
	// to contain a different class name than used for the imports
	// due to classloaderstack issues in the same JVM.
	final String MERGE_CLASS_NAME = "ClassForMerging";
	SqoopOptions options = getSqoopOptions(newConf());
	options.setTableName(TABLE_NAME);
	options.setClassName(MERGE_CLASS_NAME);

	CodeGenTool codeGen = new CodeGenTool();
	Sqoop codeGenerator = new Sqoop(codeGen, options.getConf(), options);
	int ret = Sqoop.runSqoop(codeGenerator, new String[0]);
	if (0 != ret) {
	fail("Nonzero exit from codegen: " + ret);
	}

	List<String> jars = codeGen.getGeneratedJarFiles();
	String jarFileName = jars.get(0);

	// Now do the imports.

	Path warehouse = new Path(BaseSqoopTestCase.LOCAL_WAREHOUSE_DIR);

	options = getSqoopOptions(newConf());
	options.setTableName(TABLE_NAME);
	options.setNumMappers(1);

	// Do an import of this data into the "old" dataset.
	options.setTargetDir(new Path(warehouse, "merge-old").toString());
	options.setIncrementalMode(IncrementalMode.DateLastModified);
	options.setIncrementalTestColumn("LASTMOD");

	ImportTool importTool = new ImportTool();
	Sqoop importer = new Sqoop(importTool, options.getConf(), options);
	ret = Sqoop.runSqoop(importer, new String[0]);
	if (0 != ret) {
	fail("Initial import failed with exit code " + ret);
	}

	// Check that we got records that meet our expected values.
	assertRecordStartsWith("0,0,", "merge-old");
	assertRecordStartsWith("1,42,", "merge-old");

	long prevImportEnd = System.currentTimeMillis();

	Thread.sleep(25);

	// Modify the data in the warehouse.
	PreparedStatement s = conn.prepareStatement("UPDATE \"" + TABLE_NAME + "\" SET val=43, lastmod=NOW() WHERE id=1");
	try {
	s.executeUpdate();
	conn.commit();
	} finally {
	s.close();
	}

	s = conn.prepareStatement("INSERT INTO \"" + TABLE_NAME + "\" VALUES (3,313,NOW())");
	try {
	s.executeUpdate();
	conn.commit();
	} finally {
	s.close();
	}

	Thread.sleep(25);

	// Do another import, into the "new" dir.
	options = getSqoopOptions(newConf());
	options.setTableName(TABLE_NAME);
	options.setNumMappers(1);
	options.setTargetDir(new Path(warehouse, "merge-new").toString());
	options.setIncrementalMode(IncrementalMode.DateLastModified);
	options.setIncrementalTestColumn("LASTMOD");
	options.setIncrementalLastValue(new Timestamp(prevImportEnd).toString());

	importTool = new ImportTool();
	importer = new Sqoop(importTool, options.getConf(), options);
	ret = Sqoop.runSqoop(importer, new String[0]);
	if (0 != ret) {
	fail("Second import failed with exit code " + ret);
	}

	assertRecordStartsWith("1,43,", "merge-new");
	assertRecordStartsWith("3,313,", "merge-new");

	// Now merge the results!
	ClassLoaderStack.addJarFile(jarFileName, MERGE_CLASS_NAME);

	options = getSqoopOptions(newConf());
	options.setMergeOldPath(new Path(warehouse, "merge-old").toString());
	options.setMergeNewPath(new Path(warehouse, "merge-new").toString());
	options.setMergeKeyCol("ID");
	options.setTargetDir(new Path(warehouse, "merge-final").toString());
	options.setClassName(MERGE_CLASS_NAME);

	MergeTool mergeTool = new MergeTool();
	Sqoop merger = new Sqoop(mergeTool, options.getConf(), options);
	ret = Sqoop.runSqoop(merger, new String[0]);
	if (0 != ret) {
	fail("Merge failed with exit code " + ret);
	}

	assertRecordStartsWith("0,0,", "merge-final");
	assertRecordStartsWith("1,43,", "merge-final");
	assertRecordStartsWith("3,313,", "merge-final");
	}

	/**
	* @return true if the file specified by path 'p' contains a line
	* that starts with 'prefix'
	*/
	protected boolean checkFileForLine(FileSystem fs, Path p, String prefix)
	throws IOException {
	BufferedReader r = new BufferedReader(new InputStreamReader(fs.open(p)));
	try {
	while (true) {
	String in = r.readLine();
	if (null == in) {
	break; // done with the file.
	}

	if (in.startsWith(prefix)) {
	return true;
	}
	}
	} finally {
	r.close();
	}

	return false;
	}

	/**
	* Return true if there's a file in 'dirName' with a line that starts with
	* 'prefix'.
	*/
	protected boolean recordStartsWith(String prefix, String dirName)
	throws Exception {
	Path warehousePath = new Path(LOCAL_WAREHOUSE_DIR);
	Path targetPath = new Path(warehousePath, dirName);

	FileSystem fs = FileSystem.getLocal(new Configuration());
	FileStatus [] files = fs.listStatus(targetPath);

	if (null == files \|\| files.length == 0) {
	fail("Got no import files!");
	}

	for (FileStatus stat : files) {
	Path p = stat.getPath();
	if (p.getName().startsWith("part-")) {
	if (checkFileForLine(fs, p, prefix)) {
	// We found the line. Nothing further to do.
	return true;
	}
	}
	}

	return false;
	}

	protected void assertRecordStartsWith(String prefix, String dirName)
	throws Exception {
	if (!recordStartsWith(prefix, dirName)) {
	fail("No record found that starts with " + prefix + " in " + dirName);
	}
	}
	}