jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

feat: Extend mongodb component with mTLS (#4591)

Browse source 
* Extend mongodb component with mTLS

* [autofix.ci] apply automated fixes

* fix check style errors

* [autofix.ci] apply automated fixes

* bug fix

* Update src/backend/base/langflow/components/vectorstores/mongodb_atlas.py

Co-authored-by: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>

* adjustment for MyPy

* [autofix.ci] apply automated fixes

---------

Co-authored-by: Thorsten Philipp <thorsten.philipp@msg.group>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
This commit is contained in:
Thorsten Ph. 2024-11-19 16:13:40 +01:00 • committed by GitHub
commit 67a9fff581
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 44 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

46
src/backend/base/langflow/components/vectorstores/mongodb_atlas.py
View file

@ -1,8 +1,11 @@
import tempfile
import certifi
from langchain_community.vectorstores import MongoDBAtlasVectorSearch
from langflow.base.vectorstores.model import LCVectorStoreComponent, check_cached_vector_store
from langflow.helpers.data import docs_to_data
from langflow.io import DataInput, HandleInput, IntInput, MultilineInput, SecretStrInput, StrInput
from langflow.io import BoolInput, DataInput, HandleInput, IntInput, MultilineInput, SecretStrInput, StrInput
from langflow.schema import Data
@ -15,6 +18,15 @@ class MongoVectorStoreComponent(LCVectorStoreComponent):
inputs = [
SecretStrInput(name="mongodb_atlas_cluster_uri", display_name="MongoDB Atlas Cluster URI", required=True),
BoolInput(name="enable_mtls", display_name="Enable mTLS", value=False, advanced=True, required=True),
SecretStrInput(
name="mongodb_atlas_client_cert",
display_name="MongoDB Atlas Combined Client Certificate",
required=False,
info="Client Certificate combined with the private key in the following format:\n "
"-----BEGIN PRIVATE KEY-----\n...\n -----END PRIVATE KEY-----\n-----BEGIN CERTIFICATE-----\n"
"...\n-----END CERTIFICATE-----\n",
),
StrInput(name="db_name", display_name="Database Name", required=True),
StrInput(name="collection_name", display_name="Collection Name", required=True),
StrInput(name="index_name", display_name="Index Name", required=True),
@ -42,9 +54,39 @@ class MongoVectorStoreComponent(LCVectorStoreComponent):
msg = "Please install pymongo to use MongoDB Atlas Vector Store"
raise ImportError(msg) from e
# Create temporary files for the client certificate
if self.enable_mtls:
client_cert_path = None
try:
client_cert = self.mongodb_atlas_client_cert.replace(" ", "\n")
client_cert = client_cert.replace("-----BEGIN\nPRIVATE\nKEY-----", "-----BEGIN PRIVATE KEY-----")
client_cert = client_cert.replace(
"-----END\nPRIVATE\nKEY-----\n-----BEGIN\nCERTIFICATE-----",
"-----END PRIVATE KEY-----\n-----BEGIN CERTIFICATE-----",
)
client_cert = client_cert.replace("-----END\nCERTIFICATE-----", "-----END CERTIFICATE-----")
with tempfile.NamedTemporaryFile(delete=False) as client_cert_file:
client_cert_file.write(client_cert.encode("utf-8"))
client_cert_path = client_cert_file.name
except Exception as e:
msg = f"Failed to write certificate to temporary file: {e}"
raise ValueError(msg) from e
try:
mongo_client: MongoClient = MongoClient(self.mongodb_atlas_cluster_uri)
mongo_client: MongoClient = (
MongoClient(
self.mongodb_atlas_cluster_uri,
tls=True,
tlsCertificateKeyFile=client_cert_path,
tlsCAFile=certifi.where(),
)
if self.enable_mtls
else MongoClient(self.mongodb_atlas_cluster_uri)
)
collection = mongo_client[self.db_name][self.collection_name]
collection.drop() # Drop collection to override the vector store
except Exception as e:
msg = f"Failed to connect to MongoDB Atlas: {e}"
raise ValueError(msg) from e