bench_executor.mysql

  1#!/usr/bin/env python3
  2
  3"""
  4MySQL is an open-source relational database management system developed by
  5Oracle Corporation.
  6
  7**Website**: https://www.mysql.com/<br>
  8**Repository**: https://github.com/mysql/mysql-server
  9"""
 10
 11import os
 12import pymysql
 13import tempfile
 14from csv import reader
 15from typing import List, Tuple
 16from timeout_decorator import timeout, TimeoutError  # type: ignore
 17from bench_executor.container import Container
 18from bench_executor.logger import Logger
 19
 20VERSION = '8.0'
 21HOST = 'localhost'
 22USER = 'root'
 23PASSWORD = 'root'
 24DB = 'db'
 25PORT = '3306'
 26CLEAR_TABLES_TIMEOUT = 5 * 60  # 5 minutes
 27
 28
 29class MySQL(Container):
 30    """MySQL container for executing SQL queries."""
 31    def __init__(self, data_path: str, config_path: str, directory: str,
 32                 verbose: bool):
 33        """Creates an instance of the MySQL class.
 34
 35        Parameters
 36        ----------
 37        data_path : str
 38            Path to the data directory of the case.
 39        config_path : str
 40            Path to the config directory of the case.
 41        directory : str
 42            Path to the directory to store logs.
 43        verbose : bool
 44            Enable verbose logs.
 45        """
 46        self._data_path = os.path.abspath(data_path)
 47        self._config_path = os.path.abspath(config_path)
 48        self._logger = Logger(__name__, directory, verbose)
 49        self._tables: List[str] = []
 50        tmp_dir = os.path.join(tempfile.gettempdir(), 'mysql')
 51        os.umask(0)
 52        os.makedirs(tmp_dir, exist_ok=True)
 53        os.makedirs(os.path.join(self._data_path, 'mysql'), exist_ok=True)
 54
 55        super().__init__(f'kgconstruct/mysql:v{VERSION}', 'MySQL',
 56                         self._logger,
 57                         ports={PORT: PORT},
 58                         environment={'MYSQL_ROOT_PASSWORD': 'root',
 59                                      'MYSQL_DATABASE': 'db'},
 60                         volumes=[f'{self._data_path}/shared/:/data/shared',
 61                                  f'{self._config_path}/mysql/'
 62                                  f'mysql-secure-file-prive.cnf:'
 63                                  f'/etc/mysql/conf.d/'
 64                                  f'mysql-secure-file-prive.cnf',
 65                                  f'{tmp_dir}:/var/lib/mysql'])
 66
 67    def initialization(self) -> bool:
 68        """Initialize MySQL's database.
 69
 70        Returns
 71        -------
 72        success : bool
 73            Whether the initialization was successfull or not.
 74        """
 75        # MySQL should start with a initialized database, start MySQL
 76        # if not initialized to avoid the pre-run start during benchmark
 77        # execution
 78        success = self.wait_until_ready()
 79        if not success:
 80            return False
 81        success = self.stop()
 82
 83        return success
 84
 85    @property
 86    def root_mount_directory(self) -> str:
 87        """Subdirectory in the root directory of the case for MySQL.
 88
 89        Returns
 90        -------
 91        subdirectory : str
 92            Subdirectory of the root directory for MySQL.
 93        """
 94        return __name__.lower()
 95
 96    def wait_until_ready(self, command: str = '') -> bool:
 97        """Wait until MySQL is ready to execute SQL queries.
 98
 99        Parameters
100        ----------
101        command : str
102            Command to execute in the MySQL container, optionally, defaults to
103            no command.
104
105        Returns
106        -------
107        success : bool
108            Whether the MySQL was initialized successfull or not.
109        """
110        log_line = f'port: {PORT}  MySQL Community Server - GPL.'
111        return self.run_and_wait_for_log(log_line, command=command)
112
113    def load(self, csv_file: str, table: str) -> bool:
114        """Load a single CSV file into MySQL.
115
116        Parameters
117        ----------
118        csv_file : str
119            Name of the CSV file.
120        table : str
121            Name of the table.
122
123        Returns
124        -------
125        success : bool
126            Whether the execution was successfull or not.
127        """
128        return self._load_csv(csv_file, table, True)
129
130    def load_multiple(self, csv_files: List[dict]) -> bool:
131        """Load multiple CSV files into MySQL.
132
133        Parameters
134        ----------
135        csv_files : list
136            List of CSV files to load. Each entry consist of a `file` and
137            `table` key.
138
139        Returns
140        -------
141        success : bool
142            Whether the execution was successfull or not.
143        """
144        for entry in csv_files:
145            if not self._load_csv(entry['file'], entry['table'], True):
146                return False
147        return True
148
149    def load_sql_schema(self, schema_file: str,
150                        csv_files: List[Tuple[str, str]]) -> bool:
151        """Execute SQL schema with MySQL.
152
153        Executes a .sql file with MySQL.
154        If the data is not loaded by the .sql file but only the schema is
155        provided through the .sql file, a list of CSV files can be provided to
156        load them as well.
157
158        Parameters
159        ----------
160        schema_file : str
161            Name of the .sql file.
162        csv_files : List[Tuple[str, str]]
163            List of CSV file names to load in the tables created with the .sql
164            file, may also be an empty list. Each entry contains a Tuple with
165            the name of the CSV file and the table name.
166
167        Returns
168        -------
169        success : bool
170            Whether the execution was successfull or not.
171        """
172        success = True
173
174        # Load SQL schema
175        cmd = f'/bin/sh -c \'mysql --host={HOST} --port={PORT} ' + \
176              f'--user={USER} --password={PASSWORD} --database={DB} ' + \
177              f'< /data/shared/{schema_file}\''
178        success, output = self.exec(cmd)
179
180        # Load CSVs
181        if success:
182            for csv_file, table in csv_files:
183                success = self._load_csv(csv_file, table, False)
184                if not success:
185                    break
186
187        return success
188
189    def _load_csv(self, csv_file: str, table: str, create: bool) -> bool:
190        """Load a single CSV file into MySQL.
191
192        Parameters
193        ----------
194        csv_file : str
195            Name of the CSV file.
196        table : str
197            Name of the table to store the data in.
198        create : bool
199            Whether to drop and create the table or re-use it
200
201        Returns
202        -------
203        success : bool
204            Whether the execution was successfull or not.
205        """
206        success = True
207        columns = None
208        table = table.lower()
209        path = os.path.join(self._data_path, 'shared', csv_file)
210        path2 = os.path.join(self._data_path, 'shared', f'tmp_{csv_file}')
211
212        self._tables.append(table)
213
214        # Analyze CSV for loading
215        if not os.path.exists(path):
216            self._logger.error(f'CSV file "{path}" does not exist')
217            return False
218
219        with open(path, 'r') as f:
220            csv_reader = reader(f)
221            columns = next(csv_reader)
222            columns = [x.lower() for x in columns]
223
224        # MySQL cannot set NULL as NULL keyword, use their own specific syntax
225        # for this: \N
226        with open(path, 'r') as f:
227            data = f.read()
228            data = data.replace('NULL', '\\N')
229
230            with open(path2, 'w') as f2:
231                f2.write(data)
232
233        # Load CSV
234        connection = pymysql.connect(host=HOST, user=USER, password=PASSWORD,
235                                     db=DB)
236        try:
237            cursor = connection.cursor()
238
239            if create:
240                cursor.execute(f'DROP TABLE IF EXISTS {table};')
241                c = ' TEXT , '.join(columns) + ' TEXT'
242                cursor.execute(f'CREATE TABLE {table} (k INT ZEROFILL '
243                               f'NOT NULL AUTO_INCREMENT, {c}, '
244                               f'PRIMARY KEY(k));')
245            c = ','.join(columns)
246            cursor.execute(f'LOAD DATA INFILE \'/data/shared/tmp_{csv_file}\' '
247                           f'INTO TABLE {table} FIELDS TERMINATED BY \',\' '
248                           f'ENCLOSED BY \'\\"\' LINES TERMINATED BY \'\\n\' '
249                           f'IGNORE 1 ROWS ({c});')
250            cursor.execute('COMMIT;')
251
252            header = '| ID | ' + ' | '.join(columns) + ' |'
253            self._logger.debug(header)
254            self._logger.debug('-' * len(header))
255
256            cursor.execute(f'SELECT * FROM {table};')
257            number_of_records = 0
258            for record in cursor:
259                number_of_records += 1
260                self._logger.debug(record)
261
262            if number_of_records == 0:
263                success = False
264        except Exception as e:
265            self._logger.error(f'Failed to load CSV: "{e}"')
266            success = False
267        finally:
268            connection.close()
269
270        return success
271
272    @timeout(CLEAR_TABLES_TIMEOUT)
273    def _clear_tables(self):
274        """Clears all tables with a provided timeout."""
275        connection = pymysql.connect(host=HOST, database=DB,
276                                     user=PASSWORD, password=PASSWORD)
277        cursor = connection.cursor()
278        for table in self._tables:
279            cursor.execute(f'DROP TABLE IF EXISTS {table};')
280            cursor.execute('COMMIT;')
281        self._tables = []
282        connection.close()
283
284    def stop(self) -> bool:
285        """Stop MySQL
286        Clears all tables and stops the MySQL container.
287
288        Returns
289        -------
290        success : bool
291            Whether the execution was successfull or not.
292        """
293        try:
294            self._clear_tables()
295        except TimeoutError:
296            self._logger.warning('Clearing MySQL tables timed out after '
297                                 f'{CLEAR_TABLES_TIMEOUT}s!')
298
299        return super().stop()
300
301
302if __name__ == '__main__':
303    print(f'ℹ️  Starting up MySQL v{VERSION}...')
304    m = MySQL('data', 'config', 'log', True)
305    m.wait_until_ready()
306    input('ℹ️  Press any key to stop')
307    m.stop()
308    print('ℹ️  Stopped')