def test_get_conn_with_extra_config(self, _):
try:
hook = DockerHook(docker_conn_id='docker_with_extras',
base_url='unix://var/run/docker.sock',
version='auto')
client = hook.get_conn()
self.assertIsNotNone(client)
except Exception:
self.fail('Could not get connection from Airflow')
def test_conn_with_extra_config_passes_parameters(self, _):
hook = DockerHook(docker_conn_id='docker_with_extras',
base_url='unix://var/run/docker.sock',
version='auto')
client = hook.get_conn()
client.login.assert_called_with(username='******',
password='******',
registry='some.docker.registry.com',
reauth=False,
email='*****@*****.**')
def test_conn_with_standard_config_passes_parameters(self, _):
hook = DockerHook(docker_conn_id='docker_default',
base_url='unix://var/run/docker.sock',
version='auto')
client = hook.get_conn()
client.login.assert_called_with(username='******',
password='******',
registry='some.docker.registry.com',
reauth=True,
email=None)
def test_get_conn_override_defaults(self, docker_client_mock):
hook = DockerHook(docker_conn_id='docker_default',
base_url='https://index.docker.io/v1/',
version='1.23',
tls='someconfig')
hook.get_conn()
docker_client_mock.assert_called_with(
base_url='https://index.docker.io/v1/',
version='1.23',
tls='someconfig')
def test_get_conn_with_standard_config(self, _):
try:
hook = DockerHook(
docker_conn_id='docker_default',
base_url='unix://var/run/docker.sock',
version='auto'
)
client = hook.get_conn()
self.assertIsNotNone(client)
except:
self.fail('Could not get connection from Airflow')
def test_get_conn_with_standard_config(self, _):
try:
hook = DockerHook(
docker_conn_id='docker_default',
base_url='unix://var/run/docker.sock',
version='auto'
)
client = hook.get_conn()
self.assertIsNotNone(client)
except Exception: # pylint: disable=broad-except
self.fail('Could not get connection from Airflow')
def test_get_conn_override_defaults(self, docker_client_mock):
hook = DockerHook(
docker_conn_id='docker_default',
base_url='https://index.docker.io/v1/',
version='1.23',
tls='someconfig'
)
hook.get_conn()
docker_client_mock.assert_called_with(
base_url='https://index.docker.io/v1/',
version='1.23',
tls='someconfig'
)
def test_conn_with_extra_config_passes_parameters(self, _):
hook = DockerHook(
docker_conn_id='docker_with_extras',
base_url='unix://var/run/docker.sock',
version='auto'
)
client = hook.get_conn()
client.login.assert_called_with(
username='******',
password='******',
registry='some.docker.registry.com',
reauth=False,
email='*****@*****.**'
)
def test_conn_with_standard_config_passes_parameters(self, _):
hook = DockerHook(
docker_conn_id='docker_default',
base_url='unix://var/run/docker.sock',
version='auto'
)
client = hook.get_conn()
client.login.assert_called_with(
username='******',
password='******',
registry='some.docker.registry.com',
reauth=True,
email=None
)
def get_hook(self):
return DockerHook(
docker_conn_id=self.docker_conn_id,
base_url=self.docker_url,
version=self.api_version,
tls=self.__get_tls_config(),
)
def test_init_fails_when_no_api_version_given(self, _):
with self.assertRaises(AirflowException):
DockerHook(
docker_conn_id='docker_default',
base_url='unix://var/run/docker.sock',
tls=None
)
def test_init_fails_when_no_base_url_given(self, _):
with self.assertRaises(AirflowException):
DockerHook(
docker_conn_id='docker_default',
version='auto',
tls=None
)
def test_conn_with_broken_config_missing_host_fails(self, _):
db.merge_conn(
Connection(conn_id='docker_without_host',
conn_type='docker',
login='******',
password='******'))
with self.assertRaises(AirflowException):
DockerHook(docker_conn_id='docker_without_host',
base_url='unix://var/run/docker.sock',
version='auto')
def get_hook(self) -> DockerHook:
"""
Retrieves hook for the operator.
:return: The Docker Hook
"""
return DockerHook(docker_conn_id=self.docker_conn_id,
base_url=self.docker_url,
version=self.api_version,
tls=self.__get_tls_config())
def test_conn_with_broken_config_missing_username_fails(self, _):
db.merge_conn(
Connection(conn_id='docker_without_username',
conn_type='docker',
host='some.docker.registry.com',
password='******',
extra='{"email": "*****@*****.**"}'))
with self.assertRaises(AirflowException):
DockerHook(docker_conn_id='docker_without_username',
base_url='unix://var/run/docker.sock',
version='auto')
def get_hook(self):
try:
if self.conn_type == 'mysql':
from airflow.hooks.mysql_hook import MySqlHook
return MySqlHook(mysql_conn_id=self.conn_id)
elif self.conn_type == 'google_cloud_platform':
from airflow.contrib.hooks.bigquery_hook import BigQueryHook
return BigQueryHook(bigquery_conn_id=self.conn_id)
elif self.conn_type == 'postgres':
from airflow.hooks.postgres_hook import PostgresHook
return PostgresHook(postgres_conn_id=self.conn_id)
elif self.conn_type == 'hive_cli':
from airflow.hooks.hive_hooks import HiveCliHook
return HiveCliHook(hive_cli_conn_id=self.conn_id)
elif self.conn_type == 'presto':
from airflow.hooks.presto_hook import PrestoHook
return PrestoHook(presto_conn_id=self.conn_id)
elif self.conn_type == 'hiveserver2':
from airflow.hooks.hive_hooks import HiveServer2Hook
return HiveServer2Hook(hiveserver2_conn_id=self.conn_id)
elif self.conn_type == 'sqlite':
from airflow.hooks.sqlite_hook import SqliteHook
return SqliteHook(sqlite_conn_id=self.conn_id)
elif self.conn_type == 'jdbc':
from airflow.hooks.jdbc_hook import JdbcHook
return JdbcHook(jdbc_conn_id=self.conn_id)
elif self.conn_type == 'mssql':
from airflow.hooks.mssql_hook import MsSqlHook
return MsSqlHook(mssql_conn_id=self.conn_id)
elif self.conn_type == 'oracle':
from airflow.hooks.oracle_hook import OracleHook
return OracleHook(oracle_conn_id=self.conn_id)
elif self.conn_type == 'vertica':
from airflow.contrib.hooks.vertica_hook import VerticaHook
return VerticaHook(vertica_conn_id=self.conn_id)
elif self.conn_type == 'cloudant':
from airflow.contrib.hooks.cloudant_hook import CloudantHook
return CloudantHook(cloudant_conn_id=self.conn_id)
elif self.conn_type == 'jira':
from airflow.contrib.hooks.jira_hook import JiraHook
return JiraHook(jira_conn_id=self.conn_id)
elif self.conn_type == 'redis':
from airflow.contrib.hooks.redis_hook import RedisHook
return RedisHook(redis_conn_id=self.conn_id)
elif self.conn_type == 'wasb':
from airflow.contrib.hooks.wasb_hook import WasbHook
return WasbHook(wasb_conn_id=self.conn_id)
elif self.conn_type == 'docker':
from airflow.hooks.docker_hook import DockerHook
return DockerHook(docker_conn_id=self.conn_id)
except:
pass
def get_hook(self):
if self.conn_type == 'mysql':
from airflow.hooks.mysql_hook import MySqlHook
return MySqlHook(mysql_conn_id=self.conn_id)
elif self.conn_type == 'google_cloud_platform':
from airflow.gcp.hooks.bigquery import BigQueryHook
return BigQueryHook(bigquery_conn_id=self.conn_id)
elif self.conn_type == 'postgres':
from airflow.hooks.postgres_hook import PostgresHook
return PostgresHook(postgres_conn_id=self.conn_id)
elif self.conn_type == 'pig_cli':
from airflow.hooks.pig_hook import PigCliHook
return PigCliHook(pig_cli_conn_id=self.conn_id)
elif self.conn_type == 'hive_cli':
from airflow.hooks.hive_hooks import HiveCliHook
return HiveCliHook(hive_cli_conn_id=self.conn_id)
elif self.conn_type == 'presto':
from airflow.hooks.presto_hook import PrestoHook
return PrestoHook(presto_conn_id=self.conn_id)
elif self.conn_type == 'hiveserver2':
from airflow.hooks.hive_hooks import HiveServer2Hook
return HiveServer2Hook(hiveserver2_conn_id=self.conn_id)
elif self.conn_type == 'sqlite':
from airflow.hooks.sqlite_hook import SqliteHook
return SqliteHook(sqlite_conn_id=self.conn_id)
elif self.conn_type == 'jdbc':
from airflow.hooks.jdbc_hook import JdbcHook
return JdbcHook(jdbc_conn_id=self.conn_id)
elif self.conn_type == 'mssql':
from airflow.hooks.mssql_hook import MsSqlHook
return MsSqlHook(mssql_conn_id=self.conn_id)
elif self.conn_type == 'oracle':
from airflow.hooks.oracle_hook import OracleHook
return OracleHook(oracle_conn_id=self.conn_id)
elif self.conn_type == 'vertica':
from airflow.contrib.hooks.vertica_hook import VerticaHook
return VerticaHook(vertica_conn_id=self.conn_id)
elif self.conn_type == 'cloudant':
from airflow.contrib.hooks.cloudant_hook import CloudantHook
return CloudantHook(cloudant_conn_id=self.conn_id)
elif self.conn_type == 'jira':
from airflow.contrib.hooks.jira_hook import JiraHook
return JiraHook(jira_conn_id=self.conn_id)
elif self.conn_type == 'redis':
from airflow.contrib.hooks.redis_hook import RedisHook
return RedisHook(redis_conn_id=self.conn_id)
elif self.conn_type == 'wasb':
from airflow.contrib.hooks.wasb_hook import WasbHook
return WasbHook(wasb_conn_id=self.conn_id)
elif self.conn_type == 'docker':
from airflow.hooks.docker_hook import DockerHook
return DockerHook(docker_conn_id=self.conn_id)
elif self.conn_type == 'azure_data_lake':
from airflow.contrib.hooks.azure_data_lake_hook import AzureDataLakeHook
return AzureDataLakeHook(azure_data_lake_conn_id=self.conn_id)
elif self.conn_type == 'azure_cosmos':
from airflow.contrib.hooks.azure_cosmos_hook import AzureCosmosDBHook
return AzureCosmosDBHook(azure_cosmos_conn_id=self.conn_id)
elif self.conn_type == 'cassandra':
from airflow.contrib.hooks.cassandra_hook import CassandraHook
return CassandraHook(cassandra_conn_id=self.conn_id)
elif self.conn_type == 'mongo':
from airflow.contrib.hooks.mongo_hook import MongoHook
return MongoHook(conn_id=self.conn_id)
elif self.conn_type == 'gcpcloudsql':
from airflow.gcp.hooks.cloud_sql import CloudSqlDatabaseHook
return CloudSqlDatabaseHook(gcp_cloudsql_conn_id=self.conn_id)
elif self.conn_type == 'grpc':
from airflow.contrib.hooks.grpc_hook import GrpcHook
return GrpcHook(grpc_conn_id=self.conn_id)
raise AirflowException("Unknown hook type {}".format(self.conn_type))
