| 32 | |
| 33 | |
| 34 | class MySQLDatabase(Database): |
| 35 | |
| 36 | def __init__(self, env_controller: EnvironmentController): |
| 37 | super().__init__(TYPE_MYSQL) |
| 38 | |
| 39 | self.env_controller = env_controller |
| 40 | assert isinstance(self.env_controller.delegation, DBBenchEnvironmentDelegation) |
| 41 | self.session_id: Optional[str] = None |
| 42 | self.container_id: Optional[str] = None |
| 43 | self.container_ip: Optional[str] = None |
| 44 | |
| 45 | self._conn: Optional[mysql_connector.MySQLConnectionAbstract] = None |
| 46 | self.password = self.env_controller.delegation.password |
| 47 | self.database: Optional[str] = None |
| 48 | |
| 49 | async def initialize(self): |
| 50 | session_id, container_ids, container_ips = await self.env_controller.start_session(self.type) |
| 51 | self.session_id = session_id |
| 52 | self.container_id = container_ids[self.type] |
| 53 | self.container_ip = container_ips[self.type] |
| 54 | await self._create_database() |
| 55 | |
| 56 | async def delete(self): |
| 57 | if self.database: |
| 58 | # try to connect and delete the database |
| 59 | try: |
| 60 | conn = await self._get_conn() |
| 61 | async with await conn.cursor() as cursor: |
| 62 | await cursor.execute(f'DROP DATABASE IF EXISTS {self.database}') |
| 63 | await conn.commit() |
| 64 | except: |
| 65 | self.logger.warning(f'Error dropping MySQL database {self.database}:', exc_info=True) |
| 66 | self.database = None |
| 67 | |
| 68 | if self._conn: |
| 69 | try: |
| 70 | await self._conn.close() |
| 71 | except: |
| 72 | self.logger.warning(f'Error closing MySQL connection:', exc_info=True) |
| 73 | self._conn = None |
| 74 | |
| 75 | if self.session_id: |
| 76 | try: |
| 77 | await self.env_controller.end_session(self.session_id) |
| 78 | except: |
| 79 | self.logger.warning(f'Error ending environment session {self.session_id}:', exc_info=True) |
| 80 | self.session_id = None |
| 81 | self.container_id = None |
| 82 | self.container_ip = None |
| 83 | |
| 84 | async def execute(self, sql: str, data: Union[Sequence, Dict[str, Any]] = ()) -> str: |
| 85 | conn = await self._get_conn() |
| 86 | try: |
| 87 | async with await conn.cursor() as cursor: |
| 88 | results = [] |
| 89 | await cursor.execute(sql, data) |
| 90 | if cursor.with_rows: |
| 91 | rows = await cursor.fetchall() |