def execute(self, context: 'Context') -> Optional[bool]: hook = SpannerHook( gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, ) 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.providers.google.common.hooks.base_google.GoogleBaseHook.__init__', new=mock_base_gcp_hook_no_default_project_id, ): self.spanner_hook_no_default_project_id = SpannerHook( gcp_conn_id='test')
def execute(self, context: 'Context') -> None: hook = SpannerHook( gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, ) if not hook.get_database(project_id=self.project_id, instance_id=self.instance_id, database_id=self.database_id): raise AirflowException( f"The Cloud Spanner database '{self.database_id}' in project '{self.project_id}' " f"and instance '{self.instance_id}' is missing. " f"Create the database first before you can update it.") else: SpannerDatabaseLink.persist( context=context, task_instance=self, instance_id=self.instance_id, database_id=self.database_id, project_id=self.project_id or hook.project_id, ) 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: 'Context'): hook = SpannerHook( gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, ) if isinstance(self.query, str): queries = [x.strip() for x in self.query.split(';')] self.sanitize_queries(queries) else: queries = self.query 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, ) SpannerDatabaseLink.persist( context=context, task_instance=self, instance_id=self.instance_id, database_id=self.database_id, project_id=self.project_id or hook.project_id, )
def execute(self, context: 'Context') -> None: hook = SpannerHook( gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, ) 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, ) SpannerInstanceLink.persist( context=context, task_instance=self, instance_id=self.instance_id, project_id=self.project_id or hook.project_id, )
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: 'Context') -> Optional[bool]: hook = SpannerHook( gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, ) 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) 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
class TestGcpSpannerHookNoDefaultProjectID(unittest.TestCase): def setUp(self): with mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.__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.providers.google.cloud.hooks.spanner.SpannerHook.client_info", new_callable=mock.PropertyMock ) @mock.patch( "airflow.providers.google.cloud.hooks.spanner.SpannerHook._get_credentials", return_value="CREDENTIALS", ) @mock.patch("airflow.providers.google.cloud.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.providers.google.cloud.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.providers.google.cloud.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.providers.google.cloud.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.providers.google.cloud.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.providers.google.cloud.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.providers.google.cloud.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.providers.google.cloud.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.providers.google.cloud.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.providers.google.cloud.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.providers.google.cloud.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.providers.google.cloud.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.providers.google.cloud.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)