Added bulk_copy method to Observations manager

Author: Ken Lippold

sta/management/utils.py

@@ -40,11 +40,9 @@ def generate_test_timeseries(datastream_id: UUID):
             result=result
         ))
 
-        if len(observations) >= 50000:
-            Observation.objects.bulk_create(observations)
-            print(f"Created {len(observations)} observations...")
+        if len(observations) >= 100000:
+            Observation.objects.bulk_copy(observations, batch_size=100000)
             observations.clear()
 
     if observations:
-        Observation.objects.bulk_create(observations)
-        print(f"Created {len(observations)} observations...")
+        Observation.objects.bulk_copy(observations)

sta/models/observation.py

@@ -1,7 +1,8 @@
+import io
 import uuid
 import typing
 from typing import Literal, Optional
-from django.db import models
+from django.db import models, connection
 from django.db.models import Q
 from iam.models import Workspace
 from iam.models.utils import PermissionChecker
@@ -49,6 +50,34 @@ def removable(self, user: Optional["User"]):
                   datastream__thing__workspace__collaborators__role__permissions__permission_type__in=["*", "delete"])
             )
 
+    def bulk_copy(self, observations, batch_size=100_000):
+        db_table_sql = connection.ops.quote_name(self.model._meta.db_table)  # noqa
+        db_fields = [field.column for field in self.model._meta.fields if not field.primary_key]  # noqa
+        db_fields_sql = ", ".join(connection.ops.quote_name(field) for field in db_fields)
+
+        def escape_pg_copy(value):
+            if value is None:
+                return r"\N"
+            if isinstance(value, str):
+                return value.replace("\\", "\\\\").replace("\t", "\\t").replace("\n", "\\n").replace("\r", "\\r")
+            return str(value)
+
+        with connection.cursor() as cursor:
+            with cursor.copy(f"COPY {db_table_sql} ({db_fields_sql}) FROM STDIN") as copy:
+                buffer = io.StringIO()
+                for i in range(0, len(observations), batch_size):
+                    batch = observations[i: i + batch_size]
+                    buffer.write("\n".join(
+                        "\t".join(escape_pg_copy(getattr(obs, field, None)) for field in db_fields)
+                        for obs in batch
+                    ) + "\n")
+                    buffer.seek(0)
+                    copy.write(buffer.read())
+                    buffer.truncate(0)
+                    buffer.seek(0)
+
+        return observations
+
 
 class Observation(models.Model, PermissionChecker):
     pk = models.CompositePrimaryKey("datastream_id", "phenomenon_time", "id")

sta/services/sensorthings/observation.py

@@ -180,7 +180,7 @@ def create_observations(
             #     )
 
             try:
-                new_observations_for_datastream = Observation.objects.bulk_create([
+                new_observations_for_datastream = Observation.objects.bulk_copy([
                     Observation(
                         datastream_id=observation.datastream.id,
                         phenomenon_time=observation.phenomenon_time,