bench_executor.postgresql

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