From 67a9fff581e5202a296cf3acff54c3c4c129d877 Mon Sep 17 00:00:00 2001 From: "Thorsten Ph." Date: Tue, 19 Nov 2024 16:13:40 +0100 Subject: [PATCH] feat: Extend mongodb component with mTLS (#4591) * 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 * adjustment for MyPy * [autofix.ci] apply automated fixes --------- Co-authored-by: Thorsten Philipp Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: Gabriel Luiz Freitas Almeida --- .../components/vectorstores/mongodb_atlas.py | 46 ++++++++++++++++++- 1 file changed, 44 insertions(+), 2 deletions(-) diff --git a/src/backend/base/langflow/components/vectorstores/mongodb_atlas.py b/src/backend/base/langflow/components/vectorstores/mongodb_atlas.py index 298a0a0cd..fde29063e 100644 --- a/src/backend/base/langflow/components/vectorstores/mongodb_atlas.py +++ b/src/backend/base/langflow/components/vectorstores/mongodb_atlas.py @@ -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