def configure_orm(disable_connection_pool=False): """Configure ORM using SQLAlchemy""" log.debug("Setting up DB connection pool (PID %s)", os.getpid()) global engine global Session engine_args = prepare_engine_args(disable_connection_pool) # Allow the user to specify an encoding for their DB otherwise default # to utf-8 so jobs & users with non-latin1 characters can still use us. engine_args['encoding'] = conf.get('core', 'SQL_ENGINE_ENCODING', fallback='utf-8') if conf.has_option('core', 'sql_alchemy_connect_args'): connect_args = conf.getimport('core', 'sql_alchemy_connect_args') else: connect_args = {} engine = create_engine(SQL_ALCHEMY_CONN, connect_args=connect_args, **engine_args) setup_event_handlers(engine) Session = scoped_session( sessionmaker( autocommit=False, autoflush=False, bind=engine, expire_on_commit=False, ) )
def configure_orm(disable_connection_pool=False): """Configure ORM using SQLAlchemy""" from airflow.utils.log.secrets_masker import mask_secret log.debug("Setting up DB connection pool (PID %s)", os.getpid()) global engine global Session engine_args = prepare_engine_args(disable_connection_pool) if conf.has_option('database', 'sql_alchemy_connect_args'): connect_args = conf.getimport('database', 'sql_alchemy_connect_args') else: connect_args = {} engine = create_engine(SQL_ALCHEMY_CONN, connect_args=connect_args, **engine_args) mask_secret(engine.url.password) setup_event_handlers(engine) Session = scoped_session( sessionmaker( autocommit=False, autoflush=False, bind=engine, expire_on_commit=False, )) if engine.dialect.name == 'mssql': session = Session() try: result = session.execute( sqlalchemy.text( 'SELECT is_read_committed_snapshot_on FROM sys.databases WHERE name=:database_name' ), params={"database_name": engine.url.database}, ) data = result.fetchone()[0] if data != 1: log.critical( "MSSQL database MUST have READ_COMMITTED_SNAPSHOT enabled." ) log.critical("The database %s has it disabled.", engine.url.database) log.critical( "This will cause random deadlocks, Refusing to start.") log.critical( "See https://airflow.apache.org/docs/apache-airflow/stable/howto/" "set-up-database.html#setting-up-a-mssql-database") raise Exception( "MSSQL database MUST have READ_COMMITTED_SNAPSHOT enabled." ) finally: session.close()
def configure_orm(disable_connection_pool=False): """ Configure ORM using SQLAlchemy""" log.debug("Setting up DB connection pool (PID %s)", os.getpid()) global engine global Session engine_args = {} pool_connections = conf.getboolean('core', 'SQL_ALCHEMY_POOL_ENABLED') if disable_connection_pool or not pool_connections: engine_args['poolclass'] = NullPool log.debug("settings.configure_orm(): Using NullPool") elif 'sqlite' not in SQL_ALCHEMY_CONN: # Pool size engine args not supported by sqlite. # If no config value is defined for the pool size, select a reasonable value. # 0 means no limit, which could lead to exceeding the Database connection limit. pool_size = conf.getint('core', 'SQL_ALCHEMY_POOL_SIZE', fallback=5) # The maximum overflow size of the pool. # When the number of checked-out connections reaches the size set in pool_size, # additional connections will be returned up to this limit. # When those additional connections are returned to the pool, they are disconnected and discarded. # It follows then that the total number of simultaneous connections # the pool will allow is pool_size + max_overflow, # and the total number of “sleeping” connections the pool will allow is pool_size. # max_overflow can be set to -1 to indicate no overflow limit; # no limit will be placed on the total number # of concurrent connections. Defaults to 10. max_overflow = conf.getint('core', 'SQL_ALCHEMY_MAX_OVERFLOW', fallback=10) # The DB server already has a value for wait_timeout (number of seconds after # which an idle sleeping connection should be killed). Since other DBs may # co-exist on the same server, SQLAlchemy should set its # pool_recycle to an equal or smaller value. pool_recycle = conf.getint('core', 'SQL_ALCHEMY_POOL_RECYCLE', fallback=1800) # Check connection at the start of each connection pool checkout. # Typically, this is a simple statement like “SELECT 1”, but may also make use # of some DBAPI-specific method to test the connection for liveness. # More information here: # https://docs.sqlalchemy.org/en/13/core/pooling.html#disconnect-handling-pessimistic pool_pre_ping = conf.getboolean('core', 'SQL_ALCHEMY_POOL_PRE_PING', fallback=True) log.debug( "settings.configure_orm(): Using pool settings. pool_size=%d, max_overflow=%d, " "pool_recycle=%d, pid=%d", pool_size, max_overflow, pool_recycle, os.getpid()) engine_args['pool_size'] = pool_size engine_args['pool_recycle'] = pool_recycle engine_args['pool_pre_ping'] = pool_pre_ping engine_args['max_overflow'] = max_overflow # Allow the user to specify an encoding for their DB otherwise default # to utf-8 so jobs & users with non-latin1 characters can still use us. engine_args['encoding'] = conf.get('core', 'SQL_ENGINE_ENCODING', fallback='utf-8') if conf.has_option('core', 'sql_alchemy_connect_args'): connect_args = conf.getimport('core', 'sql_alchemy_connect_args') else: connect_args = {} engine = create_engine(SQL_ALCHEMY_CONN, connect_args=connect_args, **engine_args) setup_event_handlers(engine) Session = scoped_session( sessionmaker(autocommit=False, autoflush=False, bind=engine, expire_on_commit=False))