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))