Source PostHog: add support for self-hosted instances (#6058)

airbytehq · Sep 15, 2021 · 0552b17 · 0552b17
1 parent e837048
commit 0552b17
Show file tree

Hide file tree

Showing 6 changed files with 41 additions and 18 deletions.
diff --git a/...ain/resources/config/STANDARD_SOURCE_DEFINITION/af6d50ee-dddf-4126-a8ee-7faee990774f.json b/...ain/resources/config/STANDARD_SOURCE_DEFINITION/af6d50ee-dddf-4126-a8ee-7faee990774f.json
@@ -2,6 +2,6 @@
   "sourceDefinitionId": "af6d50ee-dddf-4126-a8ee-7faee990774f",
   "name": "PostHog",
   "dockerRepository": "airbyte/source-posthog",
-  "dockerImageTag": "0.1.3",
+  "dockerImageTag": "0.1.4",
   "documentationUrl": "https://docs.airbyte.io/integrations/sources/posthog"
 }
diff --git a/airbyte-config/init/src/main/resources/seed/source_definitions.yaml b/airbyte-config/init/src/main/resources/seed/source_definitions.yaml
@@ -67,7 +67,7 @@
 - sourceDefinitionId: af6d50ee-dddf-4126-a8ee-7faee990774f
   name: PostHog
   dockerRepository: airbyte/source-posthog
-  dockerImageTag: 0.1.3
+  dockerImageTag: 0.1.4
   documentationUrl: https://docs.airbyte.io/integrations/sources/posthog
 - sourceDefinitionId: cd42861b-01fc-4658-a8ab-5d11d0510f01
   name: Recurly

diff --git a/airbyte-integrations/connectors/source-posthog/Dockerfile b/airbyte-integrations/connectors/source-posthog/Dockerfile
@@ -12,5 +12,5 @@ RUN pip install .
 ENV AIRBYTE_ENTRYPOINT "python /airbyte/integration_code/main.py"
 ENTRYPOINT ["python", "/airbyte/integration_code/main.py"]
 
-LABEL io.airbyte.version=0.1.3
+LABEL io.airbyte.version=0.1.4
 LABEL io.airbyte.name=airbyte/source-posthog
diff --git a/airbyte-integrations/connectors/source-posthog/source_posthog/source.py b/airbyte-integrations/connectors/source-posthog/source_posthog/source.py
@@ -46,13 +46,17 @@
     Trends,
 )
 
+DEFAULT_BASE_URL = "https://app.posthog.com"
+
 
 class SourcePosthog(AbstractSource):
     def check_connection(self, logger: AirbyteLogger, config: Mapping[str, Any]) -> Tuple[bool, Any]:
         try:
             _ = pendulum.parse(config["start_date"])
             authenticator = TokenAuthenticator(token=config["api_key"])
-            stream = PingMe(authenticator=authenticator)
+            base_url = config.get("base_url", DEFAULT_BASE_URL)
+
+            stream = PingMe(authenticator=authenticator, base_url=base_url)
             records = stream.read_records(sync_mode=SyncMode.full_refresh)
             _ = next(records)
             return True, None
@@ -69,15 +73,17 @@ def streams(self, config: Mapping[str, Any]) -> List[Stream]:
         This stream was requested to be removed due to this reason.
         """
         authenticator = TokenAuthenticator(token=config["api_key"])
+        base_url = config.get("base_url", DEFAULT_BASE_URL)
+
         return [
-            Annotations(authenticator=authenticator, start_date=config["start_date"]),
-            Cohorts(authenticator=authenticator),
-            Events(authenticator=authenticator, start_date=config["start_date"]),
-            EventsSessions(authenticator=authenticator),
-            FeatureFlags(authenticator=authenticator),
-            Insights(authenticator=authenticator),
-            InsightsPath(authenticator=authenticator),
-            InsightsSessions(authenticator=authenticator),
-            Persons(authenticator=authenticator),
-            Trends(authenticator=authenticator),
+            Annotations(authenticator=authenticator, start_date=config["start_date"], base_url=base_url),
+            Cohorts(authenticator=authenticator, base_url=base_url),
+            Events(authenticator=authenticator, start_date=config["start_date"], base_url=base_url),
+            EventsSessions(authenticator=authenticator, base_url=base_url),
+            FeatureFlags(authenticator=authenticator, base_url=base_url),
+            Insights(authenticator=authenticator, base_url=base_url),
+            InsightsPath(authenticator=authenticator, base_url=base_url),
+            InsightsSessions(authenticator=authenticator, base_url=base_url),
+            Persons(authenticator=authenticator, base_url=base_url),
+            Trends(authenticator=authenticator, base_url=base_url),
         ]
diff --git a/airbyte-integrations/connectors/source-posthog/source_posthog/spec.json b/airbyte-integrations/connectors/source-posthog/source_posthog/spec.json
@@ -12,12 +12,22 @@
         "type": "string",
         "description": "The date from which you'd like to replicate the data",
         "pattern": "^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}Z$",
-        "examples": "2021-01-01T00:00:00.000000Z"
+        "examples": [
+          "2021-01-01T00:00:00.000000Z"
+        ]
       },
       "api_key": {
         "type": "string",
         "airbyte_secret": true,
         "description": "API Key. See the <a href=\"https://docs.airbyte.io/integrations/sources/posthog\">docs</a> for information on how to generate this key."
+      },
+      "base_url": {
+        "type": "string",
+        "default": "https://app.posthog.com",
+        "description": "Base PostHog url. Defaults to PostHog Cloud (https://app.posthog.com).",
+        "examples": [
+          "https://posthog.example.com"
+        ]
       }
     }
   }

diff --git a/airbyte-integrations/connectors/source-posthog/source_posthog/streams.py b/airbyte-integrations/connectors/source-posthog/source_posthog/streams.py
@@ -34,10 +34,17 @@
 
 
 class PosthogStream(HttpStream, ABC):
-    url_base = "https://app.posthog.com/api/"
     primary_key = "id"
     data_field = "results"
 
+    def __init__(self, base_url: str, **kwargs):
+        super().__init__(**kwargs)
+        self._url_base = f"{base_url}/api/"
+
+    @property
+    def url_base(self) -> str:
+        return self._url_base
+
     def next_page_token(self, response: requests.Response) -> Optional[Mapping[str, Any]]:
         resp_json = response.json()
         if resp_json.get("next"):
@@ -76,8 +83,8 @@ class IncrementalPosthogStream(PosthogStream, ABC):
 
     state_checkpoint_interval = math.inf
 
-    def __init__(self, start_date: str, **kwargs):
-        super().__init__(**kwargs)
+    def __init__(self, base_url: str, start_date: str, **kwargs):
+        super().__init__(base_url=base_url, **kwargs)
         self._start_date = start_date
         self._initial_state = None  # we need to keep it here because next_page_token doesn't accept state argument