Make sqlalchemy engine connect listener recorder specific (#34908)

This commit is contained in:
J. Nick Koston 2020-04-30 02:08:40 -05:00 committed by GitHub
parent 8467b91390
commit 574d8d30a7
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -9,9 +9,7 @@ import threading
import time
from typing import Any, Dict, Optional
from sqlalchemy import create_engine, exc, select
from sqlalchemy.engine import Engine
from sqlalchemy.event import listens_for
from sqlalchemy import create_engine, event as sqlalchemy_event, exc, select
from sqlalchemy.orm import scoped_session, sessionmaker
from sqlalchemy.pool import StaticPool
import voluptuous as vol
@ -488,15 +486,13 @@ class Recorder(threading.Thread):
"""Ensure database is ready to fly."""
kwargs = {}
# pylint: disable=unused-variable
@listens_for(Engine, "connect")
def setup_connection(dbapi_connection, connection_record):
def setup_recorder_connection(dbapi_connection, connection_record):
"""Dbapi specific connection settings."""
# We do not import sqlite3 here so mysql/other
# users do not have to pay for it to be loaded in
# memory
if self.db_url == "sqlite://" or ":memory:" in self.db_url:
if self.db_url.startswith("sqlite://"):
old_isolation = dbapi_connection.isolation_level
dbapi_connection.isolation_level = None
cursor = dbapi_connection.cursor()
@ -519,6 +515,9 @@ class Recorder(threading.Thread):
self.engine.dispose()
self.engine = create_engine(self.db_url, **kwargs)
sqlalchemy_event.listen(self.engine, "connect", setup_recorder_connection)
Base.metadata.create_all(self.engine)
self.get_session = scoped_session(sessionmaker(bind=self.engine))