async def start_job_and_update_search(auth_header, subscription_search,
input_tokens, async_splunk_client,
async_kvstore_client):
if subscription_search.parent_search_key:
parent_search = await fetch_search(
auth_header, subscription_search.parent_search_key,
async_kvstore_client)
sid = parent_search.sid
LOGGER.debug("Updating sid from parent search_key=%s, sid=%s",
subscription_search.key(), sid)
else:
owner, app_name, dashboard_name = parse_dashboard_id(
subscription_search.dashboard_id)
sid = await spawn_search_job(auth_header, app_name,
subscription_search, input_tokens,
async_splunk_client)
LOGGER.debug("Updating sid from job search_key=%s, sid=%s",
subscription_search.key(), sid)
subscription_search.sid = sid
subscription_search.dispatch_state = DispatchState.NONE.value
subscription_search.done_progress = 0
if subscription_search.refresh_interval_seconds:
subscription_search.next_update_time = \
get_expiration_timestamp_str(ttl_seconds=subscription_search.refresh_interval_seconds)
LOGGER.debug("Updated subscription_search search_key=%s, sid=%s, data=%s",
subscription_search.key(), sid, subscription_search)
return True
def build_subscription_search(request_context,
dashboard_id,
search_defn,
search_query,
input_tokens,
parent_id,
search_key,
search_type,
search_type_id,
shard_id,
sid,
visualization_type=None,
trellis_enabled=False,
trellis_split_by=None):
try:
# Dashboard Definition refresh will override and refresh_interval
refresh_interval_seconds = refresh_to_seconds(search_defn.refresh)
# extract username and password
query = search_query
earliest_time = search_defn.earliest
latest_time = search_defn.latest
sample_ratio = search_defn.sample_ratio
ref = search_defn.ref
base = search_defn.base
next_update_time = get_expiration_timestamp_str(
ttl_seconds=refresh_interval_seconds)
last_update_time = get_current_timestamp_str()
# Create Search object
subscription_search = SubscriptionSearch(
_key=search_key,
dashboard_id=dashboard_id,
search_type_id=search_type_id,
search_type=search_type,
owner=request_context.current_user,
ref=ref,
base=base,
sid=sid,
query=query,
parent_search_key=parent_id,
earliest_time=earliest_time,
latest_time=latest_time,
sample_ratio=sample_ratio,
refresh_interval_seconds=refresh_interval_seconds,
next_update_time=next_update_time,
last_update_time=last_update_time,
input_tokens=input_tokens,
shard_id=shard_id,
visualization_type=visualization_type,
trellis_enabled=trellis_enabled,
trellis_split_by=trellis_split_by)
except Exception as e:
LOGGER.exception("Failed to build subscription_search")
raise e
return subscription_search
async def _batch_extend_subs(self, auth_header, subscription_ids):
LOGGER.debug("Updating subscriptions=%s", len(subscription_ids))
if not subscription_ids:
return
containedin = build_containedin_clause(constants.KEY,
list(subscription_ids))
params = {constants.QUERY: json.dumps(containedin)}
response = await self.kvstore_client.async_kvstore_get_request(
collection=constants.SUBSCRIPTIONS_COLLECTION_NAME,
params=json.dumps(params),
auth_header=auth_header)
await _ensure_valid_response(response, [HTTPStatus.OK],
"Failed to retrieve subscriptions")
subscriptions = await response.json()
subscriptions = [
Subscription.from_json(subscription_json)
for subscription_json in subscriptions
]
for subscription in subscriptions:
subscription.expired_time = get_expiration_timestamp_str(
ttl_seconds=subscription.ttl_seconds)
subscription.last_update_time = get_current_timestamp_str()
subscriptions = [
subscription.__dict__ for subscription in subscriptions
]
await self.kvstore_client.async_batch_save_request(
auth_header=auth_header,
collection=constants.SUBSCRIPTIONS_COLLECTION_NAME,
entries=subscriptions)
async def update_search_job_status(auth_header, owner, app_name, search, sid,
async_splunk_client, async_kvstore_client):
"""
Helper method to update search job status on a search collection object
:param auth_header:
:param owner:
:param app_name:
:param search:
:param sid:
:param async_splunk_client:
:param async_kvstore_client:
:return:
"""
# Loops and calls itself until job has results
loop_response = await deferred_loop(
poll_interval_seconds=DEFAULT_JOB_RESULTS_POLL_INTERVAL,
timeout_seconds=DEFAULT_JOB_RESULTS_TIMEOUT,
deferred_function=get_running_search_job_content,
# params to deferred_function
auth_header=auth_header,
owner=owner,
app_name=app_name,
sid=sid,
async_splunk_client=async_splunk_client)
# search_content check if None
if loop_response is None or loop_response.response is None:
await delete_search_job(auth_header=auth_header,
owner=owner,
app_name=app_name,
sid=sid,
async_splunk_client=async_splunk_client)
raise SpacebridgeApiRequestError(
"Search job timed out processing. sid=%s" % sid,
status_code=HTTPStatus.REQUEST_TIMEOUT)
# set valid search_job_content from loop_response
search_job_content = loop_response.response
LOGGER.info("Search Job Processed: %s", search_job_content)
# Check if search job Failed
if search_job_content.is_failed():
error_message = search_job_content.get_first_error_message(
'Search job Failed. sid=%s' % sid)
raise SpacebridgeApiRequestError(
error_message, status_code=HTTPStatus.FAILED_DEPENDENCY)
if search_job_content.is_done:
# only update the next_update_time if the refresh_interval_seconds has a value
search.next_update_time = get_expiration_timestamp_str(ttl_seconds=search.refresh_interval_seconds) \
if search.refresh_interval_seconds else ""
# update metadata on search
search.sid = sid
search.dispatch_state = search_job_content.dispatch_state
search.done_progress = search_job_content.done_progress
# Update Search Collection
response = await async_kvstore_client.async_kvstore_post_request(
collection=SEARCHES_COLLECTION_NAME,
data=jsonpickle.encode(
search, unpicklable=False), # Don't write py/object field
key_id=search.key(),
auth_header=auth_header)
if response.code != HTTPStatus.OK:
error = await response.text()
LOGGER.error(
"Unable to update sid on kvstore search. status_code=%s, error=%s, %s",
response.code, error, search)
raise SpacebridgeApiRequestError(format_splunk_error(
response.code, error),
status_code=response.code)
return search
async def create_subscription(request_context,
ttl_seconds,
search_key,
shard_id,
async_splunk_client,
async_kvstore_client,
visualization_id,
id_gen=uuid4):
"""
Create a visualization subscription object in kvstore collection [subscriptions]
:param request_context:
:param ttl_seconds:
:param search_key:
:param async_splunk_client:
:param async_kvstore_client:
:return:
"""
# extract params for subscription
subscription_id = str(id_gen())
device_id = request_context.device_id
expiration_time = get_expiration_timestamp_str(ttl_seconds=ttl_seconds)
if isinstance(request_context.auth_header, JWTAuthHeader):
LOGGER.debug("JWTAuthHeader detected. Setting session_key_type = %s",
JWT_TOKEN_TYPE)
session_type = JWT_TOKEN_TYPE
session_key = request_context.auth_header.token
else:
LOGGER.debug(
"SplunkAuthHeader detected. Setting session_key_type = %s",
SPLUNK_SESSION_TOKEN_TYPE)
session_type = SPLUNK_SESSION_TOKEN_TYPE
session_key = await get_splunk_cookie(
request_context=request_context,
async_splunk_client=async_splunk_client,
username=request_context.auth_header.username,
password=request_context.auth_header.password)
now = get_current_timestamp_str()
# Create Subscription object
subscription = Subscription(_key=subscription_id,
ttl_seconds=ttl_seconds,
device_id=device_id,
subscription_key=search_key,
expired_time=expiration_time,
shard_id=shard_id,
user=request_context.current_user,
visualization_id=visualization_id,
last_update_time=now)
auth = SubscriptionCredential(subscription_id=subscription_id,
session_key=session_key,
session_type=session_type,
shard_id=request_context.shard_id,
last_update_time=now,
_key=SUBSCRIPTION_CREDENTIAL_GLOBAL)
# TODO check if gather wworks as expected
await asyncio.gather(
_create_subscription(request_context, subscription,
async_kvstore_client),
_create_subscription_credentials(request_context, auth,
async_kvstore_client))
LOGGER.debug("Subscription created. subscription_id=%s, search_key=%s",
subscription_id, search_key)
return subscription_id
def process_subscribe_drone_mode_ipad(request_context,
client_subscription_message,
server_subscription_response,
async_client_factory):
"""
Process subscribe to drone mode ipad events for a device id
:param request_context: Used to authenticate kvstore requests
:param client_single_request: client request object protobuf
:param single_server_response: server response object protobuf
:param async_client factory: factory class used to generate kvstore and spacebridge clients
"""
LOGGER.debug('start of subscribe drone mode ipad')
async_kvstore_client = async_client_factory.kvstore_client()
LOGGER.debug('request_context=%s', request_context)
# base64 encode device id
device_id = request_context.device_id
urlsafe_b64encoded_device_id = b64_to_urlsafe_b64(device_id)
query = {
constants.QUERY: json.dumps({
constants.DEVICE_ID: device_id,
}),
constants.LIMIT: 1
}
# fetch device from registered_devices_table to get device name
response = yield async_kvstore_client.async_kvstore_get_request(
collection=constants.REGISTERED_DEVICES_COLLECTION_NAME,
params=query,
owner=request_context.current_user,
auth_header=request_context.auth_header)
yield check_and_raise_error(response, request_context, 'Fetch registered devices')
device = yield response.json()
if device:
device = device[0]
else:
raise SpacebridgeApiRequestError('Invalid device id={}'.format(device_id), status_code=http.BAD_REQUEST)
device_name = device[constants.DEVICE_NAME]
# write to ipads collection
ipad_json = {constants.DEVICE_NAME: device_name}
response = yield async_kvstore_client.async_kvstore_post_or_update_request(
constants.DRONE_MODE_IPADS_COLLECTION_NAME,
json.dumps(ipad_json),
request_context.system_auth_header,
key_id=urlsafe_b64encoded_device_id,
owner=request_context.current_user)
# construct parameters for subscription creation / updating
ttl_seconds = client_subscription_message.clientSubscribeRequest.ttlSeconds
expiration_time = get_expiration_timestamp_str(ttl_seconds=ttl_seconds)
now = get_current_timestamp_str()
params = {
constants.TTL_SECONDS: str(ttl_seconds),
constants.DEVICE_ID: device_id,
constants.LAST_UPDATE_TIME: now,
constants.EXPIRED_TIME: expiration_time,
constants.USER: request_context.current_user,
constants.SUBSCRIPTION_TYPE: constants.DRONE_MODE_IPAD
}
# delete existing subscriptions
query = {
constants.QUERY: json.dumps({
constants.DEVICE_ID: device_id,
constants.USER: request_context.current_user,
constants.SUBSCRIPTION_TYPE: constants.DRONE_MODE_IPAD
})
}
delete_response = yield async_kvstore_client.async_kvstore_delete_request(
collection=constants.SUBSCRIPTIONS_COLLECTION_NAME,
params=query,
owner=constants.NOBODY,
auth_header=request_context.auth_header)
yield check_and_raise_error(delete_response, request_context, 'Delete existing ipad subscriptions')
# create subscription
subscription_id = yield write_drone_mode_subscription(request_context,
params,
async_client_factory)
server_subscription_response.subscriptionId = subscription_id
server_subscription_response.serverSubscribeResponse.droneModeiPadSubscribe.SetInParent()
LOGGER.debug('Successfully wrote drone mode ipad subscription with id=%s', subscription_id)
def process_subscribe_drone_mode_tv(request_context,
client_subscription_message,
server_subscription_response,
async_client_factory):
"""
Process subscribe to drone mode tv events for a device id
:param request_context: Used to authenticate kvstore requests
:param client_single_request: client request object protobuf
:param single_server_response: server response object protobuf
:param async_client factory: factory class used to generate kvstore and spacebridge clients
"""
LOGGER.debug('start of subscribe drone mode tv')
async_kvstore_client = async_client_factory.kvstore_client()
# base64 encode device id
device_id = client_subscription_message.clientSubscribeRequest.droneModeTVSubscribe.deviceId
encoded_device_id = b64encode_to_str(device_id)
urlsafe_encoded_device_id = urlsafe_b64encode_to_str(device_id)
# Grab TV data
tv_data = yield get_registered_tvs(request_context.auth_header,
request_context.current_user,
async_kvstore_client)
LOGGER.debug('raw device id=%s, encoded_device_id=%s, urlsafe_encoded_device_id=%s, tv_data=%s',
device_id,
encoded_device_id,
urlsafe_encoded_device_id,
tv_data)
# validate that device is a valid apple tv registered to this user
valid_device_id = any(device.get(constants.DEVICE_ID) == encoded_device_id for device in tv_data)
if not valid_device_id:
error_message = 'Invalid device id={}'.format(encoded_device_id)
raise SpacebridgeApiRequestError(error_message, status_code=http.BAD_REQUEST)
# construct parameters for subscription creation / updating
ttl_seconds = client_subscription_message.clientSubscribeRequest.ttlSeconds
expiration_time = get_expiration_timestamp_str(ttl_seconds=ttl_seconds)
now = get_current_timestamp_str()
params = {
constants.TTL_SECONDS: str(ttl_seconds),
constants.DEVICE_ID: encoded_device_id,
constants.EXPIRED_TIME: expiration_time,
constants.USER: request_context.current_user,
constants.LAST_UPDATE_TIME: now,
constants.SUBSCRIPTION_TYPE: constants.DRONE_MODE_TV
}
# delete existing subscriptions
query = {
constants.QUERY: json.dumps({
constants.DEVICE_ID: encoded_device_id,
constants.USER: request_context.current_user,
constants.SUBSCRIPTION_TYPE: constants.DRONE_MODE_TV
})
}
delete_response = yield async_kvstore_client.async_kvstore_delete_request(
collection=constants.SUBSCRIPTIONS_COLLECTION_NAME,
params=query,
owner=constants.NOBODY,
auth_header=request_context.auth_header)
yield check_and_raise_error(delete_response, request_context, 'Delete existing tv subscriptions')
# create subscription
subscription_id = yield write_drone_mode_subscription(request_context,
params,
async_client_factory)
server_subscription_response.subscriptionId = subscription_id
server_subscription_response.serverSubscribeResponse.droneModeTVSubscribe.deviceId = device_id
LOGGER.debug('Successfully wrote drone mode tv subscription with id=%s', subscription_id)
def post(self, request):
"""
Creates/updates a subscription
Request Parameters
json version of a subscription collection object. Contains required fields
_key (required for update, must be empty for create)
ttl_seconds
subscription_key
subscription_type
device_id
shard_id
user
Example:
{
"ttl_seconds": "10",
"subscription_key": "key"
"subscription_type": "Splunk TV",
"device_id": "device_id",
"shard_id": "shard_id",
"user": "******",
}
Returns:
id of created/update subscription
Example:
{
"message": "5f04e4d1ed7a6aab3e6f2e91",
"status": 201
}
"""
session_key = request[constants.SESSION][constants.AUTHTOKEN]
payload = json.loads(request[constants.PAYLOAD])
if not payload:
return {
constants.PAYLOAD: {
'Error',
'You must provide subscription information to update/create'
},
constants.STATUS: HTTPStatus.BAD_REQUEST
}
kvstore_object = KVStore(
collection=constants.SUBSCRIPTIONS_COLLECTION_NAME,
session_key=session_key)
# Verification of provided fields
required_fields = (Subscription.TTL_SECONDS,
Subscription.SUBSCRIPTION_KEY,
Subscription.SUBSCRIPTION_TYPE,
Subscription.DEVICE_ID, Subscription.SHARD_ID,
Subscription.USER)
missing_fields = []
for field in required_fields:
if field not in payload:
missing_fields.append(field)
if missing_fields:
return {
constants.PAYLOAD: {
'Error': f'Missing field(s) {missing_fields}'
},
constants.STATUS: HTTPStatus.BAD_REQUEST
}
# Setting default fields
subscription = Subscription.from_json(payload)
subscription.expired_time = get_expiration_timestamp_str(
ttl_seconds=int(subscription.ttl_seconds))
subscription.version = constants.SUBSCRIPTION_VERSION_2
subscription.last_update_time = get_current_timestamp_str()
if subscription.key(): # update case
return update_subscription(kvstore_object=kvstore_object,
subscription=subscription)
else: # create case
return create_subscription(
kvstore_object=kvstore_object,
subscription=subscription,
session_key=session_key,
)