fix tests and imports

PostHog · Dec 20, 2023 · d94ffa6 · d94ffa6
1 parent 1bb4af7
commit d94ffa6
Show file tree

Hide file tree

Showing 6 changed files with 28 additions and 17 deletions.
diff --git a/posthog/temporal/data_imports/external_data_job.py b/posthog/temporal/data_imports/external_data_job.py
@@ -11,6 +11,7 @@
 from posthog.temporal.batch_exports.base import PostHogWorkflow
 
 from posthog.warehouse.data_load.validate_schema import validate_schema_and_update_table
+from posthog.temporal.data_imports.pipelines.schemas import PIPELINE_TYPE_SCHEMA_DEFAULT_MAPPING
 from posthog.temporal.data_imports.pipelines.pipeline import DataImportPipeline, PipelineInputs
 from posthog.warehouse.external_data_source.jobs import (
     create_external_data_job,
@@ -141,13 +142,18 @@ async def run_external_data_job(inputs: ExternalDataJobInputs) -> None:
         dataset_name=model.folder_path,
     )
 
-    if model.pipeline.source_type == 'stripe':
+    source = None
+    if model.pipeline.source_type == ExternalDataSource.Type.STRIPE:
         from posthog.temporal.data_imports.pipelines.stripe.helpers import stripe_source
-        stripe_secret_key = model.pipeline.job_inputs.get('stripe_secret_key', None)
+
+        stripe_secret_key = model.pipeline.job_inputs.get("stripe_secret_key", None)
         if not stripe_secret_key:
             raise ValueError(f"Stripe secret key not found for job {model.id}")
         source = stripe_source(api_key=stripe_secret_key, endpoints=inputs.schemas)
 
+    if not source:
+        raise ValueError(f"Source not found for job {model.id}")
+
     await DataImportPipeline(job_inputs, source, logger).run()
 
 

diff --git a/posthog/temporal/data_imports/pipelines/pipeline.py b/posthog/temporal/data_imports/pipelines/pipeline.py
@@ -11,6 +11,7 @@
 from structlog.typing import FilteringBoundLogger
 from dlt.sources import DltResource
 
+
 @dataclass
 class PipelineInputs:
     source_id: UUID
@@ -40,14 +41,14 @@ def _get_destination(self):
             credentials = {
                 "aws_access_key_id": settings.AIRBYTE_BUCKET_KEY,
                 "aws_secret_access_key": settings.AIRBYTE_BUCKET_SECRET,
-                "endpoint_url": settings.OBJECT_STORAGE_ENDPOINT
+                "endpoint_url": settings.OBJECT_STORAGE_ENDPOINT,
+            }
+        else:
+            credentials = {
+                "aws_access_key_id": settings.AIRBYTE_BUCKET_KEY,
+                "aws_secret_access_key": settings.AIRBYTE_BUCKET_SECRET,
             }
-
-        credentials = {
-            "aws_access_key_id": settings.AIRBYTE_BUCKET_KEY,
-            "aws_secret_access_key": settings.AIRBYTE_BUCKET_SECRET,
-        }
-
+
         return dlt.destinations.filesystem(
             credentials=credentials,
             bucket_url=settings.BUCKET_URL,  # type: ignore
@@ -69,9 +70,9 @@ def _get_schemas(self):
         if not self.inputs.schemas:
             self.logger.info(f"No schemas found for source id {self.inputs.source_id}")
             return None
-    
+
         return self.inputs.schemas
-        
+
     def _run(self):
         pipeline = self._create_pipeline()
         pipeline.run(self.source, loader_file_format=self.loader_file_format)
@@ -85,4 +86,4 @@ async def run(self) -> None:
             await asyncio.to_thread(self._run)
         except PipelineStepFailed:
             self.logger.error(f"Data import failed for endpoint")
-            raise
+            raise
diff --git a/posthog/temporal/data_imports/pipelines/schemas.py b/posthog/temporal/data_imports/pipelines/schemas.py
@@ -0,0 +1,4 @@
+from posthog.warehouse.models import ExternalDataSource
+from posthog.temporal.data_imports.pipelines.stripe.settings import ENDPOINTS
+
+PIPELINE_TYPE_SCHEMA_DEFAULT_MAPPING = {ExternalDataSource.Type.STRIPE: ENDPOINTS}
diff --git a/posthog/temporal/tests/test_external_data_job.py b/posthog/temporal/tests/test_external_data_job.py
@@ -28,10 +28,10 @@
     ExternalDataSchema,
 )
 
-from posthog.temporal.data_imports.pipelines.stripe.stripe_pipeline import (
-    PIPELINE_TYPE_RUN_MAPPING,
+from posthog.temporal.data_imports.pipelines.schemas import (
     PIPELINE_TYPE_SCHEMA_DEFAULT_MAPPING,
 )
+from posthog.temporal.data_imports.pipelines.pipeline import DataImportPipeline
 from temporalio.testing import WorkflowEnvironment
 from temporalio.common import RetryPolicy
 from temporalio.worker import UnsandboxedWorkflowRunner, Worker
@@ -449,7 +449,7 @@ async def mock_async_func(inputs):
 
     with mock.patch(
         "posthog.warehouse.models.table.DataWarehouseTable.get_columns", return_value={"id": "string"}
-    ), mock.patch.dict(PIPELINE_TYPE_RUN_MAPPING, {ExternalDataSource.Type.STRIPE: mock_async_func}):
+    ), mock.patch.object(DataImportPipeline, "run", mock_async_func):
         with override_settings(AIRBYTE_BUCKET_KEY="test-key", AIRBYTE_BUCKET_SECRET="test-secret"):
             async with await WorkflowEnvironment.start_time_skipping() as activity_environment:
                 async with Worker(

diff --git a/posthog/warehouse/api/external_data_source.py b/posthog/warehouse/api/external_data_source.py
@@ -21,7 +21,7 @@
 )
 from posthog.warehouse.models import ExternalDataSource, ExternalDataSchema, ExternalDataJob
 from posthog.warehouse.api.external_data_schema import ExternalDataSchemaSerializer
-from posthog.temporal.data_imports.pipelines.stripe.stripe_pipeline import (
+from posthog.temporal.data_imports.pipelines.schemas import (
     PIPELINE_TYPE_SCHEMA_DEFAULT_MAPPING,
 )
 

diff --git a/posthog/warehouse/api/test/test_external_data_source.py b/posthog/warehouse/api/test/test_external_data_source.py
@@ -2,7 +2,7 @@
 from posthog.warehouse.models import ExternalDataSource, ExternalDataSchema
 import uuid
 from unittest.mock import patch
-from posthog.temporal.data_imports.pipelines.stripe.stripe_pipeline import (
+from posthog.temporal.data_imports.pipelines.schemas import (
     PIPELINE_TYPE_SCHEMA_DEFAULT_MAPPING,
 )