Skip to content

Commit d1fe938

Browse files
[8.19] [Network Drive] Switch to usage of local connection cache (#3868) (#3871)
# Backport This will backport the following commits from `main` to `8.19`: - [[Network Drive] Switch to usage of local connection cache (#3868)](#3868) <!--- Backport version: 9.6.6 --> ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sorenlouv/backport) --------- Co-authored-by: Elastic Machine <[email protected]>
1 parent 2941816 commit d1fe938

File tree

1 file changed

+24
-6
lines changed

1 file changed

+24
-6
lines changed

connectors/sources/network_drive.py

Lines changed: 24 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -261,13 +261,14 @@ def fetch_members(self, group_name):
261261
class SMBSession:
262262
_connection = None
263263

264-
def __init__(self, server_ip, username, password, port):
264+
def __init__(self, server_ip, username, password, port, connection_cache):
265265
self.server_ip = server_ip
266266
self.username = username
267267
self.password = password
268268
self.port = port
269269
self.session = None
270270
self._logger = logger
271+
self._connection_cache = connection_cache
271272

272273
def create_connection(self):
273274
"""Creates an SMB session to the shared drive."""
@@ -277,6 +278,7 @@ def create_connection(self):
277278
username=self.username,
278279
password=self.password,
279280
port=self.port,
281+
connection_cache=self._connection_cache,
280282
)
281283
except SMBResponseException as exception:
282284
self.handle_smb_response_errors(exception=exception)
@@ -336,13 +338,20 @@ def __init__(self, configuration):
336338
self.drive_type = self.configuration["drive_type"]
337339
self.identity_mappings = self.configuration["identity_mappings"]
338340
self.security_info = SecurityInfo(self.username, self.password, self.server_ip)
341+
self._connection_cache = {}
339342

340343
def advanced_rules_validators(self):
341344
return [NetworkDriveAdvancedRulesValidator(self)]
342345

343346
@cached_property
344347
def smb_connection(self):
345-
return SMBSession(self.server_ip, self.username, self.password, self.port)
348+
return SMBSession(
349+
self.server_ip,
350+
self.username,
351+
self.password,
352+
self.port,
353+
self._connection_cache,
354+
)
346355

347356
@classmethod
348357
def get_default_configuration(cls):
@@ -464,6 +473,7 @@ async def traverse_diretory(self, path):
464473
username=self.username,
465474
password=self.password,
466475
port=self.port,
476+
connection_cache=self._connection_cache,
467477
),
468478
)
469479
)
@@ -518,6 +528,9 @@ async def traverse_directory_for_syncrule(self, path, glob_pattern, indexed_rule
518528
smbclient.scandir,
519529
path=current_path,
520530
port=self.port,
531+
username=self.username,
532+
password=self.password,
533+
connection_cache=self._connection_cache,
521534
),
522535
)
523536
for file in directory_info:
@@ -598,9 +611,9 @@ async def validate_config(self):
598611

599612
async def ping(self):
600613
"""Verify the connection with Network Drive"""
614+
if self.smb_connection.session is None:
615+
await asyncio.to_thread(self.smb_connection.create_connection)
601616

602-
await asyncio.to_thread(self.smb_connection.create_connection)
603-
await self.close()
604617
self._logger.info("Successfully connected to the Network Drive")
605618

606619
async def close(self):
@@ -611,7 +624,10 @@ async def close(self):
611624
await loop.run_in_executor(
612625
executor=None,
613626
func=partial(
614-
smbclient.delete_session, server=self.server_ip, port=self.port
627+
smbclient.delete_session,
628+
server=self.server_ip,
629+
port=self.port,
630+
connection_cache=self._connection_cache,
615631
),
616632
)
617633

@@ -876,7 +892,9 @@ async def get_docs(self, filtering=None):
876892
Yields:
877893
dictionary: Dictionary containing the Network Drive files and folders as documents
878894
"""
879-
await asyncio.to_thread(self.smb_connection.create_connection)
895+
if self.smb_connection.session is None:
896+
await asyncio.to_thread(self.smb_connection.create_connection)
897+
880898
if filtering and filtering.has_advanced_rules():
881899
advanced_rules = filtering.get_advanced_rules()
882900
async for document in self.fetch_filtered_directory(advanced_rules):

0 commit comments

Comments
 (0)