def test_instance_query_dml(self, mock_hook):
mock_hook.return_value.execute_dml.return_value = None
op = CloudSpannerInstanceDatabaseQueryOperator(project_id=PROJECT_ID,
instance_id=INSTANCE_ID,
database_id=DB_ID,
query=INSERT_QUERY,
task_id="id")
op.execute(None)
mock_hook.assert_called_once_with(gcp_conn_id="google_cloud_default")
mock_hook.return_value.execute_dml.assert_called_once_with(
PROJECT_ID, INSTANCE_ID, DB_ID, [INSERT_QUERY])
def test_instance_query_missing_project_id(self, mock_hook):
mock_hook.return_value.execute_sql.return_value = None
op = CloudSpannerInstanceDatabaseQueryOperator(instance_id=INSTANCE_ID,
database_id=DB_ID,
query=INSERT_QUERY,
task_id="id")
result = op.execute(None)
mock_hook.assert_called_once_with(gcp_conn_id="google_cloud_default")
mock_hook.return_value.execute_dml.assert_called_once_with(
project_id=None,
instance_id=INSTANCE_ID,
database_id=DB_ID,
queries=[INSERT_QUERY])
self.assertIsNone(result)
def test_instance_query_dml_list(self, mock_hook):
mock_hook.return_value.execute_dml.return_value = None
op = CloudSpannerInstanceDatabaseQueryOperator(
project_id=PROJECT_ID,
instance_id=INSTANCE_ID,
database_id=DB_ID,
query=[INSERT_QUERY, INSERT_QUERY_2],
task_id="id"
)
op.execute(None)
mock_hook.assert_called_once_with(gcp_conn_id="google_cloud_default")
mock_hook.return_value.execute_dml.assert_called_once_with(
PROJECT_ID, INSTANCE_ID, DB_ID, [INSERT_QUERY, INSERT_QUERY_2]
)
def test_instance_query_missing_project_id(self, mock_hook):
mock_hook.return_value.execute_sql.return_value = None
op = CloudSpannerInstanceDatabaseQueryOperator(
instance_id=INSTANCE_ID,
database_id=DB_ID,
query=INSERT_QUERY,
task_id="id"
)
result = op.execute(None)
mock_hook.assert_called_once_with(gcp_conn_id="google_cloud_default")
mock_hook.return_value.execute_dml.assert_called_once_with(
project_id=None, instance_id=INSTANCE_ID,
database_id=DB_ID, queries=[INSERT_QUERY]
)
self.assertIsNone(result)
def test_instance_query(self, mock_hook):
mock_hook.return_value.execute_sql.return_value = None
op = CloudSpannerInstanceDatabaseQueryOperator(
project_id=PROJECT_ID,
instance_id=INSTANCE_ID,
database_id=DB_ID,
query=INSERT_QUERY,
task_id="id"
)
result = op.execute(None) # pylint: disable=assignment-from-no-return
mock_hook.assert_called_once_with(gcp_conn_id="google_cloud_default")
mock_hook.return_value.execute_dml.assert_called_once_with(
project_id=PROJECT_ID, instance_id=INSTANCE_ID,
database_id=DB_ID,
queries=[INSERT_QUERY]
)
self.assertIsNone(result)
def test_instance_query_ex_if_param_missing(self, project_id, instance_id,
database_id, query, exp_msg,
mock_hook):
with self.assertRaises(AirflowException) as cm:
CloudSpannerInstanceDatabaseQueryOperator(project_id=project_id,
instance_id=instance_id,
database_id=database_id,
query=query,
task_id="id")
err = cm.exception
self.assertIn("The required parameter '{}' is empty".format(exp_msg),
str(err))
mock_hook.assert_not_called()
],
task_id='spanner_database_update_idempotent1_task')
spanner_database_update_idempotent2_task = CloudSpannerInstanceDatabaseUpdateOperator(
instance_id=GCP_SPANNER_INSTANCE_ID,
database_id=GCP_SPANNER_DATABASE_ID,
operation_id=OPERATION_ID,
ddl_statements=[
"CREATE TABLE my_table_unique (id INT64, name STRING(MAX)) PRIMARY KEY (id)",
],
task_id='spanner_database_update_idempotent2_task')
# [END howto_operator_spanner_database_update_idempotent]
# [START howto_operator_spanner_query]
spanner_instance_query_task = CloudSpannerInstanceDatabaseQueryOperator(
project_id=GCP_PROJECT_ID,
instance_id=GCP_SPANNER_INSTANCE_ID,
database_id=GCP_SPANNER_DATABASE_ID,
query=["DELETE FROM my_table2 WHERE true"],
task_id='spanner_instance_query_task')
spanner_instance_query_task2 = CloudSpannerInstanceDatabaseQueryOperator(
instance_id=GCP_SPANNER_INSTANCE_ID,
database_id=GCP_SPANNER_DATABASE_ID,
query=["DELETE FROM my_table2 WHERE true"],
task_id='spanner_instance_query_task2')
# [END howto_operator_spanner_query]
# [START howto_operator_spanner_database_delete]
spanner_database_delete_task = CloudSpannerInstanceDatabaseDeleteOperator(
project_id=GCP_PROJECT_ID,
instance_id=GCP_SPANNER_INSTANCE_ID,
database_id=GCP_SPANNER_DATABASE_ID,
task_id='spanner_database_delete_task')
project_id=GCP_PROJECT_ID,
instance_id=GCP_SPANNER_INSTANCE_ID,
database_id=GCP_SPANNER_DATABASE_ID,
operation_id=OPERATION_ID,
ddl_statements=[
"CREATE TABLE my_table_unique (id INT64, name STRING(MAX)) PRIMARY KEY (id)",
],
task_id='spanner_database_update_idempotent2_task'
)
# [END howto_operator_spanner_database_update_idempotent]
# [START howto_operator_spanner_query]
spanner_instance_query_task = CloudSpannerInstanceDatabaseQueryOperator(
project_id=GCP_PROJECT_ID,
instance_id=GCP_SPANNER_INSTANCE_ID,
database_id='db1',
query=["DELETE FROM my_table2 WHERE true"],
task_id='spanner_instance_query'
)
# [END howto_operator_spanner_query]
spanner_instance_query2_task = CloudSpannerInstanceDatabaseQueryOperator(
project_id=GCP_PROJECT_ID,
instance_id=GCP_SPANNER_INSTANCE_ID,
database_id='db1',
query="example_gcp_spanner.sql",
task_id='spanner_instance_query2'
)
# [START howto_operator_spanner_database_delete]
spanner_database_delete_task = CloudSpannerInstanceDatabaseDeleteOperator(
