From ebd2e2312bdd23a92bbf403a26f64194807f70e6 Mon Sep 17 00:00:00 2001
From: david-leifker <114954101+david-leifker@users.noreply.github.com>
Date: Fri, 10 Nov 2023 22:10:00 -0600
Subject: [PATCH] fix(datahub-ingestion): remove old jars, sync pyspark version
 (#9217)

---
 docker/datahub-ingestion-base/build.gradle |  7 ++++---
 docker/datahub-ingestion/Dockerfile        | 16 ++++++++++++++--
 docker/datahub-ingestion/pyspark_jars.sh   | 22 ++++++++++++++++++++++
 metadata-ingestion/setup.py                |  4 ++--
 4 files changed, 42 insertions(+), 7 deletions(-)
 create mode 100755 docker/datahub-ingestion/pyspark_jars.sh

diff --git a/docker/datahub-ingestion-base/build.gradle b/docker/datahub-ingestion-base/build.gradle
index bbd8242553cc5..64635671343ef 100644
--- a/docker/datahub-ingestion-base/build.gradle
+++ b/docker/datahub-ingestion-base/build.gradle
@@ -9,20 +9,21 @@ ext {
     docker_registry = rootProject.ext.docker_registry == 'linkedin' ? 'acryldata' : docker_registry
     docker_repo = 'datahub-ingestion-base'
     docker_dir = 'datahub-ingestion-base'
+    docker_target = project.getProperties().getOrDefault("dockerTarget", "slim")
 
     revision = 2 // increment to trigger rebuild
 }
 
 docker {
-    name "${docker_registry}/${docker_repo}:v${version}-slim"
-    version "v${version}-slim"
+    name "${docker_registry}/${docker_repo}:v${version}-${docker_target}"
+    version "v${version}-${docker_target}"
     dockerfile file("${rootProject.projectDir}/docker/${docker_dir}/Dockerfile")
     files fileTree(rootProject.projectDir) {
         include "docker/${docker_dir}/*"
     }.exclude {
         i -> i.file.isHidden() || i.file == buildDir
     }
-    buildArgs([APP_ENV: 'slim'])
+    buildArgs([APP_ENV: docker_target])
 }
 tasks.getByName('docker').dependsOn('build')
 
diff --git a/docker/datahub-ingestion/Dockerfile b/docker/datahub-ingestion/Dockerfile
index 0132ceaa9b1a9..2abd4e2f33bef 100644
--- a/docker/datahub-ingestion/Dockerfile
+++ b/docker/datahub-ingestion/Dockerfile
@@ -22,10 +22,22 @@ ENV PATH="/datahub-ingestion/.local/bin:$PATH"
 FROM base as slim-install
 RUN pip install --no-cache --user ".[base,datahub-rest,datahub-kafka,snowflake,bigquery,redshift,mysql,postgres,hive,clickhouse,glue,dbt,looker,lookml,tableau,powerbi,superset,datahub-business-glossary]"
 
-FROM base as full-install
+FROM base as full-install-build
+
+USER 0
+RUN apt-get update && apt-get install -y -qq maven
+
+USER datahub
+COPY ./docker/datahub-ingestion/pyspark_jars.sh .
+
 RUN pip install --no-cache --user ".[base]" && \
     pip install --no-cache --user "./airflow-plugin[acryl-datahub-airflow-plugin]" && \
-    pip install --no-cache --user ".[all]"
+    pip install --no-cache --user ".[all]" && \
+    ./pyspark_jars.sh
+
+FROM base as full-install
+
+COPY --from=full-install-build /datahub-ingestion/.local /datahub-ingestion/.local
 
 FROM base as dev-install
 # Dummy stage for development. Assumes code is built on your machine and mounted to this image.
diff --git a/docker/datahub-ingestion/pyspark_jars.sh b/docker/datahub-ingestion/pyspark_jars.sh
new file mode 100755
index 0000000000000..ecd24e78c4105
--- /dev/null
+++ b/docker/datahub-ingestion/pyspark_jars.sh
@@ -0,0 +1,22 @@
+#!/bin/bash
+
+set -ex
+
+HADOOP_CLIENT_DEPENDENCY="${HADOOP_CLIENT_DEPENDENCY:-org.apache.hadoop:hadoop-client:3.3.6}"
+ZOOKEEPER_DEPENDENCY="${ZOOKEEPER_DEPENDENCY:-org.apache.zookeeper:zookeeper:3.7.2}"
+PYSPARK_JARS="$(python -m site --user-site)/pyspark/jars"
+
+# Remove conflicting versions
+echo "Removing version conflicts from $PYSPARK_JARS"
+CONFLICTS="zookeeper hadoop- slf4j-"
+for jar in $CONFLICTS; do
+  rm "$PYSPARK_JARS/$jar"*.jar
+done
+
+# Fetch dependencies
+mvn dependency:get -Dtransitive=true -Dartifact="$HADOOP_CLIENT_DEPENDENCY"
+mvn dependency:get -Dtransitive=true -Dartifact="$ZOOKEEPER_DEPENDENCY"
+
+# Move to pyspark location
+echo "Moving jars to $PYSPARK_JARS"
+find "$HOME/.m2" -type f -name "*.jar" -exec mv {} "$PYSPARK_JARS/" \;
diff --git a/metadata-ingestion/setup.py b/metadata-ingestion/setup.py
index 5f44f14c3d74c..f3782abe576d3 100644
--- a/metadata-ingestion/setup.py
+++ b/metadata-ingestion/setup.py
@@ -242,7 +242,7 @@
 }
 
 data_lake_profiling = {
-    "pydeequ==1.1.0",
+    "pydeequ~=1.1.0",
     "pyspark~=3.3.0",
 }
 
@@ -256,7 +256,7 @@
 databricks = {
     # 0.1.11 appears to have authentication issues with azure databricks
     "databricks-sdk>=0.9.0",
-    "pyspark",
+    "pyspark~=3.3.0",
     "requests",
 }