Coverage for mindsdb / integrations / handlers / frappe_handler / frappe_tables.py: 0%

1import pandas as pd 

2from typing import List 

3 

4from mindsdb.integrations.libs.api_handler import APITable 

5from mindsdb.integrations.utilities.sql_utils import extract_comparison_conditions 

6from mindsdb_sql_parser import ast 

7 

8 

9class FrappeDocumentsTable(APITable): 

10 

11 def select(self, query: ast.Select) -> pd.DataFrame: 

12 """Selects data from the Frappe API and returns it as a pandas DataFrame. 

13 

14 Returns dataframe representing the Frappe API results. 

15 

16 Args: 

17 query (ast.Select): Given SQL SELECT query 

18 """ 

19 conditions = extract_comparison_conditions(query.where) 

20 

21 params = {} 

22 filters = [] 

23 for op, arg1, arg2 in conditions: 

24 if arg1 == 'doctype': 

25 if op != '=': 

26 raise NotImplementedError 

27 params['doctype'] = arg2 

28 elif arg1 == 'name': 

29 params['name'] = arg2 

30 else: 

31 filters.append([arg1, op, arg2]) 

32 

33 if 'doctype' not in params: 

34 raise ValueError('"doctype" parameter required') 

35 

36 if query.limit: 

37 params['limit'] = query.limit.value 

38 if filters: 

39 params['filters'] = filters 

40 

41 if 'name' in params: 

42 document_data = self.handler.call_frappe_api( 

43 method_name='get_document', 

44 params=params 

45 ) 

46 else: 

47 document_data = self.handler.call_frappe_api( 

48 method_name='get_documents', 

49 params=params 

50 ) 

51 

52 # Only return the columns we need to. 

53 columns = [] 

54 for target in query.targets: 

55 if isinstance(target, ast.Star): 

56 columns = document_data.columns 

57 break 

58 elif isinstance(target, ast.Identifier): 

59 columns.append(target.parts[-1]) 

60 else: 

61 raise NotImplementedError 

62 

63 if len(document_data) == 0: 

64 return pd.DataFrame([], columns=columns) 

65 

66 # Remove columns not part of select. 

67 for col in set(document_data.columns).difference(set(columns)): 

68 document_data = document_data.drop(col, axis=1) 

69 

70 return document_data 

71 

72 def insert(self, query: ast.Insert) -> pd.DataFrame: 

73 columns = [col.name for col in query.columns] 

74 

75 for row in query.values: 

76 params = dict(zip(columns, row)) 

77 

78 self.handler.call_frappe_api('create_document', params) 

79 

80 def get_columns(self) -> List: 

81 """Gets all columns to be returned in pandas DataFrame responses""" 

82 return ['doctype', 'data']