def execute(self, context): hook = CloudSpannerHook(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
class CloudSpannerInstanceDeleteOperator(BaseOperator): """ Deletes a Cloud Spanner instance. If an instance does not exist, no action is taken and the operator succeeds. .. seealso:: For more information on how to use this operator, take a look at the guide: :ref:`howto/operator:CloudSpannerInstanceDeleteOperator` :param instance_id: The Cloud Spanner instance ID. :type instance_id: str :param project_id: Optional, the ID of the project that owns the Cloud Spanner Database. If set to None or missing, the default project_id from the GCP connection is used. :type project_id: str :param gcp_conn_id: The connection ID used to connect to Google Cloud Platform. :type gcp_conn_id: str """ # [START gcp_spanner_delete_template_fields] template_fields = ('project_id', 'instance_id', 'gcp_conn_id') # [END gcp_spanner_delete_template_fields] @apply_defaults def __init__(self, instance_id, project_id=None, gcp_conn_id='google_cloud_default', *args, **kwargs): self.instance_id = instance_id self.project_id = project_id self.gcp_conn_id = gcp_conn_id self._validate_inputs() self._hook = CloudSpannerHook(gcp_conn_id=gcp_conn_id) super().__init__(*args, **kwargs) def _validate_inputs(self): if self.project_id == '': raise AirflowException( "The required parameter 'project_id' is empty") if not self.instance_id: raise AirflowException("The required parameter 'instance_id' " "is empty or None") def execute(self, context): if self._hook.get_instance(project_id=self.project_id, instance_id=self.instance_id): return self._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
class TestGcpSpannerHookNoDefaultProjectID(unittest.TestCase): def setUp(self): with mock.patch('airflow.gcp.hooks.base.GoogleCloudBaseHook.__init__', new=mock_base_gcp_hook_no_default_project_id): self.spanner_hook_no_default_project_id = CloudSpannerHook( gcp_conn_id='test') @mock.patch("airflow.gcp.hooks.spanner.CloudSpannerHook.client_info", new_callable=mock.PropertyMock) @mock.patch("airflow.gcp.hooks.spanner.CloudSpannerHook._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.GoogleCloudBaseHook.project_id', new_callable=PropertyMock, return_value=None) @mock.patch('airflow.gcp.hooks.spanner.CloudSpannerHook._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.CloudSpannerHook._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.CloudSpannerHook._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.GoogleCloudBaseHook.project_id', new_callable=PropertyMock, return_value=None) @mock.patch('airflow.gcp.hooks.spanner.CloudSpannerHook._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.CloudSpannerHook._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.GoogleCloudBaseHook.project_id', new_callable=PropertyMock, return_value=None) @mock.patch('airflow.gcp.hooks.spanner.CloudSpannerHook._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.CloudSpannerHook._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.GoogleCloudBaseHook.project_id', new_callable=PropertyMock, return_value=None) @mock.patch('airflow.gcp.hooks.spanner.CloudSpannerHook._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.CloudSpannerHook._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.GoogleCloudBaseHook.project_id', new_callable=PropertyMock, return_value=None) @mock.patch('airflow.gcp.hooks.spanner.CloudSpannerHook._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.CloudSpannerHook._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.GoogleCloudBaseHook.project_id', new_callable=PropertyMock, return_value=None) @mock.patch('airflow.gcp.hooks.spanner.CloudSpannerHook._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.CloudSpannerHook._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.GoogleCloudBaseHook.project_id', new_callable=PropertyMock, return_value=None) @mock.patch('airflow.gcp.hooks.spanner.CloudSpannerHook._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.CloudSpannerHook._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.CloudSpannerHook._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.GoogleCloudBaseHook.project_id', new_callable=PropertyMock, return_value=None) @mock.patch('airflow.gcp.hooks.spanner.CloudSpannerHook._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.CloudSpannerHook._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.CloudSpannerHook._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.GoogleCloudBaseHook.project_id', new_callable=PropertyMock, return_value=None) @mock.patch('airflow.gcp.hooks.spanner.CloudSpannerHook._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.CloudSpannerHook._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)
class TestGcpSpannerHookDefaultProjectId(unittest.TestCase): def setUp(self): with mock.patch( 'airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook.__init__', new=mock_base_gcp_hook_default_project_id): self.spanner_hook_default_project_id = CloudSpannerHook( gcp_conn_id='test') @mock.patch('airflow.gcp.hooks.spanner.CloudSpannerHook._get_client') def test_get_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 = True res = self.spanner_hook_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.CloudSpannerHook._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_default_project_id.get_instance( instance_id=SPANNER_INSTANCE, project_id='new-project') get_client.assert_called_once_with(project_id='new-project') instance_method.assert_called_once_with(instance_id='instance') self.assertIsNotNone(res) @mock.patch('airflow.gcp.hooks.spanner.CloudSpannerHook._get_client') def test_create_instance(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_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') 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.spanner.CloudSpannerHook._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_default_project_id.create_instance( project_id='new-project', instance_id=SPANNER_INSTANCE, configuration_name=SPANNER_CONFIGURATION, node_count=1, display_name=SPANNER_DATABASE) get_client.assert_called_once_with(project_id='new-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.spanner.CloudSpannerHook._get_client') def test_update_instance(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_default_project_id.update_instance( 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.spanner.CloudSpannerHook._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_default_project_id.update_instance( project_id='new-project', instance_id=SPANNER_INSTANCE, configuration_name=SPANNER_CONFIGURATION, node_count=2, display_name=SPANNER_DATABASE) get_client.assert_called_once_with(project_id='new-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.spanner.CloudSpannerHook._get_client') def test_delete_instance(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_default_project_id.delete_instance( 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.spanner.CloudSpannerHook._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_default_project_id.delete_instance( project_id='new-project', instance_id=SPANNER_INSTANCE) get_client.assert_called_once_with(project_id='new-project') instance_method.assert_called_once_with('instance') delete_method.assert_called_once_with() self.assertIsNone(res) @mock.patch('airflow.gcp.hooks.spanner.CloudSpannerHook._get_client') def test_get_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_exists_method = instance_method.return_value.exists database_exists_method.return_value = True res = self.spanner_hook_default_project_id.get_database( 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.spanner.CloudSpannerHook._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_default_project_id.get_database( project_id='new-project', instance_id=SPANNER_INSTANCE, database_id=SPANNER_DATABASE) get_client.assert_called_once_with(project_id='new-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.spanner.CloudSpannerHook._get_client') def test_create_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_create_method = database_method.return_value.create res = self.spanner_hook_default_project_id.create_database( 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.spanner.CloudSpannerHook._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_default_project_id.create_database( project_id='new-project', instance_id=SPANNER_INSTANCE, database_id=SPANNER_DATABASE, ddl_statements=[]) get_client.assert_called_once_with(project_id='new-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.spanner.CloudSpannerHook._get_client') def test_update_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_update_ddl_method = database_method.return_value.update_ddl res = self.spanner_hook_default_project_id.update_database( 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.CloudSpannerHook._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_default_project_id.update_database( project_id='new-project', instance_id=SPANNER_INSTANCE, database_id=SPANNER_DATABASE, ddl_statements=[]) get_client.assert_called_once_with(project_id='new-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.CloudSpannerHook._get_client') def test_delete_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 = True res = self.spanner_hook_default_project_id.delete_database( 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.CloudSpannerHook._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_default_project_id.delete_database( project_id='new-project', instance_id=SPANNER_INSTANCE, database_id=SPANNER_DATABASE) get_client.assert_called_once_with(project_id='new-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.CloudSpannerHook._get_client') def test_execute_dml(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_default_project_id.execute_dml( 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) @mock.patch('airflow.gcp.hooks.spanner.CloudSpannerHook._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_default_project_id.execute_dml( project_id='new-project', instance_id=SPANNER_INSTANCE, database_id=SPANNER_DATABASE, queries='') get_client.assert_called_once_with(project_id='new-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)