From 3ea2daefc6765d414d052e242d13e87f16ae272a Mon Sep 17 00:00:00 2001
From: Dustin Black <dblack@redhat.com>
Date: Wed, 13 Dec 2023 15:29:38 +0100
Subject: [PATCH] add ability to disable tls verification

---
 .../opensearch_plugin.py                      |  9 ++++-----
 .../opensearch_schema.py                      | 15 +++++++-------
 .../test_arcaflow_plugin_opensearch.py        | 20 +++++--------------
 3 files changed, 17 insertions(+), 27 deletions(-)

diff --git a/arcaflow_plugin_opensearch/opensearch_plugin.py b/arcaflow_plugin_opensearch/opensearch_plugin.py
index ffe40dc..f8b07fd 100644
--- a/arcaflow_plugin_opensearch/opensearch_plugin.py
+++ b/arcaflow_plugin_opensearch/opensearch_plugin.py
@@ -18,11 +18,12 @@
 def store(
     params: StoreDocumentRequest,
 ) -> typing.Tuple[str, typing.Union[SuccessOutput, ErrorOutput]]:
-
     try:
         if params.username:
             opensearch = OpenSearch(
-                hosts=params.url, basic_auth=[params.username, params.password]
+                hosts=params.url,
+                basic_auth=[params.username, params.password],
+                verify_certs=params.tls_verify,
             )
         # Support for servers that don't require authentication
         else:
@@ -35,9 +36,7 @@ def store(
             f"Successfully uploaded document for index {params.index}"
         )
     except Exception as ex:
-        return "error", ErrorOutput(
-            f"Failed to create OpenSearch document: {ex}"
-        )
+        return "error", ErrorOutput(f"Failed to create OpenSearch document: {ex}")
 
 
 if __name__ == "__main__":
diff --git a/arcaflow_plugin_opensearch/opensearch_schema.py b/arcaflow_plugin_opensearch/opensearch_schema.py
index d42d3ff..02e8cf9 100644
--- a/arcaflow_plugin_opensearch/opensearch_schema.py
+++ b/arcaflow_plugin_opensearch/opensearch_schema.py
@@ -5,26 +5,22 @@
 
 @dataclass
 class StoreDocumentRequest:
-
     url: typing.Annotated[
         str,
         schema.name("url"),
         schema.description("The URL for the Opensearch-compatible instance."),
     ]
-
     index: typing.Annotated[
         str,
         validation.min(1),
         schema.name("index"),
         schema.description("Name of the index that will receive the data."),
     ]
-
     data: typing.Annotated[
         typing.Dict[str, typing.Any],
         schema.name("data"),
         schema.description("Data to upload to your index"),
     ]
-
     username: typing.Annotated[
         typing.Optional[str],
         validation.min(1),
@@ -34,21 +30,26 @@ class StoreDocumentRequest:
             "Opensearch-compatible instance."
         ),
     ] = None
-
     password: typing.Annotated[
         typing.Optional[str],
         schema.name("password"),
         schema.description("The password for the given user."),
     ] = None
+    tls_verify: typing.Annotated[
+        bool,
+        schema.name("TLS verify"),
+        schema.description(
+            "For development and testing purposes, this can be set to False to disable"
+            " TLS verification for connections to Opensearch-compatible services."
+        ),
+    ] = True
 
 
 @dataclass
 class SuccessOutput:
-
     message: str
 
 
 @dataclass
 class ErrorOutput:
-
     error: str
diff --git a/tests/integration/test_arcaflow_plugin_opensearch.py b/tests/integration/test_arcaflow_plugin_opensearch.py
index e5f52c8..e8338ed 100644
--- a/tests/integration/test_arcaflow_plugin_opensearch.py
+++ b/tests/integration/test_arcaflow_plugin_opensearch.py
@@ -27,9 +27,7 @@ def test_empty_data(self) -> None:
             argv=[
                 "",
                 "-f",
-                StoreIntegrationTest.build_fixture_file_path(
-                    "empty_data.yaml"
-                ),
+                StoreIntegrationTest.build_fixture_file_path("empty_data.yaml"),
             ],
         )
 
@@ -44,9 +42,7 @@ def test_simple_data(self) -> None:
             argv=[
                 "",
                 "-f",
-                StoreIntegrationTest.build_fixture_file_path(
-                    "simple_data.yaml"
-                ),
+                StoreIntegrationTest.build_fixture_file_path("simple_data.yaml"),
             ],
         )
 
@@ -64,9 +60,7 @@ def test_nested_data(self) -> None:
             argv=[
                 "",
                 "-f",
-                StoreIntegrationTest.build_fixture_file_path(
-                    "nested_data.yaml"
-                ),
+                StoreIntegrationTest.build_fixture_file_path("nested_data.yaml"),
             ],
         )
 
@@ -94,9 +88,7 @@ def assertStoredData(self, expectedData: dict, index: str):
             if len(actualData["hits"]["hits"]) == 0:
                 time.sleep(i + 1)
                 continue
-            self.assertDictEqual(
-                expectedData, actualData["hits"]["hits"][0]["_source"]
-            )
+            self.assertDictEqual(expectedData, actualData["hits"]["hits"][0]["_source"])
             return
         self.fail(f"No documents found for index {index}")
 
@@ -113,9 +105,7 @@ def get_opensearch_data(sample: str) -> dict:
             "OPENSEARCH_PASSWORD",
         )
         elastiUrl = f"{url}/{sample}/_search"
-        with requests.get(
-            elastiUrl, auth=HTTPBasicAuth(user, password)
-        ) as resp:
+        with requests.get(elastiUrl, auth=HTTPBasicAuth(user, password)) as resp:
             return json.loads(resp.text)