feat(data-warehouse): Added support for using a different dataset for…

… temp tables in bigquery (#26506)
PostHog · Nov 29, 2024 · 88ba586 · 88ba586
1 parent b1810ea
commit 88ba586
Show file tree

Hide file tree

Showing 5 changed files with 33 additions and 1 deletion.
diff --git a/frontend/src/scenes/data-warehouse/external/forms/SourceForm.tsx b/frontend/src/scenes/data-warehouse/external/forms/SourceForm.tsx
@@ -19,6 +19,7 @@ const sourceFieldToElement = (field: SourceFieldConfig, sourceConfig: SourceConf
             <LemonField key={field.name} name={[field.name, 'enabled']} label={field.label}>
                 {({ value, onChange }) => (
                     <>
+                        {!!field.caption && <p>{field.caption}</p>}
                         <LemonSwitch
                             checked={value === undefined || value === null ? lastValue?.['enabled'] : value}
                             onChange={onChange}

diff --git a/frontend/src/scenes/data-warehouse/new/sourceWizardLogic.tsx b/frontend/src/scenes/data-warehouse/new/sourceWizardLogic.tsx
@@ -642,6 +642,23 @@ export const SOURCE_DETAILS: Record<ExternalDataSourceType, SourceConfig> = {
                 required: true,
                 placeholder: '',
             },
+            {
+                type: 'switch-group',
+                name: 'temporary-dataset',
+                label: 'Use a different dataset for the temporary tables?',
+                caption:
+                    "We have to create and delete temporary tables when querying your data, this is a requirement of querying large BigQuery tables. We can use a different dataset if you'd like to limit the permissions available to the service account provided.",
+                default: false,
+                fields: [
+                    {
+                        type: 'text',
+                        name: 'temporary_dataset_id',
+                        label: 'Dataset ID for temporary tables',
+                        required: true,
+                        placeholder: '',
+                    },
+                ],
+            },
         ],
         caption: '',
     },

diff --git a/frontend/src/types.ts b/frontend/src/types.ts
@@ -4455,6 +4455,7 @@ export interface SourceFieldSwitchGroupConfig {
     label: string
     default: string | number | boolean
     fields: SourceFieldConfig[]
+    caption?: string
 }
 
 export interface SourceFieldFileUploadConfig {

diff --git a/posthog/temporal/data_imports/workflow_activities/import_data_sync.py b/posthog/temporal/data_imports/workflow_activities/import_data_sync.py
@@ -355,7 +355,14 @@ def import_data_activity_sync(inputs: ImportDataActivityInputs):
             client_email = model.pipeline.job_inputs.get("client_email")
             token_uri = model.pipeline.job_inputs.get("token_uri")
 
-            destination_table = f"{project_id}.{dataset_id}.__posthog_import_{inputs.run_id}_{str(datetime.now().timestamp()).replace('.', '')}"
+            temporary_dataset_id = model.pipeline.job_inputs.get("temporary_dataset_id")
+            using_temporary_dataset = (
+                model.pipeline.job_inputs.get("using_temporary_dataset", False) and temporary_dataset_id is not None
+            )
+
+            destination_table_dataset_id = temporary_dataset_id if using_temporary_dataset else dataset_id
+            destination_table = f"{project_id}.{destination_table_dataset_id}.__posthog_import_{inputs.run_id}_{str(datetime.now().timestamp()).replace('.', '')}"
+
             try:
                 source = bigquery_source(
                     dataset_id=dataset_id,

diff --git a/posthog/warehouse/api/external_data_source.py b/posthog/warehouse/api/external_data_source.py
@@ -683,6 +683,10 @@ def _handle_bigquery_source(
         client_email = key_file.get("client_email")
         token_uri = key_file.get("token_uri")
 
+        temporary_dataset = request.data.get("temporary-dataset", {})
+        using_temporary_dataset = temporary_dataset.get("enabled", False)
+        temporary_dataset_id = temporary_dataset.get("temporary_dataset_id", None)
+
         new_source_model = ExternalDataSource.objects.create(
             source_id=str(uuid.uuid4()),
             connection_id=str(uuid.uuid4()),
@@ -697,6 +701,8 @@ def _handle_bigquery_source(
                 "private_key_id": private_key_id,
                 "client_email": client_email,
                 "token_uri": token_uri,
+                "using_temporary_dataset": using_temporary_dataset,
+                "temporary_dataset_id": temporary_dataset_id,
             },
             prefix=prefix,
         )