datahub/metadata-ingestion-modules/airflow-plugin/build.gradle

plugins {
    id 'base'
}

apply from: "../../gradle/coverage/python-coverage.gradle"

ext {
  python_executable = 'python3'
  venv_name = 'venv'
}

ext.venv_activate_command = "set +x && source ${venv_name}/bin/activate && set -x && "

if (!project.hasProperty("extra_pip_requirements")) {
    ext.extra_pip_requirements = ""
}
if (!project.hasProperty("extra_pip_constraints")) {
    ext.extra_pip_constraints = ""
}
if (!project.hasProperty("extra_pip_extras")) {
    ext.extra_pip_extras = ""
}
// If extra_pip_extras is non-empty, we need to add a comma to the beginning of the string.
if (extra_pip_extras != "") {
    ext.extra_pip_extras = "," + extra_pip_extras
}

def pip_install_command = "VIRTUAL_ENV=${venv_name} ${venv_name}/bin/uv pip install -e ../../metadata-ingestion ${extra_pip_requirements}"

task environmentSetup(type: Exec) {
  def sentinel_file = "${venv_name}/.venv_environment_sentinel"
  inputs.file file('setup.py')
  outputs.file(sentinel_file)
  commandLine 'bash', '-c', 
    "${python_executable} -m venv ${venv_name} && set -x && " +
    "${venv_name}/bin/python -m pip install --upgrade uv && " +
    "touch ${sentinel_file}"
}

task installPackage(type: Exec, dependsOn: [environmentSetup, ':metadata-ingestion:codegen']) {
  def sentinel_file = "${venv_name}/.build_install_package_sentinel"
  inputs.file file('setup.py')
  outputs.file(sentinel_file)
  commandLine 'bash', '-c',
    venv_activate_command +
    "${pip_install_command} -e .[ignore${extra_pip_extras}] ${extra_pip_constraints} &&" +
    "touch ${sentinel_file}"
}

task install(dependsOn: [installPackage])

task installDev(type: Exec, dependsOn: [install]) {
  def sentinel_file = "${venv_name}/.build_install_dev_sentinel"
  inputs.file file('setup.py')
  outputs.file("${sentinel_file}")
  // For some reason, the Airflow constraints files pin things like black and pytest.
  // This allows us to not respect those constraints ourselves.
  commandLine 'bash', '-c',
    venv_activate_command +
    "${pip_install_command} -e .[dev${extra_pip_extras}] && " +
    "touch ${sentinel_file}"
}
task installTest(type: Exec, dependsOn: [installDev]) {
  def sentinel_file = "${venv_name}/.build_install_test_sentinel"
  inputs.file file('setup.py')
  outputs.file(sentinel_file)
  commandLine 'bash', '-c',
    venv_activate_command +
    "${pip_install_command} -e .[integration-tests${extra_pip_extras}] ${extra_pip_constraints} && " +
    "touch ${sentinel_file}"
}

task lint(type: Exec, dependsOn: installDev) {
  /*
    The find/sed combo below is a temporary work-around for the following mypy issue with airflow 2.2.0:
   "venv/lib/python3.8/site-packages/airflow/_vendor/connexion/spec.py:169: error: invalid syntax".
   */
  commandLine 'bash', '-c',
    "find ${venv_name}/lib -path *airflow/_vendor/connexion/spec.py -exec sed -i.bak -e '169,169s/  # type: List\\[str\\]//g' {} \\; && " +
    venv_activate_command +
    "ruff check src/ tests/ && " +
    "ruff format --check src/ tests/ && " +
    "mypy --show-traceback --show-error-codes src/ tests/"
}
task lintFix(type: Exec, dependsOn: installDev) {
  commandLine 'bash', '-c',
    venv_activate_command +
    "ruff check --fix src/ tests/ && " +
    "ruff format src/ tests/ "
}

// HACK: Some of the Airflow constraint files conflict with packages that we install (e.g. black).
// I'm not sure why they constrain those in the first place, but we don't want to respect them
// when running linting. However, we do want to respect them when running integration tests.
// There's two install steps, and these constraints ensure that we run them in the right order.
installTest.shouldRunAfter lint
installTest.shouldRunAfter lintFix

task testQuick(type: Exec, dependsOn: installTest) {
  inputs.files(project.fileTree(dir: "src/", include: "**/*.py"))
  inputs.files(project.fileTree(dir: "tests/"))
  commandLine 'bash', '-c',
    venv_activate_command +
    "pytest --cov-config=setup.cfg ${get_coverage_args('quick')} -vv --continue-on-collection-errors --junit-xml=junit.quick.xml"
}


task cleanPythonCache(type: Exec) {
  commandLine 'bash', '-c',
    "find src -type f -name '*.py[co]' -delete -o -type d -name __pycache__ -delete -o -type d -empty -delete"
}
task buildWheel(type: Exec, dependsOn: [environmentSetup]) {
  commandLine 'bash', '-c',
    venv_activate_command +
    'uv pip install build && RELEASE_VERSION="\${RELEASE_VERSION:-0.0.0.dev1}" RELEASE_SKIP_INSTALL=1 RELEASE_SKIP_UPLOAD=1 ./scripts/release.sh'
}

build.dependsOn install
check.dependsOn lint
check.dependsOn testQuick

clean {
  delete venv_name
  delete 'build'
  delete 'dist'
  delete '.ruff_cache'
  delete '.mypy_cache'
  delete '.pytest_cache'
}
clean.dependsOn cleanPythonCache