def test_poke(self, mock_hook):
sensor = CassandraRecordSensor(task_id='test_task',
cassandra_conn_id='cassandra_default',
table='t',
keys={'foo': 'bar'})
sensor.poke(None)
mock_hook.return_value.record_exists.assert_called_once_with(
't', {'foo': 'bar'})
def setUp(self):
args = {'owner': 'airflow', 'start_date': DEFAULT_DATE}
self.dag = DAG('test_dag_id', default_args=args)
self.sensor = CassandraRecordSensor(
task_id='test_task',
cassandra_conn_id='cassandra_default',
dag=self.dag,
table='t',
keys={'foo': 'bar'})
def test_poke_should_not_fail_with_empty_keys(self, mock_hook):
sensor = CassandraRecordSensor(
task_id='test_task',
cassandra_conn_id=TEST_CASSANDRA_CONN_ID,
table=TEST_CASSANDRA_TABLE,
keys=None,
)
exists = sensor.poke(dict())
self.assertTrue(exists)
mock_hook.return_value.record_exists.assert_called_once_with(
TEST_CASSANDRA_TABLE, None)
mock_hook.assert_called_once_with(TEST_CASSANDRA_CONN_ID)
def test_poke(self, mock_hook):
sensor = CassandraRecordSensor(
task_id='test_task',
cassandra_conn_id=TEST_CASSANDRA_CONN_ID,
table=TEST_CASSANDRA_TABLE,
keys=TEST_CASSANDRA_KEY,
)
exists = sensor.poke(dict())
self.assertTrue(exists)
mock_hook.return_value.record_exists.assert_called_once_with(
TEST_CASSANDRA_TABLE, TEST_CASSANDRA_KEY)
mock_hook.assert_called_once_with(TEST_CASSANDRA_CONN_ID)
class TestCassandraRecordSensor(unittest.TestCase):
def setUp(self):
args = {'owner': 'airflow', 'start_date': DEFAULT_DATE}
self.dag = DAG('test_dag_id', default_args=args)
self.sensor = CassandraRecordSensor(
task_id='test_task',
cassandra_conn_id='cassandra_default',
dag=self.dag,
table='t',
keys={'foo': 'bar'})
@patch("airflow.contrib.hooks.cassandra_hook.CassandraHook.record_exists")
def test_poke(self, mock_record_exists):
self.sensor.poke(None)
mock_record_exists.assert_called_once_with('t', {'foo': 'bar'})
def test_poke_should_return_false_for_non_existing_table(self, mock_hook):
mock_hook.return_value.record_exists.return_value = False
sensor = CassandraRecordSensor(
task_id='test_task',
cassandra_conn_id=TEST_CASSANDRA_CONN_ID,
table=TEST_CASSANDRA_TABLE,
keys=TEST_CASSANDRA_KEY,
)
exists = sensor.poke(dict())
assert not exists
mock_hook.return_value.record_exists.assert_called_once_with(
TEST_CASSANDRA_TABLE, TEST_CASSANDRA_KEY)
mock_hook.assert_called_once_with(TEST_CASSANDRA_CONN_ID)
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
"""
Example Airflow DAG to check if a Cassandra Table and a Records exists
or not using `CassandraTableSensor` and `CassandraRecordSensor`.
"""
from datetime import datetime
from airflow.models import DAG
from airflow.providers.apache.cassandra.sensors.record import CassandraRecordSensor
from airflow.providers.apache.cassandra.sensors.table import CassandraTableSensor
# [START howto_operator_cassandra_sensors]
with DAG(
dag_id='example_cassandra_operator',
schedule_interval=None,
start_date=datetime(2021, 1, 1),
default_args={'table': 'keyspace_name.table_name'},
catchup=False,
tags=['example'],
) as dag:
table_sensor = CassandraTableSensor(task_id="cassandra_table_sensor")
record_sensor = CassandraRecordSensor(task_id="cassandra_record_sensor",
keys={
"p1": "v1",
"p2": "v2"
})
# [END howto_operator_cassandra_sensors]
args = {
'owner': 'Airflow',
}
with DAG(
dag_id='example_cassandra_operator',
default_args=args,
schedule_interval=None,
start_date=days_ago(2),
tags=['example'],
) as dag:
# [START howto_operator_cassandra_table_sensor]
table_sensor = CassandraTableSensor(
task_id="cassandra_table_sensor",
cassandra_conn_id="cassandra_default",
table="keyspace_name.table_name",
)
# [END howto_operator_cassandra_table_sensor]
# [START howto_operator_cassandra_record_sensor]
record_sensor = CassandraRecordSensor(
task_id="cassandra_record_sensor",
cassandra_conn_id="cassandra_default",
table="keyspace_name.table_name",
keys={
"p1": "v1",
"p2": "v2"
},
)
# [END howto_operator_cassandra_record_sensor]