Skip to content

Commit

Permalink
feat: record API usage for any type of query (#26259)
Browse files Browse the repository at this point in the history
  • Loading branch information
@raquelmsmith
raquelmsmith authored Nov 22, 2024
1 parent 0039410 commit 199e220
Show file tree
Hide file tree
Showing 3 changed files with 116 additions and 83 deletions.
6 changes: 0 additions & 6 deletions posthog/tasks/test/__snapshots__/test_usage_report.ambr
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,6 @@
WHERE (type = 'QueryFinish'
OR type = 'ExceptionWhileProcessing')
AND is_initial_query = 1
AND query_type IN (['hogql_query', 'HogQLQuery'])
AND query_start_time between '2022-01-10 00:00:00' AND '2022-01-10 23:59:59'
AND access_method = ''
GROUP BY team_id
Expand All @@ -54,7 +53,6 @@
WHERE (type = 'QueryFinish'
OR type = 'ExceptionWhileProcessing')
AND is_initial_query = 1
AND query_type IN (['hogql_query', 'HogQLQuery'])
AND query_start_time between '2022-01-10 00:00:00' AND '2022-01-10 23:59:59'
AND access_method = 'personal_api_key'
GROUP BY team_id
Expand All @@ -71,7 +69,6 @@
WHERE (type = 'QueryFinish'
OR type = 'ExceptionWhileProcessing')
AND is_initial_query = 1
AND query_type IN (['hogql_query', 'HogQLQuery'])
AND query_start_time between '2022-01-10 00:00:00' AND '2022-01-10 23:59:59'
AND access_method = 'personal_api_key'
GROUP BY team_id
Expand All @@ -88,7 +85,6 @@
WHERE (type = 'QueryFinish'
OR type = 'ExceptionWhileProcessing')
AND is_initial_query = 1
AND query_type IN (['hogql_query', 'HogQLQuery'])
AND query_start_time between '2022-01-10 00:00:00' AND '2022-01-10 23:59:59'
AND access_method = 'personal_api_key'
GROUP BY team_id
Expand Down Expand Up @@ -342,7 +338,6 @@
WHERE (type = 'QueryFinish'
OR type = 'ExceptionWhileProcessing')
AND is_initial_query = 1
AND query_type IN (['hogql_query', 'HogQLQuery'])
AND query_start_time between '2022-01-10 00:00:00' AND '2022-01-10 23:59:59'
AND access_method = ''
GROUP BY team_id
Expand All @@ -359,7 +354,6 @@
WHERE (type = 'QueryFinish'
OR type = 'ExceptionWhileProcessing')
AND is_initial_query = 1
AND query_type IN (['hogql_query', 'HogQLQuery'])
AND query_start_time between '2022-01-10 00:00:00' AND '2022-01-10 23:59:59'
AND access_method = ''
GROUP BY team_id
Expand Down
111 changes: 76 additions & 35 deletions posthog/tasks/test/test_usage_report.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
from ee.models.license import License
from ee.settings import BILLING_SERVICE_URL
from posthog.clickhouse.client import sync_execute
from posthog.clickhouse.query_tagging import tag_queries
from posthog.cloud_utils import TEST_clear_instance_license_cache
from posthog.hogql.query import execute_hogql_query
from posthog.hogql_queries.events_query_runner import EventsQueryRunner
Expand Down Expand Up @@ -529,12 +530,12 @@ def _test_usage_report(self) -> list[dict]:
"local_evaluation_requests_count_in_period": 0,
"billable_feature_flag_requests_count_in_period": 0,
"survey_responses_count_in_period": 1,
"hogql_app_bytes_read": 0,
"hogql_app_rows_read": 0,
"hogql_app_duration_ms": 0,
"hogql_api_bytes_read": 0,
"hogql_api_rows_read": 0,
"hogql_api_duration_ms": 0,
"query_app_bytes_read": 0,
"query_app_rows_read": 0,
"query_app_duration_ms": 0,
"query_api_bytes_read": 0,
"query_api_rows_read": 0,
"query_api_duration_ms": 0,
"event_explorer_app_bytes_read": 0,
"event_explorer_app_rows_read": 0,
"event_explorer_app_duration_ms": 0,
Expand Down Expand Up @@ -583,12 +584,12 @@ def _test_usage_report(self) -> list[dict]:
"local_evaluation_requests_count_in_period": 0,
"billable_feature_flag_requests_count_in_period": 0,
"survey_responses_count_in_period": 1,
"hogql_app_bytes_read": 0,
"hogql_app_rows_read": 0,
"hogql_app_duration_ms": 0,
"hogql_api_bytes_read": 0,
"hogql_api_rows_read": 0,
"hogql_api_duration_ms": 0,
"query_app_bytes_read": 0,
"query_app_rows_read": 0,
"query_app_duration_ms": 0,
"query_api_bytes_read": 0,
"query_api_rows_read": 0,
"query_api_duration_ms": 0,
"event_explorer_app_bytes_read": 0,
"event_explorer_app_rows_read": 0,
"event_explorer_app_duration_ms": 0,
Expand Down Expand Up @@ -631,12 +632,12 @@ def _test_usage_report(self) -> list[dict]:
"local_evaluation_requests_count_in_period": 0,
"billable_feature_flag_requests_count_in_period": 0,
"survey_responses_count_in_period": 0,
"hogql_app_bytes_read": 0,
"hogql_app_rows_read": 0,
"hogql_app_duration_ms": 0,
"hogql_api_bytes_read": 0,
"hogql_api_rows_read": 0,
"hogql_api_duration_ms": 0,
"query_app_bytes_read": 0,
"query_app_rows_read": 0,
"query_app_duration_ms": 0,
"query_api_bytes_read": 0,
"query_api_rows_read": 0,
"query_api_duration_ms": 0,
"event_explorer_app_bytes_read": 0,
"event_explorer_app_rows_read": 0,
"event_explorer_app_duration_ms": 0,
Expand Down Expand Up @@ -702,12 +703,12 @@ def _test_usage_report(self) -> list[dict]:
"local_evaluation_requests_count_in_period": 0,
"billable_feature_flag_requests_count_in_period": 0,
"survey_responses_count_in_period": 0,
"hogql_app_bytes_read": 0,
"hogql_app_rows_read": 0,
"hogql_app_duration_ms": 0,
"hogql_api_bytes_read": 0,
"hogql_api_rows_read": 0,
"hogql_api_duration_ms": 0,
"query_app_bytes_read": 0,
"query_app_rows_read": 0,
"query_app_duration_ms": 0,
"query_api_bytes_read": 0,
"query_api_rows_read": 0,
"query_api_duration_ms": 0,
"event_explorer_app_bytes_read": 0,
"event_explorer_app_rows_read": 0,
"event_explorer_app_duration_ms": 0,
Expand Down Expand Up @@ -756,12 +757,12 @@ def _test_usage_report(self) -> list[dict]:
"local_evaluation_requests_count_in_period": 0,
"billable_feature_flag_requests_count_in_period": 0,
"survey_responses_count_in_period": 0,
"hogql_app_bytes_read": 0,
"hogql_app_rows_read": 0,
"hogql_app_duration_ms": 0,
"hogql_api_bytes_read": 0,
"hogql_api_rows_read": 0,
"hogql_api_duration_ms": 0,
"query_app_bytes_read": 0,
"query_app_rows_read": 0,
"query_app_duration_ms": 0,
"query_api_bytes_read": 0,
"query_api_rows_read": 0,
"query_api_duration_ms": 0,
"event_explorer_app_bytes_read": 0,
"event_explorer_app_rows_read": 0,
"event_explorer_app_duration_ms": 0,
Expand Down Expand Up @@ -893,7 +894,7 @@ def test_usage_report_hogql_queries(self) -> None:
sync_execute("TRUNCATE TABLE system.query_log")

execute_hogql_query(
query="select * from events limit 200",
query="select * from events limit 400",
team=self.team,
query_type="HogQLQuery",
)
Expand All @@ -906,16 +907,56 @@ def test_usage_report_hogql_queries(self) -> None:

report = _get_team_report(all_reports, self.team)

# We selected 200 or 50 rows, but still read 100 rows to return the query
assert report.hogql_app_rows_read == 100
assert report.hogql_app_bytes_read > 0
# We selected 400 rows, but still read 200 rows to return the query
assert report.query_app_rows_read == 200
assert report.query_app_bytes_read > 0
# We selected 50 rows, but still read 100 rows to return the query
assert report.event_explorer_app_rows_read == 100
assert report.event_explorer_app_bytes_read > 0

# Nothing was read via the API
assert report.hogql_api_rows_read == 0
assert report.query_api_rows_read == 0
assert report.event_explorer_api_rows_read == 0

@also_test_with_materialized_columns(event_properties=["$lib"], verify_no_jsonextract=False)
def test_usage_report_api_queries(self) -> None:
for _ in range(0, 100):
_create_event(
distinct_id="hello",
event="$event1",
properties={"$lib": "web"},
timestamp=now() - relativedelta(hours=12),
team=self.team,
)
flush_persons_and_events()
sync_execute("SYSTEM FLUSH LOGS")
sync_execute("TRUNCATE TABLE system.query_log")
tag_queries(kind="request", id="1", access_method="personal_api_key")

execute_hogql_query(
query="select * from events limit 400",
team=self.team,
query_type="HogQLQuery",
)
EventsQueryRunner(query=EventsQuery(select=["event"], limit=50), team=self.team).calculate()
sync_execute("SYSTEM FLUSH LOGS")

period = get_previous_day(at=now() + relativedelta(days=1))
period_start, period_end = period
all_reports = _get_all_usage_data_as_team_rows(period_start, period_end)

report = _get_team_report(all_reports, self.team)

# No queries were read via the app
assert report.query_app_rows_read == 0
assert report.query_app_bytes_read == 0
assert report.event_explorer_app_rows_read == 0
assert report.event_explorer_app_bytes_read == 0

# Queries were read via the API
assert report.query_api_rows_read == 200
assert report.event_explorer_api_rows_read == 100


@freeze_time("2022-01-10T00:01:00Z")
class TestFeatureFlagsUsageReport(ClickhouseDestroyTablesMixin, TestCase, ClickhouseTestMixin):
Expand Down
Loading

0 comments on commit 199e220

Please sign in to comment.