feat(data-warehouse): Added Stripe account table (#24701)

PostHog · Sep 6, 2024 · 7bed962 · 7bed962
1 parent fa6b950
commit 7bed962
Show file tree

Hide file tree

Showing 6 changed files with 116 additions and 1 deletion.
diff --git a/latest_migrations.manifest b/latest_migrations.manifest
@@ -5,7 +5,7 @@ contenttypes: 0002_remove_content_type_name
 ee: 0016_rolemembership_organization_member
 otp_static: 0002_throttling
 otp_totp: 0002_auto_20190420_0723
-posthog: 0464_action_pinned_at
+posthog: 0465_datawarehouse_stripe_account
 sessions: 0001_initial
 social_django: 0010_uid_db_index
 two_factor: 0007_auto_20201201_1019
diff --git a/mypy-baseline.txt b/mypy-baseline.txt
@@ -780,6 +780,7 @@ posthog/temporal/data_imports/pipelines/stripe/__init__.py:0: error: Unused "typ
 posthog/temporal/data_imports/pipelines/stripe/__init__.py:0: error: Unused "type: ignore" comment  [unused-ignore]
 posthog/temporal/data_imports/pipelines/stripe/__init__.py:0: error: Unused "type: ignore" comment  [unused-ignore]
 posthog/temporal/data_imports/pipelines/stripe/__init__.py:0: error: Unused "type: ignore" comment  [unused-ignore]
+posthog/temporal/data_imports/pipelines/stripe/__init__.py:0: error: Unused "type: ignore" comment  [unused-ignore]
 posthog/tasks/test/test_email.py:0: error: Argument 1 to "send_batch_export_run_failure" has incompatible type "UUID"; expected "str"  [arg-type]
 posthog/tasks/test/test_email.py:0: error: Argument 1 to "send_batch_export_run_failure" has incompatible type "UUID"; expected "str"  [arg-type]
 posthog/tasks/test/test_email.py:0: error: Argument 1 to "send_batch_export_run_failure" has incompatible type "UUID"; expected "str"  [arg-type]

diff --git a/posthog/migrations/0465_datawarehouse_stripe_account.py b/posthog/migrations/0465_datawarehouse_stripe_account.py
@@ -0,0 +1,35 @@
+# Generated by Django 4.2.14 on 2024-08-12 12:04
+
+from django.db import connection, migrations
+
+
+def insert_stripe_account_schemas(apps, schema_editor):
+    with connection.cursor() as cursor:
+        cursor.execute("SELECT id, team_id FROM posthog_externaldatasource where source_type = 'Stripe'")
+        stripe_sources = cursor.fetchall()
+
+    ExternalDataSchema = apps.get_model("posthog", "ExternalDataSchema")
+    for source in stripe_sources:
+        schema = ExternalDataSchema.objects.create(
+            name="Account",
+            source_id=source[0],
+            team_id=source[1],
+            should_sync=False,
+            sync_type=None,
+            sync_type_config={},
+        )
+        schema.save()
+
+
+def reverse(apps, _):
+    pass
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("posthog", "0464_action_pinned_at"),
+    ]
+
+    operations = [
+        migrations.RunPython(insert_stripe_account_schemas, reverse),
+    ]
diff --git a/posthog/temporal/data_imports/pipelines/stripe/__init__.py b/posthog/temporal/data_imports/pipelines/stripe/__init__.py
@@ -10,6 +10,41 @@
 
 def get_resource(name: str, is_incremental: bool) -> EndpointResource:
     resources: dict[str, EndpointResource] = {
+        "Account": {
+            "name": "Account",
+            "table_name": "account",
+            "primary_key": "id",
+            "write_disposition": {
+                "disposition": "merge",
+                "strategy": "upsert",
+            }
+            if is_incremental
+            else "replace",
+            "columns": get_dlt_mapping_for_external_table("stripe_account"),  # type: ignore
+            "endpoint": {
+                "data_selector": "data",
+                "path": "/v1/accounts",
+                "params": {
+                    # the parameters below can optionally be configured
+                    "created[gte]": {
+                        "type": "incremental",
+                        "cursor_path": "created",
+                        "initial_value": 0,  # type: ignore
+                    }
+                    if is_incremental
+                    else None,
+                    # "currency": "OPTIONAL_CONFIG",
+                    # "ending_before": "OPTIONAL_CONFIG",
+                    # "expand": "OPTIONAL_CONFIG",
+                    "limit": 100,
+                    # "payout": "OPTIONAL_CONFIG",
+                    # "source": "OPTIONAL_CONFIG",
+                    # "starting_after": "OPTIONAL_CONFIG",
+                    # "type": "OPTIONAL_CONFIG",
+                },
+            },
+            "table_format": "delta",
+        },
         "BalanceTransaction": {
             "name": "BalanceTransaction",
             "table_name": "balance_transaction",

diff --git a/posthog/temporal/data_imports/pipelines/stripe/settings.py b/posthog/temporal/data_imports/pipelines/stripe/settings.py
@@ -12,6 +12,14 @@
 INCREMENTAL_ENDPOINTS = ("BalanceTransaction", "Subscription", "Customer", "Product", "Price", "Invoice", "Charge")
 
 INCREMENTAL_FIELDS: dict[str, list[IncrementalField]] = {
+    "Account": [
+        {
+            "label": "created_at",
+            "type": IncrementalFieldType.DateTime,
+            "field": "created",
+            "field_type": IncrementalFieldType.Integer,
+        }
+    ],
     "BalanceTransaction": [
         {
             "label": "created_at",

diff --git a/posthog/warehouse/models/external_table_definitions.py b/posthog/warehouse/models/external_table_definitions.py
@@ -17,6 +17,42 @@
         "__dlt_id": StringDatabaseField(name="_dlt_id", hidden=True),
         "__dlt_load_id": StringDatabaseField(name="_dlt_load_id", hidden=True),
     },
+    "stripe_account": {
+        "id": StringDatabaseField(name="id"),
+        "object": StringDatabaseField(name="object"),
+        "business_profile": StringJSONDatabaseField(name="business_profile"),
+        "business_type": StringDatabaseField(name="business_type"),
+        "capabilities": StringJSONDatabaseField(name="capabilities"),
+        "charges_enabled": BooleanDatabaseField(name="charges_enabled"),
+        "controller": StringJSONDatabaseField(name="controller"),
+        "country": StringDatabaseField(name="country"),
+        "__created": IntegerDatabaseField(name="created", hidden=True),
+        "created_at": ast.ExpressionField(
+            isolate_scope=True,
+            expr=ast.Call(
+                name="toDateTime",
+                args=[
+                    ast.Call(
+                        name="toString",
+                        args=[ast.Call(name="fromUnixTimestamp", args=[ast.Field(chain=["__created"])])],
+                    )
+                ],
+            ),
+            name="created_at",
+        ),
+        "default_currency": StringDatabaseField(name="default_currency"),
+        "details_submitted": BooleanDatabaseField(name="details_submitted"),
+        "email": StringDatabaseField(name="email"),
+        "external_accounts": StringJSONDatabaseField(name="external_accounts"),
+        "future_requirements": StringJSONDatabaseField(name="future_requirements"),
+        "login_links": StringJSONDatabaseField(name="login_links"),
+        "metadata": StringJSONDatabaseField(name="metadata"),
+        "payouts_enabled": BooleanDatabaseField(name="payouts_enabled"),
+        "requirements": StringJSONDatabaseField(name="requirements"),
+        "settings": StringJSONDatabaseField(name="settings"),
+        "tos_acceptance": StringJSONDatabaseField(name="tos_acceptance"),
+        "type": StringDatabaseField(name="type"),
+    },
     "stripe_customer": {
         "id": StringDatabaseField(name="id"),
         "name": StringDatabaseField(name="name"),