From 1aecdcedbdc04e93426b44f87b291fdbc6d33560 Mon Sep 17 00:00:00 2001
From: Shannon Carey <rehevkor5@gmail.com>
Date: Thu, 14 Nov 2024 10:03:04 -0600
Subject: [PATCH] feat[confluence]: Permit passing params to Confluence client

- This enables use of parameters such as backoff_and_retry which is
important for the client to react well to rate limit (429) responses,
as described in
https://developer.atlassian.com/cloud/confluence/rate-limiting/
For available parameters, refer also to:
https://github.com/atlassian-api/atlassian-python-api/blob/3.41.16/atlassian/rest_client.py#L48
---
 .../llama-index-readers-confluence/README.md  |  6 +++++-
 .../llama_index/readers/confluence/base.py    | 21 +++++++++++++++----
 .../pyproject.toml                            |  2 +-
 .../tests/test_readers_confluence.py          | 18 ++++++++++++++++
 4 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/llama-index-integrations/readers/llama-index-readers-confluence/README.md b/llama-index-integrations/readers/llama-index-readers-confluence/README.md
index 11ae8dbfed39b..618780ea07e1b 100644
--- a/llama-index-integrations/readers/llama-index-readers-confluence/README.md
+++ b/llama-index-integrations/readers/llama-index-readers-confluence/README.md
@@ -68,7 +68,11 @@ base_url = "https://yoursite.atlassian.com/wiki"
 page_ids = ["<page_id_1>", "<page_id_2>", "<page_id_3"]
 space_key = "<space_key>"
 
-reader = ConfluenceReader(base_url=base_url, oauth2=oauth2_dict)
+reader = ConfluenceReader(
+    base_url=base_url,
+    oauth2=oauth2_dict,
+    client_args={"backoff_and_retry": True},
+)
 documents = reader.load_data(
     space_key=space_key, include_attachments=True, page_status="current"
 )
diff --git a/llama-index-integrations/readers/llama-index-readers-confluence/llama_index/readers/confluence/base.py b/llama-index-integrations/readers/llama-index-readers-confluence/llama_index/readers/confluence/base.py
index b6c3278b9c1a5..973da486ea703 100644
--- a/llama-index-integrations/readers/llama-index-readers-confluence/llama_index/readers/confluence/base.py
+++ b/llama-index-integrations/readers/llama-index-readers-confluence/llama_index/readers/confluence/base.py
@@ -33,6 +33,7 @@ class ConfluenceReader(BaseReader):
         api_token (str): Confluence API token, see https://confluence.atlassian.com/cloud/api-tokens-938839638.html
         user_name (str): Confluence username, used for basic auth. Must be used with `password`.
         password (str): Confluence password, used for basic auth. Must be used with `user_name`.
+        client_args (dict): Additional keyword arguments to pass directly to the Atlassian Confluence client constructor, for example `{'backoff_and_retry': True}`.
 
     """
 
@@ -44,6 +45,7 @@ def __init__(
         api_token: Optional[str] = None,
         user_name: Optional[str] = None,
         password: Optional[str] = None,
+        client_args: Optional[dict] = None,
     ) -> None:
         if base_url is None:
             raise ValueError("Must provide `base_url`")
@@ -58,20 +60,30 @@ def __init__(
                 " atlassian-python-api`"
             )
         self.confluence: Confluence = None
+        if client_args is None:
+            client_args = {}
         if oauth2:
-            self.confluence = Confluence(url=base_url, oauth2=oauth2, cloud=cloud)
+            self.confluence = Confluence(
+                url=base_url, oauth2=oauth2, cloud=cloud, **client_args
+            )
         else:
             if api_token is not None:
-                self.confluence = Confluence(url=base_url, token=api_token, cloud=cloud)
+                self.confluence = Confluence(
+                    url=base_url, token=api_token, cloud=cloud, **client_args
+                )
             elif user_name is not None and password is not None:
                 self.confluence = Confluence(
-                    url=base_url, username=user_name, password=password, cloud=cloud
+                    url=base_url,
+                    username=user_name,
+                    password=password,
+                    cloud=cloud,
+                    **client_args,
                 )
             else:
                 api_token = os.getenv(CONFLUENCE_API_TOKEN)
                 if api_token is not None:
                     self.confluence = Confluence(
-                        url=base_url, token=api_token, cloud=cloud
+                        url=base_url, token=api_token, cloud=cloud, **client_args
                     )
                 else:
                     user_name = os.getenv(CONFLUENCE_USERNAME)
@@ -82,6 +94,7 @@ def __init__(
                             username=user_name,
                             password=password,
                             cloud=cloud,
+                            **client_args,
                         )
                     else:
                         raise ValueError(
diff --git a/llama-index-integrations/readers/llama-index-readers-confluence/pyproject.toml b/llama-index-integrations/readers/llama-index-readers-confluence/pyproject.toml
index 21cb68aed91df..a85030add9e6a 100644
--- a/llama-index-integrations/readers/llama-index-readers-confluence/pyproject.toml
+++ b/llama-index-integrations/readers/llama-index-readers-confluence/pyproject.toml
@@ -28,7 +28,7 @@ license = "MIT"
 maintainers = ["zywilliamli"]
 name = "llama-index-readers-confluence"
 readme = "README.md"
-version = "0.2.1"
+version = "0.2.2"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"
diff --git a/llama-index-integrations/readers/llama-index-readers-confluence/tests/test_readers_confluence.py b/llama-index-integrations/readers/llama-index-readers-confluence/tests/test_readers_confluence.py
index 6a8d9d2b1e23b..8040a37515970 100644
--- a/llama-index-integrations/readers/llama-index-readers-confluence/tests/test_readers_confluence.py
+++ b/llama-index-integrations/readers/llama-index-readers-confluence/tests/test_readers_confluence.py
@@ -1,3 +1,5 @@
+from unittest.mock import patch
+
 import pytest
 from llama_index.readers.confluence import ConfluenceReader
 
@@ -31,6 +33,22 @@ def test_confluence_reader_with_api_token():
     assert reader.confluence is not None
 
 
+def test_confluence_reader_with_client_args():
+    with patch("atlassian.Confluence") as MockConstructor:
+        reader = ConfluenceReader(
+            base_url="https://example.atlassian.net/wiki",
+            api_token="example_api_token",
+            client_args={"backoff_and_retry": True},
+        )
+        assert reader.confluence is not None
+        MockConstructor.assert_called_once_with(
+            url="https://example.atlassian.net/wiki",
+            token="example_api_token",
+            cloud=True,
+            backoff_and_retry=True,
+        )
+
+
 def test_confluence_reader_with_basic_auth():
     reader = ConfluenceReader(
         base_url="https://example.atlassian.net/wiki",