def execute(self, context): hook = SpannerHook(gcp_conn_id=self.gcp_conn_id) if hook.get_instance(project_id=self.project_id, instance_id=self.instance_id): return hook.delete_instance(project_id=self.project_id, instance_id=self.instance_id) else: self.log.info("Instance '%s' does not exist in project '%s'. " "Aborting delete.", self.instance_id, self.project_id) return True
def execute(self, context): hook = SpannerHook(gcp_conn_id=self.gcp_conn_id) if not hook.get_instance(project_id=self.project_id, instance_id=self.instance_id): self.log.info("Creating Cloud Spanner instance '%s'", self.instance_id) func = hook.create_instance else: self.log.info("Updating Cloud Spanner instance '%s'", self.instance_id) func = hook.update_instance func(project_id=self.project_id, instance_id=self.instance_id, configuration_name=self.configuration_name, node_count=self.node_count, display_name=self.display_name)
def execute(self, context): hook = SpannerHook(gcp_conn_id=self.gcp_conn_id) database = hook.get_database(project_id=self.project_id, instance_id=self.instance_id, database_id=self.database_id) if not database: self.log.info("The Cloud Spanner database was missing: " "'%s' in project '%s' and instance '%s'. Assuming success.", self.database_id, self.project_id, self.instance_id) return True else: return hook.delete_database(project_id=self.project_id, instance_id=self.instance_id, database_id=self.database_id)
def execute(self, context): hook = SpannerHook(gcp_conn_id=self.gcp_conn_id) queries = self.query if isinstance(self.query, str): queries = [x.strip() for x in self.query.split(';')] self.sanitize_queries(queries) self.log.info("Executing DML query(-ies) on " "projects/%s/instances/%s/databases/%s", self.project_id, self.instance_id, self.database_id) self.log.info(queries) hook.execute_dml(project_id=self.project_id, instance_id=self.instance_id, database_id=self.database_id, queries=queries)
def execute(self, context): hook = SpannerHook(gcp_conn_id=self.gcp_conn_id) if not hook.get_database(project_id=self.project_id, instance_id=self.instance_id, database_id=self.database_id): raise AirflowException("The Cloud Spanner database '{}' in project '{}' and " "instance '{}' is missing. Create the database first " "before you can update it.".format(self.database_id, self.project_id, self.instance_id)) else: return hook.update_database(project_id=self.project_id, instance_id=self.instance_id, database_id=self.database_id, ddl_statements=self.ddl_statements, operation_id=self.operation_id)
def execute(self, context): hook = SpannerHook(gcp_conn_id=self.gcp_conn_id) if not hook.get_database(project_id=self.project_id, instance_id=self.instance_id, database_id=self.database_id): self.log.info("Creating Cloud Spanner database " "'%s' in project '%s' and instance '%s'", self.database_id, self.project_id, self.instance_id) return hook.create_database(project_id=self.project_id, instance_id=self.instance_id, database_id=self.database_id, ddl_statements=self.ddl_statements) else: self.log.info("The database '%s' in project '%s' and instance '%s'" " already exists. Nothing to do. Exiting.", self.database_id, self.project_id, self.instance_id) return True
def setUp(self): with mock.patch('airflow.gcp.hooks.base.CloudBaseHook.__init__', new=mock_base_gcp_hook_no_default_project_id): self.spanner_hook_no_default_project_id = SpannerHook( gcp_conn_id='test')
class TestGcpSpannerHookNoDefaultProjectID(unittest.TestCase): def setUp(self): with mock.patch('airflow.gcp.hooks.base.CloudBaseHook.__init__', new=mock_base_gcp_hook_no_default_project_id): self.spanner_hook_no_default_project_id = SpannerHook( gcp_conn_id='test') @mock.patch("airflow.gcp.hooks.spanner.SpannerHook.client_info", new_callable=mock.PropertyMock) @mock.patch("airflow.gcp.hooks.spanner.SpannerHook._get_credentials", return_value="CREDENTIALS") @mock.patch("airflow.gcp.hooks.spanner.Client") def test_spanner_client_creation(self, mock_client, mock_get_creds, mock_client_info): result = self.spanner_hook_no_default_project_id._get_client( GCP_PROJECT_ID_HOOK_UNIT_TEST) mock_client.assert_called_once_with( project=GCP_PROJECT_ID_HOOK_UNIT_TEST, credentials=mock_get_creds.return_value, client_info=mock_client_info.return_value) self.assertEqual(mock_client.return_value, result) self.assertEqual(self.spanner_hook_no_default_project_id._client, result) @mock.patch('airflow.gcp.hooks.base.CloudBaseHook.project_id', new_callable=PropertyMock, return_value=None) @mock.patch('airflow.gcp.hooks.spanner.SpannerHook._get_client') def test_get_existing_instance_missing_project_id(self, get_client, mock_project_id): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True with self.assertRaises(AirflowException) as cm: self.spanner_hook_no_default_project_id.get_instance( instance_id=SPANNER_INSTANCE) get_client.assert_not_called() instance_method.assert_not_called() instance_exists_method.assert_not_called() err = cm.exception self.assertIn("The project id must be passed", str(err)) @mock.patch('airflow.gcp.hooks.spanner.SpannerHook._get_client') def test_get_existing_instance_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True res = self.spanner_hook_no_default_project_id.get_instance( instance_id=SPANNER_INSTANCE, project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST) get_client.assert_called_once_with(project_id='example-project') instance_method.assert_called_once_with(instance_id='instance') self.assertIsNotNone(res) @mock.patch('airflow.gcp.hooks.spanner.SpannerHook._get_client') def test_get_non_existing_instance(self, get_client): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = False res = self.spanner_hook_no_default_project_id.get_instance( instance_id=SPANNER_INSTANCE, project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST) get_client.assert_called_once_with(project_id='example-project') instance_method.assert_called_once_with(instance_id='instance') self.assertIsNone(res) @mock.patch('airflow.gcp.hooks.base.CloudBaseHook.project_id', new_callable=PropertyMock, return_value=None) @mock.patch('airflow.gcp.hooks.spanner.SpannerHook._get_client') def test_create_instance_missing_project_id(self, get_client, mock_project_id): instance_method = get_client.return_value.instance create_method = instance_method.return_value.create create_method.return_value = False with self.assertRaises(AirflowException) as cm: self.spanner_hook_no_default_project_id.create_instance( instance_id=SPANNER_INSTANCE, configuration_name=SPANNER_CONFIGURATION, node_count=1, display_name=SPANNER_DATABASE) get_client.assert_called_once_with(project_id='example-project') get_client.assert_not_called() instance_method.assert_not_called() create_method.assert_not_called() err = cm.exception self.assertIn("The project id must be passed", str(err)) @mock.patch('airflow.gcp.hooks.spanner.SpannerHook._get_client') def test_create_instance_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance create_method = instance_method.return_value.create create_method.return_value = False res = self.spanner_hook_no_default_project_id.create_instance( project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, instance_id=SPANNER_INSTANCE, configuration_name=SPANNER_CONFIGURATION, node_count=1, display_name=SPANNER_DATABASE) get_client.assert_called_once_with(project_id='example-project') instance_method.assert_called_once_with( instance_id='instance', configuration_name='configuration', display_name='database-name', node_count=1) self.assertIsNone(res) @mock.patch('airflow.gcp.hooks.base.CloudBaseHook.project_id', new_callable=PropertyMock, return_value=None) @mock.patch('airflow.gcp.hooks.spanner.SpannerHook._get_client') def test_update_instance_missing_project_id(self, get_client, mock_project_id): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True update_method = instance_method.return_value.update update_method.return_value = False with self.assertRaises(AirflowException) as cm: self.spanner_hook_no_default_project_id.update_instance( instance_id=SPANNER_INSTANCE, configuration_name=SPANNER_CONFIGURATION, node_count=2, display_name=SPANNER_DATABASE) get_client.assert_not_called() instance_method.assert_not_called() update_method.assert_not_called() err = cm.exception self.assertIn("The project id must be passed", str(err)) @mock.patch('airflow.gcp.hooks.spanner.SpannerHook._get_client') def test_update_instance_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True update_method = instance_method.return_value.update update_method.return_value = False res = self.spanner_hook_no_default_project_id.update_instance( project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, instance_id=SPANNER_INSTANCE, configuration_name=SPANNER_CONFIGURATION, node_count=2, display_name=SPANNER_DATABASE) get_client.assert_called_once_with(project_id='example-project') instance_method.assert_called_once_with( instance_id='instance', configuration_name='configuration', display_name='database-name', node_count=2) update_method.assert_called_once_with() self.assertIsNone(res) @mock.patch('airflow.gcp.hooks.base.CloudBaseHook.project_id', new_callable=PropertyMock, return_value=None) @mock.patch('airflow.gcp.hooks.spanner.SpannerHook._get_client') def test_delete_instance_missing_project_id(self, get_client, mock_project_id): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True delete_method = instance_method.return_value.delete delete_method.return_value = False with self.assertRaises(AirflowException) as cm: self.spanner_hook_no_default_project_id.delete_instance( instance_id=SPANNER_INSTANCE) get_client.assert_not_called() instance_method.assert_not_called() delete_method.assert_not_called() err = cm.exception self.assertIn("The project id must be passed", str(err)) @mock.patch('airflow.gcp.hooks.spanner.SpannerHook._get_client') def test_delete_instance_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True delete_method = instance_method.return_value.delete delete_method.return_value = False res = self.spanner_hook_no_default_project_id.delete_instance( project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, instance_id=SPANNER_INSTANCE) get_client.assert_called_once_with(project_id='example-project') instance_method.assert_called_once_with('instance') delete_method.assert_called_once_with() self.assertIsNone(res) @mock.patch('airflow.gcp.hooks.base.CloudBaseHook.project_id', new_callable=PropertyMock, return_value=None) @mock.patch('airflow.gcp.hooks.spanner.SpannerHook._get_client') def test_get_database_missing_project_id(self, get_client, mock_project_id): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True database_method = instance_method.return_value.database database_exists_method = instance_method.return_value.exists database_exists_method.return_value = True with self.assertRaises(AirflowException) as cm: self.spanner_hook_no_default_project_id.get_database( instance_id=SPANNER_INSTANCE, database_id=SPANNER_DATABASE) get_client.assert_not_called() instance_method.assert_not_called() database_method.assert_not_called() err = cm.exception self.assertIn("The project id must be passed", str(err)) @mock.patch('airflow.gcp.hooks.spanner.SpannerHook._get_client') def test_get_database_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True database_method = instance_method.return_value.database database_exists_method = instance_method.return_value.exists database_exists_method.return_value = True res = self.spanner_hook_no_default_project_id.get_database( project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, instance_id=SPANNER_INSTANCE, database_id=SPANNER_DATABASE) get_client.assert_called_once_with(project_id='example-project') instance_method.assert_called_once_with(instance_id='instance') database_method.assert_called_once_with(database_id='database-name') database_exists_method.assert_called_once_with() self.assertIsNotNone(res) @mock.patch('airflow.gcp.hooks.base.CloudBaseHook.project_id', new_callable=PropertyMock, return_value=None) @mock.patch('airflow.gcp.hooks.spanner.SpannerHook._get_client') def test_create_database_missing_project_id(self, get_client, mock_project_id): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True database_method = instance_method.return_value.database database_create_method = database_method.return_value.create with self.assertRaises(AirflowException) as cm: self.spanner_hook_no_default_project_id.create_database( instance_id=SPANNER_INSTANCE, database_id=SPANNER_DATABASE, ddl_statements=[]) get_client.assert_not_called() instance_method.assert_not_called() database_method.assert_not_called() database_create_method.assert_not_called() err = cm.exception self.assertIn("The project id must be passed", str(err)) @mock.patch('airflow.gcp.hooks.spanner.SpannerHook._get_client') def test_create_database_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True database_method = instance_method.return_value.database database_create_method = database_method.return_value.create res = self.spanner_hook_no_default_project_id.create_database( project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, instance_id=SPANNER_INSTANCE, database_id=SPANNER_DATABASE, ddl_statements=[]) get_client.assert_called_once_with(project_id='example-project') instance_method.assert_called_once_with(instance_id='instance') database_method.assert_called_once_with(database_id='database-name', ddl_statements=[]) database_create_method.assert_called_once_with() self.assertIsNone(res) @mock.patch('airflow.gcp.hooks.base.CloudBaseHook.project_id', new_callable=PropertyMock, return_value=None) @mock.patch('airflow.gcp.hooks.spanner.SpannerHook._get_client') def test_update_database_missing_project_id(self, get_client, mock_project_id): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True database_method = instance_method.return_value.database database_update_method = database_method.return_value.update with self.assertRaises(AirflowException) as cm: self.spanner_hook_no_default_project_id.update_database( instance_id=SPANNER_INSTANCE, database_id=SPANNER_DATABASE, ddl_statements=[]) get_client.assert_not_called() instance_method.assert_not_called() database_method.assert_not_called() database_update_method.assert_not_called() err = cm.exception self.assertIn("The project id must be passed", str(err)) @mock.patch('airflow.gcp.hooks.spanner.SpannerHook._get_client') def test_update_database_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True database_method = instance_method.return_value.database database_update_ddl_method = database_method.return_value.update_ddl res = self.spanner_hook_no_default_project_id.update_database( project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, instance_id=SPANNER_INSTANCE, database_id=SPANNER_DATABASE, ddl_statements=[]) get_client.assert_called_once_with(project_id='example-project') instance_method.assert_called_once_with(instance_id='instance') database_method.assert_called_once_with(database_id='database-name') database_update_ddl_method.assert_called_once_with(ddl_statements=[], operation_id=None) self.assertIsNone(res) @mock.patch('airflow.gcp.hooks.spanner.SpannerHook._get_client') def test_update_database_overridden_project_id_and_operation( self, get_client): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True database_method = instance_method.return_value.database database_update_ddl_method = database_method.return_value.update_ddl res = self.spanner_hook_no_default_project_id.update_database( project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, instance_id=SPANNER_INSTANCE, database_id=SPANNER_DATABASE, operation_id="operation", ddl_statements=[]) get_client.assert_called_once_with(project_id='example-project') instance_method.assert_called_once_with(instance_id='instance') database_method.assert_called_once_with(database_id='database-name') database_update_ddl_method.assert_called_once_with( ddl_statements=[], operation_id="operation") self.assertIsNone(res) @mock.patch('airflow.gcp.hooks.base.CloudBaseHook.project_id', new_callable=PropertyMock, return_value=None) @mock.patch('airflow.gcp.hooks.spanner.SpannerHook._get_client') def test_delete_database_missing_project_id(self, get_client, mock_project_id): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True database_method = instance_method.return_value.database database_drop_method = database_method.return_value.drop database_exists_method = database_method.return_value.exists database_exists_method.return_value = True with self.assertRaises(AirflowException) as cm: self.spanner_hook_no_default_project_id.delete_database( instance_id=SPANNER_INSTANCE, database_id=SPANNER_DATABASE) get_client.assert_not_called() instance_method.assert_not_called() database_method.assert_not_called() database_exists_method.assert_not_called() database_drop_method.assert_not_called() err = cm.exception self.assertIn("The project id must be passed", str(err)) @mock.patch('airflow.gcp.hooks.spanner.SpannerHook._get_client') def test_delete_database_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True database_method = instance_method.return_value.database database_drop_method = database_method.return_value.drop database_exists_method = database_method.return_value.exists database_exists_method.return_value = True res = self.spanner_hook_no_default_project_id.delete_database( project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, instance_id=SPANNER_INSTANCE, database_id=SPANNER_DATABASE) get_client.assert_called_once_with(project_id='example-project') instance_method.assert_called_once_with(instance_id='instance') database_method.assert_called_once_with(database_id='database-name') database_exists_method.assert_called_once_with() database_drop_method.assert_called_once_with() self.assertTrue(res) @mock.patch('airflow.gcp.hooks.spanner.SpannerHook._get_client') def test_delete_database_missing_database(self, get_client): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True database_method = instance_method.return_value.database database_drop_method = database_method.return_value.drop database_exists_method = database_method.return_value.exists database_exists_method.return_value = False self.spanner_hook_no_default_project_id.delete_database( project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, instance_id=SPANNER_INSTANCE, database_id=SPANNER_DATABASE) get_client.assert_called_once_with(project_id='example-project') instance_method.assert_called_once_with(instance_id='instance') database_method.assert_called_once_with(database_id='database-name') database_exists_method.assert_called_once_with() database_drop_method.assert_not_called() @mock.patch('airflow.gcp.hooks.base.CloudBaseHook.project_id', new_callable=PropertyMock, return_value=None) @mock.patch('airflow.gcp.hooks.spanner.SpannerHook._get_client') def test_execute_dml_missing_project_id(self, get_client, mock_project_id): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True database_method = instance_method.return_value.database run_in_transaction_method = database_method.return_value.run_in_transaction with self.assertRaises(AirflowException) as cm: self.spanner_hook_no_default_project_id.execute_dml( instance_id=SPANNER_INSTANCE, database_id=SPANNER_DATABASE, queries='') get_client.assert_not_called() instance_method.assert_not_called() database_method.assert_not_called() run_in_transaction_method.assert_not_called() err = cm.exception self.assertIn("The project id must be passed", str(err)) @mock.patch('airflow.gcp.hooks.spanner.SpannerHook._get_client') def test_execute_dml_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True database_method = instance_method.return_value.database run_in_transaction_method = database_method.return_value.run_in_transaction res = self.spanner_hook_no_default_project_id.execute_dml( project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, instance_id=SPANNER_INSTANCE, database_id=SPANNER_DATABASE, queries='') get_client.assert_called_once_with(project_id='example-project') instance_method.assert_called_once_with(instance_id='instance') database_method.assert_called_once_with(database_id='database-name') run_in_transaction_method.assert_called_once_with(mock.ANY) self.assertIsNone(res)