def setUp(self): if sys.version_info[0] == 3: raise unittest.SkipTest('TestSparkSubmitHook won\'t work with ' 'python3. No need to test anything here') configuration.load_test_config() db.merge_conn( models.Connection( conn_id='spark_yarn_cluster', conn_type='spark', host='yarn://yarn-master', extra='{"queue": "root.etl", "deploy-mode": "cluster"}')) db.merge_conn( models.Connection(conn_id='spark_default_mesos', conn_type='spark', host='mesos://host', port=5050)) db.merge_conn( models.Connection(conn_id='spark_home_set', conn_type='spark', host='yarn://yarn-master', extra='{"spark-home": "/opt/myspark"}')) db.merge_conn( models.Connection(conn_id='spark_home_not_set', conn_type='spark', host='yarn://yarn-master'))
def setUp(self): configuration.load_test_config() db.merge_conn( models.Connection( conn_id='cassandra_test', conn_type='cassandra', host='host-1,host-2', port='9042', schema='test_keyspace', extra='{"load_balancing_policy":"TokenAwarePolicy"}')) db.merge_conn( models.Connection(conn_id='cassandra_default_with_schema', conn_type='cassandra', host='localhost', port='9042', schema='s')) hook = CassandraHook("cassandra_default") session = hook.get_conn() cqls = [ "DROP SCHEMA IF EXISTS s", """ CREATE SCHEMA s WITH REPLICATION = { 'class' : 'SimpleStrategy', 'replication_factor' : 1 } """, ] for cql in cqls: session.execute(cql) session.shutdown() hook.shutdown_cluster()
def setUp(self): configuration.conf.load_test_config() db.merge_conn( models.Connection( conn_id='s3_test', conn_type='s3', schema='test', extra= '{"aws_access_key_id": "aws_access_key_id", "aws_secret_access_key":' ' "aws_secret_access_key"}')) db.merge_conn( models.Connection(conn_id='mysql_test', conn_type='mysql', host='some.host.com', schema='test_db', login='******', password='******')) self.s3_to_mysql_transfer_kwargs = { 'aws_conn_id': 's3_test', 'mysql_conn_id': 'mysql_test', 's3_source_key': 'test/s3_to_mysql_test.csv', 'mysql_table': 'mysql_table', 'mysql_duplicate_key_handling': 'IGNORE', 'mysql_extra_options': """ FIELDS TERMINATED BY ',' IGNORE 1 LINES """, 'task_id': 'task_id', 'dag': None }
def setUp(self): load_test_config() db.merge_conn( models.Connection(conn_id='google_test', host='google', conn_type="google_cloud_platform", schema='refresh_token', login='******', password='******')) db.merge_conn( models.Connection( conn_id='s3_test', conn_type='s3', schema='test', extra= '{"aws_access_key_id": "aws_access_key_id", "aws_secret_access_key":' ' "aws_secret_access_key"}')) self.kwargs = { 'gcp_conn_id': 'google_test', 'google_api_service_name': 'test_service', 'google_api_service_version': 'v3', 'google_api_endpoint_path': 'analyticsreporting.reports.batchGet', 'google_api_endpoint_params': {}, 'google_api_pagination': False, 'google_api_num_retries': 0, 'aws_conn_id': 's3_test', 's3_destination_key': 'test/google_api_to_s3_test.csv', 's3_overwrite': True, 'task_id': 'task_id', 'dag': None }
def setUp(self): configuration.load_test_config() db.merge_conn( models.Connection( conn_id='spark_yarn_cluster', conn_type='spark', host='yarn://yarn-master', extra='{"queue": "root.etl", "deploy-mode": "cluster"}')) db.merge_conn( models.Connection(conn_id='spark_k8s_cluster', conn_type='spark', host='k8s://https://k8s-master', extra='{"spark-home": "/opt/spark", ' + '"deploy-mode": "cluster", ' + '"namespace": "mynamespace"}')) db.merge_conn( models.Connection(conn_id='spark_default_mesos', conn_type='spark', host='mesos://host', port=5050)) db.merge_conn( models.Connection(conn_id='spark_home_set', conn_type='spark', host='yarn://yarn-master', extra='{"spark-home": "/opt/myspark"}')) db.merge_conn( models.Connection(conn_id='spark_home_not_set', conn_type='spark', host='yarn://yarn-master')) db.merge_conn( models.Connection(conn_id='spark_binary_set', conn_type='spark', host='yarn', extra='{"spark-binary": "custom-spark-submit"}')) db.merge_conn( models.Connection(conn_id='spark_binary_and_home_set', conn_type='spark', host='yarn', extra='{"spark-home": "/path/to/spark_home", ' + '"spark-binary": "custom-spark-submit"}')) db.merge_conn( models.Connection( conn_id='spark_standalone_cluster', conn_type='spark', host='spark://spark-standalone-master:6066', extra= '{"spark-home": "/path/to/spark_home", "deploy-mode": "cluster"}' )) db.merge_conn( models.Connection( conn_id='spark_standalone_cluster_client_mode', conn_type='spark', host='spark://spark-standalone-master:6066', extra= '{"spark-home": "/path/to/spark_home", "deploy-mode": "client"}' ))
def setUp(self): configuration.load_test_config() db.merge_conn( models.Connection(conn_id='wasb_test_key', conn_type='wasb', login='******', password='******')) db.merge_conn( models.Connection(conn_id='wasb_test_sas_token', conn_type='wasb', login='******', extra=json.dumps({'sas_token': 'token'})))
def setUp(self): configuration.load_test_config() db.merge_conn( models.Connection( conn_id='spark_yarn_cluster', conn_type='spark', host='yarn://yarn-mater', extra='{"queue": "root.etl", "deploy-mode": "cluster"}')) db.merge_conn( models.Connection(conn_id='spark_default_mesos', conn_type='spark', host='mesos://host', port=5050))
def airflow_connections(): db.merge_conn(models.Connection( conn_id='https_conn', conn_type='https', host='www.google.com', port=443, extra=json_dumps({'param': 1, 'option': 2}), )) db.merge_conn(models.Connection( conn_id='postgres_conn', uri='postgresql+psycopg2://postgres:postgres@postgres:5432/postgres', ))
def get_airflow_connection(conn_id=None): return models.Connection( conn_id='http_default', conn_type='http', host='test:8080/', extra='{"bareer": "test"}' )
def test_host_encoded_https_connection(self, mock_get_connection): c = models.Connection(conn_id='http_default', conn_type='http', host='https://localhost') mock_get_connection.return_value = c hook = HttpHook() hook.get_conn({}) self.assertEqual(hook.base_url, 'https://localhost')
def get_airflow_connection_with_port(conn_id=None): return models.Connection( conn_id='http_default', conn_type='http', host='test.com', port=1234 )
def setUp(self): configuration.load_test_config() db.merge_conn( models.Connection( conn_id='jira_default', conn_type='jira', host='https://localhost/jira/', port=443, extra='{"verify": "False", "project": "AIRFLOW"}'))
def setUp(self): configuration.load_test_config() db.merge_conn( models.Connection( conn_id='slack-webhook-default', extra='{"webhook_token": "your_token_here"}') )
def setUp(self): configuration.load_test_config() db.merge_conn( models.Connection(conn_id='docker_default', conn_type='docker', host='some.docker.registry.com', login='******', password='******')) db.merge_conn( models.Connection( conn_id='docker_with_extras', conn_type='docker', host='some.docker.registry.com', login='******', password='******', extra='{"email": "*****@*****.**", "reauth": "no"}'))
def setUp(self): configuration.load_test_config() db.merge_conn( models.Connection(conn_id='spark_default', conn_type='spark', host='yarn://yarn-master'))
def setUp(self): configuration.load_test_config() db.merge_conn( models.Connection(conn_id='valohai_default', conn_type='HTTP', host='app.valohai.com', password='******'))
def setUp(self): configuration.load_test_config() db.merge_conn( models.Connection( conn_id='spark-default', conn_type='spark', host='yarn://yarn-master', extra='{"queue": "root.etl", "deploy-mode": "cluster"}') ) db.merge_conn( models.Connection( conn_id='jdbc-default', conn_type='postgres', host='localhost', schema='default', port=5432, login='******', password='******', extra='{"conn_prefix":"jdbc:postgresql://"}' ) )
def add_connections(self): merge_conn( models.Connection(conn_id="process_report", conn_type="http", host="localhost", port="3070", extra="{\"endpoint\":\"/airflow/\"}"))
def setUp(self): configuration.load_test_config() db.merge_conn( models.Connection( conn_id='cassandra_test', conn_type='cassandra', host='host-1,host-2', port='9042', schema='test_keyspace', extra='{"load_balancing_policy":"TokenAwarePolicy"}'))
def setUp(self): configuration.load_test_config() db.merge_conn( models.Connection( conn_id='jdbc_default', conn_type='jdbc', host='jdbc://localhost/', port=443, extra='{"extra__jdbc__drv_path": "/path1/test.jar,/path2/t.jar2", "extra__jdbc__drv_clsname": "com.driver.main"}'))
def update_connections(connections): """ Add or update Airflow connections. :param connections: Connections as dictionary. :return: Job done. """ print('Updating connections...') session = settings.Session() try: for conn_id, conn in connections.items(): existing = session.query( models.Connection).filter_by(conn_id=conn_id).first() if existing: existing.host = conn['host'] existing.conn_type = conn['conn_type'] existing.port = conn.get('port') existing.login = conn.get('login') existing.password = conn.get('password') existing.schema = conn.get('schema') existing.extra = conn.get('extra') session.merge(existing) print('\tUPDATED: connection %s' % conn_id) else: c = models.Connection(conn_id=conn_id, **conn) session.add(c) print('\tADDED: connection %s' % conn_id) session.commit() print('Changes commited.') finally: session.close()
def create_or_replace_connection(connection_id, connection_type, ip, port, login, password, schema, extra): """ :param connection_id: Airflow connection_id to be created or replaced :param connection_type: Airflow connection type :param ip: Airflow connection ip :param port: Airflow connection port :param login: Airflow connection login/user :param password: Airflow connection password :param schema: Airflow connection schema :param extra: Airflow connection extra parameter supplied as json string """ session = settings.Session() conn = models.Connection(conn_id=connection_id, conn_type=connection_type, host=ip, port=port, login=login, password=password, schema=schema, extra=extra) ExtendedEmrCreateJobFlowOperator.delete_connection(connection_id) session.add(conn) session.commit()
def get_or_update_conn(name, **kwargs): """Returns a connection by id """ from airflow import settings, models session = settings.Session() con = get_connection(name) if not con: con = models.Connection(name, **kwargs) session.add(con) session.commit() else: for key, value in kwargs.items(): if key == "extra": con.set_extra(value) else: setattr(con, key, value) session.commit() return con
def setUp(self): configuration.load_test_config() db.merge_conn( models.Connection(conn_id='imap_test', host='base_url', login='******', password='******'))
def test_no_host_key_check_no_ignore(self, get_connection): connection = models.Connection( login='******', host='host', extra='{"ignore_hostkey_verification": false}') get_connection.return_value = connection hook = SFTPHook() self.assertEqual(hook.no_host_key_check, False)
def setUp(self): configuration.load_test_config() db.merge_conn( models.Connection(conn_id='imap_default', host='imap_server_address', login='******', password='******'))
def test_no_host_key_check_disabled_for_all_but_true(self, get_connection): connection = models.Connection(login='******', host='host', extra='{"no_host_key_check": "foo"}') get_connection.return_value = connection hook = SFTPHook() self.assertEqual(hook.no_host_key_check, False)
def setUp(self): configuration.load_test_config() db.merge_conn( models.Connection(conn_id='sagemaker_test_conn_id', conn_type='sagemaker', login='******', password='******', extra=json.dumps(db_config)))
def setUp(self): configuration.load_test_config() db.merge_conn( models.Connection( conn_id='sqoop_test', conn_type='sqoop', schema='schema', host='rmdbs', port=5050, extra=json.dumps(self._config_json) ) )
def test_no_host_key_check_enabled(self, get_connection): connection = models.Connection(login='******', host='host', extra='{"no_host_key_check": true}') get_connection.return_value = connection hook = SFTPHook() self.assertEqual(hook.no_host_key_check, True)