connection sync DAG with secrets support

dags/datacat_connection_sync.py

+
+from typing import Dict
+from airflow.decorators import dag, task
+from airflow.models.connection import Connection
+from airflow.utils.dates import days_ago
+from airflow import  settings
+import logging
+from sqlalchemy.orm.session import Session as SASession
+from datacat_integration.secrets import DataCatConnectionWithSecrets
+
+from datacat_integration.hooks import DataCatalogHook
+
+default_args = {
+    'owner': 'airflow',
+}
+
+connections_type = "airflow_connections"
+substring_start = len(connections_type) + 1
+substring_end = substring_start + 36 # length of a UUID4
+
+log = logging.getLogger(__name__)
+
+def get_conn_name(datacat_type: str, oid: str):
+    return "{}/{}-connection".format(datacat_type, oid)
+
+def get_normal_or_secret_property(key: str, props: Dict[str,str], secrets: Dict[str, str], default_value = None):
+    return props.get(key, secrets.get(key, default_value))
+
+
+def get_connection(hook: DataCatalogHook, datacat_type: str, oid: str):
+    conn_id = get_conn_name(datacat_type, oid)
+    secrets_connection = DataCatConnectionWithSecrets(hook.connection.url, hook.connection.user, hook.connection._password)
+    datacat_entry: Dict[str,str] = secrets_connection.get_object(datacat_type, oid)['metadata']
+    datacat_entry_secrets = secrets_connection.get_all_secret_key_value(datacat_type, oid)
+    extra={}
+    predefined_keys = ['conn_type', 'description', 'host', 'login', 'password', 'schema', 'port']
+    # build extra from non-predefined keys
+    for key in datacat_entry:
+        if key not in predefined_keys:
+            extra[key] = datacat_entry[key]
+    
+    for key in datacat_entry_secrets:
+        if key not in predefined_keys:
+            extra[key] = datacat_entry_secrets[key]
+
+    
+    return Connection(
+        conn_id=conn_id,
+        conn_type=get_normal_or_secret_property('conn_type', datacat_entry, datacat_entry_secrets),
+        description=get_normal_or_secret_property('description', datacat_entry, datacat_entry_secrets, 'Automatically generated Connection from the datacatalog object {}/{}'.format(connections_type, oid)),
+        host=get_normal_or_secret_property('host', datacat_entry, datacat_entry_secrets),
+        login=get_normal_or_secret_property('login', datacat_entry, datacat_entry_secrets),
+        password=get_normal_or_secret_property('password', datacat_entry, datacat_entry_secrets),
+        schema=get_normal_or_secret_property('schema', datacat_entry, datacat_entry_secrets),
+        port=int(get_normal_or_secret_property('port', datacat_entry, datacat_entry_secrets)),
+        extra=extra
+    )
+
+
+@dag(default_args=default_args, schedule_interval='@hourly', start_date=days_ago(1), tags=['dls-service-dag'])
+def sync_connections():
+
+    @task
+    def list_catalog_connections(**kwargs):
+        hook = DataCatalogHook("datacatalog")
+        objects = hook.list_type(connections_type)
+        oid_list = [element[1] for element in objects]
+        return oid_list
+
+    @task
+    def list_airflow_connections(**kwargs):
+        session : SASession = settings.Session()
+        conns = session.query(Connection).filter(Connection.conn_id.like("{}/%-connection".format(connections_type)))
+        oid_list = [conn.conn_id[substring_start:substring_end] for conn in conns]
+        return oid_list
+
+    @task 
+    def get_add_list(catalog_connections, airflow_connections, **kwargs):
+        return list(set(catalog_connections).difference(airflow_connections))
+    
+    @task 
+    def get_remove_list(catalog_connections, airflow_connections, **kwargs):
+        return list(set(airflow_connections).difference(catalog_connections))
+
+    @task
+    def remove_connections(oid_list, **kwargs):
+        log.info("Going to remove from conections: " + ','.join(oid_list))
+        session : SASession = settings.Session()
+        for oid in oid_list:
+            session.query(Connection).filter(Connection.conn_id == get_conn_name(connections_type, oid)).delete()
+        session.commit()
+    
+    @task
+    def add_connections(oid_list, **kwargs):
+        log.info("Going to add to conections: " + ','.join(oid_list))
+        hook = DataCatalogHook("datacatalog")
+        connections = []
+        for oid in oid_list:
+            connections.append(get_connection(hook, connections_type, oid))
+        
+        session = settings.Session()
+        # no check for existsnce necessary, since it is handled by get_add_list()
+        for conn in connections: 
+            session.add(conn)
+
+        session.commit()
+
+    cat_conn = list_catalog_connections()
+    air_conn = list_airflow_connections()
+
+    add_list = get_add_list(cat_conn, air_conn)
+    remove_list = get_remove_list(cat_conn, air_conn)
+
+    add_connections(add_list)
+
+    remove_connections(remove_list)
+
+
+dag = sync_connections()