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)
