fix(data-warehouse): Check table name before creation (#23947)

PostHog · Jul 24, 2024 · 0a015ed · 0a015ed
1 parent af8a881
commit 0a015ed
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 3 deletions.
diff --git a/posthog/warehouse/api/table.py b/posthog/warehouse/api/table.py
@@ -1,6 +1,6 @@
 from typing import Any
 
-from rest_framework import filters, request, response, serializers, status, viewsets
+from rest_framework import exceptions, filters, request, response, serializers, status, viewsets
 from rest_framework.decorators import action
 
 from posthog.api.routing import TeamAndOrgViewSetMixin
@@ -89,11 +89,17 @@ def get_external_schema(self, instance: DataWarehouseTable):
         return SimpleExternalDataSchemaSerializer(instance.externaldataschema_set.first(), read_only=True).data or None
 
     def create(self, validated_data):
-        validated_data["team_id"] = self.context["team_id"]
+        team_id = self.context["team_id"]
+
+        table_name_exists = DataWarehouseTable.objects.filter(team_id=team_id, name=validated_data["name"]).exists()
+        if table_name_exists:
+            raise exceptions.ValidationError("Table name already exists.")
+
+        validated_data["team_id"] = team_id
         validated_data["created_by"] = self.context["request"].user
         if validated_data.get("credential"):
             validated_data["credential"] = DataWarehouseCredential.objects.create(
-                team_id=self.context["team_id"],
+                team_id=team_id,
                 access_key=validated_data["credential"]["access_key"],
                 access_secret=validated_data["credential"]["access_secret"],
             )

diff --git a/posthog/warehouse/api/test/test_table.py b/posthog/warehouse/api/test/test_table.py
@@ -232,3 +232,49 @@ def test_update_schema_400_with_invalid_type(self):
         assert response.status_code == 400
         assert response.json()["message"] == "Can not parse type another_type for column id - type does not exist"
         assert table.columns == columns
+
+    @patch(
+        "posthog.warehouse.models.table.DataWarehouseTable.get_columns",
+        return_value={
+            "id": {"clickhouse": "Nullable(String)", "hogql": "StringDatabaseField", "valid": True},
+            "a_column": {"clickhouse": "Nullable(String)", "hogql": "StringDatabaseField", "valid": True},
+        },
+    )
+    @patch(
+        "posthog.warehouse.models.table.DataWarehouseTable.validate_column_type",
+        return_value=True,
+    )
+    def test_table_name_duplicate(self, patch_get_columns, patch_validate_column_type):
+        response = self.client.post(
+            f"/api/projects/{self.team.id}/warehouse_tables/",
+            {
+                "name": "whatever",
+                "url_pattern": "https://your-org.s3.amazonaws.com/bucket/whatever.pqt",
+                "credential": {
+                    "access_key": "_accesskey",
+                    "access_secret": "_accesssecret",
+                },
+                "format": "Parquet",
+            },
+        )
+        assert response.status_code == 201
+        data: dict[str, Any] = response.json()
+
+        table = DataWarehouseTable.objects.get(id=data["id"])
+
+        assert table is not None
+
+        response = self.client.post(
+            f"/api/projects/{self.team.id}/warehouse_tables/",
+            {
+                "name": "whatever",
+                "url_pattern": "https://your-org.s3.amazonaws.com/bucket/whatever.pqt",
+                "credential": {
+                    "access_key": "_accesskey",
+                    "access_secret": "_accesssecret",
+                },
+                "format": "Parquet",
+            },
+        )
+        assert response.status_code == 400
+        assert DataWarehouseTable.objects.count() == 1