Ubiq Encryption for Airflow

Included is a set of two Airflow Operators for use in your DAGs as a simple interface for performing structured encrypt and decrypt operations.

Setup

In your Airflow Python environment, run

pip install ubiq_security

In the Airflow Documentation about Custom Operators there is an important note:

Note: For imports to work, you should place the file in a directory that is present in the PYTHONPATH env. Airflow adds dags/, plugins/, and config/ directories in the Airflow home to PYTHONPATH by default.

To add the Ubiq Operators, we set it up by creating an operators/ folder within plugins/ (plugins/operators/). This way on startup, the operators will automatically be ready for importing into any custom DAG via

from operators.ubiq_decrypt_operator import UbiqDecryptOperator
from operators.ubiq_encrypt_operator import UbiqEncryptOperator

Usage

Credentials

In your DAG, you will need to pass the Operator a Ubiq Credentials object. This will define what datasets the DAG can access & what permissions it has for encrypting and decrypting.

For example:

import ubiq_security as ubiq

from operators.ubiq_decrypt_operator import UbiqDecryptOperator
from operators.ubiq_encrypt_operator import UbiqEncryptOperator

credentials = ubiq.configCredentials(config_file = "some-credential-file", profile = "some-profile")

with DAG(
    'dag_name',
    schedule="@once",
    start_date=datetime(2021,1,1),
    catchup=False,
    tags=["example", "ubiq"]
) as dag:
        
    ... operations ...

For more options or detail about initializing a credentials object, see the Ubiq Library for Python

Operators

There are two operators, UbiqEncryptOperator and UbiqDecryptOperator. Both expect data to come in as a list of lists. (This is how the BigQueryGetDataOperator handles data.)

Example of data:

# id, name, ssn, birthdate
[
    [1,'Joe Richards', '000-00-0000', '01-01-1970'],
    [1,'Susan Carmichael', '000-00-0001', '05-17-1970']
]

Operator Paramters:

• Parameters from Airflow BaseOperator (Documentation)
• credentials: The Ubiq credentials object
• data_task: The task_id from a prior task that set up the data to be encrypted/decrypted
• dataset_names: A list of dataset names, each corresponding to a column in the data. For data that won't be encrypted/decrypted, simply pass an empty string. The length of the dataset_names array should match the length of a row of the processing data.

...
) as dag:
     get_data = BigQueryGetDataOperator(
        task_id = "get_data",
        dataset_id = DATASET_NAME,
        table_id = TABLE_NAME,
        max_results=10,
        selected_fields="name,email,ssn,birthday",
    )
    
    enc_task = UbiqEncryptOperator(
        task_id="enc_task",
        credentials=creds,
        data_task=get_data.task_id,
        dataset_names=["", "", "SSN", "Birthdate"]
    )

    dec_task = UbiqDecryptOperator(
        task_id = "ubiq_dec_task",
        credentials=creds,
        data_task=enc_task.task_id,
        dataset_names=["", "", "SSN", "Birthdate"]
    )

    (
        get_data
        >> enc_task
        >> dec_task
    )