From 5b5a5dfe71e3714593050140731839165d156ffc Mon Sep 17 00:00:00 2001
From: Jesse S <jschmidt@aerospike.com>
Date: Tue, 13 Aug 2024 09:27:15 -0700
Subject: [PATCH] feat: Support tls app -> avs connection (#54)

---
 .../quote-search/avs_client.py                |  6 ++++++
 quote-semantic-search/quote-search/config.py  | 20 ++++++++++++++++++-
 .../quote-search/requirements.txt             |  2 +-
 3 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/quote-semantic-search/quote-search/avs_client.py b/quote-semantic-search/quote-search/avs_client.py
index 6323a84..22bfbba 100644
--- a/quote-semantic-search/quote-search/avs_client.py
+++ b/quote-semantic-search/quote-search/avs_client.py
@@ -7,6 +7,9 @@
         host=Config.AVS_HOST,
         port=Config.AVS_PORT,
     ),
+    root_certificate=Config.AVS_TLS_CA,
+    certificate_chain=Config.AVS_TLS_CERT,
+    private_key=Config.AVS_TLS_KEY,
     listener_name=Config.AVS_ADVERTISED_LISTENER,
     is_loadbalancer=Config.AVS_IS_LOADBALANCER,
 )
@@ -17,6 +20,9 @@
         host=Config.AVS_HOST,
         port=Config.AVS_PORT,
     ),
+    root_certificate=Config.AVS_TLS_CA,
+    certificate_chain=Config.AVS_TLS_CERT,
+    private_key=Config.AVS_TLS_KEY,
     listener_name=Config.AVS_ADVERTISED_LISTENER,
     is_loadbalancer=Config.AVS_IS_LOADBALANCER,
 )
diff --git a/quote-semantic-search/quote-search/config.py b/quote-semantic-search/quote-search/config.py
index a50fb25..63050fa 100644
--- a/quote-semantic-search/quote-search/config.py
+++ b/quote-semantic-search/quote-search/config.py
@@ -5,7 +5,7 @@ def get_bool_env(name, default):
     env = os.environ.get(name)
     if env is None:
         return default
-    env = env.lower()    
+    env = env.lower()
 
     if env in ["true", "1"]:
         return True
@@ -20,6 +20,12 @@ class Config(object):
     AVS_HOST = os.environ.get("AVS_HOST") or "localhost"
     AVS_PORT = int(os.environ.get("AVS_PORT") or 5000)
     AVS_ADVERTISED_LISTENER = os.environ.get("AVS_ADVERTISED_LISTENER") or None
+    AVS_TLS_CA_FILE = os.environ.get("AVS_TLS_CA_FILE") or None
+    AVS_TLS_CA = None
+    AVS_TLS_CERT_FILE = os.environ.get("AVS_TLS_CERT_FILE") or None
+    AVS_TLS_CERT = None
+    AVS_TLS_KEY_FILE = os.environ.get("AVS_TLS_KEY_FILE") or None
+    AVS_TLS_KEY = None
     AVS_INDEX_NAME = os.environ.get("AVS_INDEX_NAME") or "quote-semantic-search"
     AVS_NAMESPACE = os.environ.get("AVS_NAMESPACE") or "test"
     AVS_SET = os.environ.get("AVS_SET") or "quote-data"
@@ -37,3 +43,15 @@ class Config(object):
 
     if NUM_QUOTES > 100000:
         NUM_QUOTES = 100000
+
+    if AVS_TLS_CA_FILE:
+        with open(AVS_TLS_CA_FILE, "rb") as f:
+            AVS_TLS_CA = f.read()
+
+    if AVS_TLS_CERT_FILE:
+        with open(AVS_TLS_CERT_FILE, "rb") as f:
+            AVS_TLS_CERT = f.read()
+
+    if AVS_TLS_KEY_FILE:
+        with open(AVS_TLS_KEY_FILE, "rb") as f:
+            AVS_TLS_KEY = f.read()
diff --git a/quote-semantic-search/quote-search/requirements.txt b/quote-semantic-search/quote-search/requirements.txt
index 9f21f22..dd7e06b 100644
--- a/quote-semantic-search/quote-search/requirements.txt
+++ b/quote-semantic-search/quote-search/requirements.txt
@@ -1,6 +1,6 @@
 # TODO: Include exact versions
 # Aerospike dependencies
-aerospike-vector-search==1.0.0
+aerospike-vector-search==1.0.1
 
 
 # Flask framework