MG-3470 Update to Airflow 2.4.1 and Pandas 1.5.0 (#61)

* feat: try to upgrade to airflow 2.4
* fix: install package by pip because python-poetry/poetry#1214
* refactor: update docker images
* refactor: update aws connections
* feat: add XCOM_WORK_KEY_PREFIX constant
* fix: copy_from don't work with schema.table for psycopg > 2.9 psycopg/psycopg2#1294
* refactor: use DagRunState class instead of str
* fix: use right TaskInstance
* feat: use new Xcom logic
* refactor: use schedule instead of schedule_interval
* refactor: remove create dagrun from fixture
* feat: add create_dag_run to dag_generator tests
* feat: add create_dag_run to operators tests
* feat: updata pandas to 1.5.0
* fix: size of empty DataFrame changed
* fix: ports in docker-compose after review
* fix: down version to 2.1.0

Co-authored-by: Ivan Kanashov <i.kanashov@tinkoff.ru>

Author: ikanashov

data-detective-airflow/Dockerfile

@@ -1,6 +1,6 @@
 # Main docker building
 ARG PYTHON_VERSION=${PYTHON_VERSION:-3.9}
-FROM python:${PYTHON_VERSION}-slim-buster as base
+FROM python:${PYTHON_VERSION}-slim-bullseye as base
 
 ARG AIRFLOW_USER_UID=50000
 ENV DEBIAN_FRONTEND=noninteractive \
@@ -22,10 +22,10 @@ ENV PYTHONPATH=$PATH:${AIRFLOW_HOME} \
 
 
 # Устанавливаем pip из более полного образа python-image (в основном - slim не достаточно утилит для pip install)
-FROM python:${PYTHON_VERSION}-buster as builder
+FROM python:${PYTHON_VERSION}-bullseye as builder
 
 ARG AIRFLOW_USER_UID=50000
-ARG POETRY_VERSION=${POETRY_VERSION:-1.1.13}
+ARG POETRY_VERSION=${POETRY_VERSION:-1.2.2}
 
 ENV AIRFLOW_HOME=/usr/local/airflow/ \
     AIRFLOW_USER_UID=${AIRFLOW_USER_UID}
@@ -43,9 +43,11 @@ ENV POETRY_VERSION=${POETRY_VERSION} \
 
 WORKDIR ${AIRFLOW_HOME}
 COPY pyproject.toml poetry.lock ${AIRFLOW_HOME}/
+# poetry not work without virtualenv https://github.com/python-poetry/poetry/issues/1214
 RUN pip install --no-cache-dir "poetry==${POETRY_VERSION}" \
     && poetry config virtualenvs.create false \
-    && poetry install --no-interaction --no-ansi
+    && poetry export -f requirements.txt --with dev --output requirements.txt \
+    && pip install -r requirements.txt
 
 
 FROM base as dev
@@ -60,7 +62,7 @@ WORKDIR ${AIRFLOW_HOME}
 USER root
 RUN echo "airflow ALL = NOPASSWD: /usr/sbin/sshd" >> /etc/sudoers.d/airflow \
  && apt-get update --allow-releaseinfo-change\
- && apt-get install -y openssh-server netcat make\
+ && apt-get install -y openssh-server netcat make libpq5\
  && apt-get clean \
  && mkdir /var/run/sshd \
  && sed -i 's/\#PermitRootLogin prohibit-password/PermitRootLogin yes/' /etc/ssh/sshd_config \

data-detective-airflow/data_detective_airflow/constants.py

@@ -13,6 +13,7 @@
 WORK_S3_PREFIX: str = 'dd_airflow'
 WORK_FILE_PREFIX: str = 'wrk'
 WORK_PG_SCHEMA_PREFIX: str = 'wrk'
+XCOM_WORK_KEY_PREFIX = 'work_'
 
 # DAG
 DEFAULT_START_DATE = datetime(2020, 4, 8)

data-detective-airflow/data_detective_airflow/dag_generator/dags/python_dag.py

@@ -25,7 +25,7 @@ def __init__(self, dag_dir: str, config: None):
             dag_id=Path(dag_dir).name,
             factory=self.config['factory'],
             start_date=DEFAULT_START_DATE,
-            schedule_interval=self.config['schedule_interval'],
+            schedule=self.config['schedule_interval'],
             description=self.config.get('description', ''),
             default_args=self.config['default_args'],
             template_searchpath=dag_dir,

data-detective-airflow/data_detective_airflow/dag_generator/dags/tdag.py

@@ -1,12 +1,13 @@
 # -*- coding: utf-8 -*-
+import re
 from typing import Callable
 
 from airflow import DAG
 from airflow.models.xcom import XCom
 from airflow.utils.context import Context
 from airflow.utils.module_loading import import_string
 
-from data_detective_airflow.constants import PG_CONN_ID, S3_CONN_ID, SFTP_CONN_ID
+from data_detective_airflow.constants import PG_CONN_ID, S3_CONN_ID, SFTP_CONN_ID, XCOM_WORK_KEY_PREFIX
 from data_detective_airflow.dag_generator.results import PgResult, PickleResult
 from data_detective_airflow.dag_generator.results.base_result import BaseResult, ResultType
 from data_detective_airflow.dag_generator.works import FileWork, PgWork, S3Work, SFTPWork
@@ -93,11 +94,14 @@ def clear_all_works(self, context: Context):
     def get_all_works(self, context: Context):
         """Clearing all work on completion of execution"""
         dag_id = self.dag_id
-        execution_date = context['logical_date']
-        xcoms = XCom.get_many(task_ids='work', dag_ids=dag_id, execution_date=execution_date)
+        run_id = context['run_id']
+        xcoms = [xcom for xcom in XCom.get_many(dag_ids=dag_id, run_id=run_id)
+                 if xcom.key.startswith(XCOM_WORK_KEY_PREFIX)]
         works = set()
         for xcom in xcoms:
-            work = self.get_work(work_type=xcom.key.split('-')[0], work_conn_id=xcom.key.split('-')[1])
+            work_key = re.sub(XCOM_WORK_KEY_PREFIX, '', xcom.key).split('-', 2)
+            work_type, work_conn_id = work_key[0], work_key[1]
+            work = self.get_work(work_type=work_type, work_conn_id=work_conn_id)
             work.set_params(params=dict(xcom.value))
             works.add(work)
 

data-detective-airflow/data_detective_airflow/dag_generator/dags/yaml_dag.py

@@ -27,7 +27,7 @@ def __init__(self, dag_dir: str, config: dict[str, Any]):
             dag_id=Path(dag_dir).name,
             factory=self.config['factory'],
             start_date=constants.DEFAULT_START_DATE,
-            schedule_interval=self.config.get('schedule_interval'),
+            schedule=self.config.get('schedule_interval'),
             description=self.config.get('description', ''),
             default_args=self.config['default_args'],
             template_searchpath=dag_dir,

data-detective-airflow/data_detective_airflow/dag_generator/works/base_work.py

@@ -10,12 +10,12 @@
 from collections.abc import Callable
 from enum import Enum
 from functools import wraps
-from typing import Any, Dict, Optional, Text, Union
+from typing import Any, Optional, Union
 
 from airflow.models.xcom import XCom
 from airflow.utils.log.logging_mixin import LoggingMixin
 
-from data_detective_airflow.constants import CLEAR_WORK_KEY, DAG_ID_KEY, TASK_ID_KEY
+from data_detective_airflow.constants import CLEAR_WORK_KEY, DAG_ID_KEY, TASK_ID_KEY, XCOM_WORK_KEY_PREFIX
 
 
 class WorkType(Enum):
@@ -57,8 +57,6 @@ class BaseWork(ABC, LoggingMixin):
     Note: each worker will have its own instance of this class created, so one worker will not know that a work has already been created in another worker
     """
 
-    XCOM_TASK_KEY = 'work'
-
     def __init__(self, dag, work_type: str, conn_id: str = None, **kwargs):  # pylint: disable=super-init-not-called
         del kwargs
 
@@ -78,7 +76,7 @@ def __eq__(self, other):
     def __ne__(self, other):
         return not self.__eq__(other)
 
-    def get_path(self, context: Dict, prefix: str = 'work_airflow'):
+    def get_path(self, context: dict, prefix: str = 'work_airflow'):
         """Return a unique work name for this context
         (for dag_run)
         Important note! The name of the work is uniquely determined by dag_run
@@ -119,13 +117,13 @@ def clear(self, context: dict):
             self._clear_logic(context)
             self._exists = False
 
-    def get_xcom_key(self, context: Dict):
+    def get_xcom_key(self, context: dict):
         """Return the key for XCom for the current work and task"""
         task = context.get('task')
-        base_xcom_key = f'{self.work_type}-{self.conn_id}'
+        base_xcom_key = f'{XCOM_WORK_KEY_PREFIX}{self.work_type}-{self.conn_id}'
         return f'{base_xcom_key}-{task.task_id}' if task else base_xcom_key
 
-    def get_xcom_params(self, context: Dict) -> Dict:
+    def get_xcom_params(self, context: dict) -> dict:
         """Serialize work to a dictionary for writing to XCom"""
         return {
             'key': self.get_xcom_key(context),
@@ -134,12 +132,12 @@ def get_xcom_params(self, context: Dict) -> Dict:
             'value': {
                 '_exists': self._exists,
             },
-            'execution_date': context['execution_date'],
-            TASK_ID_KEY: BaseWork.XCOM_TASK_KEY,
+            'run_id': context['run_id'],
+            TASK_ID_KEY: context.get('task').task_id,
             DAG_ID_KEY: self.dag.dag_id,
         }
 
-    def _get_attrs_from_xcom(self, context: Dict) -> Dict:
+    def _get_attrs_from_xcom(self, context: dict) -> dict:
         """Вернуть из XCom атрибуты для текущего таска"""
         to_get = self.get_xcom_params(context)
         to_get.pop('value')
@@ -153,7 +151,7 @@ def set_params(self, params: Optional[dict[str, Any]] = None):
             setattr(self, key, value)
 
     @abstractmethod
-    def _create_logic(self, context: Dict):
+    def _create_logic(self, context: dict):
         """Create a work unique for the context.
         Important note! The inheritors need to make this method idempotent, i.e. the method can be called
         as many times as needed.
@@ -163,7 +161,7 @@ def _create_logic(self, context: Dict):
         """
 
     @abstractmethod
-    def _clear_logic(self, context: Dict):
+    def _clear_logic(self, context: dict):
         """Deinitialization"""
 
     @abstractmethod

data-detective-airflow/data_detective_airflow/test_utilities/airflow.py

@@ -3,6 +3,7 @@
 
 from airflow.models import DagRun, TaskInstance
 from airflow.utils.context import Context
+from airflow.utils.state import DagRunState
 from airflow.utils.types import DagRunType
 
 
@@ -18,7 +19,7 @@ def create_or_get_dagrun(dag, task) -> DagRun:
         return dag.create_dagrun(
             execution_date=datetime.now(timezone(offset=timedelta(hours=0))),
             run_type=DagRunType.MANUAL,
-            state='queued',
+            state=DagRunState.QUEUED,
             data_interval=(
                 datetime.now(timezone(offset=timedelta(hours=0))),
                 datetime.now(timezone(offset=timedelta(hours=0))) + timedelta(hours=1)

data-detective-airflow/data_detective_airflow/test_utilities/generate_df.py

@@ -72,7 +72,7 @@ def fill_table_from_dataframe(conn: psycopg2.extensions.connection, dframe: Data
         logging.info(f'Run sql: {trunc_cmd}')
         cur.execute(trunc_cmd)
         logging.info('Fill data')
-        cur.copy_from(file=buffer, table=full_table_name, sep=sep)
+        cur.copy_expert(f"COPY {full_table_name} FROM STDIN DELIMITER '{sep}';", file=buffer)
         conn.commit()
         return True
     except psycopg2.DatabaseError as err:

data-detective-airflow/data_detective_airflow/utils/logging_thread.py

@@ -7,7 +7,7 @@
 
 class LoggingThread(LoggingMixin):
     def __init__(self, context, interval=300):
-        super().__init__(context=context)
+        super().__init__(context=context['ti'])
         self.interval = interval
         self.dag_id = context['dag'].dag_id
         self.task_id = context['task'].task_id

data-detective-airflow/docker-compose.CeleryExecutor.yml

@@ -9,9 +9,9 @@ x-airflow: &airflow
     - AIRFLOW__CELERY__WORKER_AUTOSCALE=16,12
     - AIRFLOW__CELERY__BROKER_URL=redis://redis:6379/1
     - AIRFLOW__CELERY__RESULT_BACKEND=db+postgresql://${META_USER}:${META_PASS}@metadb:5432/${META_USER}
-    - AIRFLOW__CORE__SQL_ALCHEMY_CONN=postgresql+psycopg2://${META_USER}:${META_PASS}@metadb:5432/${META_USER}
+    - AIRFLOW__DATABASE__SQL_ALCHEMY_CONN=postgresql+psycopg2://${META_USER}:${META_PASS}@metadb:5432/${META_USER}
     - AIRFLOW__WEBSERVER__SECRET_KEY=${SECRET_KEY}
-    - AIRFLOW__CORE__LOAD_DEFAULT_CONNECTIONS=False
+    - AIRFLOW__DATABASE__LOAD_DEFAULT_CONNECTIONS=False
     - AIRFLOW__CORE__LOAD_EXAMPLES=False
 
 services:
@@ -54,7 +54,7 @@ services:
     ports:
       - "5005:5432"
     volumes:
-      - ./docker/init-pgwork.sql:/docker-entrypoint-initdb.d/init-pgwork.sql
+      - ./docker/init-pg.sql:/docker-entrypoint-initdb.d/init-pg.sql
     healthcheck:
       test: ["CMD-SHELL", "pg_isready -U airflow"]
       interval: 2s
@@ -63,14 +63,16 @@ services:
     restart: always
 
   s3:
-    image: localstack/localstack:0.11.5
+    image: localstack/localstack:1.0.4
     ports:
-      - 4566:4566
-      - 8055:8080
+      - "4566:4566"
+      - "8055:8080"
     environment:
+      - EAGER_SERVICE_LOADING=1
       - SERVICES=s3
       - DEBUG=0
-      - DATA_DIR=/tmp/localstack/data
+      # Disable LocalStack Event Analytics
+      - DISABLE_EVENTS=1
     volumes:
       - ./docker/init-s3.sh:/docker-entrypoint-initaws.d/init-s3.sh
     healthcheck:
@@ -91,11 +93,11 @@ services:
         condition: service_healthy
     volumes:
       - ./dags:${AIRFLOW_HOME}/dags
-      - ./data_detectivedata_detectivedata_detectivedata_detective_airflow:${AIRFLOW_HOME}/data_detectivedata_detectivedata_detectivedata_detective_airflow
+      - ./data_detective_airflow:${AIRFLOW_HOME}/data_detective_airflow
       - ./tests:${AIRFLOW_HOME}/tests
       - ./tests_data:${AIRFLOW_HOME}/tests_data
     ports:
-      - 127.0.0.1:9922:22
+      - "127.0.0.1:9922:22"
     command: scheduler-ssh
     restart: always
 
@@ -107,7 +109,7 @@ services:
       - ./dags:${AIRFLOW_HOME}/dags
       - ./data_detective_airflow:${AIRFLOW_HOME}/data_detective_airflow
     ports:
-      - 8080:8080
+      - "8080:8080"
     command: webserver
     healthcheck:
       test: curl -s -I -o /dev/null -w '%{http_code}' webserver:8080 | grep -qE '302|200'
@@ -121,7 +123,7 @@ services:
     depends_on:
       - scheduler
     ports:
-      - 5555:5555
+      - "5555:5555"
     command: flower
     healthcheck:
       test: ["CMD", "curl", "--fail", "http://flower:5555/"]

data-detective-airflow/docker-compose.tests.yml

@@ -5,10 +5,10 @@ x-airflow: &airflow
     context: .
     target: dev
   environment:
-    - AIRFLOW__CORE__SQL_ALCHEMY_CONN=postgresql+psycopg2://${META_USER}:${META_PASS}@metadb:5432/${META_USER}
+    - AIRFLOW__DATABASE__SQL_ALCHEMY_CONN=postgresql+psycopg2://${META_USER}:${META_PASS}@metadb:5432/${META_USER}
     - AIRFLOW__CORE__EXECUTOR=LocalExecutor
     - AIRFLOW__WEBSERVER__SECRET_KEY=${SECRET_KEY}
-    - AIRFLOW__CORE__LOAD_DEFAULT_CONNECTIONS=False
+    - AIRFLOW__DATABASE__LOAD_DEFAULT_CONNECTIONS=False
     - AIRFLOW__CORE__LOAD_EXAMPLES=False
     - PYTHON_VERSION=${PYTHON_VERSION}
     - POETRY_VERSION=${POETRY_VERSION}
@@ -53,11 +53,13 @@ services:
       start_period: 20s
 
   s3:
-    image: localstack/localstack:0.11.5
+    image: localstack/localstack:1.0.4
     environment:
+      - EAGER_SERVICE_LOADING=1
       - SERVICES=s3
       - DEBUG=0
-      - DATA_DIR=/tmp/localstack/data
+      # Disable LocalStack Event Analytics
+      - DISABLE_EVENTS=1
     volumes:
       - ./docker/init-s3.sh:/docker-entrypoint-initaws.d/init-s3.sh
     healthcheck:

data-detective-airflow/docker-compose.yml

@@ -5,10 +5,10 @@ x-airflow: &airflow
     context: .
     target: dev
   environment:
-    - AIRFLOW__CORE__SQL_ALCHEMY_CONN=postgresql+psycopg2://${META_USER}:${META_PASS}@metadb:5432/${META_USER}
+    - AIRFLOW__DATABASE__SQL_ALCHEMY_CONN=postgresql+psycopg2://${META_USER}:${META_PASS}@metadb:5432/${META_USER}
     - AIRFLOW__CORE__EXECUTOR=LocalExecutor
     - AIRFLOW__WEBSERVER__SECRET_KEY=${SECRET_KEY}
-    - AIRFLOW__CORE__LOAD_DEFAULT_CONNECTIONS=False
+    - AIRFLOW__DATABASE__LOAD_DEFAULT_CONNECTIONS=False
     - AIRFLOW__CORE__LOAD_EXAMPLES=False
 
 services:
@@ -37,17 +37,19 @@ services:
     ports:
       - "5005:5432"
     volumes:
-      - ./docker/init-pgwork.sql:/docker-entrypoint-initdb.d/init-pgwork.sql
+      - ./docker/init-pg.sql:/docker-entrypoint-initdb.d/init-pg.sql
 
   s3:
-    image: localstack/localstack:0.11.5
+    image: localstack/localstack:1.0.4
     ports:
       - "4566:4566"
       - "8055:8080"
     environment:
+      - EAGER_SERVICE_LOADING=1
       - SERVICES=s3
       - DEBUG=0
-      - DATA_DIR=/tmp/localstack/data
+      # Disable LocalStack Event Analytics
+      - DISABLE_EVENTS=1
     volumes:
       - ./docker/init-s3.sh:/docker-entrypoint-initaws.d/init-s3.sh
 

data-detective-airflow/docker/init-connections.sh

@@ -3,4 +3,4 @@
 # Here it is necessary to set all connections to test/dev servers
 airflow connections add --conn-uri 'postgres://airflow:airflow@pg:5432' pg
 airflow connections add --conn-uri 'ftp://airflow:airflow@ssh_service:22' ssh_service
-airflow connections add --conn-type s3 --conn-extra "{\"aws_access_key_id\": \"accessKey1\", \"aws_secret_access_key\":\"verySecretKey1\", \"host\":\"http://s3:4566\"}" s3
+airflow connections add --conn-type aws --conn-extra "{\"aws_access_key_id\": \"accessKey1\", \"aws_secret_access_key\":\"verySecretKey1\", \"endpoint_url\":\"http://s3:4566\"}" s3