code-infrastructure-lab/modules/airflow/airflow.tf

resource "kubernetes_namespace" "spark_jobs_namespace" {
  metadata {
    name = "${var.namespace_prefix}spark-jobs"
  }
}


# resource "kubernetes_service_account_v1" "spark_sa" {
#   depends_on = [kubernetes_namespace.spark_jobs_namespace]
#   metadata {
#     name      = "spark"
#     namespace = "${var.namespace_prefix}spark-jobs"
#   }
# }
#
resource "kubernetes_role" "airflow_spark_role" {
  depends_on = [kubernetes_namespace.spark_jobs_namespace]
  metadata {
    name      = "airflow-spark-role"
    namespace = "${var.namespace_prefix}spark-jobs"
  }

  rule {
    api_groups = ["sparkoperator.k8s.io"]
    resources = [
      "sparkapplications", "sparkapplications/status",
      "scheduledsparkapplications", "scheduledsparkapplications/status"
    ]
    verbs = ["*"]
  }

  rule {
    api_groups = [""]
    resources = ["pods", "pods/log"]
    verbs = ["*"]
  }
}

resource "kubernetes_role_binding_v1" "airflow_spark_role_binding" {
  depends_on = [kubernetes_namespace.spark_jobs_namespace]
  metadata {
    name      = "airflow-spark-role-binding"
    namespace = "${var.namespace_prefix}spark-jobs"
  }

  subject {
    kind      = "ServiceAccount"
    name      = "airflow-worker"
    namespace = "${var.namespace_prefix}airflow"
  }

  role_ref {
    api_group = "rbac.authorization.k8s.io"
    kind      = "Role"
    name      = "airflow-spark-role"
  }
}

resource "kubernetes_role_binding_v1" "airflow_spark_role_binding2" {
  depends_on = [kubernetes_namespace.spark_jobs_namespace]
  metadata {
    name      = "airflow-spark-role-binding2"
    namespace = "${var.namespace_prefix}spark-jobs"
  }

  subject {
    kind      = "ServiceAccount"
    name      = "airflow-worker"
    namespace = "${var.namespace_prefix}airflow"
  }

  role_ref {
    api_group = "rbac.authorization.k8s.io"
    kind      = "Role"
    name      = "spark-role"
  }
}

resource "helm_release" "gcp_spark_operator" {
  depends_on = [kubernetes_namespace.spark_jobs_namespace]
  name              = "gcp-spark-operator"
  chart             = "spark-operator"
  repository        = "https://kubeflow.github.io/spark-operator"
  create_namespace  = "true"
  namespace         = "${var.namespace_prefix}gcp-spark-operator"
  dependency_update = "true"
  version           = "2.0.2"

  set {
    name  = "image.repository"
    value = "spark-operator"
  }

  set {
    name  = "image.tag"
    value = "2.0.2"
  }

  set {
    name  = "spark.jobNamespaces"
    value = "{${var.namespace_prefix}spark-jobs}"
  }

  set {
    name  = "spark.serviceAccount.create"
    value = "true"
  }

  set {
    name  = "spark.serviceAccount.name"
    value = "spark"
  }

  set {
    name  = "controller.serviceAccount.create"
    value = "true"
  }

  set {
    name  = "controller.serviceAccount.name"
    value = "spark"
  }

  set {
    name  = "webhook.enabled"
    value = "true"
  }

  set {
    name  = "driver.ingressUrlFormat"
    value = "\\{\\{$appName\\}\\}.\\{\\{$appNamespace\\}\\}.${var.domain}"
    type  = "string"
  }
}

resource "kubernetes_namespace" "airflow" {
  metadata {
    name = "${var.namespace_prefix}airflow"
  }
}

resource "kubernetes_secret" "s3_conn_secrets" {
  depends_on = [kubernetes_namespace.airflow]
  metadata {
    name      = "s3-conn-secrets"
    namespace = "${var.namespace_prefix}airflow"
  }

  data = {
    username             = var.s3_key
    password             = var.s3_secret
    AIRFLOW_CONN_S3_CONN = <<EOT
{
  "conn_type": "aws",
  "extra": {
    "aws_access_key_id": "${var.s3_key}",
    "aws_secret_access_key": "${var.s3_secret}",
    "endpoint_url": "${var.s3_endpoint}",
    "verify": false
  }
}
EOT
  }

  type = "Opaque"
}


resource "helm_release" "airflow" {
  depends_on = [kubernetes_secret.s3_conn_secrets]

  name              = "airflow"
  chart             = "airflow"
  repository        = "https://airflow.apache.org"
  namespace         = "${var.namespace_prefix}airflow"
  dependency_update = "true"
  version           = "1.15.0"

  values = [
    file("./envs/${var.env}/airflow.yaml")
  ]

  set {
    name  = "fernetkey"
    value = "TG9mVjJvVEpoREVYdmdTRWlHdENXQ05zOU5OU2VGY0U="
  }

  set {
    name  = "webserver.defaultUser.password"
    value = var.admin_password
  }

  set {
    name = "spec.values.env"
    value = yamlencode([
      {
        name  = "AIRFLOW__WEBSERVER__BASE_URL",
        value = "https://airflow.${var.domain}"
      },
      {
        name  = "AIRFLOW__WEBSERVER__ENABLE_PROXY_FIX",
        value = "True"
      }
    ])
  }

  set {
    name  = "dags.gitSync.repo"
    value = var.repo_url
  }

  set {
    name  = "dags.gitSync.branch"
    value = var.branch_name
  }

  set {
    name  = "dags.gitSync.subPath"
    value = var.dag_path
  }

  # set {
  #   name  = "images.airflow.repository"
  #   value = "gbloisi/airflow"
  # }

  set {
    name  = "images.airflow.tag"
    value = "2.9.3-python3.11"
  }

  set {
    name  = "ingress.web.host"
    value = "airflow.${var.domain}"
  }
  set {
    name  = "ingress.flower.host"
    value = "airflow.${var.domain}"
  }
}