hive-lineage: address comments

Signed-off-by: Pahulpreet Singh <[email protected]>
GoogleCloudDataproc · Feb 3, 2025 · 3d250a2 · 3d250a2
1 parent 3bb025e
commit 3d250a2
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 11 deletions.
diff --git a/BUILD b/BUILD
@@ -138,14 +138,16 @@ py_library(
 
 py_test(
     name = "test_hive_lineage",
-    size = "small",
+    size = "enormous",
     srcs = ["hive-lineage/test_hive_lineage.py"],
     data = [
         "hive-lineage/hive-lineage.sh",
         "hive-lineage/hivetest.hive",
     ],
     local = True,
+    shard_count = 3,
     deps = [
         "//integration_tests:dataproc_test_case",
+        "@io_abseil_py//absl/testing:parameterized",
     ],
 )
diff --git a/hive-lineage/hive-lineage.sh b/hive-lineage/hive-lineage.sh
@@ -1,5 +1,7 @@
 #!/bin/bash
 #
+# Copyright 2025 Google LLC and contributors
+#
 # Licensed 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
@@ -17,14 +19,15 @@
 
 set -euxo pipefail
 
-export HIVE_HOME="/usr/lib/hive"
-export HIVE_CONF_DIR="/etc/hive/conf"
-export HIVE_CONF_FILE="$HIVE_CONF_DIR/hive-site.xml"
-export HIVE_LIB_DIR="/usr/lib/hive/lib"
-export INSTALLATION_SOURCE="gs://hadoop-lib/hive-lineage"
-export HIVE_OL_HOOK_VERSION="1.0.0-preview"
-export HIVE_OL_HOOK="io.openlineage.hive.hooks.HiveOpenLineageHook"
-
+function prepare_env() {
+  export HIVE_HOME="/usr/lib/hive"
+  export HIVE_CONF_DIR="/etc/hive/conf"
+  export HIVE_CONF_FILE="$HIVE_CONF_DIR/hive-site.xml"
+  export HIVE_LIB_DIR="$HIVE_HOME/lib"
+  export INSTALLATION_SOURCE="gs://hadoop-lib/hive-lineage"
+  export HIVE_OL_HOOK_VERSION="1.0.0-preview"
+  export HIVE_OL_HOOK="io.openlineage.hive.hooks.HiveOpenLineageHook"
+}
 
 function set_hive_lineage_conf() {
   declare -A properties=(
@@ -44,7 +47,7 @@ function set_hive_lineage_conf() {
 
 function install_jars() {
   echo "Installing openlineage-hive hook"
-  gsutil cp -P "$INSTALLATION_SOURCE/hive-openlineage-hook-$HIVE_HIVE_OL_HOOK_VERSION.jar" "$HIVE_LIB_DIR/hive-openlineage-hook.jar"
+  gsutil cp -P "$INSTALLATION_SOURCE/hive-openlineage-hook-$HIVE_OL_HOOK_VERSION.jar" "$HIVE_LIB_DIR/hive-openlineage-hook.jar"
 }
 
 function restart_hive_server2_master() {
@@ -55,6 +58,7 @@ function restart_hive_server2_master() {
   fi
 }
 
+prepare_env
 install_jars
 set_hive_lineage_conf
 restart_hive_server2_master
diff --git a/hive-lineage/test_hive_lineage.py b/hive-lineage/test_hive_lineage.py
@@ -1,8 +1,10 @@
 from absl.testing import absltest
+from absl.testing import parameterized
 
 from integration_tests.dataproc_test_case import DataprocTestCase
 
 class HiveLineageTestCase(DataprocTestCase):
+  COMPONENT = "hive-lineage"
   INIT_ACTIONS = ["hive-lineage/hive-lineage.sh"]
   TEST_SCRIPT_FILE = "hive-lineage/hivetest.hive"
 
@@ -13,8 +15,14 @@ def __submit_hive_job(self, cluster_name):
   def verify_cluster(self, name):
     self.__submit_hive_job(name)
 
+  @parameterized.parameters(
+      'STANDARD',
+      'HA',
+  )
   def test_hive_job_success(self, configuration):
-    self.createCluster(configuration, self.INIT_ACTIONS, scopes='cloud-platform')
+    self.createCluster(configuration,
+                       self.INIT_ACTIONS,
+                       scopes='cloud-platform')
     self.verify_cluster(self.getClusterName())