feat: airflow version upgrade

README.md

@@ -7,7 +7,7 @@
 
 ## Compatibility
 
-Apache Airflow 2.0.1+
+Apache Airflow 3.1.5+
 
 ## Usage
 

sample/dags/example.py

@@ -2,9 +2,9 @@
 from random import choice
 
 from airflow import DAG
-from airflow.contrib.sensors.aws_sqs_sensor import SQSSensor
-from airflow.operators.bash_operator import BashOperator
-from airflow.operators.dummy_operator import DummyOperator
+from airflow.providers.amazon.aws.sensors.sqs import SqsSensor
+from airflow.operators.bash import BashOperator
+from airflow.operators.dummy import DummyOperator
 from airflow.operators.ergo import ErgoTaskProducerOperator
 from airflow.sensors.ergo import ErgoJobResultSensor
 from airflow.utils.dates import days_ago

src/dags/dag_job_collector.py

@@ -1,7 +1,7 @@
 from datetime import timedelta
 
 from airflow import DAG
-from airflow.contrib.sensors.aws_sqs_sensor import SQSSensor
+from airflow.providers.amazon.aws.sensors.sqs import SqsSensor
 from airflow.utils import timezone
 from airflow.utils.dates import days_ago
 
@@ -32,9 +32,9 @@
     dagrun_timeout=timedelta(minutes=15),
     max_active_runs=Config.max_runs_dag_job_collector
 ) as dag:
-    sqs_collector = SQSSensor(
+    sqs_collector = SqsSensor(
         task_id=TASK_ID_SQS_COLLECTOR,
-        sqs_queue=sqs_queue_url,
+        queue_url=sqs_queue_url,
         max_messages=10,
         wait_time_seconds=10,
         poke_interval=poke_interval_collector,

src/operators/deferred_job_result.py

@@ -1,8 +1,7 @@
 from datetime import datetime, timedelta
 from airflow.utils.db import provide_session
-from airflow.utils.decorators import apply_defaults
 from airflow.utils.state import State
-from airflow.models import BaseOperator
+from airflow.models.baseoperator import BaseOperator
 from airflow.sensors.base import BaseSensorOperator
 from airflow.triggers.temporal import TimeDeltaTrigger
 from ergo.exceptions import ErgoFailedResultException
@@ -14,7 +13,6 @@
 
 class ErgoDeferredJobResult(BaseOperator):
 
-    @apply_defaults
     def __init__(
             self,
             pusher_task_id: str,

src/operators/ergo_task_producer.py

@@ -1,10 +1,9 @@
 import json
 from typing import Union, List, Tuple
-from airflow.contrib.hooks.aws_sqs_hook import SQSHook
-from airflow.models import BaseOperator
+from airflow.providers.amazon.aws.hooks.sqs import SqsHook
+from airflow.models.baseoperator import BaseOperator
 from ergo.links.ergo_task_detail import ErgoTaskDetailLink
 from airflow.utils.db import provide_session
-from airflow.utils.decorators import apply_defaults
 from airflow.utils.state import State
 from ergo.models import ErgoJob, ErgoTask
 from ergo.config import Config
@@ -17,7 +16,6 @@ class ErgoTaskQueuerOperator(BaseOperator):
 
     operator_extra_links = (ErgoTaskDetailLink(),)
 
-    @apply_defaults
     def __init__(
             self,
             ergo_task_callable: callable = None,
@@ -126,7 +124,7 @@ def execute(self, context, session=None):
         session.commit()
 
     def _send_to_sqs(self, queue_url, task) -> Tuple[List, List]:
-        sqs_client = SQSHook(aws_conn_id=self.aws_conn_id).get_conn()
+        sqs_client = SqsHook(aws_conn_id=self.aws_conn_id).get_conn()
         self.log.info('Trying to push a message on queue: %s\n', queue_url)
         self.log.info('Request task: %s', task.task_id)
         entries = [

src/operators/sqs/result_from_messages.py

@@ -1,9 +1,8 @@
 import json
 
-from airflow.models import BaseOperator
+from airflow.models.baseoperator import BaseOperator
 from airflow.utils import timezone
 from airflow.utils.db import provide_session
-from airflow.utils.decorators import apply_defaults
 from airflow.utils.state import State
 from sqlalchemy.orm import joinedload
 
@@ -12,7 +11,6 @@
 
 
 class JobResultFromMessagesOperator(BaseOperator):
-    @apply_defaults
     def __init__(
         self,
         sqs_sensor_task_id: str,

src/operators/sqs/sqs_task_pusher.py

@@ -1,16 +1,14 @@
 from typing import List, Tuple
 
-from airflow.contrib.hooks.aws_sqs_hook import SQSHook
-from airflow.models import BaseOperator
+from airflow.providers.amazon.aws.hooks.sqs import SqsHook
+from airflow.models.baseoperator import BaseOperator
 from airflow.utils.db import provide_session
-from airflow.utils.decorators import apply_defaults
 from airflow.utils.state import State
 from ergo.models import ErgoJob, ErgoTask
 
 class SqsTaskPusherOperator(BaseOperator):
     filter_ergo_task = ErgoTask.state.in_([State.SCHEDULED, State.UP_FOR_RESCHEDULE])
 
-    @apply_defaults
     def __init__(
         self,
         task_id_collector: str,
@@ -66,7 +64,7 @@ def execute(self, context, session=None):
 
     def _send_to_sqs(self, queue_url, query) -> Tuple[List, List]:
         tasks = list(query)
-        sqs_client = SQSHook(aws_conn_id=self.aws_conn_id).get_conn()
+        sqs_client = SqsHook(aws_conn_id=self.aws_conn_id).get_conn()
         self.log.info('Trying to push %d messages on queue: %s\n',
                       len(tasks), queue_url)
         self.log.info('Request tasks: ' + '\n'.join([str(task) for task in tasks]))

src/operators/task_producer.py

@@ -1,10 +1,9 @@
 import json
 from typing import Union
 
-from airflow.contrib.hooks.aws_sqs_hook import SQSHook
-from airflow.models import BaseOperator
+from airflow.providers.amazon.aws.hooks.sqs import SqsHook
+from airflow.models.baseoperator import BaseOperator
 from airflow.utils.db import provide_session
-from airflow.utils.decorators import apply_defaults
 
 from ergo.config import Config
 from ergo.links.ergo_task_detail import ErgoTaskDetailLink
@@ -16,7 +15,6 @@ class ErgoTaskProducerOperator(BaseOperator):
 
     operator_extra_links = (ErgoTaskDetailLink(),)
 
-    @apply_defaults
     def __init__(
         self,
         ergo_task_callable: callable = None,

src/sensors/job_result_sensor.py

@@ -2,7 +2,6 @@
 
 from airflow.sensors.base_sensor_operator import BaseSensorOperator
 from airflow.utils.db import provide_session
-from airflow.utils.decorators import apply_defaults
 from airflow.utils.state import State
 from ergo.exceptions import ErgoFailedResultException
 from ergo.models import ErgoJob, ErgoTask
@@ -12,7 +11,6 @@
 class ErgoJobResultSensor(BaseSensorOperator):
     poke_context_fields = ('pusher_task_id', 'wait_for_state')
 
-    @apply_defaults
     def __init__(
         self,
         pusher_task_id: str,

src/sensors/task_requests_batcher.py

@@ -3,7 +3,6 @@
 from airflow.sensors.base_sensor_operator import BaseSensorOperator
 from airflow.utils import timezone
 from airflow.utils.db import provide_session
-from airflow.utils.decorators import apply_defaults
 from airflow.utils.state import State
 from sqlalchemy import func, text
 
@@ -32,7 +31,6 @@ class TaskRequestBatchSensor(BaseSensorOperator):
     filter_ergo_task = ErgoTask.state.in_(
         [State.SCHEDULED, State.UP_FOR_RESCHEDULE])
 
-    @apply_defaults
     def __init__(
         self,
         max_requests: int,