def __exit__(
self,
exc_type: Type[BaseException],
exc_val: BaseException,
exc_tb: TracebackType,
) -> None:
try:
if exc_val:
self.c.rollback()
raise exc_val
elif not self.commit:
self.c.rollback()
else:
self.c.commit()
except psycopg2.InterfaceError as e:
self.c.close(timeout=0)
raise InterfaceError(e)
except psycopg2.DataError as e:
raise DataError(e)
except psycopg2.OperationalError as e:
raise OperationalError(e)
except psycopg2.IntegrityError as e:
raise IntegrityError(e)
except psycopg2.InternalError as e:
raise InternalError(e)
except psycopg2.ProgrammingError as e:
raise ProgrammingError(e)
except psycopg2.NotSupportedError as e:
raise NotSupportedError(e)
except psycopg2.DatabaseError as e:
raise DatabaseError(e)
except psycopg2.Error as e:
raise PersistenceError(e)
finally:
self.c.is_idle.set()
def _create_connection(self) -> SQLiteConnection:
# Make a connection to an SQLite database.
try:
c = sqlite3.connect(
database=self.db_name,
uri=True,
check_same_thread=False,
isolation_level=None, # Auto-commit mode.
cached_statements=True,
timeout=self.lock_timeout or SQLITE3_DEFAULT_LOCK_TIMEOUT,
)
except (sqlite3.Error, TypeError) as e:
raise InterfaceError(e)
# Use WAL (write-ahead log) mode if file-based database.
if not self.is_sqlite_memory_mode:
if not self.is_journal_mode_wal:
cursor = c.cursor()
cursor.execute("PRAGMA journal_mode;")
mode = cursor.fetchone()[0]
if mode.lower() == "wal":
self.is_journal_mode_wal = True
else:
cursor.execute("PRAGMA journal_mode=WAL;")
self.is_journal_mode_wal = True
self.journal_mode_was_changed_to_wal = True
# Set the row factory.
c.row_factory = sqlite3.Row
# Return the connection.
return SQLiteConnection(sqlite_conn=c, max_age=self.max_age)
def __exit__(
self,
exc_type: Type[BaseException],
exc_val: BaseException,
exc_tb: TracebackType,
) -> None:
try:
if exc_val:
# Roll back all changes
# if an exception occurs.
self.connection.rollback()
raise exc_val
elif not self.commit:
self.connection.rollback()
else:
self.connection.commit()
except sqlite3.InterfaceError as e:
raise InterfaceError(e) from e
except sqlite3.DataError as e:
raise DataError(e) from e
except sqlite3.OperationalError as e:
raise OperationalError(e) from e
except sqlite3.IntegrityError as e:
raise IntegrityError(e) from e
except sqlite3.InternalError as e:
raise InternalError(e) from e
except sqlite3.ProgrammingError as e:
raise ProgrammingError(e) from e
except sqlite3.NotSupportedError as e:
raise NotSupportedError(e) from e
except sqlite3.DatabaseError as e:
raise DatabaseError(e) from e
except sqlite3.Error as e:
raise PersistenceError(e) from e
def _create_connection(self, thread_id: int) -> Connection:
# Make a connection to a Postgres database.
try:
psycopg_c = psycopg2.connect(
dbname=self.dbname,
host=self.host,
port=self.port,
user=self.user,
password=self.password,
connect_timeout=5,
)
psycopg_c.cursor().execute(
f"SET idle_in_transaction_session_timeout = "
f"'{self.idle_in_transaction_session_timeout}s'"
)
except psycopg2.Error as e:
raise InterfaceError(e)
else:
c = Connection(
psycopg_c,
max_age=self.conn_max_age,
)
self._connections[thread_id] = c
return c