0
0
mirror of https://github.com/PostHog/posthog.git synced 2024-11-28 18:26:15 +01:00
posthog/ee/tasks/subscriptions/subscription_utils.py
Paul D'Ambra 18c5927f57
fix: celery autoimport was ignoring CSV exports (#10586)
* launch celery with debug logging

* autoimport a single task which decides what type of export to run

* still need to manually inject root folder so tests can clear up

* fix mock
2022-06-30 17:42:28 +01:00

54 lines
1.6 KiB
Python

from time import sleep
from typing import List, Tuple
import structlog
from celery import group
from posthog.models.dashboard_tile import get_tiles_ordered_by_position
from posthog.models.exported_asset import ExportedAsset
from posthog.models.insight import Insight
from posthog.models.subscription import Subscription
from posthog.tasks import exporter
logger = structlog.get_logger(__name__)
UTM_TAGS_BASE = "utm_source=posthog&utm_campaign=subscription_report"
DEFAULT_MAX_ASSET_COUNT = 6
def generate_assets(
subscription: Subscription, max_asset_count: int = DEFAULT_MAX_ASSET_COUNT
) -> Tuple[List[Insight], List[ExportedAsset]]:
insights = []
if subscription.dashboard:
tiles = get_tiles_ordered_by_position(subscription.dashboard)
insights = [tile.insight for tile in tiles]
elif subscription.insight:
insights = [subscription.insight]
else:
raise Exception("There are no insights to be sent for this Subscription")
# Create all the assets we need
assets = [
ExportedAsset(
team=subscription.team, export_format="image/png", insight=insight, dashboard=subscription.dashboard
)
for insight in insights[:max_asset_count]
]
ExportedAsset.objects.bulk_create(assets)
# Wait for all assets to be exported
tasks = [exporter.export_asset.s(asset.id) for asset in assets]
parallel_job = group(tasks).apply_async()
max_wait = 30
while not parallel_job.ready():
max_wait = max_wait - 1
sleep(1)
if max_wait < 0:
raise Exception("Timed out waiting for exports")
return insights, assets