Coverage for mindsdb / integrations / handlers / documentdb_handler / documentdb_handler.py: 0%

1from pymongo import MongoClient 

2 

3from mindsdb.integrations.handlers.mongodb_handler import Handler as MongoDBHandler 

4 

5 

6class DocumentDBHandler(MongoDBHandler): 

7 """ 

8 This handler handles connection and execution of the DocumentDB statements. 

9 """ 

10 

11 name = 'documentdb' 

12 

13 def __init__(self, name, **kwargs): 

14 super().__init__(name, **kwargs) 

15 connection_data = kwargs.get('connection_data', {}) 

16 self.host = connection_data.get("host") 

17 self.port = int(connection_data.get("port") or 27017) 

18 self.user = connection_data.get("username") 

19 self.password = connection_data.get("password") 

20 self.database = connection_data.get('database') 

21 self.flatten_level = connection_data.get('flatten_level', 0) 

22 self.mykwargs = connection_data.get('kwargs', {}) 

23 

24 self.connection = None 

25 self.is_connected = False 

26 

27 def connect(self): 

28 kwargs = {} 

29 if isinstance(self.user, str) and len(self.user) > 0: 

30 kwargs['username'] = self.user 

31 

32 if isinstance(self.password, str) and len(self.password) > 0: 

33 kwargs['password'] = self.password 

34 

35 connection = MongoClient( 

36 host=self.host, port=self.port, **{**kwargs, **self.mykwargs} 

37 ) 

38 

39 self.is_connected = True 

40 self.connection = connection 

41 return self.connection