| 1 | # -*- coding: utf-8 -*- |
| 2 | # Description: |
| 3 | # Author: Ilya Mashchenko (ilyam8) |
| 4 | # SPDX-License-Identifier: GPL-3.0-or-later |
| 5 | |
| 6 | from sys import exc_info |
| 7 | |
| 8 | try: |
| 9 | import MySQLdb |
| 10 | |
| 11 | PY_MYSQL = True |
| 12 | except ImportError: |
| 13 | try: |
| 14 | import pymysql as MySQLdb |
| 15 | |
| 16 | PY_MYSQL = True |
| 17 | except ImportError: |
| 18 | PY_MYSQL = False |
| 19 | |
| 20 | from bases.FrameworkServices.SimpleService import SimpleService |
| 21 | |
| 22 | |
| 23 | class MySQLService(SimpleService): |
| 24 | def __init__(self, configuration=None, name=None): |
| 25 | SimpleService.__init__(self, configuration=configuration, name=name) |
| 26 | self.__connection = None |
| 27 | self.__conn_properties = dict() |
| 28 | self.extra_conn_properties = dict() |
| 29 | self.__queries = self.configuration.get('queries', dict()) |
| 30 | self.queries = dict() |
| 31 | |
| 32 | def __connect(self): |
| 33 | try: |
| 34 | connection = MySQLdb.connect(connect_timeout=self.update_every, **self.__conn_properties) |
| 35 | except (MySQLdb.MySQLError, TypeError, AttributeError) as error: |
| 36 | return None, str(error) |
| 37 | else: |
| 38 | return connection, None |
| 39 | |
| 40 | def check(self): |
| 41 | def get_connection_properties(conf, extra_conf): |
| 42 | properties = dict() |
| 43 | if conf.get('user'): |
| 44 | properties['user'] = conf['user'] |
| 45 | if conf.get('pass'): |
| 46 | properties['passwd'] = conf['pass'] |
| 47 | |
| 48 | if conf.get('socket'): |
| 49 | properties['unix_socket'] = conf['socket'] |
| 50 | elif conf.get('host'): |
| 51 | properties['host'] = conf['host'] |
| 52 | properties['port'] = int(conf.get('port', 3306)) |
| 53 | elif conf.get('my.cnf'): |
| 54 | properties['read_default_file'] = conf['my.cnf'] |
| 55 | |
| 56 | if conf.get('ssl'): |
| 57 | properties['ssl'] = conf['ssl'] |
| 58 | |
| 59 | if isinstance(extra_conf, dict) and extra_conf: |
| 60 | properties.update(extra_conf) |
| 61 | |
| 62 | return properties or None |
| 63 | |
| 64 | def is_valid_queries_dict(raw_queries, log_error): |
| 65 | """ |
| 66 | :param raw_queries: dict: |
| 67 | :param log_error: function: |
| 68 | :return: dict or None |
| 69 | |
| 70 | raw_queries is valid when: type <dict> and not empty after is_valid_query(for all queries) |
| 71 | """ |
| 72 | |
| 73 | def is_valid_query(query): |
| 74 | return all([isinstance(query, str), |
| 75 | query.startswith(('SELECT', 'select', 'SHOW', 'show'))]) |
| 76 | |
| 77 | if hasattr(raw_queries, 'keys') and raw_queries: |
| 78 | valid_queries = dict([(n, q) for n, q in raw_queries.items() if is_valid_query(q)]) |
| 79 | bad_queries = set(raw_queries) - set(valid_queries) |
| 80 | |
| 81 | if bad_queries: |
| 82 | log_error('Removed query(s): {queries}'.format(queries=bad_queries)) |
| 83 | return valid_queries |
| 84 | else: |
| 85 | log_error('Unsupported "queries" format. Must be not empty <dict>') |
| 86 | return None |
| 87 | |
| 88 | if not PY_MYSQL: |
| 89 | self.error('MySQLdb or PyMySQL module is needed to use mysql.chart.py plugin') |
| 90 | return False |
| 91 | |
| 92 | # Preference: 1. "queries" from the configuration file 2. "queries" from the module |
| 93 | self.queries = self.__queries or self.queries |
| 94 | # Check if "self.queries" exist, not empty and all queries are in valid format |
| 95 | self.queries = is_valid_queries_dict(self.queries, self.error) |
| 96 | if not self.queries: |
| 97 | return None |
| 98 | |
| 99 | # Get connection properties |
| 100 | self.__conn_properties = get_connection_properties(self.configuration, self.extra_conn_properties) |
| 101 | if not self.__conn_properties: |
| 102 | self.error('Connection properties are missing') |
| 103 | return False |
| 104 | |
| 105 | # Create connection to the database |
| 106 | self.__connection, error = self.__connect() |
| 107 | if error: |
| 108 | self.error('Can\'t establish connection to MySQL: {error}'.format(error=error)) |
| 109 | return False |
| 110 | |
| 111 | try: |
| 112 | data = self._get_data() |
| 113 | except Exception as error: |
| 114 | self.error('_get_data() failed. Error: {error}'.format(error=error)) |
| 115 | return False |
| 116 | |
| 117 | if isinstance(data, dict) and data: |
| 118 | return True |
| 119 | self.error("_get_data() returned no data or type is not <dict>") |
| 120 | return False |
| 121 | |
| 122 | def _get_raw_data(self, description=None): |
| 123 | """ |
| 124 | Get raw data from MySQL server |
| 125 | :return: dict: fetchall() or (fetchall(), description) |
| 126 | """ |
| 127 | |
| 128 | if not self.__connection: |
| 129 | self.__connection, error = self.__connect() |
| 130 | if error: |
| 131 | return None |
| 132 | |
| 133 | raw_data = dict() |
| 134 | queries = dict(self.queries) |
| 135 | try: |
| 136 | cursor = self.__connection.cursor() |
| 137 | for name, query in queries.items(): |
| 138 | try: |
| 139 | cursor.execute(query) |
| 140 | except (MySQLdb.ProgrammingError, MySQLdb.OperationalError) as error: |
| 141 | if self.__is_error_critical(err_class=exc_info()[0], err_text=str(error)): |
| 142 | cursor.close() |
| 143 | raise RuntimeError |
| 144 | self.error('Removed query: {name}[{query}]. Error: error'.format(name=name, |
| 145 | query=query, |
| 146 | error=error)) |
| 147 | self.queries.pop(name) |
| 148 | continue |
| 149 | else: |
| 150 | raw_data[name] = (cursor.fetchall(), cursor.description) if description else cursor.fetchall() |
| 151 | cursor.close() |
| 152 | self.__connection.commit() |
| 153 | except (MySQLdb.MySQLError, RuntimeError, TypeError, AttributeError): |
| 154 | self.__connection.close() |
| 155 | self.__connection = None |
| 156 | return None |
| 157 | else: |
| 158 | return raw_data or None |
| 159 | |
| 160 | @staticmethod |
| 161 | def __is_error_critical(err_class, err_text): |
| 162 | return err_class == MySQLdb.OperationalError and all(['denied' not in err_text, |
| 163 | 'Unknown column' not in err_text]) |