Python CloudSpannerHook.execute_dml方法代码示例

# 需要导入模块: from airflow.contrib.hooks.gcp_spanner_hook import CloudSpannerHook [as 别名]
# 或者: from airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook import execute_dml [as 别名]

#.........这里部分代码省略.........
        database_method.assert_called_with(database_id='database-name')
        database_update_ddl_method.assert_called_with(ddl_statements=[], operation_id="operation")
        self.assertIsNone(res)

    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')
    def test_delete_database_missing_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
        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.contrib.hooks.gcp_spanner_hook.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_with(database_id='database-name')
        database_exists_method.assert_called_with()
        database_drop_method.assert_called_with()
        self.assertIsNone(res)

    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.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_with(database_id='database-name')
        database_exists_method.assert_called_with()
        database_drop_method.assert_not_called()

    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')
    def test_execute_dml_missing_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
        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.contrib.hooks.gcp_spanner_hook.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_with(database_id='database-name')
        run_in_transaction_method.assert_called_with(mock.ANY)
        self.assertIsNone(res)

# 需要导入模块: from airflow.contrib.hooks.gcp_spanner_hook import CloudSpannerHook [as 别名]
# 或者: from airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook import execute_dml [as 别名]
class CloudSpannerInstanceDatabaseQueryOperator(BaseOperator):
    """
    Executes an arbitrary DML query (INSERT, UPDATE, DELETE).

    .. seealso::
        For more information on how to use this operator, take a look at the guide:
        :ref:`howto/operator:CloudSpannerInstanceDatabaseQueryOperator`

    :param instance_id: The Cloud Spanner instance ID.
    :type instance_id: str
    :param database_id: The Cloud Spanner database ID.
    :type database_id: str
    :param query: The query or list of queries to be executed. Can be a path to a SQL
       file.
    :type query: str or list
    :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_query_template_fields]
    template_fields = ('project_id', 'instance_id', 'database_id', 'query', 'gcp_conn_id')
    template_ext = ('.sql',)
    # [END gcp_spanner_query_template_fields]

    @apply_defaults
    def __init__(self,
                 instance_id,
                 database_id,
                 query,
                 project_id=None,
                 gcp_conn_id='google_cloud_default',
                 *args, **kwargs):
        self.instance_id = instance_id
        self.project_id = project_id
        self.database_id = database_id
        self.query = query
        self.gcp_conn_id = gcp_conn_id
        self._validate_inputs()
        self._hook = CloudSpannerHook(gcp_conn_id=gcp_conn_id)
        super(CloudSpannerInstanceDatabaseQueryOperator, self).__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")
        if not self.database_id:
            raise AirflowException("The required parameter 'database_id' "
                                   "is empty or None")
        if not self.query:
            raise AirflowException("The required parameter 'query' is empty")

    def execute(self, context):
        queries = self.query
        if isinstance(self.query, six.string_types):
            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)
        self._hook.execute_dml(project_id=self.project_id,
                               instance_id=self.instance_id,
                               database_id=self.database_id,
                               queries=queries)

    @staticmethod
    def sanitize_queries(queries):
        if len(queries) and queries[-1] == '':
            del queries[-1]