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

1from typing import Optional 

2 

3from mindsdb_sql_parser.ast.base import ASTNode 

4 

5from mindsdb.integrations.libs.base import DatabaseHandler 

6from mindsdb.integrations.libs.response import ( 

7 HandlerStatusResponse as StatusResponse 

8) 

9 

10from mindsdb.integrations.handlers.mysql_handler.mysql_handler import MySQLHandler 

11from mindsdb.integrations.handlers.postgres_handler.postgres_handler import PostgresHandler 

12from mindsdb.integrations.handlers.mssql_handler.mssql_handler import SqlServerHandler 

13 

14 

15class CloudSQLHandler(DatabaseHandler): 

16 """ 

17 This handler handles connection and execution of the Google Cloud SQL statements. 

18 """ 

19 name = 'cloud_sql' 

20 

21 def __init__(self, name: str, connection_data: Optional[dict], **kwargs): 

22 """ 

23 Initialize the handler. 

24 Args: 

25 name (str): name of particular handler instance 

26 connection_data (dict): parameters for connecting to the database 

27 **kwargs: arbitrary keyword arguments. 

28 """ 

29 super().__init__(name) 

30 

31 self.dialect = 'cloud_sql' 

32 self.connection_data = connection_data 

33 self.kwargs = kwargs 

34 

35 if self.connection_data['db_engine'] == 'mysql': 

36 self.db = MySQLHandler( 

37 name=name + 'mysql', 

38 connection_data={key: self.connection_data[key] for key in self.connection_data if key != 'db_engine'} 

39 ) 

40 elif self.connection_data['db_engine'] == 'postgresql': 

41 self.db = PostgresHandler( 

42 name=name + 'postgresql', 

43 connection_data={key: self.connection_data[key] for key in self.connection_data if key != 'db_engine'} 

44 ) 

45 elif self.connection_data['db_engine'] == 'mssql': 

46 self.db = SqlServerHandler( 

47 name=name + 'mssql', 

48 connection_data={key: self.connection_data[key] for key in self.connection_data if key != 'db_engine'} 

49 ) 

50 else: 

51 raise Exception("The database engine should be either MySQL, PostgreSQL or SQL Server!") 

52 

53 def __del__(self): 

54 self.db.__del__() 

55 

56 def connect(self) -> StatusResponse: 

57 """ 

58 Set up the connection required by the handler. 

59 Returns: 

60 HandlerStatusResponse 

61 """ 

62 

63 return self.db.connect() 

64 

65 def disconnect(self): 

66 """ 

67 Close any existing connections. 

68 """ 

69 

70 return self.db.disconnect() 

71 

72 def check_connection(self) -> StatusResponse: 

73 """ 

74 Check connection to the handler. 

75 Returns: 

76 HandlerStatusResponse 

77 """ 

78 

79 return self.db.check_connection() 

80 

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

82 """ 

83 Receive raw query and act upon it somehow. 

84 Args: 

85 query (str): query in native format 

86 Returns: 

87 HandlerResponse 

88 """ 

89 

90 return self.db.native_query(query) 

91 

92 def query(self, query: ASTNode) -> StatusResponse: 

93 """ 

94 Receive query as AST (abstract syntax tree) and act upon it somehow. 

95 Args: 

96 query (ASTNode): sql query represented as AST. May be any kind 

97 of query: SELECT, INTSERT, DELETE, etc 

98 Returns: 

99 HandlerResponse 

100 """ 

101 

102 return self.db.query(query) 

103 

104 def get_tables(self) -> StatusResponse: 

105 """ 

106 Return list of entities that will be accessible as tables. 

107 Returns: 

108 HandlerResponse 

109 """ 

110 

111 return self.db.get_tables() 

112 

113 def get_columns(self, table_name: str) -> StatusResponse: 

114 """ 

115 Returns a list of entity columns. 

116 Args: 

117 table_name (str): name of one of tables returned by self.get_tables() 

118 Returns: 

119 HandlerResponse 

120 """ 

121 

122 return self.db.get_columns(table_name)