Changing atlassian JIRA SDK to official atlassian-python-api SDK (apa…

…che#27633) * Changing atlassian JIRA SDK to official atlassian-python-api SDK Co-authored-by: eladkal <[email protected]>
Jorricks · Dec 7, 2022 · f3c68d7 · f3c68d7
1 parent b5338b5
commit f3c68d7
Show file tree

Hide file tree

Showing 9 changed files with 103 additions and 121 deletions.
diff --git a/airflow/providers/atlassian/jira/CHANGELOG.rst b/airflow/providers/atlassian/jira/CHANGELOG.rst
@@ -24,6 +24,23 @@
 Changelog
 ---------
 
+2.0.0
+.....
+
+Breaking changes
+~~~~~~~~~~~~~~~~
+
+* ``Changing atlassian JIRA SDK to official atlassian-python-api SDK (#27633)``
+
+Migrated ``Jira`` provider from Atlassian ``Jira`` SDK to ``atlassian-python-api`` SDK.
+``Jira`` provider doesn't support ``validate`` and ``get_server_info`` in connection extra dict.
+Changed the return type of ``JiraHook.get_conn`` to return an ``atlassian.Jira`` object instead of a ``jira.Jira`` object.
+
+.. warning:: Due to the underlying SDK change, the ``JiraOperator`` now requires ``jira_method`` and ``jira_method_args``
+             arguments as per ``atlassian-python-api``.
+
+             Please refer `Atlassian Python API Documentation <https://atlassian-python-api.readthedocs.io/jira.html>`__
+
 1.1.0
 .....
 

diff --git a/airflow/providers/atlassian/jira/hooks/jira.py b/airflow/providers/atlassian/jira/hooks/jira.py
@@ -20,16 +20,15 @@
 
 from typing import Any
 
-from jira import JIRA
-from jira.exceptions import JIRAError
+from atlassian import Jira
 
 from airflow.exceptions import AirflowException
 from airflow.hooks.base import BaseHook
 
 
 class JiraHook(BaseHook):
     """
-    Jira interaction hook, a Wrapper around JIRA Python SDK.
+    Jira interaction hook, a Wrapper around Atlassian Jira Python SDK.
 
     :param jira_conn_id: reference to a pre-defined Jira Connection
     """
@@ -43,48 +42,33 @@ def __init__(self, jira_conn_id: str = default_conn_name, proxies: Any | None =
         super().__init__()
         self.jira_conn_id = jira_conn_id
         self.proxies = proxies
-        self.client: JIRA | None = None
+        self.client: Jira | None = None
         self.get_conn()
 
-    def get_conn(self) -> JIRA:
+    def get_conn(self) -> Jira:
         if not self.client:
             self.log.debug("Creating Jira client for conn_id: %s", self.jira_conn_id)
 
-            get_server_info = True
-            validate = True
-            extra_options = {}
+            verify = True
             if not self.jira_conn_id:
                 raise AirflowException("Failed to create jira client. no jira_conn_id provided")
 
             conn = self.get_connection(self.jira_conn_id)
             if conn.extra is not None:
                 extra_options = conn.extra_dejson
                 # only required attributes are taken for now,
-                # more can be added ex: async, logging, max_retries
+                # more can be added ex: timeout, cloud, session
 
                 # verify
                 if "verify" in extra_options and extra_options["verify"].lower() == "false":
-                    extra_options["verify"] = False
-
-                # validate
-                if "validate" in extra_options and extra_options["validate"].lower() == "false":
-                    validate = False
-
-                if "get_server_info" in extra_options and extra_options["get_server_info"].lower() == "false":
-                    get_server_info = False
-
-            try:
-                self.client = JIRA(
-                    conn.host,
-                    options=extra_options,
-                    basic_auth=(conn.login, conn.password),
-                    get_server_info=get_server_info,
-                    validate=validate,
-                    proxies=self.proxies,
-                )
-            except JIRAError as jira_error:
-                raise AirflowException(f"Failed to create jira client, jira error: {str(jira_error)}")
-            except Exception as e:
-                raise AirflowException(f"Failed to create jira client, error: {str(e)}")
+                    verify = False
+
+            self.client = Jira(
+                url=conn.host,
+                username=conn.login,
+                password=conn.password,
+                verify_ssl=verify,
+                proxies=self.proxies,
+            )
 
         return self.client
diff --git a/airflow/providers/atlassian/jira/operators/jira.py b/airflow/providers/atlassian/jira/operators/jira.py
@@ -19,9 +19,8 @@
 
 from typing import TYPE_CHECKING, Any, Callable, Sequence
 
-from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator
-from airflow.providers.atlassian.jira.hooks.jira import JIRAError, JiraHook
+from airflow.providers.atlassian.jira.hooks.jira import JiraHook
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -30,10 +29,10 @@
 class JiraOperator(BaseOperator):
     """
     JiraOperator to interact and perform action on Jira issue tracking system.
-    This operator is designed to use Jira Python SDK: http://jira.readthedocs.io
+    This operator is designed to use Atlassian Jira SDK: https://atlassian-python-api.readthedocs.io/jira.html
 
     :param jira_conn_id: reference to a pre-defined Jira Connection
-    :param jira_method: method name from Jira Python SDK to be called
+    :param jira_method: method name from Atlassian Jira Python SDK to be called
     :param jira_method_args: required method parameters for the jira_method. (templated)
     :param result_processor: function to further process the response from Jira
     :param get_jira_resource_method: function or operator to get jira resource
@@ -60,32 +59,27 @@ def __init__(
         self.get_jira_resource_method = get_jira_resource_method
 
     def execute(self, context: Context) -> Any:
-        try:
-            if self.get_jira_resource_method is not None:
-                # if get_jira_resource_method is provided, jira_method will be executed on
-                # resource returned by executing the get_jira_resource_method.
-                # This makes all the provided methods of JIRA sdk accessible and usable
-                # directly at the JiraOperator without additional wrappers.
-                # ref: http://jira.readthedocs.io/en/latest/api.html
-                if isinstance(self.get_jira_resource_method, JiraOperator):
-                    resource = self.get_jira_resource_method.execute(**context)
-                else:
-                    resource = self.get_jira_resource_method(**context)
+        if self.get_jira_resource_method is not None:
+            # if get_jira_resource_method is provided, jira_method will be executed on
+            # resource returned by executing the get_jira_resource_method.
+            # This makes all the provided methods of atlassian-python-api JIRA sdk accessible and usable
+            # directly at the JiraOperator without additional wrappers.
+            # ref: https://atlassian-python-api.readthedocs.io/jira.html
+            if isinstance(self.get_jira_resource_method, JiraOperator):
+                resource = self.get_jira_resource_method.execute(**context)
             else:
-                # Default method execution is on the top level jira client resource
-                hook = JiraHook(jira_conn_id=self.jira_conn_id)
-                resource = hook.client
+                resource = self.get_jira_resource_method(**context)
+        else:
+            # Default method execution is on the top level jira client resource
+            hook = JiraHook(jira_conn_id=self.jira_conn_id)
+            resource = hook.client
 
-            # Current Jira-Python SDK (1.0.7) has issue with pickling the jira response.
-            # ex: self.xcom_push(context, key='operator_response', value=jira_response)
-            # This could potentially throw error if jira_result is not picklable
-            jira_result = getattr(resource, self.method_name)(**self.jira_method_args)
-            if self.result_processor:
-                return self.result_processor(context, jira_result)
+        jira_result = getattr(resource, self.method_name)(**self.jira_method_args)
 
-            return jira_result
+        output = jira_result.get("id", None) if jira_result is not None else None
+        self.xcom_push(context, key="id", value=output)
 
-        except JIRAError as jira_error:
-            raise AirflowException(f"Failed to execute jiraOperator, error: {str(jira_error)}")
-        except Exception as e:
-            raise AirflowException(f"Jira operator error: {str(e)}")
+        if self.result_processor:
+            return self.result_processor(context, jira_result)
+
+        return jira_result
diff --git a/airflow/providers/atlassian/jira/provider.yaml b/airflow/providers/atlassian/jira/provider.yaml
@@ -22,12 +22,13 @@ description: |
     `Atlassian Jira <https://www.atlassian.com/>`__
 
 versions:
+  - 2.0.0
   - 1.1.0
   - 1.0.0
 
 dependencies:
   - apache-airflow>=2.3.0
-  - JIRA>1.0.7
+  - atlassian-python-api>=1.14.2
 
 integrations:
   - integration-name: Atlassian Jira

diff --git a/airflow/providers/atlassian/jira/sensors/jira.py b/airflow/providers/atlassian/jira/sensors/jira.py
@@ -19,10 +19,7 @@
 
 from typing import TYPE_CHECKING, Any, Callable, Sequence
 
-from jira.resources import Issue, Resource
-
 from airflow.providers.atlassian.jira.hooks.jira import JiraHook
-from airflow.providers.atlassian.jira.operators.jira import JIRAError
 from airflow.sensors.base import BaseSensorOperator
 
 if TYPE_CHECKING:
@@ -34,7 +31,7 @@ class JiraSensor(BaseSensorOperator):
     Monitors a jira ticket for any change.
 
     :param jira_conn_id: reference to a pre-defined Jira Connection
-    :param method_name: method name from jira-python-sdk to be execute
+    :param method_name: method name from atlassian-python-api JIRA sdk to execute
     :param method_params: parameters for the method method_name
     :param result_processor: function that return boolean and act as a sensor response
     """
@@ -96,40 +93,37 @@ def __init__(
         if field_checker_func is None:
             field_checker_func = self.issue_field_checker
 
-        super().__init__(jira_conn_id=jira_conn_id, result_processor=field_checker_func, **kwargs)
+        super().__init__(
+            jira_conn_id=jira_conn_id, method_name="issue", result_processor=field_checker_func, **kwargs
+        )
 
     def poke(self, context: Context) -> Any:
         self.log.info("Jira Sensor checking for change in ticket: %s", self.ticket_id)
 
         self.method_name = "issue"
-        self.method_params = {"id": self.ticket_id, "fields": self.field}
+        self.method_params = {"key": self.ticket_id, "fields": self.field}
         return JiraSensor.poke(self, context=context)
 
-    def issue_field_checker(self, issue: Issue) -> bool | None:
+    def issue_field_checker(self, jira_result: dict) -> bool | None:
         """Check issue using different conditions to prepare to evaluate sensor."""
         result = None
-        try:
-            if issue is not None and self.field is not None and self.expected_value is not None:
-
-                field_val = getattr(issue.fields, self.field)
-                if field_val is not None:
-                    if isinstance(field_val, list):
-                        result = self.expected_value in field_val
-                    elif isinstance(field_val, str):
-                        result = self.expected_value.lower() == field_val.lower()
-                    elif isinstance(field_val, Resource) and getattr(field_val, "name"):
-                        result = self.expected_value.lower() == field_val.name.lower()
-                    else:
-                        self.log.warning(
-                            "Not implemented checker for issue field %s which "
-                            "is neither string nor list nor Jira Resource",
-                            self.field,
-                        )
-
-        except JIRAError as jira_error:
-            self.log.error("Jira error while checking with expected value: %s", jira_error)
-        except Exception:
-            self.log.exception("Error while checking with expected value %s:", self.expected_value)
+        if jira_result is not None and self.field is not None and self.expected_value is not None:
+
+            field_val = jira_result.get("fields", {}).get(self.field, None)
+            if field_val is not None:
+                if isinstance(field_val, list):
+                    result = self.expected_value in field_val
+                elif isinstance(field_val, str):
+                    result = self.expected_value.lower() == field_val.lower()
+                elif isinstance(field_val, dict) and field_val.get("name", None):
+                    result = self.expected_value.lower() == field_val.get("name", "").lower()
+                else:
+                    self.log.warning(
+                        "Not implemented checker for issue field %s which "
+                        "is neither string nor list nor Jira Resource",
+                        self.field,
+                    )
+
         if result is True:
             self.log.info(
                 "Issue field %s has expected value %s, returning success", self.field, self.expected_value

diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json
@@ -168,8 +168,8 @@
   },
   "atlassian.jira": {
     "deps": [
-      "JIRA>1.0.7",
-      "apache-airflow>=2.3.0"
+      "apache-airflow>=2.3.0",
+      "atlassian-python-api>=1.14.2"
     ],
     "cross-providers-deps": []
   },

diff --git a/tests/providers/atlassian/jira/hooks/test_jira.py b/tests/providers/atlassian/jira/hooks/test_jira.py
@@ -38,7 +38,7 @@ def setup_method(self):
             )
         )
 
-    @patch("airflow.providers.atlassian.jira.hooks.jira.JIRA", autospec=True, return_value=jira_client_mock)
+    @patch("airflow.providers.atlassian.jira.hooks.jira.Jira", autospec=True, return_value=jira_client_mock)
     def test_jira_client_connection(self, jira_mock):
         jira_hook = JiraHook()
 

diff --git a/tests/providers/atlassian/jira/operators/test_jira.py b/tests/providers/atlassian/jira/operators/test_jira.py
@@ -53,35 +53,35 @@ def setup_method(self):
             )
         )
 
-    @patch("airflow.providers.atlassian.jira.hooks.jira.JIRA", autospec=True, return_value=jira_client_mock)
+    @patch("airflow.providers.atlassian.jira.hooks.jira.Jira", autospec=True, return_value=jira_client_mock)
     def test_issue_search(self, jira_mock):
         jql_str = "issuekey=TEST-1226"
-        jira_mock.return_value.search_issues.return_value = minimal_test_ticket
+        jira_mock.return_value.jql_get_list_of_tickets.return_value = minimal_test_ticket
 
         jira_ticket_search_operator = JiraOperator(
             task_id="search-ticket-test",
-            jira_method="search_issues",
-            jira_method_args={"jql_str": jql_str, "maxResults": "1"},
+            jira_method="jql_get_list_of_tickets",
+            jira_method_args={"jql": jql_str, "limit": "1"},
             dag=self.dag,
         )
 
         jira_ticket_search_operator.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
 
         assert jira_mock.called
-        assert jira_mock.return_value.search_issues.called
+        assert jira_mock.return_value.jql_get_list_of_tickets.called
 
-    @patch("airflow.providers.atlassian.jira.hooks.jira.JIRA", autospec=True, return_value=jira_client_mock)
+    @patch("airflow.providers.atlassian.jira.hooks.jira.Jira", autospec=True, return_value=jira_client_mock)
     def test_update_issue(self, jira_mock):
-        jira_mock.return_value.add_comment.return_value = True
+        jira_mock.return_value.issue_add_comment.return_value = minimal_test_ticket
 
         add_comment_operator = JiraOperator(
             task_id="add_comment_test",
-            jira_method="add_comment",
-            jira_method_args={"issue": minimal_test_ticket.get("key"), "body": "this is test comment"},
+            jira_method="issue_add_comment",
+            jira_method_args={"issue_key": minimal_test_ticket.get("key"), "comment": "this is test comment"},
             dag=self.dag,
         )
 
         add_comment_operator.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
 
         assert jira_mock.called
-        assert jira_mock.return_value.add_comment.called
+        assert jira_mock.return_value.issue_add_comment.called
diff --git a/tests/providers/atlassian/jira/sensors/test_jira.py b/tests/providers/atlassian/jira/sensors/test_jira.py
@@ -27,22 +27,15 @@
 DEFAULT_DATE = timezone.datetime(2017, 1, 1)
 jira_client_mock = Mock(name="jira_client_for_test")
 
-
-class _MockJiraTicket(dict):
-    class _TicketFields:
-        labels = ["test-label-1", "test-label-2"]
-        description = "this is a test description"
-
-    fields = _TicketFields
-
-
-minimal_test_ticket = _MockJiraTicket(
-    {
-        "id": "911539",
-        "self": "https://sandbox.localhost/jira/rest/api/2/issue/911539",
-        "key": "TEST-1226",
-    }
-)
+minimal_test_ticket = {
+    "id": "911539",
+    "self": "https://sandbox.localhost/jira/rest/api/2/issue/911539",
+    "key": "TEST-1226",
+    "fields": {
+        "labels": ["test-label-1", "test-label-2"],
+        "description": "this is a test description",
+    },
+}
 
 
 class TestJiraSensor:
@@ -60,12 +53,11 @@ def setup_method(self):
             )
         )
 
-    @patch("airflow.providers.atlassian.jira.hooks.jira.JIRA", autospec=True, return_value=jira_client_mock)
+    @patch("airflow.providers.atlassian.jira.hooks.jira.Jira", autospec=True, return_value=jira_client_mock)
     def test_issue_label_set(self, jira_mock):
         jira_mock.return_value.issue.return_value = minimal_test_ticket
 
         ticket_label_sensor = JiraTicketSensor(
-            method_name="issue",
             task_id="search-ticket-test",
             ticket_id="TEST-1226",
             field="labels",