Coverage for mindsdb / api / executor / sql_query / steps / prepare_steps.py: 35%

1from mindsdb_sql_parser.ast import ( 

2 Identifier, 

3 Constant, 

4 Select, 

5 Star, 

6) 

7from mindsdb.api.executor.planner.steps import ( 

8 GetPredictorColumns, 

9 GetTableColumns, 

10) 

11 

12from mindsdb.api.executor.sql_query.result_set import ResultSet, Column 

13from mindsdb.utilities.config import config 

14 

15from .base import BaseStepCall 

16 

17 

18class GetPredictorColumnsCall(BaseStepCall): 

19 

20 bind = GetPredictorColumns 

21 

22 def call(self, step): 

23 mindsdb_database_name = config.get('default_project') 

24 

25 predictor_name = step.predictor.parts[-1] 

26 dn = self.session.datahub.get(mindsdb_database_name) 

27 columns_names = dn.get_table_columns_names(predictor_name) 

28 

29 data = ResultSet() 

30 for column_name in columns_names: 

31 data.add_column(Column( 

32 name=column_name, 

33 table_name=predictor_name, 

34 database=mindsdb_database_name 

35 )) 

36 return data 

37 

38 

39class GetTableColumnsCall(BaseStepCall): 

40 

41 bind = GetTableColumns 

42 

43 def call(self, step): 

44 

45 table = step.table 

46 dn = self.session.datahub.get(step.namespace) 

47 ds_query = Select(from_table=Identifier(table), targets=[Star()], limit=Constant(0)) 

48 

49 response = dn.query(ds_query, session=self.session) 

50 

51 data = ResultSet() 

52 for column in response.columns: 

53 data.add_column(Column( 

54 name=column['name'], 

55 type=column.get('type'), 

56 table_name=table, 

57 database=self.context.get('database') 

58 )) 

59 return data