diff --git a/.eslintrc.js b/.eslintrc.js index 59dae5ce57f..341f903da4c 100644 --- a/.eslintrc.js +++ b/.eslintrc.js @@ -50,7 +50,7 @@ module.exports = { 'posthog', 'simple-import-sort', 'import', - "unused-imports" + 'unused-imports', ], rules: { // PyCharm always adds curly braces, I guess vscode doesn't, PR reviewers often complain they are present on props that don't need them @@ -74,7 +74,7 @@ module.exports = { html: true, }, ], - "unused-imports/no-unused-imports": "error", + 'unused-imports/no-unused-imports': 'error', 'no-unused-vars': 'off', '@typescript-eslint/no-unused-vars': [ 'error', @@ -292,7 +292,7 @@ module.exports = { ], }, ], - 'no-else-return': 'warn' + 'no-else-return': 'warn', }, overrides: [ { diff --git a/frontend/__snapshots__/scenes-app-insights-error-empty-states--long-loading--dark.png b/frontend/__snapshots__/scenes-app-insights-error-empty-states--long-loading--dark.png index 5b7f3caaa2a..56d27b14f84 100644 Binary files a/frontend/__snapshots__/scenes-app-insights-error-empty-states--long-loading--dark.png and b/frontend/__snapshots__/scenes-app-insights-error-empty-states--long-loading--dark.png differ diff --git a/frontend/__snapshots__/scenes-app-insights-error-empty-states--long-loading--light.png b/frontend/__snapshots__/scenes-app-insights-error-empty-states--long-loading--light.png index 97dfc2a3b50..035408ca972 100644 Binary files a/frontend/__snapshots__/scenes-app-insights-error-empty-states--long-loading--light.png and b/frontend/__snapshots__/scenes-app-insights-error-empty-states--long-loading--light.png differ diff --git a/frontend/src/lib/api.ts b/frontend/src/lib/api.ts index e38e83faea3..189e6068501 100644 --- a/frontend/src/lib/api.ts +++ b/frontend/src/lib/api.ts @@ -694,8 +694,12 @@ class ApiRequest { return this.projectsDetail(teamId).addPathComponent('query') } - public queryStatus(queryId: string, teamId?: TeamType['id']): ApiRequest { - return this.query(teamId).addPathComponent(queryId) + public queryStatus(queryId: string, showProgress: boolean, teamId?: TeamType['id']): ApiRequest { + const apiRequest = this.query(teamId).addPathComponent(queryId) + if (showProgress) { + return apiRequest.withQueryString('showProgress=true') + } + return apiRequest } // Notebooks @@ -2084,8 +2088,8 @@ const api = { }, queryStatus: { - async get(queryId: string): Promise { - return await new ApiRequest().queryStatus(queryId).get() + async get(queryId: string, showProgress: boolean): Promise { + return await new ApiRequest().queryStatus(queryId, showProgress).get() }, }, diff --git a/frontend/src/lib/components/Cards/InsightCard/InsightCard.tsx b/frontend/src/lib/components/Cards/InsightCard/InsightCard.tsx index 15eea38dd7a..83e5c159a21 100644 --- a/frontend/src/lib/components/Cards/InsightCard/InsightCard.tsx +++ b/frontend/src/lib/components/Cards/InsightCard/InsightCard.tsx @@ -222,7 +222,7 @@ export function FilterBasedCardContent({ ) : empty ? ( ) : !loading && timedOut ? ( - + ) : apiErrored && !loading ? ( ) : ( diff --git a/frontend/src/lib/constants.tsx b/frontend/src/lib/constants.tsx index 02a80d85aa8..383c5a82a39 100644 --- a/frontend/src/lib/constants.tsx +++ b/frontend/src/lib/constants.tsx @@ -210,6 +210,7 @@ export const FEATURE_FLAGS = { THEME: 'theme', // owner: @aprilfools SESSION_TABLE_PROPERTY_FILTERS: 'session-table-property-filters', // owner: @robbie-c SESSION_REPLAY_HOG_QL_FILTERING: 'session-replay-hogql-filtering', // owner: #team-replay + INSIGHT_LOADING_BAR: 'insight-loading-bar', // owner: @aspicer SESSION_REPLAY_ARTIFICIAL_LAG: 'artificial-lag-query-performance', // owner: #team-replay PROXY_AS_A_SERVICE: 'proxy-as-a-service', // owner: #team-infrastructure } as const diff --git a/frontend/src/lib/lemon-ui/LoadingBar/LoadingBar.scss b/frontend/src/lib/lemon-ui/LoadingBar/LoadingBar.scss new file mode 100644 index 00000000000..e7bffa602d0 --- /dev/null +++ b/frontend/src/lib/lemon-ui/LoadingBar/LoadingBar.scss @@ -0,0 +1,26 @@ +.progress-outer { + display: flex; + align-items: center; +} + +.progress-info { + width: 30px; + padding-left: 5px; +} + +.progress { + position: relative; + width: 100%; + height: 10px; + min-height: 1px; + overflow: hidden; + background: #eee; + border-radius: 3px; +} + +.progress-bar { + position: absolute; + left: 0; + height: 100%; + background: var(--primary-3000-active); +} diff --git a/frontend/src/lib/lemon-ui/LoadingBar/LoadingBar.tsx b/frontend/src/lib/lemon-ui/LoadingBar/LoadingBar.tsx new file mode 100644 index 00000000000..fb364f12b5d --- /dev/null +++ b/frontend/src/lib/lemon-ui/LoadingBar/LoadingBar.tsx @@ -0,0 +1,43 @@ +import './LoadingBar.scss' + +import { useEffect, useState } from 'react' +import { twMerge } from 'tailwind-merge' + +export interface SpinnerProps { + textColored?: boolean + className?: string +} + +/** Smoothly animated spinner for loading states. It does not indicate progress, only that something's happening. */ +export function LoadingBar({ className }: SpinnerProps): JSX.Element { + const [progress, setProgress] = useState(0) + + useEffect(() => { + const interval = setInterval(() => { + setProgress((prevProgress) => { + let newProgress = prevProgress + 0.005 + if (newProgress >= 70) { + newProgress = prevProgress + 0.0025 + } + if (newProgress >= 85) { + newProgress = prevProgress + 0.001 + } + return newProgress + }) + }, 50) + + return () => clearInterval(interval) + }, []) // Empty dependency array ensures this effect runs only once + + return ( +
+
+
+
+
+ ) +} diff --git a/frontend/src/lib/lemon-ui/LoadingBar/index.ts b/frontend/src/lib/lemon-ui/LoadingBar/index.ts new file mode 100644 index 00000000000..1a60869fd58 --- /dev/null +++ b/frontend/src/lib/lemon-ui/LoadingBar/index.ts @@ -0,0 +1 @@ +export * from './LoadingBar' diff --git a/frontend/src/queries/nodes/DataNode/dataNodeLogic.ts b/frontend/src/queries/nodes/DataNode/dataNodeLogic.ts index adaf49c500e..c7d6b2830b2 100644 --- a/frontend/src/queries/nodes/DataNode/dataNodeLogic.ts +++ b/frontend/src/queries/nodes/DataNode/dataNodeLogic.ts @@ -29,6 +29,7 @@ import { userLogic } from 'scenes/userLogic' import { dataNodeCollectionLogic, DataNodeCollectionProps } from '~/queries/nodes/DataNode/dataNodeCollectionLogic' import { removeExpressionComment } from '~/queries/nodes/DataTable/utils' import { query } from '~/queries/query' +import { QueryStatus } from '~/queries/schema' import { ActorsQuery, ActorsQueryResponse, @@ -151,6 +152,7 @@ export const dataNodeLogic = kea([ toggleAutoLoad: true, highlightRows: (rows: any[]) => ({ rows }), setElapsedTime: (elapsedTime: number) => ({ elapsedTime }), + setPollResponse: (status: QueryStatus | null) => ({ status }), }), loaders(({ actions, cache, values, props }) => ({ response: [ @@ -184,6 +186,7 @@ export const dataNodeLogic = kea([ } actions.abortAnyRunningQuery() + actions.setPollResponse(null) const abortController = new AbortController() cache.abortController = abortController const methodOptions: ApiMethodOptions = { @@ -204,7 +207,9 @@ export const dataNodeLogic = kea([ addModifiers(props.query, props.modifiers), methodOptions, refresh, - queryId + queryId, + undefined, + actions.setPollResponse )) ?? null const duration = performance.now() - now return { data, duration } @@ -299,6 +304,12 @@ export const dataNodeLogic = kea([ loadDataFailure: () => false, }, ], + queryId: [ + null as null | string, + { + loadData: (_, { queryId }) => queryId, + }, + ], newDataLoading: [ false, { @@ -324,6 +335,14 @@ export const dataNodeLogic = kea([ cancelQuery: () => true, }, ], + pollResponse: [ + null as null | Record, + { + setPollResponse: (state, { status }) => { + return { status, previousStatus: state && state.status } + }, + }, + ], autoLoadToggled: [ false, // store the 'autoload toggle' state in localstorage, separately for each data node kind diff --git a/frontend/src/queries/nodes/DataTable/dataTableLogic.test.ts b/frontend/src/queries/nodes/DataTable/dataTableLogic.test.ts index bd3693de1ab..cc397c7873f 100644 --- a/frontend/src/queries/nodes/DataTable/dataTableLogic.test.ts +++ b/frontend/src/queries/nodes/DataTable/dataTableLogic.test.ts @@ -60,7 +60,14 @@ describe('dataTableLogic', () => { response: randomResponse, }) - expect(query).toHaveBeenCalledWith(dataTableQuery.source, expect.anything(), false, expect.any(String)) + expect(query).toHaveBeenCalledWith( + dataTableQuery.source, + expect.anything(), + false, + expect.any(String), + undefined, + expect.any(Function) + ) expect(query).toHaveBeenCalledTimes(2) // TODO: Should be 1 }) diff --git a/frontend/src/queries/nodes/InsightViz/InsightVizDisplay.tsx b/frontend/src/queries/nodes/InsightViz/InsightVizDisplay.tsx index 52bab0544ea..830dc4bd4a3 100644 --- a/frontend/src/queries/nodes/InsightViz/InsightVizDisplay.tsx +++ b/frontend/src/queries/nodes/InsightViz/InsightVizDisplay.tsx @@ -1,7 +1,5 @@ import clsx from 'clsx' import { useValues } from 'kea' -import { AnimationType } from 'lib/animations/animations' -import { Animation } from 'lib/components/Animation/Animation' import { ExportButton } from 'lib/components/ExportButton/ExportButton' import { InsightLegend } from 'lib/components/InsightLegend/InsightLegend' import { Tooltip } from 'lib/lemon-ui/Tooltip' @@ -12,6 +10,7 @@ import { FunnelSingleStepState, InsightEmptyState, InsightErrorState, + InsightLoadingState, InsightTimeoutState, InsightValidationError, } from 'scenes/insights/EmptyStates' @@ -77,17 +76,14 @@ export function InsightVizDisplay({ vizSpecificOptions, query, } = useValues(insightVizDataLogic(insightProps)) - const { exportContext } = useValues(insightDataLogic(insightProps)) + const { exportContext, queryId } = useValues(insightDataLogic(insightProps)) // Empty states that completely replace the graph const BlockingEmptyState = (() => { if (insightDataLoading) { return (
- - {!!timedOutQueryId && ( - - )} +
) } @@ -111,13 +107,7 @@ export function InsightVizDisplay({ return } if (timedOutQueryId) { - return ( - - ) + return } return null diff --git a/frontend/src/queries/query.ts b/frontend/src/queries/query.ts index 49a3af52eda..eb6e466a64e 100644 --- a/frontend/src/queries/query.ts +++ b/frontend/src/queries/query.ts @@ -19,7 +19,7 @@ import { import { AnyPartialFilterType, OnlineExportContext, QueryExportContext } from '~/types' import { queryNodeToFilter } from './nodes/InsightQuery/utils/queryNodeToFilter' -import { DataNode, HogQLQuery, HogQLQueryResponse, NodeKind, PersonsNode } from './schema' +import { DataNode, HogQLQuery, HogQLQueryResponse, NodeKind, PersonsNode, QueryStatus } from './schema' import { isActorsQuery, isDataTableNode, @@ -107,12 +107,15 @@ async function executeQuery( queryNode: N, methodOptions?: ApiMethodOptions, refresh?: boolean, - queryId?: string + queryId?: string, + setPollResponse?: (response: QueryStatus) => void ): Promise> { const isAsyncQuery = !SYNC_ONLY_QUERY_KINDS.includes(queryNode.kind) && !!featureFlagLogic.findMounted()?.values.featureFlags?.[FEATURE_FLAGS.QUERY_ASYNC] + const showProgress = !!featureFlagLogic.findMounted()?.values.featureFlags?.[FEATURE_FLAGS.INSIGHT_LOADING_BAR] + const response = await api.query(queryNode, methodOptions, queryId, refresh, isAsyncQuery) if (!isAsyncQuery || !response.query_async) { @@ -126,11 +129,14 @@ async function executeQuery( await delay(currentDelay, methodOptions?.signal) currentDelay = Math.min(currentDelay * 2, QUERY_ASYNC_MAX_INTERVAL_SECONDS * 1000) - const statusResponse = await api.queryStatus.get(response.id) + const statusResponse = await api.queryStatus.get(response.id, showProgress) if (statusResponse.complete || statusResponse.error) { return statusResponse.results } + if (setPollResponse) { + setPollResponse(statusResponse) + } } throw new Error('Query timed out') } @@ -141,7 +147,8 @@ export async function query( methodOptions?: ApiMethodOptions, refresh?: boolean, queryId?: string, - legacyUrl?: string + legacyUrl?: string, + setPollResponse?: (status: QueryStatus) => void ): Promise> { if (isTimeToSeeDataSessionsNode(queryNode)) { return query(queryNode.source) @@ -364,13 +371,13 @@ export async function query( : {}), }) } else { - response = await executeQuery(queryNode, methodOptions, refresh, queryId) + response = await executeQuery(queryNode, methodOptions, refresh, queryId, setPollResponse) } } else { response = await fetchLegacyInsights() } } else { - response = await executeQuery(queryNode, methodOptions, refresh, queryId) + response = await executeQuery(queryNode, methodOptions, refresh, queryId, setPollResponse) if (isHogQLQuery(queryNode) && response && typeof response === 'object') { logParams.clickhouse_sql = (response as HogQLQueryResponse)?.clickhouse } diff --git a/frontend/src/queries/schema.json b/frontend/src/queries/schema.json index dfbe4254ca9..62e1aa5939f 100644 --- a/frontend/src/queries/schema.json +++ b/frontend/src/queries/schema.json @@ -1539,6 +1539,25 @@ ], "type": "string" }, + "ClickhouseQueryStatus": { + "additionalProperties": false, + "properties": { + "bytes_read": { + "type": "integer" + }, + "estimated_rows_total": { + "type": "integer" + }, + "rows_read": { + "type": "integer" + }, + "time_elapsed": { + "type": "integer" + } + }, + "required": ["bytes_read", "rows_read", "estimated_rows_total", "time_elapsed"], + "type": "object" + }, "CohortPropertyFilter": { "additionalProperties": false, "description": "Sync with plugin-server/src/types.ts", @@ -6476,6 +6495,9 @@ "default": true, "type": "boolean" }, + "query_progress": { + "$ref": "#/definitions/ClickhouseQueryStatus" + }, "results": {}, "start_time": { "format": "date-time", diff --git a/frontend/src/queries/schema.ts b/frontend/src/queries/schema.ts index 8a9d36a5a5e..a82db7ede7a 100644 --- a/frontend/src/queries/schema.ts +++ b/frontend/src/queries/schema.ts @@ -927,6 +927,13 @@ export interface CacheMissResponse { cache_key: string | null } +export type ClickhouseQueryStatus = { + bytes_read: integer + rows_read: integer + estimated_rows_total: integer + time_elapsed: integer +} + export type QueryStatus = { id: string /** @default true */ @@ -946,6 +953,7 @@ export type QueryStatus = { /** @format date-time */ expiration_time?: string task_id?: string + query_progress?: ClickhouseQueryStatus } export interface LifecycleQueryResponse extends AnalyticsQueryResponseBase[]> {} diff --git a/frontend/src/scenes/insights/EmptyStates/EmptyStates.tsx b/frontend/src/scenes/insights/EmptyStates/EmptyStates.tsx index 2c24bf32f01..b715ac0568c 100644 --- a/frontend/src/scenes/insights/EmptyStates/EmptyStates.tsx +++ b/frontend/src/scenes/insights/EmptyStates/EmptyStates.tsx @@ -6,12 +6,20 @@ import { IconInfo, IconPlus, IconWarning } from '@posthog/icons' import { LemonButton } from '@posthog/lemon-ui' import { Empty } from 'antd' import { useActions, useValues } from 'kea' +import { AnimationType } from 'lib/animations/animations' +import { Animation } from 'lib/components/Animation/Animation' import { BuilderHog3 } from 'lib/components/hedgehogs' import { supportLogic } from 'lib/components/Support/supportLogic' +import { FEATURE_FLAGS } from 'lib/constants' +import { dayjs } from 'lib/dayjs' import { IconErrorOutline, IconOpenInNew } from 'lib/lemon-ui/icons' import { Link } from 'lib/lemon-ui/Link' +import { LoadingBar } from 'lib/lemon-ui/LoadingBar' import { Tooltip } from 'lib/lemon-ui/Tooltip' +import { featureFlagLogic } from 'lib/logic/featureFlagLogic' +import { humanFriendlyNumber } from 'lib/utils' import posthog from 'posthog-js' +import { useEffect, useState } from 'react' import { funnelDataLogic } from 'scenes/funnels/funnelDataLogic' import { entityFilterLogic } from 'scenes/insights/filters/ActionFilter/entityFilterLogic' import { insightLogic } from 'scenes/insights/insightLogic' @@ -26,6 +34,7 @@ import { FilterType, InsightLogicProps, SavedInsightsTabs } from '~/types' import { samplingFilterLogic } from '../EditorFilters/samplingFilterLogic' import { MathAvailability } from '../filters/ActionFilter/ActionFilterRow/ActionFilterRow' +import { insightDataLogic } from '../insightDataLogic' export function InsightEmptyState({ heading = 'There are no matching events for this query', @@ -68,63 +77,82 @@ function SamplingLink({ insightProps }: { insightProps: InsightLogicProps }): JS ) } +function humanFileSize(size: number): string { + const i = size == 0 ? 0 : Math.floor(Math.log(size) / Math.log(1024)) + return +(size / Math.pow(1024, i)).toFixed(2) * 1 + ' ' + ['B', 'kB', 'MB', 'GB', 'TB'][i] +} -export function InsightTimeoutState({ - isLoading, +export function InsightLoadingStateWithLoadingBar({ queryId, insightProps, }: { - isLoading: boolean queryId?: string | null insightProps: InsightLogicProps }): JSX.Element { const { suggestedSamplingPercentage, samplingPercentage } = useValues(samplingFilterLogic(insightProps)) - const { openSupportForm } = useActions(supportLogic) + const { insightPollResponse } = useValues(insightDataLogic(insightProps)) + + const [rowsRead, setRowsRead] = useState(0) + const [bytesRead, setBytesRead] = useState(0) + const [secondsElapsed, setSecondsElapsed] = useState(0) + + useEffect(() => { + setRowsRead(insightPollResponse?.previousStatus?.query_progress?.rows_read || 0) + setBytesRead(insightPollResponse?.previousStatus?.query_progress?.bytes_read || 0) + const interval = setInterval(() => { + setRowsRead((rowsRead) => { + const diff = + (insightPollResponse?.status?.query_progress?.rows_read || 0) - + (insightPollResponse?.previousStatus?.query_progress?.rows_read || 0) + return rowsRead + diff / 30 + }) + setBytesRead((bytesRead) => { + const diff = + (insightPollResponse?.status?.query_progress?.bytes_read || 0) - + (insightPollResponse?.previousStatus?.query_progress?.bytes_read || 0) + return bytesRead + diff / 30 + }) + setSecondsElapsed(() => { + return dayjs().diff(dayjs(insightPollResponse?.status?.start_time), 'second') + }) + }, 100) + + return () => clearInterval(interval) + }, [insightPollResponse]) + const bytesPerSecond = bytesRead / (secondsElapsed || 1) return (
- {!isLoading ? ( - <> -
- -
-

Your query took too long to complete

- - ) : ( -

Crunching through hogloads of data...

- )} +

Crunching through hogloads of data...

+ +

+ {rowsRead > 0 && bytesRead > 0 && ( + <> + {humanFriendlyNumber(rowsRead || 0)} rows +
+ {humanFileSize(bytesRead || 0)} ({humanFileSize(bytesPerSecond || 0)}/s) + + )} +

- {isLoading && suggestedSamplingPercentage && !samplingPercentage ? ( + {suggestedSamplingPercentage && !samplingPercentage ? ( Need to speed things up? Try reducing the date range, removing breakdowns, or turning on{' '} . - ) : isLoading && suggestedSamplingPercentage && samplingPercentage ? ( + ) : suggestedSamplingPercentage && samplingPercentage ? ( <> Still waiting around? You must have lots of data! Kick it up a notch with{' '} . Or try reducing the date range and removing breakdowns. - ) : isLoading ? ( - <>Need to speed things up? Try reducing the date range or removing breakdowns. ) : ( - <> - Sometimes this happens. Try refreshing the page, reducing the date range, or removing - breakdowns. If you're still having issues,{' '} - { - openSupportForm({ kind: 'bug', target_area: 'analytics' }) - }} - > - let us know - - . - + <>Need to speed things up? Try reducing the date range or removing breakdowns. )}

@@ -136,6 +164,92 @@ export function InsightTimeoutState({ ) } +export function InsightLoadingState({ + queryId, + insightProps, +}: { + queryId?: string | null + insightProps: InsightLogicProps +}): JSX.Element { + const { featureFlags } = useValues(featureFlagLogic) + if (featureFlags[FEATURE_FLAGS.INSIGHT_LOADING_BAR]) { + return + } + + const { suggestedSamplingPercentage, samplingPercentage } = useValues(samplingFilterLogic(insightProps)) + return ( +
+ +
+

Crunching through hogloads of data...

+
+
+ +
+

+ {suggestedSamplingPercentage && !samplingPercentage ? ( + + Need to speed things up? Try reducing the date range, removing breakdowns, or turning on{' '} + . + + ) : suggestedSamplingPercentage && samplingPercentage ? ( + <> + Still waiting around? You must have lots of data! Kick it up a notch with{' '} + . Or try reducing the date range and + removing breakdowns. + + ) : ( + <>Need to speed things up? Try reducing the date range or removing breakdowns. + )} +

+
+ {queryId ? ( +
Query ID: {queryId}
+ ) : null} +
+
+ ) +} + +export function InsightTimeoutState({ queryId }: { queryId?: string | null }): JSX.Element { + const { openSupportForm } = useActions(supportLogic) + + return ( +
+
+ <> +
+ +
+

Your query took too long to complete

+ +
+
+ +
+

+ <> + Sometimes this happens. Try refreshing the page, reducing the date range, or removing + breakdowns. If you're still having issues,{' '} + { + openSupportForm({ kind: 'bug', target_area: 'analytics' }) + }} + > + let us know + + . + +

+
+ {queryId ? ( +
Query ID: {queryId}
+ ) : null} +
+
+ ) +} + export interface InsightErrorStateProps { excludeDetail?: boolean title?: string diff --git a/frontend/src/scenes/insights/insightDataLogic.tsx b/frontend/src/scenes/insights/insightDataLogic.tsx index fd05ff7da6e..30a29a71397 100644 --- a/frontend/src/scenes/insights/insightDataLogic.tsx +++ b/frontend/src/scenes/insights/insightDataLogic.tsx @@ -61,6 +61,8 @@ export const insightDataLogic = kea([ 'dataLoading as insightDataLoading', 'responseErrorObject as insightDataError', 'getInsightRefreshButtonDisabledReason', + 'pollResponse as insightPollResponse', + 'queryId', ], filterTestAccountsDefaultsLogic, ['filterTestAccountsDefault'], diff --git a/posthog/api/query.py b/posthog/api/query.py index 197fe79f18e..355303007ec 100644 --- a/posthog/api/query.py +++ b/posthog/api/query.py @@ -98,7 +98,9 @@ class QueryViewSet(TeamAndOrgViewSetMixin, PydanticModelMixin, viewsets.ViewSet) }, ) def retrieve(self, request: Request, pk=None, *args, **kwargs) -> JsonResponse: - query_status = get_query_status(team_id=self.team.pk, query_id=pk) + query_status = get_query_status( + team_id=self.team.pk, query_id=pk, show_progress=request.query_params.get("showProgress", False) + ) http_code: int = status.HTTP_202_ACCEPTED if query_status.error: diff --git a/posthog/api/test/test_query.py b/posthog/api/test/test_query.py index bb60c836df5..bfd7a8ad3b9 100644 --- a/posthog/api/test/test_query.py +++ b/posthog/api/test/test_query.py @@ -843,6 +843,7 @@ class TestQuery(ClickhouseTestMixin, APIBaseTest): "start_time": "2020-01-10T12:14:00Z", "task_id": mock.ANY, "team_id": mock.ANY, + "query_progress": None, }, ) @@ -1003,4 +1004,4 @@ class TestQueryRetrieve(APIBaseTest): ).encode() response = self.client.delete(f"/api/projects/{self.team.id}/query/{self.valid_query_id}/") self.assertEqual(response.status_code, 204) - self.redis_client_mock.delete.assert_called_once() + self.assertEqual(self.redis_client_mock.delete.call_count, 2) diff --git a/posthog/clickhouse/client/execute_async.py b/posthog/clickhouse/client/execute_async.py index 4beb4ca0dc2..d4376744ca3 100644 --- a/posthog/clickhouse/client/execute_async.py +++ b/posthog/clickhouse/client/execute_async.py @@ -1,5 +1,7 @@ import datetime -import json + +import orjson as json +import math from functools import partial from typing import Optional import uuid @@ -11,12 +13,13 @@ from rest_framework.exceptions import NotFound from django.db import transaction from posthog import celery, redis +from posthog.clickhouse.client import sync_execute from posthog.clickhouse.query_tagging import tag_queries from posthog.errors import ExposedCHQueryError from posthog.hogql.constants import LimitContext from posthog.hogql.errors import ExposedHogQLError from posthog.renderers import SafeJSONRenderer -from posthog.schema import QueryStatus +from posthog.schema import QueryStatus, ClickhouseQueryStatus from posthog.tasks.tasks import process_query_task logger = structlog.get_logger(__name__) @@ -46,10 +49,18 @@ class QueryStatusManager: def results_key(self) -> str: return f"{self.KEY_PREFIX_ASYNC_RESULTS}:{self.team_id}:{self.query_id}" + @property + def clickhouse_query_status_key(self) -> str: + return f"{self.KEY_PREFIX_ASYNC_RESULTS}:{self.team_id}:{self.query_id}:status" + def store_query_status(self, query_status: QueryStatus): - value = SafeJSONRenderer().render(query_status.model_dump()) + value = SafeJSONRenderer().render(query_status.model_dump(exclude={"clickhouse_query_progress"})) self.redis_client.set(self.results_key, value, ex=self.STATUS_TTL_SECONDS) + def store_clickhouse_query_status(self, query_statuses): + value = json.dumps(query_statuses) + self.redis_client.set(self.clickhouse_query_status_key, value, ex=self.STATUS_TTL_SECONDS) + def _get_results(self): try: byte_results = self.redis_client.get(self.results_key) @@ -58,20 +69,82 @@ class QueryStatusManager: return byte_results + def _get_clickhouse_query_status(self): + try: + byte_results = self.redis_client.get(self.clickhouse_query_status_key) + except Exception: + # Don't fail because of progress checking + return {} + + return json.loads(byte_results) if byte_results is not None else {} + def has_results(self): return self._get_results() is not None - def get_query_status(self) -> QueryStatus: + def get_query_status(self, show_progress=False) -> QueryStatus: byte_results = self._get_results() if not byte_results: raise QueryNotFoundError(f"Query {self.query_id} not found for team {self.team_id}") - return QueryStatus(**json.loads(byte_results)) + query_status = QueryStatus(**json.loads(byte_results)) + + if show_progress and not query_status.complete: + CLICKHOUSE_SQL = """ + SELECT + query_id, + initial_query_id, + read_rows, + read_bytes, + total_rows_approx, + elapsed + FROM clusterAllReplicas(posthog, system.processes) + WHERE query_id like %(query_id)s + """ + + clickhouse_query_progress_dict = self._get_clickhouse_query_status() + try: + results, types = sync_execute( + CLICKHOUSE_SQL, {"query_id": f"%{self.query_id}%"}, with_column_types=True + ) + + noNaNInt = lambda num: 0 if math.isnan(num) else int(num) + + new_clickhouse_query_progress = { + result[0]: { + "bytes_read": noNaNInt(result[3]), + "rows_read": noNaNInt(result[2]), + "estimated_rows_total": noNaNInt(result[4]), + "time_elapsed": noNaNInt(result[5]), + } + for result in results + } + clickhouse_query_progress_dict.update(new_clickhouse_query_progress) + self.store_clickhouse_query_status(clickhouse_query_progress_dict) + + query_progress = { + "bytes_read": 0, + "rows_read": 0, + "estimated_rows_total": 0, + "time_elapsed": 0, + } + for single_query_progress in clickhouse_query_progress_dict.values(): + query_progress["bytes_read"] += single_query_progress["bytes_read"] + query_progress["rows_read"] += single_query_progress["rows_read"] + query_progress["estimated_rows_total"] += single_query_progress["estimated_rows_total"] + query_progress["time_elapsed"] += single_query_progress["time_elapsed"] + query_status.query_progress = ClickhouseQueryStatus(**query_progress) + + except Exception as e: + logger.error("Clickhouse Status Check Failed", e) + pass + + return query_status def delete_query_status(self): logger.info("Deleting redis query key %s", self.results_key) self.redis_client.delete(self.results_key) + self.redis_client.delete(self.clickhouse_query_status_key) def execute_process_query( @@ -200,12 +273,12 @@ def enqueue_process_query_task( return query_status -def get_query_status(team_id, query_id) -> QueryStatus: +def get_query_status(team_id, query_id, show_progress=False) -> QueryStatus: """ Abstracts away the manager for any caller and returns a QueryStatus object """ manager = QueryStatusManager(query_id, team_id) - return manager.get_query_status() + return manager.get_query_status(show_progress=show_progress) def cancel_query(team_id, query_id): diff --git a/posthog/hogql_queries/query_runner.py b/posthog/hogql_queries/query_runner.py index 364f9d71aa1..9062ec6a8ef 100644 --- a/posthog/hogql_queries/query_runner.py +++ b/posthog/hogql_queries/query_runner.py @@ -333,9 +333,8 @@ class QueryRunner(ABC, Generic[Q, R, CR]): return isinstance(data, self.query_type) def is_cached_response(self, data) -> TypeGuard[dict]: - return ( - hasattr(data, "is_cached") # Duck typing for backwards compatibility with `CachedQueryResponse` - or (isinstance(data, dict) and "is_cached" in data) + return hasattr(data, "is_cached") or ( # Duck typing for backwards compatibility with `CachedQueryResponse` + isinstance(data, dict) and "is_cached" in data ) @abstractmethod diff --git a/posthog/schema.py b/posthog/schema.py index 528664753b3..e02e7ef8839 100644 --- a/posthog/schema.py +++ b/posthog/schema.py @@ -178,6 +178,16 @@ class ChartDisplayType(str, Enum): WorldMap = "WorldMap" +class ClickhouseQueryStatus(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + bytes_read: int + estimated_rows_total: int + rows_read: int + time_elapsed: int + + class CohortPropertyFilter(BaseModel): model_config = ConfigDict( extra="forbid", @@ -752,6 +762,7 @@ class QueryStatus(BaseModel): expiration_time: Optional[AwareDatetime] = None id: str query_async: Optional[bool] = True + query_progress: Optional[ClickhouseQueryStatus] = None results: Optional[Any] = None start_time: Optional[AwareDatetime] = None task_id: Optional[str] = None diff --git a/test-runner-jest-environment.js b/test-runner-jest-environment.js index e7e8a419a12..c5aa2479732 100644 --- a/test-runner-jest-environment.js +++ b/test-runner-jest-environment.js @@ -16,9 +16,12 @@ class CustomEnvironment extends PlaywrightEnvironment { // Take screenshots on test failures - these become Actions artifacts const parentName = event.test.parent.parent.name.replace(/\W/g, '-').toLowerCase() const specName = event.test.parent.name.replace(/\W/g, '-').toLowerCase() - await this.global.page.locator('body, main').last().screenshot({ - path: `frontend/__snapshots__/__failures__/${parentName}--${specName}.png`, - }) + await this.global.page + .locator('body, main') + .last() + .screenshot({ + path: `frontend/__snapshots__/__failures__/${parentName}--${specName}.png`, + }) } await super.handleTestEvent(event) }