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

1import os 

2from pyzotero import zotero 

3from mindsdb.utilities import log 

4from mindsdb.utilities.config import Config 

5from mindsdb.integrations.libs.api_handler import APIHandler, FuncParser 

6 

7from mindsdb.integrations.handlers.zotero_handler.zotero_tables import AnnotationsTable 

8 

9from mindsdb.integrations.libs.response import ( 

10 HandlerStatusResponse as StatusResponse, 

11 HandlerResponse as Response, 

12 RESPONSE_TYPE 

13) 

14 

15logger = log.getLogger(__name__) 

16 

17 

18class ZoteroHandler(APIHandler): 

19 """Handles communication with the Zotero API.""" 

20 

21 def __init__(self, name=None, **kwargs): 

22 """Initialize the Zotero handler. 

23 

24 Parameters 

25 ---------- 

26 name : str 

27 Name of the handler instance. 

28 

29 Other Parameters 

30 ---------------- 

31 connection_data : dict 

32 Dictionary containing connection data such as 'library_id', 'library_type', and 'api_key'. 

33 If not provided, will attempt to fetch from environment variables or configuration file. 

34 """ 

35 super().__init__(name) 

36 self.connection_args = self._get_connection_args(kwargs.get('connection_data', {})) 

37 self.is_connected = False 

38 self.api = None 

39 self._register_table('annotations', AnnotationsTable(self)) 

40 

41 def _get_connection_args(self, args): 

42 """Fetch connection arguments from parameters, environment variables, or configuration. 

43 

44 Parameters 

45 ---------- 

46 args 

47 Dictionary containing connection data. 

48 

49 Returns 

50 ------- 

51 connection_args 

52 Connection data list 

53 """ 

54 handler_config = Config().get('zotero_handler', {}) 

55 connection_args = {} 

56 for k in ['library_id', 'library_type', 'api_key']: 

57 connection_args[k] = args.get(k) or os.getenv(f'ZOTERO_{k.upper()}') or handler_config.get(k) 

58 return connection_args 

59 

60 def connect(self) -> StatusResponse: 

61 """Connect to the Zotero API. 

62 

63 Returns 

64 ------- 

65 StatusResponse 

66 Status of the connection attempt. 

67 """ 

68 if not self.is_connected: 

69 self.api = zotero.Zotero( 

70 self.connection_args['library_id'], 

71 self.connection_args['library_type'], 

72 self.connection_args['api_key'] 

73 ) 

74 self.is_connected = True 

75 return StatusResponse(True) 

76 

77 def check_connection(self) -> StatusResponse: 

78 """Check the connection status to the Zotero API. 

79 

80 Returns 

81 ------- 

82 StatusResponse 

83 Status of the connection. 

84 """ 

85 try: 

86 self.connect() 

87 return StatusResponse(True) 

88 except Exception as e: 

89 error_message = f'Error connecting to Zotero API: {str(e)}. Check credentials.' 

90 logger.error(error_message) 

91 self.is_connected = False 

92 return StatusResponse(False, error_message=error_message) 

93 

94 def native_query(self, query_string: str = None): 

95 """Execute a native query against the Zotero API. 

96 

97 Parameters 

98 ---------- 

99 query_string : str 

100 The query string to execute, formatted as required by the Zotero API. 

101 

102 Returns 

103 ------- 

104 Response 

105 Response object containing the result of the query. 

106 """ 

107 method_name, params = FuncParser().from_string(query_string) 

108 df = self._call_find_annotations_zotero_api(method_name, params) 

109 return Response(RESPONSE_TYPE.TABLE, data_frame=df)