def connect(self):
engine = create_engine(self._conn_string, poolclass=NullPool)
conn = engine.connect()
try:
with handle_schema_errors(conn, get_alembic_config(__file__)):
yield conn
finally:
conn.close()
def connect(self):
engine = create_engine(self._conn_string, poolclass=NullPool)
conn = engine.connect()
try:
with handle_schema_errors(
conn, get_alembic_config(__file__), msg="Sqlite run storage requires migration",
):
yield conn
finally:
conn.close()
def connect(self, _run_id=None): # pylint: disable=arguments-differ
try:
conn = self._engine.connect()
with handle_schema_errors(
conn,
get_alembic_config(__file__),
msg='Postgres run storage requires migration',
):
yield self._engine
finally:
conn.close()
def connect(self):
try:
conn = self._engine.connect()
with handle_schema_errors(
conn,
get_alembic_config(__file__),
msg='Postgres schedule storage requires migration',
):
yield self._engine
finally:
conn.close()
def connect(self, run_id=None):
engine = create_engine(self._conn_string, poolclass=NullPool)
conn = engine.connect()
try:
with handle_schema_errors(
conn,
get_alembic_config(__file__),
msg='ConsolidatedSqliteEventLogStorage requires migration',
):
yield conn
finally:
conn.close()
def create_mysql_connection(engine, dunder_file, storage_type_desc=None):
check.inst_param(engine, "engine", db.engine.Engine)
check.str_param(dunder_file, "dunder_file")
check.opt_str_param(storage_type_desc, "storage_type_desc", "")
if storage_type_desc:
storage_type_desc += " "
else:
storage_type_desc = ""
conn = None
try:
# Retry connection to gracefully handle transient connection issues
conn = retry_mysql_connection_fn(engine.connect)
with handle_schema_errors(conn, mysql_alembic_config(dunder_file)):
yield conn
finally:
if conn:
conn.close()
def _connect(self, shard):
with self._db_lock:
check.str_param(shard, "shard")
conn_string = self.conn_string_for_shard(shard)
engine = create_engine(conn_string, poolclass=NullPool)
if not shard in self._initialized_dbs:
self._initdb(engine)
self._initialized_dbs.add(shard)
conn = engine.connect()
try:
with handle_schema_errors(conn, get_alembic_config(__file__)):
yield conn
finally:
conn.close()
engine.dispose()
def connect(self, run_id=None):
check.str_param(run_id, "run_id")
conn_string = self.conn_string_for_run_id(run_id)
engine = create_engine(conn_string, poolclass=NullPool)
if not os.path.exists(self.path_for_run_id(run_id)):
self._initdb(engine)
conn = engine.connect()
try:
with handle_schema_errors(
conn,
get_alembic_config(__file__),
msg="SqliteEventLogStorage for run {run_id}".format(run_id=run_id),
):
yield conn
finally:
conn.close()
engine.dispose()
def create_pg_connection(engine, dunder_file, storage_type_desc=None):
check.inst_param(engine, "engine", sqlalchemy.engine.Engine)
check.str_param(dunder_file, "dunder_file")
check.opt_str_param(storage_type_desc, "storage_type_desc", "")
if storage_type_desc:
storage_type_desc += " "
else:
storage_type_desc = ""
conn = None
try:
# Retry connection to gracefully handle transient connection issues
conn = retry_pg_connection_fn(engine.connect)
with handle_schema_errors(
conn,
get_alembic_config(dunder_file),
msg="Postgres {}storage requires migration".format(storage_type_desc),
):
yield conn
finally:
if conn:
conn.close()
