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

1from collections import OrderedDict 

2from mindsdb_sql_parser import parse_sql 

3 

4from mindsdb.integrations.libs.const import HANDLER_CONNECTION_ARG_TYPE as ARG_TYPE 

5from mindsdb.integrations.handlers.lightdash_handler.api import Lightdash 

6from mindsdb.integrations.handlers.lightdash_handler.lightdash_tables import ( 

7 UserTable, 

8 UserAbilityTable, 

9 OrgTable, 

10 OrgProjectsTable, 

11 OrgMembersTable, 

12 ProjectTable, 

13 WarehouseConnectionTable, 

14 DBTConnectionTable, 

15 DBTEnvironmentVarsTable, 

16 ChartsTable, 

17 SpacesTable, 

18 AccessTable, 

19 ValidationTable, 

20 DashboardsTable, 

21 QueriesTable, 

22 ChartHistoryTable, 

23 ChartConfigTable, 

24 ChartAdditionalMetricsTable, 

25 ChartTableCalculationsTable, 

26 SchedulerLogsTable, 

27 SchedulerTable, 

28 SchedulerJobsTable, 

29 SchedulerJobStatus, 

30) 

31from mindsdb.integrations.libs.api_handler import APIHandler 

32from mindsdb.integrations.libs.response import HandlerStatusResponse as StatusResponse 

33 

34 

35class LightdashHandler(APIHandler): 

36 

37 def __init__(self, name: str, **kwargs) -> None: 

38 super().__init__(name) 

39 self.connection = None 

40 self.is_connected = False 

41 self.api_key = kwargs.get("connection_data", {}).get("api_key", "") 

42 self.base_url = kwargs.get("connection_data", {}).get("base_url", "") 

43 _tables = [ 

44 UserTable, 

45 UserAbilityTable, 

46 OrgTable, 

47 OrgProjectsTable, 

48 OrgMembersTable, 

49 ProjectTable, 

50 WarehouseConnectionTable, 

51 DBTConnectionTable, 

52 DBTEnvironmentVarsTable, 

53 ChartsTable, 

54 SpacesTable, 

55 AccessTable, 

56 ValidationTable, 

57 DashboardsTable, 

58 QueriesTable, 

59 ChartHistoryTable, 

60 ChartConfigTable, 

61 ChartAdditionalMetricsTable, 

62 ChartTableCalculationsTable, 

63 SchedulerLogsTable, 

64 SchedulerTable, 

65 SchedulerJobsTable, 

66 SchedulerJobStatus, 

67 ] 

68 for Table in _tables: 

69 self._register_table(Table.name, Table(self)) 

70 self.connect() 

71 

72 def connect(self) -> Lightdash: 

73 self.connection = Lightdash(self.base_url, self.api_key) 

74 return self.connection 

75 

76 def check_connection(self) -> StatusResponse: 

77 resp = StatusResponse(False) 

78 if self.connection and not self.connection.is_connected(): 

79 resp.error = "Client not connected" 

80 else: 

81 resp.success = True 

82 return resp 

83 

84 def native_query(self, query: str) -> StatusResponse: 

85 ast = parse_sql(query) 

86 return self.query(ast) 

87 

88 

89connection_args = OrderedDict( 

90 api_key={ 

91 'type': ARG_TYPE.STR, 

92 'description': 'API Token for accessing Lightdash instance', 

93 'required': True, 

94 'label': 'API Key', 

95 }, 

96 base_url={ 

97 'type': ARG_TYPE.STR, 

98 'description': 'Base URL of Lightdash instance', 

99 'required': True, 

100 'label': 'Base URL', 

101 } 

102) 

103 

104 

105connection_args_example = OrderedDict( 

106 api_key='23d6b9e0c2fab7eba2e8b7e452cead47', 

107 base_url='http://localhost:8080/' 

108)