AWS DataZone에서 OpenLineage 기반의 Airflow 데이터 계보 그리기

배경

Airflow는 데이터 마트(Data Mart)를 포함한 데이터 파이프라인 구축 및 관리에서 매우 널리 사용되는 도구입니다. 이러한 Airflow에서 데이터 계보가 중요한 이유는 데이터의 출처와 변환 과정을 명확히 추적할 수 있어 데이터의 신뢰성을 보장하고, 문제 발생 시 원인을 빠르게 파악할 수 있기 때문입니다. 또한, 데이터 계보는 규제 준수와 감사 요구사항을 충족시키는 데 도움을 주며, 데이터 파이프라인의 변경이 미치는 영향을 사전에 분석하여 위험을 최소화할 수 있게 합니다. 이를 통해 데이터 품질 문제를 쉽게 식별하고 해결할 수 있으며, 오류 발생 시 문제의 근원을 빠르게 찾아 다운타임을 줄이고 시스템 안정성을 높일 수 있습니다. 나아가 데이터 계보는 조직 내 데이터 거버넌스 정책을 강화하여 데이터 소유권, 사용 권한, 보안 정책 등을 효과적으로 관리할 수 있도록 지원합니다. 결국 데이터 계보를 통해 복잡한 데이터 파이프라인을 효율적으로 운영하며, 데이터 기반 의사결정의 신뢰성과 조직의 데이터 활용 능력을 향상시킬 수 있습니다.

솔루션 개요

이번 포스팅에서는 Amazon DataZone의 데이터 계보 기능을 활용하여 Airflow에 대한 데이터 계보를 그리는 방법에 대해서 설명합니다. Amazon DataZone 서비스에서는 DataZone 라이프 사이클 내에서 발생한 게시와 구독 이벤트에 대한 데이터 계보만 표현되며, 그 외의 경우에는 사용자가 데이터 계보를 직접 작성해야 합니다. 이를 위해 AWS Lambda를 사용하여 Airflow에 대한 데이터 계보를 OpenLineage 표준에 맞게 Amazon DataZone에 업데이트하는 방법에 대해서 설명합니다.

Airflow에 대한 데이터 계보를 그리면 아래와 같습니다. 이 그림을 보면 오른쪽 노드에 위치한 Output Node은 왼쪽 Input Node 3개를 사용해서 만들어 진 것을 알 수 있습니다. 그리고 이러한 Airflow Job은 COMPLETE 상태이며 이는 데이터 마트 생성 Job이 성공하였다는 것을 의미합니다.

이외에도 Airflow Job에 대한 상세 정보도 아래와 같이 알 수 있습니다.

솔루션 아키텍처

전체적인 아키텍처는 Prodcution 계정과 Governance 계정으로 구성되며, 두 계정은 Transit Gateway를 통해 연결되어 네트워크 통신이 가능하도록 설정됩니다. 먼저, Prodcution계정에서는 데이터의 저장 및 메타데이터 생성 작업이 이루어집니다. 관계형 데이터베이스(RDS)는 데이터를 저장하고, 사용자의 요청에 따라 뷰 테이블(View Tables)을 생성합니다. 이후 AWS Glue Crawler가 RDS 데이터를 스캔하여 메타데이터를 자동으로 생성하며, 이 메타데이터는 Glue Data Catalog에 저장됩니다.

Amazon Managed Workflows for Apache Airflow(MWAA)에서 매일 View 테이블을 갱신하는 Airflow Job이 배치 형태로 동작하게 됩니다. Airflow Job이 동작할 때 마다 로그 정보를 CloudWatch에 저장하게 됩니다. 이 로그 정보는 Assume Role을 사용해서 Governance 계정에서 조회될 수 있도록 구성합니다. Governance계정에서는 Airflow에 대한 데이터 계보 작성과 사용자 접근 제어가 이루어집니다. Lambda 함수는 매일 24시간 동안 Airflow Job에 대한 로그를 가져옵니다. 이 로그에서 프로시저 함수를 가져오고 함수 이름을 이용해 프로시저 정의 요청 쿼리문을 이용해 Input Table과 Output Table 이름을 조회합니다. 조회한 정보와 Airflow Job 정보를 조합한 후 데이터 계보를 Amazon DataZone에 업데이트 합니다.

사용자는 Data Portal을 통해 이러한 메타데이터와 계보 정보를 검색하고 필요한 데이터 자산에 접근할 수 있습니다. Data Portal은 SSO(Single Sign-On) 기능을 제공하는 Identity Center와 연동되어 사용자 인증을 처리하며, 이를 통해 안전한 데이터 접근이 가능합니다. Governance계정 내에서는 Step Functions가 워크플로우를 관리하며, Secrets Manager를 통해 민감한 정보를 안전하게 저장하고 관리합니다. 또한 Lambda 및 기타 서비스에서 발생하는 로그는 CloudWatch Logs에 기록되어 모니터링과 분석에 활용됩니다.

사전 준비 사항

다음과 같은 사항이 사전에 준비되어야 합니다.

Amazon RDS for PostgreSQL에 스키마 및 테이블 생성되어 있어야 함
Amazon DataZone에 Table과 View Table이 자산으로 등록되어 있어야 함
Amazon Managed Workflows for Apache Airflow(MWAA)에 Airflow Job이 구성되어 있어야 함
Daily Batch로 프로시저 함수(View Table 갱신 함수)가 실행되고 있어야 함

RDBMS View Table 계보 작성 단계 요약

단계 1: Amazon Managed Workflows for Apache Airflow(MWAA)에 OpenLineage구성
단계 2: Airflow 로그 Assume Role을 이용한 로그 공유
단계 3: CloudWatch 로그 쿼리
단계 4: 프로시저 정의 요청 쿼리
단계 5: OpenLineage SQL Parser을 이용해 Input/Output 테이블 분석
단계 6: 데이터 계보 작성 및 업데이트

단계 1: Amazon Managed Workflows for Apache Airflow(MWAA)에 OpenLineage 구성

Amazon Managed Workflows for Apache Airflow(MWAA)에 OpenLineage 구성을 위해서는 MWAA 내 DAG 스크립트가 실행되는 이벤트가 1) OpenLineage 형태로 이벤트가 생성되고, 2) 이 이벤트가 CloudWatch Logs에 기록됩니다. 이를 특정 3) 람다 함수(Extract Lineage Events)가 Parsing 해서 DataZone 데이터 계보 이벤트로 전달하게 됩니다. 이를 통해서 데이터 포털에서 MWAA에 대한 데이터 계보를 표현할 수 있습니다. 해당 구성은 Airflow 버전 2.9.2에서 수행하였습니다.

Amazon Managed Workflows for Apache Airflow(MWAA) 플러그인 구성

MWAA에서 OpenLineage 구성을 위해서는 OpenLineage에 대한 플로그인 구성을 아래와 같이 진행합니다.

1. AWS 콘솔 → MWAA → 환경 → 구성할 MWAA 이름 클릭합니다.

2. MWAA 콘솔 화면에서 Amazon S3의 DAG 코드 구성 시 사용한 S3 버킷을 확인합니다.

3. requirements.txt 파일을 아래와 같이 작성합니다. OpenLineage 구성 시 필요한 패키지를 지정해 줍니다. 지정 된 패키지가 플러그인 구성 시 MWAA에 설치 되게 됩니다.

apache-airflow-providers-openlineage==1.8.0
openapi-python-client==0.21.7

4. env_var_plugin.py 파일을 아래와 같이 작성합니다. OpenLineage 이벤트를 어디로 보낼 지 설정하게 됩니다.

5. os.environ[“AIRFLOW__OPENLINEAGE__NAMESPACE”] 값은 MWAA 환경 이름과 같은 값을 사용합니다. 이 이름은 데이터 포털 리니지에서 네임스페이스 이름과 같아야 합니다.

6. “type”: “console” 으로 설정해야 CloudWatch Logs로 OpenLineage 이벤트가 저장됩니다.

from airflow.plugins_manager import AirflowPlugin
import os

os.environ["AIRFLOW__OPENLINEAGE__NAMESPACE"] = "ssts-test-environment"
os.environ["AIRFLOW__OPENLINEAGE__TRANSPORT"] = '''{
  "type": "console"
}'''
os.environ["AIRFLOW__OPENLINEAGE__CONFIG_PATH"] = ""
os.environ["AIRFLOW__OPENLINEAGE__DISABLED_FOR_OPERATORS"] = ""


class EnvVarPlugin(AirflowPlugin):
    name = "env_var_plugin"

7. 이벤트가 저장되는 CloudWatch Logs 그룹은 “Airflow 태스크 로그 그룹”에서 확인할 수 있습니다.

8. env_var_plugin.py 파일을 zip 파일로 압축합니다. 압축 파일명은 plugins.zip으로 압축합니다. Mac 환경에서 압축하게 되면 플러그인 로딩 시 문제가 발생합니다. Linux PC에서 압축하는 것을 추천합니다.

9. plugins.zip 파일과 requirements.txt 파일을 DAG 코드 구성 시 사용한 S3 버킷에 업로드 합니다.

MWAA 구성 편집

1. AWS 콘솔 → MWAA → 환경 → 구성할 MWAA 이름 클릭 → 편집 클릭합니다.

2. 세부 정보 지정 > Amazon S3의 DAG 코드 설정에서 아래와 같이 설정합니다.

3. 플러그인 파일: 위에서 S3에 추가한 plugins.zip 파일을 선택한 후 현재 버전을 선택합니다.

4. 요구 사항 파일: 위에서 S3에 추가한 requirements.txt 파일을 선택한 후 현재 버전을 선택합니다.

5. 다음을 클릭합니다.

6. 고급 설정 구성 > 모니터링에서 아래와 같이 설정합니다.

Airflow 태스크 로그 enable, 로그 수준 INFO
Airflow 웹 서버 로그 enable, 로그 수준 INFO
Airflow 스케줄러 로그 enable, 로그 수준 INFO
Airflow 작업자 로그 enable, 로그 수준 INFO
Airflow DAG 프로세싱 로그 enable, 로그 수준 INFO

7. 다음을 클릭한 후 저장을 클릭합니다.

8. MWAA가 사용 가능한 상태까지 기달립니다.

9. 플러그인 설정이 성공하게 되면 Airflow UI에서 구성 된 정보를 확인할 수 있습니다.

10. Airflow UI > Admin > Plugins 화면에서 아래와 같이 설정 된 정보를 확인합니다.

CloudWatch Logs 파일 확인

1. Airflow 태스크 로그 그룹을 클릭합니다.

2. 로그 그룹 검색을 클릭합니다.

3. console.py로 검색을 수행합니다.

4. OpenLineage 형태로 CloudWatch에 저장 된 로그를 확인할 수 있습니다.

단계 2: Airflow 로그 Assume Role을 이용한 로그 공유

Airflow 로그를 Governance Account에 공유할 수 있는 AssumeRole을 생성합니다. 이 작업은 Production 계정(Airflow가 동작하고 있는)에서 수행합니다.

1. IAM → 역할 → 역할 생성

2. “사용자 지정 신뢰 정책” 선택

3. “사용자 지정 신뢰 정책” 아래와 같이 작성

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "AWS": "arn:aws:iam::{Governance 계정 ID}:root"
            },
            "Action": "sts:AssumeRole"
        }
    ]
}

4. 권한 추가에서 “CloudWatchReadOnlyAccess” 선택

5. 역할 이름에 “{Customer Name}-{Project Name}-iam-role-cloudwatch-share-log” 입력 후 생성

단계 3: CloudWatch 로그 쿼리

단계 2에서 공유한 Airflow 관련 CloudWatch 로그를 쿼리하기 위해 Assume Role을 얻어오는 코드를 아래와 같이 구성합니다.

import os
import time
import json
import boto3
import botocore


# AWS 클라이언트 생성
secrets = boto3.client('secretsmanager', region_name=config.region, config=client_config)

def get_assume_role_params(secret_name):
    """
    AWS Secrets Manager에서 Role ARN을 가져오는 함수
    Args:
        secret_name: Secrets Manager에 저장된 시크릿 이름
    Returns:
        role_arn: 역할의 ARN 문자열
    """
    response = secrets.get_secret_value(SecretId=secret_name)
    secret = json.loads(response['SecretString'])
    
    return secret['role_arn']

def get_assumed_role_session(
    role_arn,
    session_name='AssumedRoleSession'
):
    """
    지정된 IAM Role을 사용하여 임시 보안 자격 증명을 생성하는 함수
    Args:
        role_arn: 위임받을 IAM Role의 ARN
        session_name: 생성할 세션의 이름 (기본값: 'AssumedRoleSession')
    Returns:
        boto3.Session: 임시 자격 증명으로 생성된 새로운 boto3 세션
    """
    # 기본 boto3 세션 생성
    session = boto3.Session()
    # STS 클라이언트 생성
    sts_client = session.client('sts')
    
    # Role Assume 수행
    assumed_role = sts_client.assume_role(
        RoleArn=role_arn,
        RoleSessionName=session_name
    )
    
    # 임시 자격 증명으로 새로운 세션 생성 및 반환
    return boto3.Session(
        aws_access_key_id=assumed_role['Credentials']['AccessKeyId'],
        aws_secret_access_key=assumed_role['Credentials']['SecretAccessKey'],
        aws_session_token=assumed_role['Credentials']['SessionToken']
    )

# Secrets Manager에서 Role ARN을 가져와서 임시 세션 생성
assumed_session = get_assumed_role_session(get_assume_role_params(config.cloudwatch_secret_name))

# 생성된 임시 세션으로 CloudWatch 클라이언트 생성
cloudwatch = assumed_session.client('logs')

CloudWatch에서 Airflow 로그를 가져오는 코드를 아래와 같이 구성합니다.

import os
import boto3
import json

# 리니지 이벤트 시간 설정 (모든 리니지에서 고정된 timestamp 활용)
config.EVENT_TIMESTAMP = event.get('current_time', get_utc_current_time_iso())
asset_results = []
    
# 1. CloudWatch에서 Airflow 로그 가져오기
# 현재 시간과 24시간 전 시간 계산 (필요에 따라 조정 가능)
end_time = int(datetime.now().timestamp() * 1000)
start_time = int((datetime.now() - timedelta(hours=24)).timestamp() * 1000)

# 쿼리 실행
# asc 순서대로 sort 한 후 cloudwatch 이벤트 처리
query_response = cloudwatch.start_query(
    logGroupName= config.log_group_name,
    startTime= start_time,
    endTime= end_time,
    queryString="fields @timestamp, @message | filter @message like /console.py/",
    limit=1000  # 필요에 따라 조정 가능
)
    
query_id = query_response['queryId']


# 쿼리 결과 가져오기
while True:
   cloudwatch_logs = cloudwatch.get_query_results(queryId=query_id)
    if cloudwatch_logs['status'] in ['Complete', 'Failed', 'Cancelled', 'Timeout', 'Unknown']:
      break
            
# 결과 출력
print(f"[INFO] cloudwatch airflow 전체 로그 갯수: {len(cloudwatch_logs['results'])}")

단계 4: 프로시저 정의 요청 쿼리

단계 3에서 가져온 CloudWatch 로그는 아래와 같은 형태입니다.

{
  'eventTime': '2025-01-20T21:00:36.793568+00:00',
  'eventType': 'START',
  'inputs': [
    
  ],
  'job': {
    'facets': {
      'documentation': {
        '_producer': 'https://github.com/apache/airflow/tree/providers-openlineage/1.8.0',
        '_schemaURL': 'https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/DocumentationJobFacet',
        'description': ‘Daily Procedure at 7 AM'
      },
      'jobType': {
        '_producer': 'https://github.com/apache/airflow/tree/providers-openlineage/1.8.0',
        '_schemaURL': 'https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/JobTypeJobFacet',
        'integration': 'AIRFLOW',
        'jobType': 'TASK',
        'processingType': 'BATCH'
      },
      'ownership': {
        '_producer': 'https://github.com/apache/airflow/tree/providers-openlineage/1.8.0',
        '_schemaURL': 'https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/OwnershipJobFacet',
        'owners': [
          {
            'name': '{owner name}'
          }
        ]
      },
      'sql': {
        '_producer': 'https://github.com/apache/airflow/tree/providers-openlineage/1.8.0',
        '_schemaURL': 'https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/SqlJobFacet',
        'query': "CALL {Procedure Name}('20250120')"
      }
    },
    'name': ‘{MWAA Name}’,
    'namespace': '{MWAA NameSpace}'
  },
  'outputs': [
    
  ],
  'producer': 'https://github.com/apache/airflow/tree/providers-openlineage/1.8.0',
  'run': {
    'facets': {
      'airflow': {
        '_producer': 'https://github.com/apache/airflow/tree/providers-openlineage/1.8.0',
        '_schemaURL': 'https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/BaseFacet',
        'dag': {
          'dag_id': '{dag id}',
          'schedule_interval': '0 6 * * *',
          'timetable': {
            'expression': '0 6 * * *',
            'timezone': 'Asia/Seoul'
          }
        },
        'dagRun': {
          'conf': {
            
          },
          'dag_id': '{dag id}',
          'data_interval_end': '2025-01-20T21:00:00+00:00',
          'data_interval_start': '2025-01-19T21:00:00+00:00',
          'external_trigger': False,
          'run_id': 'scheduled__2025-01-19T21:00:00+00:00',
          'run_type': 'scheduled',
          'start_date': '2025-01-20T21:00:00.468808+00:00'
        },
        'task': {
          'depends_on_past': False,
          'downstream_task_ids': '[]',
          'executor_config': {
            
          },
          'ignore_first_depends_on_past': True,
          'is_setup': False,
          'is_teardown': False,
          'mapped': False,
          'multiple_outputs': False,
          'operator_class': 'PostgresOperator',
          'owner': '{owner name}',
          'priority_weight': 1,
          'queue': 'airflow-celery-101fe4a3-d5c5-4794-9311-cc1f3ab788e5',
          'retries': 1,
          'retry_exponential_backoff': False,
          'task_id': '{task id}',
          'trigger_rule': 'all_success',
          'upstream_task_ids': "['{upstream task id}']",
          'wait_for_downstream': False,
          'wait_for_past_depends_before_skipping': False,
          'weight_rule': '<<non-serializable: _DownstreamPriorityWeightStrategy>>'
        },
        'taskInstance': {
          'pool': 'default_pool',
          'try_number': 1
        },
        'taskUuid': '75840854-1401-3adf-b341-8535b1bf9997'
      },
      'nominalTime': {
        '_producer': 'https://github.com/apache/airflow/tree/providers-openlineage/1.8.0',
        '_schemaURL': 'https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/NominalTimeRunFacet',
        'nominalEndTime': '2025-01-20T21:00:00+00:00',
        'nominalStartTime': '2025-01-19T21:00:00+00:00'
      },
      'parent': {
        '_producer': 'https://github.com/apache/airflow/tree/providers-openlineage/1.8.0',
        '_schemaURL': 'https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/ParentRunFacet',
        'job': {
          'name': '{job name}',
          'namespace': '{job namespace}'
        },
        'run': {
          'runId': '3f233364-57e6-3fc9-86bb-aa083018a5fc'
        }
      },
      'parentRun': {
        '_producer': 'https://github.com/apache/airflow/tree/providers-openlineage/1.8.0',
        '_schemaURL': 'https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/ParentRunFacet',
        'job': {
          'name': '{job name}',
          'namespace': '{job namespace}'
        },
        'run': {
          'runId': '3f233364-57e6-3fc9-86bb-aa083018a5fc'
        }
      },
      'processing_engine': {
        '_producer': 'https://github.com/apache/airflow/tree/providers-openlineage/1.8.0',
        '_schemaURL': 'https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/ProcessingEngineRunFacet',
        'name': 'Airflow',
        'openlineageAdapterVersion': '1.8.0',
        'version': '2.9.2'
      }
    },
    'runId': '75840854-1401-3adf-b341-8535b1bf9997'
  },
  'schemaURL': 'https://openlineage.io/spec/1-0-5/OpenLineage.json#/definitions/RunEvent'
}

단계 5: OpenLineage SQL Parser을 이용해 Input/Output 테이블 분석

‘query’에 있는 프로시저 이름을 가져와 프로시저 정의를 요청하는 PostgreSQL 시스템 함수를 호출합니다. 얻어온 프로시저 생성 정의문에서 Input Table 이름과 Output Table 이름을 얻어와야 합니다. 이를 위해서 OpenLineage SQL Parser 파이썬 패키지를 사용합니다.

import os
import boto3
import json
import re
from openlineage_sql import parse
from psycopg2.sql import SQL, Identifier, Literal

procedure_name = re.search(r'\b([a-zA-Z_]+)\b(?=\()', query).group(1)

# 프로시저에 대한 정의를 요청하는 쿼리 요청
sql_query = """
    SELECT PROSRC
    FROM PG_CATALOG.PG_PROC
    WHERE PRONAME = {};
"""
procedure_definition = run_query(config.db_secret_name, 
    SQL(sql_query).format(
        Literal(procedure_name)
    )             
)

if not procedure_definition:
    print(f"[ERROR] {procedure_name} 프로시저에 대한 정의가 존재하지 않습니다.")
    return 'error'
else:
    dml_pattern = r'(DELETE.*?;|INSERT.*?;|UPDATE.*?;|MERGE.*?;)'
    dml_statements = re.findall(dml_pattern, procedure_definition[0]['prosrc'], re.DOTALL | re.IGNORECASE)

    sql_query = ''.join(dml_statements)

    # openlineage_sql을 이용해 파싱한다.
    meta = parse([sql_query])

    in_tables = meta.in_tables
    out_table = meta.out_tables[0]

OpenLineage SQL Parser을 이용해 Input/Output 테이블 분석 코드를 실행하게 되면 아래와 같이 Input Table 이름과 Output Table 이름을 얻을 수 있다.

input table: [wms.wms_inventory, wms.wms_location, wms.wms_movement, wms.wms_stock, wms.wms_supplier, wms.wms_warehouse]
output table: [wms.v_wms_stock_out_r]

단계 6: 데이터 계보 작성 및 업데이트

Amazon DataZone에서는 OpenLineage 호환 이벤트를 사용하여 데이터 계보를 그릴 수 있습니다. Amazon DataZone 데이터 계보는 아래와 같이 표현 됩니다. Airflow 데이터 계보는 MWAA에서 생성 된 OpenLineage 데이터 계보 이벤트에 Input node와 output node만 추가하면 됩니다.

OpenLineage 이벤트 패킷 구성은 아래와 같습니다.

eventTime: OpenLineage 이벤트가 발생한 시점을 나타내는 필드입니다. 이는 ISO 8601 형식의 타임스탬프로 기록되며, 이벤트가 생성된 정확한 시간을 포함합니다.
eventType: 특정 이벤트의 유형을 나타냅니다. 예를 들어, 데이터 처리 작업이 시작되었는지, 완료되었는지, 실패했는지를 정의합니다. 주요 이벤트 유형에는 START, COMPLETE, FAIL, ABORT 등이 있으며, 이는 작업 실행의 상태를 나타냅니다
job: 데이터 계보 처리 작업을 정의하며, 고유한 이름과 네임스페이스로 식별됩니다.
run: 특정 Job이 실행되는 인스턴스를 나타내며, 시작 및 완료 시간과 같은 정보를 포함합니다. 각 Run은 고유한 ID(UUID)로 식별되며, 이는 Job의 동적 실행을 추적하는 데 사용됩니다.
Inputs/Outputs: Inputs는 작업 실행 중 사용된 데이터셋(소스)을 나타내며, Outputs는 작업 결과로 생성된 데이터셋(대상)을 나타냅니다. 각 데이터셋은 네임스페이스와 이름으로 식별되며, 데이터 흐름과 변환 과정을 추적하는 데 중요한 역할을 합니다
facets: Job, Run, Dataset 등의 엔티티에 부가적인 메타데이터를 추가할 수 있는 확장 가능한 필드입니다. 특정 모델이나 프로세스를 더 세부적으로 표현하기 위해 사용되며, 커스터마이징이 가능합니다. 예를 들어, 데이터셋의 스키마 정보나 변환 방식 등을 포함할 수 있습니다.

AWS DataZone에서 Input 테이블과 output 테이블로 등록 된 자산 정보를 조회하는 코드를 아래와 같이 사용합니다.

import boto3
import botocore
from botocore.exceptions import ClientError

# 리니지 노드의 Source Identifier를 사용하여 해당하는 자산(Asset)의 ID를 찾는 함수
def get_asset_id_by_source_identifier(source_identifier):
    # source_identifier에서 첫번째 '/' 이후의 경로만 추출
    name = '/'.join(source_identifier.split('/')[1:])
    
    # 모든 프로젝트 목록을 가져옴
    projects = list_projects()
    for project in projects:
        project_id = project['id']
        # 각 프로젝트 내에서 name과 일치하는 자산을 검색
        assets = search_asset_in_project(project_id, name)
        if assets:
            for asset in assets:
                # 자산의 식별자(identifier) 추출
                asset_id = asset['assetItem']['identifier']
                # 자산의 상세 정보 조회
                asset_info = get_asset(asset_id)
                # AssetCommonDetailsForm 양식에서 정보 필터링
                filtered_forms = [form for form in asset_info['formsOutput'] if form['formName'] == 'AssetCommonDetailsForm']
                content = json.loads(filtered_forms[0]['content'])
                # sourceIdentifier가 일치하는 경우 해당 asset_id 반환
                if content.get('sourceIdentifier') == source_identifier:
                    return asset_id
    # 일치하는 자산을 찾지 못한 경우 None 반환
    return None

# 자산 ID를 사용하여 자산 정보와 컬럼 설명을 조회하는 함수
def get_asset_with_description(asset_id):
    # Amazon DataZone API를 사용하여 자산 정보 조회
    response = datazone.get_asset(
        domainIdentifier=config.domainId,
        identifier=asset_id
    )

    # Glue 관련 양식(GlueViewForm 또는 GlueTableForm) 찾기
    glue_form_index = next((index for (index, d) in enumerate(response['formsOutput']) 
                        if d['formName'] in ['GlueViewForm', 'GlueTableForm']), None)
    glue_data = response['formsOutput'][glue_form_index]['content']
    
    # 컬럼 비즈니스 메타데이터 양식 찾기
    business_metadata_form_index = next((index for (index, d) in enumerate(response['formsOutput']) 
                        if d['formName'] == 'ColumnBusinessMetadataForm'), None)

    # 비즈니스 메타데이터 양식이 없는 경우
    if business_metadata_form_index is None:
        print("ColumnBusinessMetadataForm가 없습니다.")
        data_dict = json.loads(glue_data)
        return data_dict['columns']
    
    # 비즈니스 메타데이터 정보 추출
    business_metadata = response['formsOutput'][business_metadata_form_index]['content']
    # Glue 데이터와 비즈니스 메타데이터를 결합하여 반환
    return enhance_columns(glue_data, business_metadata)

위 OpenLineage 이벤트 패킷 구성을 기반으로 Inputs 테이블 정보를 아래와 같이 추가합니다.

# 입력 테이블들의 정보를 lineage_event_packet에 추가하는 로직
for input_table in input_tables:
    # 소스 식별자 형식 변환
    source_identifier = f"[{config.lineage_node_namespace}] odw.{input_table}"
    
    # OpenLineage 형식에 맞는 입력 테이블 메타데이터 구성
    input_name = {
        "namespace": "postgresql",           # 데이터베이스 타입
        "name": source_identifier,           # 테이블 식별자
        "facets": {                         # OpenLineage 스키마 정보
            "schema": {
                "_producer": "https://github.com/OpenLineage/OpenLineage/tree/0.10.0/integration/airflow",
                "_schemaURL": "https://openlineage.io/spec/1-0-2/OpenLineage.json#/definitions/SchemaDatasetFacet",
                "fields": []                 # 컬럼 정보를 담을 빈 배열
            }
        }
    }

    # 소스 식별자를 이용하여 에셋 식별자 조회
    asset_identifier = get_asset_id_by_source_identifier(f"{config.lineage_node_namespace}" + "/" + source_identifier)
    
    # 자산이 포털에 등록되어 있지 않은 경우
    if asset_identifier is None:
        print(f"[INFO] {source_identifier} Asset이 포털에 등록되지 않았습니다. Default 값으로 node 정보를 업데이트 합니다.")
        lineage_event_packet['inputs'].append(input_name)
    # 자산이 포털에 등록되어 있는 경우
    else:
        # 자산의 상세 컬럼 정보 조회
        asset_columns = get_asset_with_description(asset_identifier)

        # OpenLineage 형식에 맞게 컬럼 속성명 변경
        # dataType -> type, columnName -> name으로 변경
        for column in asset_columns:
            column['type'] = column.pop('dataType')
            column['name'] = column.pop('columnName')

        # 변환된 컬럼 정보를 스키마에 추가
        input_name['facets']['schema']['fields'] = asset_columns
        # 완성된 입력 테이블 정보를 lineage_event_packet에 추가
        lineage_event_packet['inputs'].append(input_name)

위 OpenLineage 이벤트 패킷 구성을 기반으로 Outputs 테이블 정보를 아래와 같이 추가합니다.

# 출력 테이블 정보를 lineage_event_packet에 추가하는 로직
source_identifier = f"[{config.lineage_node_namespace}] odw.{output_table}"

# OpenLineage 형식에 맞는 출력 테이블 메타데이터 구성
output_name = {
    "namespace": "postgresql",           # 데이터베이스 타입
    "name": source_identifier,           # 테이블 식별자
    "facets": {                         # OpenLineage 스키마 정보
        "schema": {
            "_producer": "https://github.com/OpenLineage/OpenLineage/tree/0.10.0/integration/airflow",
            "_schemaURL": "https://openlineage.io/spec/1-0-2/OpenLineage.json#/definitions/SchemaDatasetFacet",
            "fields": []                 # 컬럼 정보를 담을 빈 배열
        }
    }
}

# 소스 식별자를 이용하여 에셋 식별자 조회
asset_identifier = get_asset_id_by_source_identifier(f"{config.lineage_node_namespace}" + "/" + source_identifier)

# 출력 테이블이 포털에 등록되어 있지 않은 경우
if asset_identifier is None:
    print(f"[INFO] {source_identifier} output 노드가 포털에 등록되지 않았습니다. 자산에 대한 계보 표현을 Skip 합니다!")
    return 'skip_log'
# 출력 테이블이 포털에 등록되어 있는 경우
else:
    # 에셋의 상세 컬럼 정보 조회
    asset_columns = get_asset_with_description(asset_identifier)

    # OpenLineage 형식에 맞게 컬럼 속성명 변경
    # dataType -> type, columnName -> name으로 변경
    for column in asset_columns:
        column['type'] = column.pop('dataType')
        column['name'] = column.pop('columnName')

    # 변환된 컬럼 정보를 스키마에 추가
    output_name['facets']['schema']['fields'] = asset_columns
    # 완성된 출력 테이블 정보를 lineage_event_packet에 추가
    lineage_event_packet['outputs'].append(output_name)

OpenLineage 이벤트 패킷 구성이 완료되면 DataZone API을 이용해 데이터 계보를 업데이트 합니다.
try:
    # DataZone API를 호출하여 계보 이벤트 포스팅
    datazone.post_lineage_event(
        domainIdentifier=config.domainId,                  # DataZone 도메인 ID
        event=json.dumps(lineage_event_packet)             # 계보 이벤트 데이터를 JSON 문자열로 변환
    )
    
    # 성공 시 로그 출력
    # 출력 테이블의 namespace와 name을 포함하여 성공 메시지 표시
    print(f"[SUCCESS] '{lineage_event_packet['outputs'][0]['namespace']}/{lineage_event_packet['outputs'][0]['name']}' 계보 이벤트가 성공적으로 포스팅되었습니다.")

except Exception as e:
    # 오류 발생 시 로그 출력
    # 실패한 출력 테이블 정보와 오류 내용을 포함하여 에러 메시지 표시
    print(f"[ERROR] '{lineage_event_packet['outputs'][0]['namespace']}/{lineage_event_packet['outputs'][0]['name']}' 계보 이벤트 포스팅 중 오류 발생: {str(e)}")

작성된 Airflow 데이터 계보를 확인합니다.

결론

이번 포스팅에서는 Amazon DataZone의 데이터 계보 기능을 활용하여 데이터 마트에서 많이 사용하는 Airflow에 대한 데이터 계보를 그리는 방법에 대해서 알아보았습니다. Airflow에서 데이터 계보를 사용하면 데이터 파이프라인의 신뢰성과 효율성을 보장하는 데 핵심적인 역할을 합니다. 데이터 계보를 통해 데이터의 출처와 변환 과정을 명확히 추적할 수 있어 데이터 품질 문제를 빠르게 식별하고 해결할 수 있습니다. 이는 시스템 다운타임을 줄이고 안정성을 높이는 데 기여합니다. 또한, 데이터 계보는 규제 준수 및 감사 요구사항을 충족하며, 데이터 파이프라인 변경 시 영향을 사전에 분석하여 위험을 최소화할 수 있도록 도와줄 수 있습니다. 이를 통해 조직은 데이터 거버넌스 정책을 강화하고, 데이터 소유권, 사용 권한, 보안 정책 등을 효과적으로 관리할 수 있습니다. 궁극적으로 데이터 계보는 복잡한 데이터 파이프라인 운영을 최적화하고, 데이터 기반 의사결정의 신뢰성과 조직의 데이터 활용 능력을 향상시킬 수 있습니다.

쿠키 기본 설정 선택

AWS 기술 블로그