feat: add support for specifying FAB roles (airflow-helm#282)

kmehkeri · Mar 2, 2022 · 205511d · 205511d
1 parent df99a33
commit 205511d
Show file tree

Hide file tree

Showing 6 changed files with 411 additions and 0 deletions.
diff --git a/charts/airflow/README.md b/charts/airflow/README.md
@@ -119,6 +119,7 @@ Review the FAQ to understand how the chart functions, here are some good startin
 - ["How to set airflow configs?"](#how-to-set-airflow-configs)
 - ["How to create airflow users?"](#how-to-create-airflow-users)
 - ["How to authenticate airflow users with LDAP/OAUTH?"](#how-to-authenticate-airflow-users-with-ldapoauth)
+- ["How to create airflow roles?"](#how-to-create-airflow-roles)
 - ["How to create airflow connections?"](#how-to-create-airflow-connections)
 - ["How to use an external database?"](#how-to-use-an-external-database)
 - ["How to persist airflow logs?"](#how-to-persist-airflow-logs)
@@ -596,6 +597,31 @@ web:
 <hr>
 </details>
 
+### How to create airflow roles? 
+<details>
+<summary>Expand</summary>
+<hr>
+
+You can use the `airflow.roles` value to create airflow roles in a declarative way.
+
+Example values to create roles `RoleA` and `RoleB` with some permissions:
+```yaml
+airflow:
+  roles:
+    - name: RoleA
+      permissions: [['can_read', 'My Profile'], ['can_read', 'Website']]
+    - name: RoleB
+      permissions: [['can_read', 'My Profile']]
+
+  ## if we create a Deployment to perpetually sync `airflow.roles`
+  rolesUpdate: true
+```
+
+Note: sync process will not remove DAG-level permissions, in order not to override permissions assigned by `access_control` attribute in DAG definition.
+
+<hr>
+</details>
+
 ### How to set a custom fernet encryption key? 
 <details>
 <summary>Expand</summary>

diff --git a/charts/airflow/templates/sync/_helpers/sync_roles.tpl b/charts/airflow/templates/sync/_helpers/sync_roles.tpl
@@ -0,0 +1,136 @@
+{{/*
+The python sync script for roles.
+*/}}
+{{- define "airflow.sync.sync_roles.py" }}
+############################
+#### BEGIN: GLOBAL CODE ####
+############################
+{{- include "airflow.sync.global_code" . }}
+##########################
+#### END: GLOBAL CODE ####
+##########################
+
+
+#############
+## Imports ##
+#############
+import sys
+from typing import List, Tuple, Dict
+from flask_appbuilder.security.sqla.models import Role
+{{- if .Values.airflow.legacyCommands }}
+import airflow.www_rbac.app as www_app
+flask_app, flask_appbuilder = www_app.create_app()
+{{- else }}
+import airflow.www.app as www_app
+flask_app = www_app.create_app()
+flask_appbuilder = flask_app.appbuilder
+{{- end }}
+
+
+#############
+## Classes ##
+#############
+class RoleWrapper(object):
+    def __init__(
+            self,
+            name: str,
+            permissions: List[Tuple[str, str]] = []
+    ):
+        self.name = name
+        self.permissions = permissions
+
+    def as_dict(self) -> Dict[str, any]:
+        return {
+            "name": self.name,
+            "permissions": self.permissions
+        }
+
+
+###############
+## Variables ##
+###############
+VAR__TEMPLATE_NAMES = []
+VAR__TEMPLATE_MTIME_CACHE = {}
+VAR__TEMPLATE_VALUE_CACHE = {}
+VAR__ROLE_WRAPPERS = {
+  {{- range .Values.airflow.roles }}
+  {{ .name | quote }}: RoleWrapper(
+    name={{ (required "each `name` in `airflow.roles` must be non-empty!" .name) | quote }},
+    permissions=[
+      {{- range .permissions }}
+        ( {{ index . 0 | quote }}, {{ index . 1 | quote }} ),
+      {{- end }}
+    ]
+  ),
+  {{- end }}
+}
+
+
+def sync_role(role_wrapper: RoleWrapper) -> None:
+    """
+    Sync the Role defined by a provided RoleWrapper into the FAB DB.
+    """
+    name = role_wrapper.name
+    r_new = role_wrapper.as_dict()
+    r_old = flask_appbuilder.sm.find_role(name=name)
+
+    if r_old:
+        role = r_old
+    else:
+        logging.info(f"Role=`{name}` is missing, adding...")
+        role = flask_appbuilder.sm.add_role(name=r_new["name"])
+        if role:
+            logging.info(f"Role=`{name}` was successfully added.")
+        else:
+            logging.error(f"Failed to add Role=`{name}`")
+            sys.exit(1)
+
+    p_old = set([(p.permission.name, p.view_menu.name) for p in role.permissions])
+    p_new = set(r_new["permissions"])
+
+    for p in (p_old - p_new):
+        # Not deleting DAG-level permissions, as they are assigned using `access_control` attribute in DAG code
+        if not p[1].startswith('DAG:'):
+            perm_view = flask_appbuilder.sm.find_permission_view_menu(p[0], p[1])
+            flask_appbuilder.sm.del_permission_role(role, perm_view)
+            logging.info(f"Deleted permission `{perm_view}` from role=`{role.name}`")
+
+    for p in (p_new - p_old):
+        perm_view = flask_appbuilder.sm.find_permission_view_menu(p[0], p[1])
+        if perm_view is None:
+            logging.error(f"Failed to add permission `{p[0]} {p[1]}` to role=`{role.name}` - no such permission")
+            sys.exit(1)
+        flask_appbuilder.sm.add_permission_role(role, perm_view)
+        logging.info(f"Added permission `{perm_view}` to role=`{role.name}`")
+
+
+def sync_all_roles(role_wrappers: Dict[str, RoleWrapper]) -> None:
+    """
+    Sync all roles in provided `role_wrappers`.
+    """
+    logging.info("BEGIN: airflow roles sync")
+    for role_wrapper in role_wrappers.values():
+        sync_role(role_wrapper)
+    logging.info("END: airflow roles sync")
+
+    # ensures than any SQLAlchemy sessions are closed (so we don't hold a connection to the database)
+    flask_app.do_teardown_appcontext()
+
+
+def sync_with_airflow() -> None:
+    """
+    Preform a sync of all objects with airflow (note, `sync_with_airflow()` is called in `main()` template).
+    """
+    sync_all_roles(role_wrappers=VAR__ROLE_WRAPPERS)
+
+
+##############
+## Run Main ##
+##############
+{{- if .Values.airflow.rolesUpdate }}
+main(sync_forever=True)
+{{- else }}
+main(sync_forever=False)
+{{- end }}
+
+{{- end }}
diff --git a/charts/airflow/templates/sync/sync-roles-deployment.yaml b/charts/airflow/templates/sync/sync-roles-deployment.yaml
@@ -0,0 +1,117 @@
+{{- if and (.Values.airflow.roles) (.Values.airflow.rolesUpdate) }}
+{{- $podNodeSelector := include "airflow.podNodeSelector" (dict "Release" .Release "Values" .Values "nodeSelector" .Values.airflow.sync.nodeSelector) }}
+{{- $podAffinity := include "airflow.podAffinity" (dict "Release" .Release "Values" .Values "affinity" .Values.airflow.sync.affinity) }}
+{{- $podTolerations := include "airflow.podTolerations" (dict "Release" .Release "Values" .Values "tolerations" .Values.airflow.sync.tolerations) }}
+{{- $podSecurityContext := include "airflow.podSecurityContext" (dict "Release" .Release "Values" .Values "securityContext" .Values.airflow.sync.securityContext) }}
+{{- $extraPipPackages := .Values.airflow.extraPipPackages }}
+{{- $volumeMounts := include "airflow.volumeMounts" (dict "Release" .Release "Values" .Values "extraPipPackages" $extraPipPackages) }}
+{{- $volumes := include "airflow.volumes" (dict "Release" .Release "Values" .Values "extraPipPackages" $extraPipPackages) }}
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: {{ include "airflow.fullname" . }}-sync-roles
+  labels:
+    app: {{ include "airflow.labels.app" . }}
+    component: sync-roles
+    chart: {{ include "airflow.labels.chart" . }}
+    release: {{ .Release.Name }}
+    heritage: {{ .Release.Service }}
+  {{- if .Values.airflow.sync.annotations }}
+  annotations:
+    {{- toYaml .Values.airflow.sync.annotations | nindent 4 }}
+  {{- end }}
+spec:
+  replicas: 1
+  strategy:
+    ## only 1 replica should run at a time
+    type: Recreate
+  selector:
+    matchLabels:
+      app: {{ include "airflow.labels.app" . }}
+      component: sync-roles
+      release: {{ .Release.Name }}
+  template:
+    metadata:
+      annotations:
+        checksum/secret-config-envs: {{ include (print $.Template.BasePath "/config/secret-config-envs.yaml") . | sha256sum }}
+        checksum/secret-local-settings: {{ include (print $.Template.BasePath "/config/secret-local-settings.yaml") . | sha256sum }}
+        checksum/sync-roles-script: {{ include "airflow.sync.sync_roles.py" . | sha256sum }}
+        {{- if .Values.airflow.podAnnotations }}
+        {{- toYaml .Values.airflow.podAnnotations | nindent 8 }}
+        {{- end }}
+        {{- if .Values.airflow.sync.podAnnotations }}
+        {{- toYaml .Values.airflow.sync.podAnnotations | nindent 8 }}
+        {{- end }}
+        {{- if .Values.airflow.sync.safeToEvict }}
+        cluster-autoscaler.kubernetes.io/safe-to-evict: "true"
+        {{- end }}
+      labels:
+        app: {{ include "airflow.labels.app" . }}
+        component: sync-roles
+        release: {{ .Release.Name }}
+        {{- if .Values.airflow.sync.podLabels }}
+        {{- toYaml .Values.airflow.sync.podLabels | nindent 8 }}
+        {{- end }}
+    spec:
+      restartPolicy: Always
+      {{- if .Values.airflow.image.pullSecret }}
+      imagePullSecrets:
+        - name: {{ .Values.airflow.image.pullSecret }}
+      {{- end }}
+      {{- if $podNodeSelector }}
+      nodeSelector:
+        {{- $podNodeSelector | nindent 8 }}
+      {{- end }}
+      {{- if $podAffinity }}
+      affinity:
+        {{- $podAffinity | nindent 8 }}
+      {{- end }}
+      {{- if $podTolerations }}
+      tolerations:
+        {{- $podTolerations | nindent 8 }}
+      {{- end }}
+      {{- if $podSecurityContext }}
+      securityContext:
+        {{- $podSecurityContext | nindent 8 }}
+      {{- end }}
+      serviceAccountName: {{ include "airflow.serviceAccountName" . }}
+      initContainers:
+        {{- if $extraPipPackages }}
+        {{- include "airflow.init_container.install_pip_packages" (dict "Release" .Release "Values" .Values "extraPipPackages" $extraPipPackages) | indent 8 }}
+        {{- end }}
+        {{- if .Values.dags.gitSync.enabled }}
+        ## git-sync is included so "airflow plugins" & "python packages" can be stored in the dags repo
+        {{- include "airflow.container.git_sync" (dict "Release" .Release "Values" .Values "sync_one_time" "true") | indent 8 }}
+        {{- end }}
+        {{- include "airflow.init_container.check_db" (dict "Release" .Release "Values" .Values "volumeMounts" $volumeMounts) | indent 8 }}
+        {{- include "airflow.init_container.wait_for_db_migrations" (dict "Release" .Release "Values" .Values "volumeMounts" $volumeMounts) | indent 8 }}
+      containers:
+        - name: sync-airflow-roles
+          {{- include "airflow.image" . | indent 10 }}
+          resources:
+            {{- toYaml .Values.airflow.sync.resources | nindent 12 }}
+          envFrom:
+            {{- include "airflow.envFrom" . | indent 12 }}
+          env:
+            {{- include "airflow.env" . | indent 12 }}
+          command:
+            {{- include "airflow.command" . | indent 12 }}
+          args:
+            - "python"
+            - "-u"
+            - "/mnt/scripts/sync_roles.py"
+          volumeMounts:
+            {{- $volumeMounts | indent 12 }}
+            - name: scripts
+              mountPath: /mnt/scripts
+              readOnly: true
+        {{- if .Values.dags.gitSync.enabled }}
+        ## git-sync is included so "airflow plugins" & "python packages" can be stored in the dags repo
+        {{- include "airflow.container.git_sync" . | indent 8 }}
+        {{- end }}
+      volumes:
+        {{- $volumes | indent 8 }}
+        - name: scripts
+          secret:
+            secretName: {{ include "airflow.fullname" . }}-sync-roles
+{{- end }}