def test_connection_type_not_supported(self, mock_get_connection):
conn = get_airflow_connection("NOT_SUPPORT")
mock_get_connection.return_value = conn
hook = GrpcHook("grpc_default")
with self.assertRaises(AirflowConfigException):
hook.get_conn()
def test_custom_connection_with_no_connection_func(self,
mock_get_connection):
conn = get_airflow_connection("CUSTOM")
mock_get_connection.return_value = conn
hook = GrpcHook("grpc_default")
with self.assertRaises(AirflowConfigException):
hook.get_conn()
def test_custom_connection(self, mock_get_connection):
conn = get_airflow_connection("CUSTOM")
mock_get_connection.return_value = conn
mocked_channel = self.channel_mock.return_value
hook = GrpcHook("grpc_default",
custom_connection_func=self.custom_conn_func)
channel = hook.get_conn()
self.assertEqual(channel, mocked_channel)
def test_simple_run(self, mock_get_conn, mock_get_connection):
conn = get_airflow_connection()
mock_get_connection.return_value = conn
mocked_channel = mock.Mock()
mocked_channel.__enter__ = mock.Mock(return_value=(mock.Mock(), None))
mocked_channel.__exit__ = mock.Mock(return_value=None)
hook = GrpcHook("grpc_default")
mock_get_conn.return_value = mocked_channel
response = hook.run(StubClass, "single_call", data={'data': 'hello'})
self.assertEqual(next(response), "hello")
def test_connection_with_port(self, mock_get_connection,
mock_insecure_channel):
conn = get_airflow_connection_with_port()
mock_get_connection.return_value = conn
hook = GrpcHook("grpc_default")
mocked_channel = self.channel_mock.return_value
mock_insecure_channel.return_value = mocked_channel
channel = hook.get_conn()
expected_url = "test.com:1234"
mock_insecure_channel.assert_called_once_with(expected_url)
self.assertEqual(channel, mocked_channel)
def test_no_auth_connection(self, mock_get_connection,
mock_insecure_channel):
conn = get_airflow_connection()
mock_get_connection.return_value = conn
hook = GrpcHook("grpc_default")
mocked_channel = self.channel_mock.return_value
mock_insecure_channel.return_value = mocked_channel
channel = hook.get_conn()
expected_url = "test:8080"
mock_insecure_channel.assert_called_once_with(expected_url)
assert channel == mocked_channel
def test_connection_with_interceptors(self, mock_insecure_channel,
mock_get_connection,
mock_intercept_channel):
conn = get_airflow_connection()
mock_get_connection.return_value = conn
mocked_channel = self.channel_mock.return_value
hook = GrpcHook("grpc_default", interceptors=["test1"])
mock_insecure_channel.return_value = mocked_channel
mock_intercept_channel.return_value = mocked_channel
channel = hook.get_conn()
self.assertEqual(channel, mocked_channel)
mock_intercept_channel.assert_called_once_with(mocked_channel, "test1")
def test_stream_run(self, mock_get_conn, mock_get_connection):
conn = get_airflow_connection()
mock_get_connection.return_value = conn
mocked_channel = mock.Mock()
mocked_channel.__enter__ = mock.Mock(return_value=(mock.Mock(), None))
mocked_channel.__exit__ = mock.Mock(return_value=None)
hook = GrpcHook("grpc_default")
mock_get_conn.return_value = mocked_channel
response = hook.run(StubClass,
"stream_call",
data={'data': ['hello!', "hi"]})
assert next(response) == ["streaming", "call"]
def test_connection_with_jwt(self, mock_secure_channel,
mock_google_default_auth, mock_google_cred,
mock_get_connection):
conn = get_airflow_connection(auth_type="JWT_GOOGLE")
mock_get_connection.return_value = conn
hook = GrpcHook("grpc_default")
mocked_channel = self.channel_mock.return_value
mock_secure_channel.return_value = mocked_channel
mock_credential_object = "test_credential_object"
mock_google_default_auth.return_value = (mock_credential_object, "")
mock_google_cred.return_value = mock_credential_object
channel = hook.get_conn()
expected_url = "test:8080"
mock_google_cred.assert_called_once_with(mock_credential_object)
mock_secure_channel.assert_called_once_with(mock_credential_object,
None, expected_url)
self.assertEqual(channel, mocked_channel)
def test_connection_with_tls(self, mock_secure_channel,
mock_channel_credentials, mock_get_connection,
mock_open):
conn = get_airflow_connection(auth_type="TLS",
credential_pem_file="pem")
mock_get_connection.return_value = conn
mock_open.return_value = StringIO('credential')
hook = GrpcHook("grpc_default")
mocked_channel = self.channel_mock.return_value
mock_secure_channel.return_value = mocked_channel
mock_credential_object = "test_credential_object"
mock_channel_credentials.return_value = mock_credential_object
channel = hook.get_conn()
expected_url = "test:8080"
mock_open.assert_called_once_with("pem")
mock_channel_credentials.assert_called_once_with('credential')
mock_secure_channel.assert_called_once_with(expected_url,
mock_credential_object)
self.assertEqual(channel, mocked_channel)
def test_connection_with_google_oauth(self, mock_secure_channel,
mock_google_default_auth,
mock_google_auth_request,
mock_get_connection):
conn = get_airflow_connection(auth_type="OATH_GOOGLE",
scopes="grpc,gcs")
mock_get_connection.return_value = conn
hook = GrpcHook("grpc_default")
mocked_channel = self.channel_mock.return_value
mock_secure_channel.return_value = mocked_channel
mock_credential_object = "test_credential_object"
mock_google_default_auth.return_value = (mock_credential_object, "")
mock_google_auth_request.return_value = "request"
channel = hook.get_conn()
expected_url = "test:8080"
mock_google_default_auth.assert_called_once_with(
scopes=["grpc", "gcs"])
mock_secure_channel.assert_called_once_with(mock_credential_object,
"request", expected_url)
self.assertEqual(channel, mocked_channel)
def _get_grpc_hook(self) -> GrpcHook:
return GrpcHook(
self.grpc_conn_id,
interceptors=self.interceptors,
custom_connection_func=self.custom_connection_func,
)
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))
