def compare_conditions(
self,
subscription: SubscriptionData,
exception: Optional[Type[Exception]],
aggregate: str,
value: Union[int, float],
) -> None:
timer = Timer("test")
if exception is not None:
with pytest.raises(exception):
request = subscription.build_request(
self.dataset,
datetime.utcnow(),
100,
timer,
)
parse_and_run_query(self.dataset, request, timer)
return
request = subscription.build_request(
self.dataset,
datetime.utcnow(),
100,
timer,
)
result = parse_and_run_query(self.dataset, request, timer)
assert result.result["data"][0][aggregate] == value
def test_invalid_time_window(self):
creator = SubscriptionCreator(self.dataset)
with raises(InvalidSubscriptionError):
creator.create(
SubscriptionData(
123,
[["platfo", "IN", ["a"]]],
[["count()", "", "count"]],
timedelta(),
timedelta(minutes=1),
),
self.timer,
)
with raises(InvalidSubscriptionError):
creator.create(
SubscriptionData(
123,
[["platfo", "IN", ["a"]]],
[["count()", "", "count"]],
timedelta(hours=48),
timedelta(minutes=1),
),
self.timer,
)
def test_subscription_task_result_encoder() -> None:
codec = SubscriptionTaskResultEncoder()
timestamp = datetime.now()
entity_subscription = EventsSubscription(data_dict={})
subscription_data = SubscriptionData(
project_id=1,
query="MATCH (events) SELECT count() AS count",
time_window_sec=60,
resolution_sec=60,
entity_subscription=entity_subscription,
)
# XXX: This seems way too coupled to the dataset.
request = subscription_data.build_request(get_dataset("events"), timestamp,
None, Timer("timer"))
result: Result = {
"meta": [{
"type": "UInt64",
"name": "count"
}],
"data": [{
"count": 1
}],
}
task_result = SubscriptionTaskResult(
ScheduledSubscriptionTask(
timestamp,
SubscriptionWithMetadata(
EntityKey.EVENTS,
Subscription(
SubscriptionIdentifier(PartitionId(1), uuid.uuid1()),
subscription_data,
),
5,
),
),
(request, result),
)
message = codec.encode(task_result)
data = json.loads(message.value.decode("utf-8"))
assert data["version"] == 3
payload = data["payload"]
assert payload["subscription_id"] == str(
task_result.task.task.subscription.identifier)
assert payload["request"] == request.original_body
assert payload["result"] == result
assert payload["timestamp"] == task_result.task.timestamp.isoformat()
assert payload["entity"] == EntityKey.EVENTS.value
def create(self, data: SubscriptionData, timer: Timer) -> SubscriptionIdentifier:
data.validate()
self._test_request(data, timer)
identifier = SubscriptionIdentifier(
self.__partitioner.build_partition_id(data),
uuid1(),
)
RedisSubscriptionDataStore(
redis_client, self.entity_key, identifier.partition
).create(
identifier.uuid,
data,
)
return identifier
def test(self) -> None:
creator = SubscriptionCreator(self.dataset, EntityKey.EVENTS)
subscription = SubscriptionData(
project_id=1,
query="MATCH (events) SELECT count() AS count",
time_window_sec=10 * 60,
resolution_sec=60,
entity_subscription=create_entity_subscription(),
)
identifier = creator.create(subscription, Timer("test"))
assert (cast(
List[Tuple[UUID, SubscriptionData]],
RedisSubscriptionDataStore(
redis_client,
self.entity_key,
identifier.partition,
).all(),
)[0][1] == subscription)
SubscriptionDeleter(self.entity_key,
identifier.partition).delete(identifier.uuid)
assert (RedisSubscriptionDataStore(
redis_client,
self.entity_key,
identifier.partition,
).all() == [])
def decode(self, value: bytes) -> SubscriptionData:
try:
data = json.loads(value.decode("utf-8"))
except json.JSONDecodeError:
raise InvalidQueryException("Invalid JSON")
return SubscriptionData.from_dict(data, self.entity_key)
def test_encode(self):
result = SubscriptionResult(
ScheduledTask(
datetime.now(),
Subscription(
SubscriptionIdentifier(PartitionId(1), uuid.uuid1()),
SubscriptionData(
1,
[],
[["count()", "", "count"]],
timedelta(minutes=1),
timedelta(minutes=1),
),
),
),
{"data": {
"count": 100
}},
)
codec = SubscriptionResultCodec()
message = codec.encode(result)
data = json.loads(message.value.decode("utf-8"))
assert data["version"] == 1
payload = data["payload"]
assert payload["subscription_id"] == str(result.task.task.identifier)
assert payload["values"] == result.result
assert payload["timestamp"] == result.task.timestamp.isoformat()
def test(self):
executor = SubscriptionExecutor(self.dataset, ThreadPoolExecutor(),
DummyMetricsBackend(strict=True))
subscription = Subscription(
SubscriptionIdentifier(PartitionId(0), uuid1()),
SubscriptionData(
project_id=self.project_id,
conditions=[["platform", "IN", ["a"]]],
aggregations=[["count()", "", "count"]],
time_window=timedelta(minutes=500),
resolution=timedelta(minutes=1),
),
)
now = datetime.utcnow()
tick = Tick(
offsets=Interval(1, 2),
timestamps=Interval(now - timedelta(minutes=1), now),
)
result = executor.execute(ScheduledTask(now, subscription),
tick).result()
assert result["data"][0]["count"] == 10
result = executor.execute(
ScheduledTask(
now + timedelta(minutes=self.minutes) +
subscription.data.time_window,
subscription,
),
tick,
).result()
assert result["data"][0]["count"] == 0
def test_subscription_task_result_encoder() -> None:
codec = SubscriptionTaskResultEncoder()
timestamp = datetime.now()
subscription_data = SubscriptionData(
1,
[],
[["count()", "", "count"]],
timedelta(minutes=1),
timedelta(minutes=1),
)
# XXX: This seems way too coupled to the dataset.
request = subscription_data.build_request(get_dataset("events"), timestamp,
None, Timer("timer"))
result: Result = {
"meta": [{
"type": "UInt64",
"name": "count"
}],
"data": [{
"count": 1
}],
}
task_result = SubscriptionTaskResult(
ScheduledTask(
timestamp,
Subscription(
SubscriptionIdentifier(PartitionId(1), uuid.uuid1()),
subscription_data,
),
),
(request, result),
)
message = codec.encode(task_result)
data = json.loads(message.value.decode("utf-8"))
assert data["version"] == 2
payload = data["payload"]
assert payload["subscription_id"] == str(task_result.task.task.identifier)
assert payload["request"] == request.body
assert payload["result"] == result
assert payload["timestamp"] == task_result.task.timestamp.isoformat()
def subscription(self) -> SubscriptionData:
return SubscriptionData(
project_id=self.project_id,
conditions=[["platform", "IN", ["a"]]],
aggregations=[["count()", "", "count"]],
time_window=timedelta(minutes=500),
resolution=timedelta(minutes=1),
)
def test_conditions(self) -> None:
subscription = SubscriptionData(
project_id=self.project_id,
conditions=[["platform", "IN", ["a"]]],
aggregations=[["count()", "", "count"]],
time_window=timedelta(minutes=500),
resolution=timedelta(minutes=1),
)
timer = Timer("test")
request = subscription.build_request(
self.dataset,
datetime.utcnow(),
100,
timer,
)
result = parse_and_run_query(self.dataset, request, timer)
assert result.result["data"][0]["count"] == 10
def subscription(self) -> Sequence[SubscriptionData]:
return [
SubscriptionData(
project_id=self.project_id,
query="MATCH (events) SELECT count() WHERE in(platform, 'a')",
time_window_sec=500 * 60,
resolution_sec=60,
entity_subscription=create_entity_subscription(),
),
SubscriptionData(
project_id=self.project_id,
time_window_sec=500 * 60,
resolution_sec=60,
query="MATCH (events) SELECT count() WHERE in(platform, 'a')",
entity_subscription=create_entity_subscription(),
),
]
def create(self, data: SubscriptionData, timer: Timer) -> SubscriptionIdentifier:
# We want to test the query out here to make sure it's valid and can run
# If there is a delegate subscription, we need to run both the SnQL and Legacy validator
if isinstance(data, DelegateSubscriptionData):
self._test_request(data.to_snql(), timer)
self._test_request(data.to_legacy(), timer)
else:
self._test_request(data, timer)
identifier = SubscriptionIdentifier(
self.__partitioner.build_partition_id(data), uuid1(),
)
RedisSubscriptionDataStore(
redis_client, self.dataset, identifier.partition
).create(
identifier.uuid, data,
)
return identifier
def decode(self, value: bytes) -> SubscriptionData:
data = json.loads(value.decode("utf-8"))
return SubscriptionData(
project_id=data["project_id"],
conditions=data["conditions"],
aggregations=data["aggregations"],
time_window=timedelta(seconds=data["time_window"]),
resolution=timedelta(seconds=data["resolution"]),
)
def test_invalid_time_window(self) -> None:
creator = SubscriptionCreator(self.dataset, EntityKey.EVENTS)
with raises(InvalidSubscriptionError):
creator.create(
SubscriptionData(
project_id=123,
time_window_sec=0,
resolution_sec=60,
query=
"MATCH (events) SELECT count() AS count WHERE platfo IN tuple('a')",
entity_subscription=create_entity_subscription(),
),
self.timer,
)
with raises(InvalidSubscriptionError):
creator.create(
SubscriptionData(
project_id=123,
query=("MATCH (events) "
"SELECT count() AS count BY time "
"WHERE "
"platform IN tuple('a') "),
time_window_sec=0,
resolution_sec=60,
entity_subscription=create_entity_subscription(),
),
self.timer,
)
with raises(InvalidSubscriptionError):
creator.create(
SubscriptionData(
project_id=123,
time_window_sec=48 * 60 * 60,
resolution_sec=60,
query=
"MATCH (events) SELECT count() AS count WHERE platfo IN tuple('a')",
entity_subscription=create_entity_subscription(),
),
self.timer,
)
def build_subscription(self, resolution: timedelta) -> Subscription:
return Subscription(
SubscriptionIdentifier(self.partition_id, uuid.uuid4()),
SubscriptionData(
project_id=1,
query="MATCH (events) SELECT count() AS count",
time_window_sec=60,
resolution_sec=int(resolution.total_seconds()),
entity_subscription=create_entity_subscription(),
),
)
def build_subscription(resolution: timedelta, sequence: int) -> Subscription:
entity_subscription = EventsSubscription(data_dict={})
return Subscription(
SubscriptionIdentifier(PartitionId(1), UUIDS[sequence]),
SubscriptionData(
project_id=1,
time_window_sec=int(timedelta(minutes=5).total_seconds()),
resolution_sec=int(resolution.total_seconds()),
query="MATCH events SELECT count()",
entity_subscription=entity_subscription,
),
)
def test_invalid_aggregation(self):
creator = SubscriptionCreator(self.dataset)
with raises(QueryException):
creator.create(
SubscriptionData(
123,
[["platform", "IN", ["a"]]],
[["cout()", "", "count"]],
timedelta(minutes=10),
timedelta(minutes=1),
),
self.timer,
)
def test_conditions(self, subscription: SubscriptionData,
exception: Optional[Exception]) -> None:
timer = Timer("test")
if exception is not None:
with pytest.raises(exception):
request = subscription.build_request(
self.dataset,
datetime.utcnow(),
100,
timer,
)
result = parse_and_run_query(self.dataset, request, timer)
return
request = subscription.build_request(
self.dataset,
datetime.utcnow(),
100,
timer,
)
result = parse_and_run_query(self.dataset, request, timer)
assert result.result["data"][0]["count"] == 10
def test(self):
creator = SubscriptionCreator(self.dataset)
subscription = SubscriptionData(
project_id=123,
conditions=[["platform", "IN", ["a"]]],
aggregations=[["count()", "", "count"]],
time_window=timedelta(minutes=10),
resolution=timedelta(minutes=1),
)
identifier = creator.create(subscription, self.timer)
RedisSubscriptionDataStore(
redis_client, self.dataset, identifier.partition,
).all()[0][1] == subscription
def build_snql_subscription_data(
entity_key: EntityKey,
organization: Optional[int] = None,
) -> SubscriptionData:
return SubscriptionData(
project_id=5,
time_window_sec=500 * 60,
resolution_sec=60,
query="MATCH events SELECT count() WHERE in(platform, 'a')",
entity_subscription=create_entity_subscription(entity_key,
organization),
)
def create_subscription() -> None:
store = RedisSubscriptionDataStore(redis_client, EntityKey.EVENTS,
PartitionId(0))
store.create(
uuid.uuid4(),
SubscriptionData(
project_id=1,
time_window_sec=60,
resolution_sec=60,
query="MATCH (events) SELECT count()",
entity_subscription=EventsSubscription(data_dict={}),
),
)
def test_invalid_resolution(self):
creator = SubscriptionCreator(self.dataset)
with raises(InvalidSubscriptionError):
creator.create(
SubscriptionData(
123,
[["platfo", "IN", ["a"]]],
[["count()", "", "count"]],
timedelta(minutes=1),
timedelta(),
),
Mock(),
)
def test_invalid_condition_column(self):
creator = SubscriptionCreator(self.dataset)
with raises(RawQueryException):
creator.create(
SubscriptionData(
123,
[["platfo", "IN", ["a"]]],
[["count()", "", "count"]],
timedelta(minutes=10),
timedelta(minutes=1),
),
Mock(),
)
def test_invalid_aggregation(self) -> None:
creator = SubscriptionCreator(self.dataset, EntityKey.EVENTS)
with raises(QueryException):
creator.create(
SubscriptionData(
project_id=123,
time_window_sec=10 * 60,
resolution_sec=60,
query=
"MATCH (events) SELECT cout() AS count WHERE platform IN tuple('a')",
entity_subscription=create_entity_subscription(),
),
self.timer,
)
def create(self, data: SubscriptionData,
timer: Timer) -> SubscriptionIdentifier:
# We want to test the query out here to make sure it's valid and can run
request = data.build_request(self.dataset, datetime.utcnow(), None,
timer)
parse_and_run_query(self.dataset, request, timer)
identifier = SubscriptionIdentifier(
self.__partitioner.build_partition_id(data),
uuid1(),
)
RedisSubscriptionDataStore(redis_client, self.dataset,
identifier.partition).create(
identifier.uuid,
data,
)
return identifier
def test(self):
creator = SubscriptionCreator(self.dataset)
subscription = SubscriptionData(
project_id=1,
conditions=[],
aggregations=[["count()", "", "count"]],
time_window=timedelta(minutes=10),
resolution=timedelta(minutes=1),
)
identifier = creator.create(subscription, Timer("test"))
RedisSubscriptionDataStore(
redis_client, self.dataset, identifier.partition,
).all()[0][1] == subscription
SubscriptionDeleter(self.dataset, identifier.partition).delete(identifier.uuid)
RedisSubscriptionDataStore(
redis_client, self.dataset, identifier.partition,
).all() == []
def test_all(self):
store = self.build_store()
assert store.all() == []
subscription_id = uuid1()
store.create(subscription_id, self.subscription)
assert store.all() == [(subscription_id, self.subscription)]
new_subscription = SubscriptionData(
project_id=self.project_id,
conditions=[["platform", "IN", ["b"]]],
aggregations=[["count()", "", "something"]],
time_window=timedelta(minutes=400),
resolution=timedelta(minutes=2),
)
new_subscription_id = uuid1()
store.create(new_subscription_id, new_subscription)
assert sorted(store.all(), key=lambda row: row[0]) == [
(subscription_id, self.subscription),
(new_subscription_id, new_subscription),
]
def test_partitions(self):
store_1 = self.build_store("1")
store_2 = self.build_store("2")
subscription_id = uuid1()
store_1.create(subscription_id, self.subscription)
assert store_2.all() == []
assert store_1.all() == [(subscription_id, self.subscription)]
new_subscription = SubscriptionData(
project_id=self.project_id,
conditions=[["platform", "IN", ["b"]]],
aggregations=[["count()", "", "something"]],
time_window=timedelta(minutes=400),
resolution=timedelta(minutes=2),
)
new_subscription_id = uuid1()
store_2.create(new_subscription_id, new_subscription)
assert store_1.all() == [(subscription_id, self.subscription)]
assert store_2.all() == [(new_subscription_id, new_subscription)]
def test_subscription_task_encoder() -> None:
encoder = SubscriptionScheduledTaskEncoder()
subscription_data = SubscriptionData(
project_id=1,
query="MATCH events SELECT count()",
time_window_sec=60,
resolution_sec=60,
entity_subscription=EventsSubscription(data_dict={}),
)
subscription_id = uuid.UUID("91b46cb6224f11ecb2ddacde48001122")
epoch = datetime(1970, 1, 1)
tick_upper_offset = 5
subscription_with_metadata = SubscriptionWithMetadata(
EntityKey.EVENTS,
Subscription(SubscriptionIdentifier(PartitionId(1), subscription_id),
subscription_data),
tick_upper_offset,
)
task = ScheduledSubscriptionTask(timestamp=epoch,
task=subscription_with_metadata)
encoded = encoder.encode(task)
assert encoded.key == b"1/91b46cb6224f11ecb2ddacde48001122"
assert encoded.value == (
b"{"
b'"timestamp":"1970-01-01T00:00:00",'
b'"entity":"events",'
b'"task":{'
b'"data":{"project_id":1,"time_window":60,"resolution":60,"query":"MATCH events SELECT count()"}},'
b'"tick_upper_offset":5'
b"}")
decoded = encoder.decode(encoded)
assert decoded == task
