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='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='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='cassandra', 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.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 setUp(self):
configuration.load_test_config()
db.merge_conn(
Connection(
conn_id='default-discord-webhook',
host='https://discordapp.com/api/',
extra='{"webhook_endpoint": "webhooks/00000/some-discord-token_000"}')
)
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(
Connection(
conn_id=self.conn_id,
conn_type='http',
host='https://oapi.dingtalk.com',
password='******'))
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 setUp(self):
configuration.load_test_config()
db.merge_conn(
Connection(
conn_id='datadog_default', conn_type='datadog',
login='******', password='******',
extra=json.dumps({'api_key': 'api_key', 'app_key': 'app_key'})
)
)
def setUp(self):
configuration.load_test_config()
db.merge_conn(
Connection(
conn_id=self.conn_id,
host='https://api.opsgenie.com/',
password='******'
)
)
def setUp(self):
configuration.load_test_config()
db.merge_conn(
models.Connection(
conn_id='imap_test',
host='base_url',
login='******',
password='******'
)
)
def setUp(self):
configuration.load_test_config()
db.merge_conn(
Connection(
conn_id='imap_default',
host='imap_server_address',
login='******',
password='******'
)
)
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 test_conn_with_extra_parameters(self):
db.merge_conn(
models.Connection(conn_id='ssh_with_extra',
host='localhost',
conn_type='ssh',
extra='{"compress" : true, "no_host_key_check" : "true"}'
)
)
ssh_hook = SSHHook(ssh_conn_id='ssh_with_extra')
self.assertEqual(ssh_hook.compress, True)
self.assertEqual(ssh_hook.no_host_key_check, True)
def setUp(self):
configuration.load_test_config()
db.merge_conn(
Connection(
conn_id='adl_test_key',
conn_type='azure_data_lake',
login='******',
password='******',
extra=json.dumps({"tenant": "tenant",
"account_name": "accountname"})
)
)
def test_conn_with_extra_parameters(self):
from airflow.contrib.hooks.ssh_hook import SSHHook
db.merge_conn(
models.Connection(conn_id='ssh_with_extra',
host='localhost',
conn_type='ssh',
extra='{"compress" : true, "no_host_key_check" : "true"}'
)
)
ssh_hook = SSHHook(ssh_conn_id='ssh_with_extra', keepalive_interval=10)
ssh_hook.get_conn()
self.assertEqual(ssh_hook.compress, True)
self.assertEqual(ssh_hook.no_host_key_check, True)
def setUp(self):
configuration.load_test_config()
args = {
'owner': 'airflow',
'start_date': DEFAULT_DATE
}
dag = DAG('test_dag_id', default_args=args)
self.dag = dag
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='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):
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 test_conn_with_broken_config_missing_host_fails(self, _):
db.merge_conn(
models.Connection(
conn_id='docker_without_host',
conn_type='docker',
login='******',
password='******'
)
)
with self.assertRaises(AirflowException):
hook = DockerHook(
docker_conn_id='docker_without_host',
base_url='unix://var/run/docker.sock',
version='auto'
)
def test_get_conn(self):
configuration.load_test_config()
db.merge_conn(
Connection(
conn_id='azure_container_registry',
login='******',
password='******',
host='test.cr',
)
)
hook = AzureContainerRegistryHook(conn_id='azure_container_registry')
self.assertIsNotNone(hook.connection)
self.assertEqual(hook.connection.username, 'myuser')
self.assertEqual(hook.connection.password, 'password')
self.assertEqual(hook.connection.server, 'test.cr')
def setUp(self):
super(TestIntegrationFTPHook, self).setUp()
from airflow import configuration
from airflow.utils import db
from airflow.models import Connection
configuration.load_test_config()
db.merge_conn(
Connection(
conn_id='ftp_passive', conn_type='ftp',
host='localhost', extra='{"passive": true}'))
db.merge_conn(
Connection(
conn_id='ftp_active', conn_type='ftp',
host='localhost', extra='{"passive": false}'))
def test_conn_with_broken_config_missing_username_fails(self, _):
db.merge_conn(
models.Connection(
conn_id='docker_without_username',
conn_type='docker',
host='some.docker.registry.com',
password='******',
extra='{"email": "*****@*****.**"}'
)
)
with self.assertRaises(AirflowException):
hook = DockerHook(
docker_conn_id='docker_without_username',
base_url='unix://var/run/docker.sock',
version='auto'
)
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 setUp(self):
# set up some test variables
self.test_end_point = 'https://test_endpoint:443'
self.test_master_key = 'magic_test_key'
self.test_database_name = 'test_database_name'
self.test_collection_name = 'test_collection_name'
configuration.load_test_config()
db.merge_conn(
Connection(
conn_id='azure_cosmos_test_key_id',
conn_type='azure_cosmos',
login=self.test_end_point,
password=self.test_master_key,
extra=json.dumps({'database_name': self.test_database_name,
'collection_name': self.test_collection_name})
)
)
def setUp(self):
configuration.load_test_config()
db.merge_conn(
Connection(
conn_id='azure_container_instance_test',
conn_type='azure_container_instances',
login='******',
password='******',
extra=json.dumps({'tenantId': 'tenant_id',
'subscriptionId': 'subscription_id'})
)
)
self.resources = ResourceRequirements(requests=ResourceRequests(
memory_in_gb='4',
cpu='1'))
with patch('azure.common.credentials.ServicePrincipalCredentials.__init__',
autospec=True, return_value=None):
with patch('azure.mgmt.containerinstance.ContainerInstanceManagementClient'):
self.testHook = AzureContainerInstanceHook(conn_id='azure_container_instance_test')
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 test_get_file_volume(self):
configuration.load_test_config()
db.merge_conn(
Connection(
conn_id='wasb_test_key',
conn_type='wasb',
login='******',
password='******'
)
)
hook = AzureContainerVolumeHook(wasb_conn_id='wasb_test_key')
volume = hook.get_file_volume(mount_name='mount',
share_name='share',
storage_account_name='storage',
read_only=True)
self.assertIsNotNone(volume)
self.assertEqual(volume.name, 'mount')
self.assertEqual(volume.azure_file.share_name, 'share')
self.assertEqual(volume.azure_file.storage_account_key, 'key')
self.assertEqual(volume.azure_file.storage_account_name, 'storage')
self.assertEqual(volume.azure_file.read_only, True)
def setUp(self):
configuration.load_test_config()
db.merge_conn(
Connection(
conn_id='mongo_test', conn_type='mongo',
host='mongo', port='27017', schema='test'))
args = {
'owner': 'airflow',
'start_date': DEFAULT_DATE
}
self.dag = DAG('test_dag_id', default_args=args)
hook = MongoHook('mongo_test')
hook.insert_one('foo', {'bar': 'baz'})
self.sensor = MongoSensor(
task_id='test_task',
mongo_conn_id='mongo_test',
dag=self.dag,
collection='foo',
query={'bar': 'baz'}
)
def setUp(self):
db.merge_conn(
Connection(conn_id='slack-webhook-default',
extra='{"webhook_token": "your_token_here"}'))
db.merge_conn(
Connection(conn_id='slack-webhook-url',
host='https://hooks.slack.com/services/T000/B000/XXX'))
db.merge_conn(
Connection(conn_id='slack-webhook-host',
host='https://hooks.slack.com/services/T000/'))
def setUpClass(cls) -> None:
db.merge_conn(
Connection(
conn_id=cls.CONN_SSH_WITH_EXTRA,
host='localhost',
conn_type='ssh',
extra='{"compress" : true, "no_host_key_check" : "true", "allow_host_key_change": false}',
)
)
db.merge_conn(
Connection(
conn_id=cls.CONN_SSH_WITH_EXTRA_FALSE_LOOK_FOR_KEYS,
host='localhost',
conn_type='ssh',
extra='{"compress" : true, "no_host_key_check" : "true", '
'"allow_host_key_change": false, "look_for_keys": false}',
)
)
db.merge_conn(
Connection(
conn_id=cls.CONN_SSH_WITH_PRIVATE_KEY_EXTRA,
host='localhost',
conn_type='ssh',
extra=json.dumps(
{
"private_key": TEST_PRIVATE_KEY,
}
),
)
)
db.merge_conn(
Connection(
conn_id=cls.CONN_SSH_WITH_PRIVATE_KEY_PASSPHRASE_EXTRA,
host='localhost',
conn_type='ssh',
extra=json.dumps(
{"private_key": TEST_ENCRYPTED_PRIVATE_KEY, "private_key_passphrase": PASSPHRASE}
),
)
)
def setUp(self):
db.merge_conn(
Connection(conn_id='kubernetes_in_cluster',
conn_type='kubernetes',
extra=json.dumps(
{'extra__kubernetes__in_cluster': True})))
db.merge_conn(
Connection(conn_id='kubernetes_kube_config',
conn_type='kubernetes',
extra=json.dumps({
'extra__kubernetes__kube_config':
'{"test": "kube"}'
})))
db.merge_conn(
Connection(conn_id='kubernetes_default_kube_config',
conn_type='kubernetes',
extra=json.dumps({})))
db.merge_conn(
Connection(conn_id='kubernetes_with_namespace',
conn_type='kubernetes',
extra=json.dumps(
{'extra__kubernetes__namespace':
'mock_namespace'})))
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_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):
db.merge_conn(Connection(conn_id='kubernetes_default', conn_type='kubernetes', extra=json.dumps({})))
db.merge_conn(
Connection(
conn_id="kubernetes_default",
conn_type="kubernetes",
extra=json.dumps({}),
)
)
db.merge_conn(
Connection(
conn_id="kubernetes_with_namespace",
conn_type="kubernetes",
extra=json.dumps({"extra__kubernetes__namespace": "mock_namespace"}),
)
)
args = {"owner": "airflow", "start_date": timezone.datetime(2020, 2, 1)}
self.dag = DAG("test_dag_id", default_args=args)
def setUp(self):
db.merge_conn(
Connection(
conn_id='telegram-webhook-without-token',
conn_type='http',
))
db.merge_conn(
Connection(
conn_id='telegram_default',
conn_type='http',
password=TELEGRAM_TOKEN,
))
db.merge_conn(
Connection(
conn_id='telegram-webhook-with-chat_id',
conn_type='http',
password=TELEGRAM_TOKEN,
host="-420913222",
))
def setUp(self):
db.merge_conn(
Connection(conn_id=self.conn_id,
host='https://api.opsgenie.com/',
password='******'))
def setUpClass(cls):
db.merge_conn(
Connection(conn_id='livy_default',
host='host',
schema='http',
port='8998'))
db.merge_conn(Connection(conn_id='default_port', host='http://host'))
db.merge_conn(Connection(conn_id='default_protocol', host='host'))
db.merge_conn(Connection(conn_id='port_set', host='host', port=1234))
db.merge_conn(
Connection(conn_id='schema_set', host='host', schema='zzz'))
db.merge_conn(
Connection(conn_id='dont_override_schema',
host='http://host',
schema='zzz'))
db.merge_conn(Connection(conn_id='missing_host', port=1234))
db.merge_conn(
Connection(conn_id='invalid_uri', uri='http://invalid_uri:4321'))
def setUp(self):
db.merge_conn(
Connection(conn_id='spark_default',
conn_type='spark',
host='yarn://yarn-master'))
def setUp(self):
db.merge_conn(Connection(conn_id='wasb_test_key', conn_type='wasb', login='******', password='******'))
self.connection_type = 'wasb'
self.connection_string_id = 'azure_test_connection_string'
self.shared_key_conn_id = 'azure_shared_key_test'
self.ad_conn_id = 'azure_AD_test'
self.sas_conn_id = 'sas_token_id'
self.public_read_conn_id = 'pub_read_id'
db.merge_conn(
Connection(
conn_id=self.public_read_conn_id,
conn_type=self.connection_type,
host='https://accountname.blob.core.windows.net',
)
)
db.merge_conn(
Connection(
conn_id=self.connection_string_id,
conn_type=self.connection_type,
extra=json.dumps({'connection_string': CONN_STRING}),
)
)
db.merge_conn(
Connection(
conn_id=self.shared_key_conn_id,
conn_type=self.connection_type,
host='https://accountname.blob.core.windows.net',
extra=json.dumps({'shared_access_key': 'token'}),
)
)
db.merge_conn(
Connection(
conn_id=self.ad_conn_id,
conn_type=self.connection_type,
extra=json.dumps(
{'tenant_id': 'token', 'application_id': 'appID', 'application_secret': "appsecret"}
),
)
)
db.merge_conn(
Connection(
conn_id=self.sas_conn_id,
conn_type=self.connection_type,
extra=json.dumps({'sas_token': 'token'}),
)
)
def test_cli_connections_import_should_not_overwrite_existing_connections(
self, mock_exists, mock_load_connections_dict, session=None):
mock_exists.return_value = True
# Add a pre-existing connection "new1"
merge_conn(
Connection(
conn_id="new1",
conn_type="mysql",
description="mysql description",
host="mysql",
login="******",
password="",
schema="airflow",
),
session=session,
)
# Sample connections to import, including a collision with "new1"
expected_connections = {
"new0": {
"conn_type": "postgres",
"description": "new0 description",
"host": "host",
"is_encrypted": False,
"is_extra_encrypted": False,
"login": "******",
"port": 5432,
"schema": "airflow",
},
"new1": {
"conn_type": "mysql",
"description": "new1 description",
"host": "host",
"is_encrypted": False,
"is_extra_encrypted": False,
"login": "******",
"port": 3306,
"schema": "airflow",
},
}
# We're not testing the behavior of load_connections_dict, assume successfully reads JSON, YAML or env
mock_load_connections_dict.return_value = expected_connections
with redirect_stdout(io.StringIO()) as stdout:
connection_command.connections_import(
self.parser.parse_args(
["connections", "import", 'sample.json']))
assert 'Could not import connection new1: connection already exists.' in stdout.getvalue(
)
# Verify that the imported connections match the expected, sample connections
current_conns = session.query(Connection).all()
comparable_attrs = [
"conn_type",
"description",
"host",
"is_encrypted",
"is_extra_encrypted",
"login",
"port",
"schema",
]
current_conns_as_dicts = {
current_conn.conn_id:
{attr: getattr(current_conn, attr)
for attr in comparable_attrs}
for current_conn in current_conns
}
assert current_conns_as_dicts['new0'] == expected_connections['new0']
# The existing connection's description should not have changed
assert current_conns_as_dicts['new1'][
'description'] == 'new1 description'
def setUp(self):
db.merge_conn(
Connection(
conn_id='jira_default', conn_type='jira',
host='https://localhost/jira/', port=443,
extra='{"verify": "False", "project": "AIRFLOW"}'))
def setUp(self, mock_batch, mock_hook):
# set up the test variable
self.test_vm_conn_id = "test_azure_batch_vm2"
self.test_cloud_conn_id = "test_azure_batch_cloud2"
self.test_account_name = "test_account_name"
self.test_account_key = 'test_account_key'
self.test_account_url = "http://test-endpoint:29000"
self.test_vm_size = "test-vm-size"
self.test_vm_publisher = "test.vm.publisher"
self.test_vm_offer = "test.vm.offer"
self.test_vm_sku = "test-sku"
self.test_cloud_os_family = "test-family"
self.test_cloud_os_version = "test-version"
self.test_node_agent_sku = "test-node-agent-sku"
# connect with vm configuration
db.merge_conn(
Connection(conn_id=self.test_vm_conn_id,
conn_type="azure_batch",
extra=json.dumps({
"account_name":
self.test_account_name,
"account_key":
self.test_account_key,
"account_url":
self.test_account_url,
"vm_publisher":
self.test_vm_publisher,
"vm_offer":
self.test_vm_offer,
"vm_sku":
self.test_vm_sku,
"node_agent_sku_id":
self.test_node_agent_sku
})))
# connect with cloud service
db.merge_conn(
Connection(conn_id=self.test_cloud_conn_id,
conn_type="azure_batch",
extra=json.dumps({
"account_name":
self.test_account_name,
"account_key":
self.test_account_key,
"account_url":
self.test_account_url,
"os_family":
self.test_cloud_os_family,
"os_version":
self.test_cloud_os_version,
"node_agent_sku_id":
self.test_node_agent_sku
})))
self.operator = AzureBatchOperator(
task_id=TASK_ID,
batch_pool_id=BATCH_POOL_ID,
batch_pool_vm_size=BATCH_VM_SIZE,
batch_job_id=BATCH_JOB_ID,
batch_task_id=BATCH_TASK_ID,
batch_task_command_line="echo hello",
azure_batch_conn_id=self.test_vm_conn_id,
timeout=2)
self.batch_client = mock_batch.return_value
self.assertEqual(self.batch_client, self.operator.hook.connection)
def setUp(self):
db.merge_conn(
Connection(conn_id='imap_default',
host='imap_server_address',
login='******',
password='******'))
def setUp(self):
db.merge_conn(
Connection(conn_id=self.conn_id,
conn_type='http',
host='https://oapi.dingtalk.com',
password='******'))
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):
db.merge_conn(Connection(conn_id=DEFAULT_CONN, conn_type='HTTP'))
def set_connection(conn: dict):
merge_conn(models.Connection(**conn))
def setUp(self):
db.merge_conn(
Connection(conn_id='imap_test',
host='base_url',
login='******',
password='******'))
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):
db.merge_conn(
Connection(
conn_id='spark_yarn_cluster',
conn_type='spark',
host='yarn://yarn-master',
extra='{"queue": "root.etl", "deploy-mode": "cluster"}',
))
db.merge_conn(
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(
Connection(conn_id='spark_default_mesos',
conn_type='spark',
host='mesos://host',
port=5050))
db.merge_conn(
Connection(
conn_id='spark_home_set',
conn_type='spark',
host='yarn://yarn-master',
extra='{"spark-home": "/opt/myspark"}',
))
db.merge_conn(
Connection(conn_id='spark_home_not_set',
conn_type='spark',
host='yarn://yarn-master'))
db.merge_conn(
Connection(
conn_id='spark_binary_set',
conn_type='spark',
host='yarn',
extra='{"spark-binary": "custom-spark-submit"}',
))
db.merge_conn(
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(
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(
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 generate_biowardrobe_workflow():
_settings.cursor.execute("select * from experimenttype limit 1")
field_names = [i[0] for i in _settings.cursor.description]
if 'workflow' not in field_names:
apply_sql_patch("labdata_alter.sql")
apply_sql_patch("experimenttype_alter.sql")
apply_sql_patch("experimenttype_patch.sql")
_template = u"""#!/usr/bin/env python3
from airflow import DAG
from biowardrobe_airflow_analysis import biowardrobe_workflow
dag = biowardrobe_workflow("{}")
"""
_settings.cursor.execute("select workflow from experimenttype")
for (workflow, ) in _settings.cursor.fetchall():
if not workflow:
continue
_filename = os.path.abspath(
os.path.join(
DAGS_FOLDER,
os.path.basename(os.path.splitext(workflow)[0]) + '.py'))
with open(_filename, 'w') as generated_workflow_stream:
generated_workflow_stream.write(_template.format(workflow))
_template = u"""#!/usr/bin/env python3
from airflow import DAG
from biowardrobe_airflow_analysis import BioWardrobeDownloadDAG, BioWardrobeDownloadTriggerDAG
d = BioWardrobeDownloadDAG
dt= BioWardrobeDownloadTriggerDAG
"""
with open(os.path.join(DAGS_FOLDER, 'biowardrobe_download.py'),
'w') as generated_workflow_stream:
generated_workflow_stream.write(_template)
_template = u"""#!/usr/bin/env python3
from airflow import DAG
from biowardrobe_airflow_analysis import BioWardrobeForceRunDAG
d = BioWardrobeForceRunDAG
"""
with open(os.path.join(DAGS_FOLDER, 'biowardrobe_force_run.py'),
'w') as generated_workflow_stream:
generated_workflow_stream.write(_template)
merge_conn(
models.Connection(conn_id='biowardrobe',
conn_type='mysql',
host=_settings.config[0],
login=_settings.config[1],
password=_settings.config[2],
schema=_settings.config[3],
extra="{\"cursor\":\"dictcursor\"}"))
try:
api_client.get_pool(name='biowardrobe_download')
except Exception as e:
api_client.create_pool(name='biowardrobe_download',
slots=5,
description="pool to download files")
try:
api_client.get_pool(name='biowardrobe_basic_analysis')
except Exception as e:
api_client.create_pool(name='biowardrobe_basic_analysis',
slots=1,
description="pool to run basic analysis")
if not conf.has_option('cwl', 'tmp_folder'):
if not os.path.exists(conf.AIRFLOW_CONFIG + '.orig'):
copyfile(conf.AIRFLOW_CONFIG, conf.AIRFLOW_CONFIG + '.orig')
with open(conf.AIRFLOW_CONFIG, 'w') as fp:
# for s in ['mesos', 'kerberos', 'celery', 'smtp', 'email', 'dask', 'ldap']:
# conf.conf.remove_section(s)
conf.conf.add_section('cwl')
conf.set('cwl', 'tmp_folder', os.path.join(AIRFLOW_HOME, 'tmp'))
conf.set('core', 'logging_level', 'WARNING')
conf.set('core', 'load_examples', 'False')
conf.set('webserver', 'dag_default_view', 'graph')
conf.set('webserver', 'dag_orientation', 'TB')
conf.set('webserver', 'web_server_worker_timeout', '120')
conf.set('scheduler', 'job_heartbeat_sec', '20')
conf.set('scheduler', 'scheduler_heartbeat_sec', '20')
conf.set('scheduler', 'min_file_process_interval', '30')
conf.conf.write(fp)
startup_scripts = [
'com.datirium.airflow-scheduler.plist',
'com.datirium.airflow-webserver.plist'
]
if platform == "darwin":
_sys_dir = os.path.expanduser('~/Library/LaunchAgents')
for scripts in startup_scripts:
with open(os.path.join(system_folder, 'macosx', scripts),
'r') as s:
data = s.read()
# OS X
dst = os.path.join(_sys_dir, scripts)
if os.path.exists(dst):
with open(dst + '.new', 'w') as w:
w.write(data.format(AIRFLOW_HOME=AIRFLOW_HOME))
else:
with open(dst, 'w') as w:
w.write(data.format(AIRFLOW_HOME=AIRFLOW_HOME))
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')
)
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"}')
)
def setUp(self):
db.merge_conn(Connection(conn_id=DEFAULT_CONN, conn_type='HTTP'))
db.merge_conn(
Connection(conn_id=TEST_CONN,
conn_type='HTTP',
host='http://localhost/api'))
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_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"}')
)
def setUp(self, mock_batch, mock_hook):
# set up the test variable
self.test_vm_conn_id = "test_azure_batch_vm2"
self.test_cloud_conn_id = "test_azure_batch_cloud2"
self.test_account_name = "test_account_name"
self.test_account_key = 'test_account_key'
self.test_account_url = "http://test-endpoint:29000"
self.test_vm_size = "test-vm-size"
self.test_vm_publisher = "test.vm.publisher"
self.test_vm_offer = "test.vm.offer"
self.test_vm_sku = "test-sku"
self.test_cloud_os_family = "test-family"
self.test_cloud_os_version = "test-version"
self.test_node_agent_sku = "test-node-agent-sku"
# connect with vm configuration
db.merge_conn(
Connection(
conn_id=self.test_vm_conn_id,
conn_type="azure_batch",
extra=json.dumps({
"account_name": self.test_account_name,
"account_key": self.test_account_key,
"account_url": self.test_account_url,
"vm_publisher": self.test_vm_publisher,
"vm_offer": self.test_vm_offer,
"vm_sku": self.test_vm_sku,
"node_agent_sku_id": self.test_node_agent_sku,
}),
))
# connect with cloud service
db.merge_conn(
Connection(
conn_id=self.test_cloud_conn_id,
conn_type="azure_batch",
extra=json.dumps({
"account_name": self.test_account_name,
"account_key": self.test_account_key,
"account_url": self.test_account_url,
"os_family": self.test_cloud_os_family,
"os_version": self.test_cloud_os_version,
"node_agent_sku_id": self.test_node_agent_sku,
}),
))
self.operator = AzureBatchOperator(
task_id=TASK_ID,
batch_pool_id=BATCH_POOL_ID,
batch_pool_vm_size=BATCH_VM_SIZE,
batch_job_id=BATCH_JOB_ID,
batch_task_id=BATCH_TASK_ID,
vm_publisher=self.test_vm_publisher,
vm_offer=self.test_vm_offer,
vm_sku=self.test_vm_sku,
vm_node_agent_sku_id=self.test_node_agent_sku,
sku_starts_with=self.test_vm_sku,
batch_task_command_line="echo hello",
azure_batch_conn_id=self.test_vm_conn_id,
target_dedicated_nodes=1,
timeout=2,
)
self.operator2_pass = AzureBatchOperator(
task_id=TASK_ID,
batch_pool_id=BATCH_POOL_ID,
batch_pool_vm_size=BATCH_VM_SIZE,
batch_job_id=BATCH_JOB_ID,
batch_task_id=BATCH_TASK_ID,
os_family="4",
batch_task_command_line="echo hello",
azure_batch_conn_id=self.test_vm_conn_id,
enable_auto_scale=True,
auto_scale_formula=FORMULA,
timeout=2,
)
self.operator2_no_formula = AzureBatchOperator(
task_id=TASK_ID,
batch_pool_id=BATCH_POOL_ID,
batch_pool_vm_size=BATCH_VM_SIZE,
batch_job_id=BATCH_JOB_ID,
batch_task_id=BATCH_TASK_ID,
os_family='4',
batch_task_command_line="echo hello",
azure_batch_conn_id=self.test_vm_conn_id,
enable_auto_scale=True,
timeout=2,
)
self.operator_fail = AzureBatchOperator(
task_id=TASK_ID,
batch_pool_id=BATCH_POOL_ID,
batch_pool_vm_size=BATCH_VM_SIZE,
batch_job_id=BATCH_JOB_ID,
batch_task_id=BATCH_TASK_ID,
os_family='4',
batch_task_command_line="echo hello",
azure_batch_conn_id=self.test_vm_conn_id,
timeout=2,
)
self.operator_mutual_exclusive = AzureBatchOperator(
task_id=TASK_ID,
batch_pool_id=BATCH_POOL_ID,
batch_pool_vm_size=BATCH_VM_SIZE,
batch_job_id=BATCH_JOB_ID,
batch_task_id=BATCH_TASK_ID,
vm_publisher=self.test_vm_publisher,
vm_offer=self.test_vm_offer,
vm_sku=self.test_vm_sku,
vm_node_agent_sku_id=self.test_node_agent_sku,
os_family="5",
sku_starts_with=self.test_vm_sku,
batch_task_command_line="echo hello",
azure_batch_conn_id=self.test_vm_conn_id,
target_dedicated_nodes=1,
timeout=2,
)
self.operator_invalid = AzureBatchOperator(
task_id=TASK_ID,
batch_pool_id=BATCH_POOL_ID,
batch_pool_vm_size=BATCH_VM_SIZE,
batch_job_id=BATCH_JOB_ID,
batch_task_id=BATCH_TASK_ID,
batch_task_command_line="echo hello",
azure_batch_conn_id=self.test_vm_conn_id,
target_dedicated_nodes=1,
timeout=2,
)
self.batch_client = mock_batch.return_value
self.mock_instance = mock_hook.return_value
self.assertEqual(self.batch_client, self.operator.hook.connection)
