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) # 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) mask_secret(engine.url.password) setup_event_handlers(engine) Session = scoped_session( sessionmaker( autocommit=False, autoflush=False, bind=engine, expire_on_commit=False, ))
def get( cls, key: str, default_var: Any = __NO_DEFAULT_SENTINEL, deserialize_json: bool = False, ) -> Any: """ Gets a value for an Airflow Variable Key :param key: Variable Key :param default_var: Default value of the Variable if the Variable doesn't exists :param deserialize_json: Deserialize the value to a Python dict """ var_val = Variable.get_variable_from_secrets(key=key) if var_val is None: if default_var is not cls.__NO_DEFAULT_SENTINEL: return default_var else: raise KeyError(f'Variable {key} does not exist') else: if deserialize_json: obj = json.loads(var_val) mask_secret(var_val, key) return obj else: mask_secret(var_val, key) return var_val
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 extra_dejson(self) -> Dict: """Returns the extra property by deserializing json.""" obj = {} if self.extra: try: obj = json.loads(self.extra) except JSONDecodeError: self.log.exception("Failed parsing the json for conn_id %s", self.conn_id) # Mask sensitive keys from this list mask_secret(obj) return obj
def __init__( self, conn_id: Optional[str] = None, conn_type: Optional[str] = None, description: Optional[str] = None, host: Optional[str] = None, login: Optional[str] = None, password: Optional[str] = None, schema: Optional[str] = None, port: Optional[int] = None, extra: Optional[Union[str, dict]] = None, uri: Optional[str] = None, ): super().__init__() self.conn_id = conn_id self.description = description if extra and not isinstance(extra, str): extra = json.dumps(extra) if uri and (conn_type or host or login or password or schema or port or extra): raise AirflowException( "You must create an object using the URI or individual values " "(conn_type, host, login, password, schema, port or extra)." "You can't mix these two ways to create this object.") if uri: self._parse_from_uri(uri) else: self.conn_type = conn_type self.host = host self.login = login self.password = password self.schema = schema self.port = port self.extra = extra if self.extra: self._validate_extra(self.extra, self.conn_id) if self.password: mask_secret(self.password)
def __init__( # pylint: disable=too-many-arguments self, conn_id: Optional[str] = None, conn_type: Optional[str] = None, description: Optional[str] = None, host: Optional[str] = None, login: Optional[str] = None, password: Optional[str] = None, schema: Optional[str] = None, port: Optional[int] = None, extra: Optional[str] = None, uri: Optional[str] = None, ): super().__init__() self.conn_id = conn_id self.description = description if uri and ( # pylint: disable=too-many-boolean-expressions conn_type or host or login or password or schema or port or extra): raise AirflowException( "You must create an object using the URI or individual values " "(conn_type, host, login, password, schema, port or extra)." "You can't mix these two ways to create this object.") if uri: self._parse_from_uri(uri) else: self.conn_type = conn_type self.host = host self.login = login self.password = password self.schema = schema self.port = port self.extra = extra if self.password: mask_secret(self.password)
def on_db_load(self): if self.password: mask_secret(self.password)
def on_db_load(self): if self._val: mask_secret(self.val, self.key)
def on_db_load(self): # pylint: disable=missing-function-docstring if self.password: mask_secret(self.password)
def on_db_load(self): # pylint: disable=missing-function-docstring if self._val: mask_secret(self.val, self.key)