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()
assert client is not None
except Exception: # pylint: disable=broad-except
self.fail('Could not get connection from Airflow')
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: # 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_once_with(
base_url='https://index.docker.io/v1/',
version='1.23',
tls='someconfig')
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_once_with( # pylint: disable=no-member
username='******',
password='******',
registry='some.docker.registry.com',
reauth=True,
email=None)
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_once_with( # pylint: disable=no-member
username='******',
password='******',
registry='another.docker.registry.com:9876',
reauth=False,
email='*****@*****.**')
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 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_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 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):
if self.conn_type == 'mysql':
from airflow.providers.mysql.hooks.mysql 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.providers.postgres.hooks.postgres import PostgresHook
return PostgresHook(postgres_conn_id=self.conn_id)
elif self.conn_type == 'pig_cli':
from airflow.providers.apache.pig.hooks.pig import PigCliHook
return PigCliHook(pig_cli_conn_id=self.conn_id)
elif self.conn_type == 'hive_cli':
from airflow.providers.apache.hive.hooks.hive import HiveCliHook
return HiveCliHook(hive_cli_conn_id=self.conn_id)
elif self.conn_type == 'presto':
from airflow.providers.presto.hooks.presto import PrestoHook
return PrestoHook(presto_conn_id=self.conn_id)
elif self.conn_type == 'hiveserver2':
from airflow.providers.apache.hive.hooks.hive import HiveServer2Hook
return HiveServer2Hook(hiveserver2_conn_id=self.conn_id)
elif self.conn_type == 'sqlite':
from airflow.providers.sqlite.hooks.sqlite import SqliteHook
return SqliteHook(sqlite_conn_id=self.conn_id)
elif self.conn_type == 'jdbc':
from airflow.providers.jdbc.hooks.jdbc import JdbcHook
return JdbcHook(jdbc_conn_id=self.conn_id)
elif self.conn_type == 'mssql':
from airflow.providers.microsoft.mssql.hooks.mssql import MsSqlHook
return MsSqlHook(mssql_conn_id=self.conn_id)
elif self.conn_type == 'odbc':
from airflow.providers.odbc.hooks.odbc import OdbcHook
return OdbcHook(odbc_conn_id=self.conn_id)
elif self.conn_type == 'oracle':
from airflow.providers.oracle.hooks.oracle import OracleHook
return OracleHook(oracle_conn_id=self.conn_id)
elif self.conn_type == 'vertica':
from airflow.providers.vertica.hooks.vertica import VerticaHook
return VerticaHook(vertica_conn_id=self.conn_id)
elif self.conn_type == 'cloudant':
from airflow.providers.cloudant.hooks.cloudant import CloudantHook
return CloudantHook(cloudant_conn_id=self.conn_id)
elif self.conn_type == 'jira':
from airflow.providers.jira.hooks.jira import JiraHook
return JiraHook(jira_conn_id=self.conn_id)
elif self.conn_type == 'redis':
from airflow.providers.redis.hooks.redis import RedisHook
return RedisHook(redis_conn_id=self.conn_id)
elif self.conn_type == 'wasb':
from airflow.providers.microsoft.azure.hooks.wasb import WasbHook
return WasbHook(wasb_conn_id=self.conn_id)
elif self.conn_type == 'docker':
from airflow.providers.docker.hooks.docker import DockerHook
return DockerHook(docker_conn_id=self.conn_id)
elif self.conn_type == 'azure_data_lake':
from airflow.providers.microsoft.azure.hooks.azure_data_lake import AzureDataLakeHook
return AzureDataLakeHook(azure_data_lake_conn_id=self.conn_id)
elif self.conn_type == 'azure_cosmos':
from airflow.providers.microsoft.azure.hooks.azure_cosmos import AzureCosmosDBHook
return AzureCosmosDBHook(azure_cosmos_conn_id=self.conn_id)
elif self.conn_type == 'cassandra':
from airflow.providers.apache.cassandra.hooks.cassandra import CassandraHook
return CassandraHook(cassandra_conn_id=self.conn_id)
elif self.conn_type == 'mongo':
from airflow.providers.mongo.hooks.mongo 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.providers.grpc.hooks.grpc import GrpcHook
return GrpcHook(grpc_conn_id=self.conn_id)
raise AirflowException("Unknown hook type {}".format(self.conn_type))