def snapshot(**kwargs):
"""
Query the TAP service and snapshot the OMM data.
#FIXME: The query should have some conditions to limit the data.
"""
logging.info('Populating inputs.')
query = Variable.get('omm_input_uri_query')
redis = RedisHook(redis_conn_id='redis_default')
data = {'QUERY': query, 'REQUEST': 'doQuery',
'LANG': 'ADQL', 'FORMAT': 'csv'}
http_connection = HttpHook(method='GET', http_conn_id='tap_service_host')
count = -1
with http_connection.run('/tap/sync?', parse.urlencode(data)) as response:
arr = response.text.split('\n')
count = len(arr)
logging.info('Found {} items.'.format(count))
sanitized_uris = []
for uri in arr[1:]:
if uri:
artifact_uri = uri.split('/')[1].strip()
sanitized_artifact_uri = artifact_uri.replace(
'+', '_').replace('%', '__')
logging.info('Output is {}'.format(sanitized_artifact_uri))
sanitized_uris.append(sanitized_artifact_uri)
redis.get_conn().rpush(redis_key, *sanitized_uris)
redis.get_conn().persist(redis_key)
return 'Extracted {} items'.format(len(sanitized_uris))
def query_and_extract(**context):
http_conn = HttpHook('GET', http_conn_id)
redis_conn = RedisHook(redis_conn_id)
prev_exec_date = context.get('prev_execution_date')
next_exec_date = context.get('next_execution_date')
query_meta = "SELECT fileName FROM archive_files WHERE archiveName = '{}'" \
" AND ingestDate > '{}' and ingestDate <= '{}' ORDER BY ingestDate".format(collection,
prev_exec_date.strftime(
datetime_format),
next_exec_date.strftime(datetime_format))
logging.info('Query: {}'.format(query_meta))
data = {
'QUERY': query_meta,
'LANG': 'ADQL',
'FORMAT': '{}'.format(output_format)
}
with http_conn.run('/ad/auth-sync?{}'.format(
parse.urlencode(data))) as response:
artifact_files_list = response.text.split()[1:]
if artifact_files_list:
redis_key = '{}_{}_{}.{}'.format(collection,
_to_milliseconds(prev_exec_date),
_to_milliseconds(next_exec_date),
output_format)
redis_conn.get_conn().rpush(redis_key, artifact_files_list)
return redis_key
def test_get_conn(self):
hook = RedisHook(redis_conn_id='redis_default')
self.assertEqual(hook.redis, None)
self.assertEqual(hook.host, None, 'host initialised as None.')
self.assertEqual(hook.port, None, 'port initialised as None.')
self.assertEqual(hook.password, None, 'password initialised as None.')
self.assertEqual(hook.db, None, 'db initialised as None.')
self.assertIs(hook.get_conn(), hook.get_conn(),
'Connection initialized only if None.')
def test_poke(self):
hook = RedisHook(redis_conn_id='redis_default')
redis = hook.get_conn()
redis.set('test_key', 'test_value')
self.assertTrue(self.sensor.poke(None), "Key exists on first call.")
redis.delete('test_key')
self.assertFalse(self.sensor.poke(None), "Key does NOT exists on second call.")
def test_poke_true(self):
sensor = RedisPubSubSensor(task_id='test_task',
dag=self.dag,
channels='test',
redis_conn_id='redis_default')
hook = RedisHook(redis_conn_id='redis_default')
redis = hook.get_conn()
redis.publish('test', 'message')
result = sensor.poke(self.mock_context)
self.assertFalse(result)
result = sensor.poke(self.mock_context)
self.assertTrue(result)
context_calls = [
call.xcom_push(key='message',
value={
'type': 'message',
'pattern': None,
'channel': b'test',
'data': b'message'
})
]
self.assertTrue(self.mock_context['ti'].method_calls == context_calls,
"context calls should be same")
result = sensor.poke(self.mock_context)
self.assertFalse(result)
def test_execute_hello(self):
operator = RedisPublishOperator(
task_id='test_task',
dag=self.dag,
message='hello',
channel=self.channel,
redis_conn_id='redis_default'
)
hook = RedisHook(redis_conn_id='redis_default')
pubsub = hook.get_conn().pubsub()
pubsub.subscribe(self.channel)
operator.execute(self.mock_context)
context_calls = []
self.assertTrue(self.mock_context['ti'].method_calls == context_calls, "context calls should be same")
message = pubsub.get_message()
self.assertEquals(message['type'], 'subscribe')
message = pubsub.get_message()
self.assertEquals(message['type'], 'message')
self.assertEquals(message['data'], b'hello')
pubsub.unsubscribe(self.channel)
def load_into_redis(file_path):
print('checking {}'.format(file_path))
lines = []
with open('/tmp/downloader_test1', 'rb') as f:
lines = f.readlines()
file_md5 = ''.join([line.decode('latin-1') for line in lines])
hook = RedisHook(redis_conn_id='redis_default')
hook.host = 'redis'
hook.port = 6379
redis = hook.get_conn()
try:
redis.ping()
except Exception as e:
print('could not ping redis')
print(str(e))
return
redis.set('google', file_md5)
md5_return_value = redis.get('google')
print('set redis with keyvalue pair <google> {}'.format(md5_return_value))
def test_poke(self):
hook = RedisHook(redis_conn_id='redis_default')
redis = hook.get_conn()
redis.set('test_key', 'test_value')
self.assertTrue(self.sensor.poke(None), "Key exists on first call.")
redis.delete('test_key')
self.assertFalse(self.sensor.poke(None), "Key does NOT exists on second call.")
def set_redis(key, value, **context):
redis_hook = RedisHook(redis_conn_id='redis_default')
r = redis_hook.get_conn()
r.set(key, value)
context['ti'].xcom_push('redis-test', value)
context['ti'].xcom_push('redis-branch-test', True)
def create_transform_task(redis_key):
redis_conn = RedisHook(redis_conn_id)
input_file_names = redis_conn.get_conn().lrange(redis_key, 0, -1)
child_dag_id = '_files_{}'.format(redis_key)
return SubDagOperator(subdag=sub_dag(child_dag_id, input_file_names,
redis_key),
task_id=child_dag_id,
dag=vlass_dag)
def test_real_get_and_set(self):
hook = RedisHook(redis_conn_id='redis_default')
redis = hook.get_conn()
self.assertTrue(redis.set('test_key', 'test_value'),
'Connection to Redis with SET works.')
self.assertEqual(redis.get('test_key'), b'test_value',
'Connection to Redis with GET works.')
self.assertEqual(redis.delete('test_key'), 1,
'Connection to Redis with DELETE works.')
def test_get_conn(self):
hook = RedisHook(redis_conn_id='redis_default')
self.assertEqual(hook.client, None)
self.assertEqual(
repr(hook.get_conn()),
(
'StrictRedis<ConnectionPool'
'<Connection<host=localhost,port=6379,db=0>>>'
)
)
def test_get_conn(self):
hook = RedisHook(redis_conn_id='redis_default')
self.assertEqual(hook.client, None)
self.assertEqual(
repr(hook.get_conn()),
(
'StrictRedis<ConnectionPool'
'<Connection<host=localhost,port=6379,db=0>>>'
)
)
def execute(self, context):
"""
Publish the message to Redis channel
:param context: the context object
:type context: dict
"""
redis_hook = RedisHook(redis_conn_id=self.redis_conn_id)
self.log.info('Sending messsage %s to Redis on channel %s', self.message, self.channel)
result = redis_hook.get_conn().publish(channel=self.channel, message=self.message)
self.log.info('Result of publishing %s', result)
def test_poke_true(self):
sensor = RedisPubSubSensor(
task_id='test_task',
dag=self.dag,
channels='test',
redis_conn_id='redis_default'
)
hook = RedisHook(redis_conn_id='redis_default')
redis = hook.get_conn()
redis.publish('test', 'message')
result = sensor.poke(self.mock_context)
self.assertFalse(result)
result = sensor.poke(self.mock_context)
self.assertTrue(result)
context_calls = [
call.xcom_push(
key='message',
value={'type': 'message', 'pattern': None, 'channel': b'test', 'data': b'message'})]
self.assertTrue(self.mock_context['ti'].method_calls == context_calls, "context calls should be same")
result = sensor.poke(self.mock_context)
self.assertFalse(result)
def test_real_ping(self):
hook = RedisHook(redis_conn_id='redis_default')
redis = hook.get_conn()
self.assertTrue(redis.ping(), 'Connection to Redis with PING works.')
def test_get_conn_password_stays_none(self):
hook = RedisHook(redis_conn_id='redis_default')
hook.get_conn()
self.assertEqual(hook.password, None)
def get_redis(key, **context):
redis_hook = RedisHook(redis_conn_id='redis_default')
r = redis_hook.get_conn()
return r.get(key)