chore(data-warehouse): make sure counts are updated at end of import (#…

…20223)

make sure counts are updated at end of import

posthog/temporal/data_imports/pipelines/stripe/helpers.py

@@ -8,7 +8,7 @@
 from dlt.sources import DltResource
 from pendulum import DateTime
 from asgiref.sync import sync_to_async
-from posthog.temporal.data_imports.pipelines.helpers import check_limit
+from posthog.temporal.data_imports.pipelines.helpers import check_limit, aupdate_job_count
 from posthog.warehouse.models import ExternalDataJob
 
 stripe.api_version = "2022-11-15"
@@ -92,6 +92,8 @@ async def stripe_pagination(
         if not response["has_more"] or status == ExternalDataJob.Status.CANCELLED:
             break
 
+    await aupdate_job_count(job_id, team_id, count)
+
 
 @dlt.source(max_table_nesting=0)
 def stripe_source(

posthog/temporal/tests/external_data/test_external_data_job.py

@@ -387,6 +387,9 @@ async def setup_job_1():
         # if job was not canceled, this job would run indefinitely
         assert len(job_1_customer_objects["Contents"]) == 1
 
+        await sync_to_async(job_1.refresh_from_db)()
+        assert job_1.rows_synced == 1
+
 
 @pytest.mark.django_db(transaction=True)
 @pytest.mark.asyncio