Feat background task nf

.gitignore

@@ -5,6 +5,9 @@ venv
 k8s/*
 !k8s/sample_app.yaml
 
+#for mac and idea
+*.DS_Store
+.idea/
 # TODO: remove
 credentials
 .terraform*

config/settings/base.py

@@ -378,6 +378,9 @@
 
 # Datasets config
 WORKSPACE_DATASETS_BUCKET = os.environ.get("WORKSPACE_DATASETS_BUCKET")
+WORKSPACE_DATASETS_FILE_SNAPSHOT_SIZE = os.environ.get(
+    "WORKSPACE_DATASETS_FILE_SNAPSHOT_SIZE", 50
+)
 
 # Base64 encoded service account key
 # To generate a service account key, follow the instructions here:

hexa/datasets/admin.py

@@ -1,6 +1,11 @@
 from django.contrib import admin
 
-from .models import Dataset, DatasetLink, DatasetVersion, DatasetVersionFile
+from .models import (
+    Dataset,
+    DatasetLink,
+    DatasetVersion,
+    DatasetVersionFile,
+)
 
 
 @admin.register(Dataset)

hexa/datasets/graphql/schema.graphql

@@ -91,6 +91,23 @@ type DatasetVersionPermissions {
   download: Boolean!
 }
 
+"""
+Statuses that can occur when generating file metadata
+"""
+enum FileMetadataStatus{
+    PROCESSING,
+    FAILED,
+    FINISHED
+}
+
+"""
+Metadata for dataset file
+"""
+type DatasetFileMetadata {
+    sample: JSON!
+    status: FileMetadataStatus!
+}
+
 """
 A file in a dataset version.
 """
@@ -101,6 +118,7 @@ type DatasetVersionFile {
   createdAt: DateTime!
   createdBy: User
   contentType: String!
+  fileMetadata: DatasetFileMetadata
 }
 
 """
@@ -315,7 +333,7 @@ type CreateDatasetVersionFileResult {
   "The created file object"
   file: DatasetVersionFile
   "The URL to upload the file to"
-  uploadUrl: String @deprecated(reason: "moved to dedicated generateDatasetUploadUrl mutation")
+  uploadUrl: String! @deprecated(reason: "moved to dedicated generateDatasetUploadUrl mutation")
   success: Boolean!
   errors: [CreateDatasetVersionFileError!]!
 }
@@ -429,6 +447,8 @@ extend type Query {
   dataset(id: ID!): Dataset
   "Get a dataset by its slug."
   datasetVersion(id: ID!): DatasetVersion
+  "Get a dataset file by its id "
+  datasetVersionFile(id: ID!): DatasetVersionFile
   "Get a dataset link by its id."
   datasetLink(id: ID!): DatasetLink
   "Get a dataset link by its slug."
@@ -437,6 +457,7 @@ extend type Query {
   datasets(query: String, page: Int = 1, perPage: Int = 15): DatasetPage!
 }
 
+
 extend type Mutation {
   "Create a new dataset."
   createDataset(input: CreateDatasetInput!): CreateDatasetResult! @loginRequired
@@ -460,4 +481,5 @@ extend type Mutation {
   deleteDatasetLink(input: DeleteDatasetLinkInput!): DeleteDatasetLinkResult! @loginRequired
   "Pin or unpin a dataset for a workspace."
   pinDataset(input: PinDatasetInput!): PinDatasetResult! @loginRequired
-}
+}
+

hexa/datasets/migrations/0006_datasetfilemetadata.py

@@ -0,0 +1,56 @@
+# Generated by Django 5.0.7 on 2024-08-07 09:02
+
+import uuid
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("datasets", "0005_datasetfilemetadatajob"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="DatasetFileMetadata",
+            fields=[
+                (
+                    "id",
+                    models.UUIDField(
+                        default=uuid.uuid4,
+                        editable=False,
+                        primary_key=True,
+                        serialize=False,
+                    ),
+                ),
+                ("created_at", models.DateTimeField(auto_now_add=True)),
+                ("updated_at", models.DateTimeField(auto_now=True)),
+                ("sample", models.JSONField(blank=True, default=list, null=True)),
+                (
+                    "status",
+                    models.CharField(
+                        choices=[
+                            ("PROCESSING", "Processing"),
+                            ("FAILED", "Failed"),
+                            ("FINISHED", "Finished"),
+                        ],
+                        default="PROCESSING",
+                        max_length=10,
+                    ),
+                ),
+                ("status_reason", models.TextField(blank=True, null=True)),
+                (
+                    "dataset_version_file",
+                    models.ForeignKey(
+                        on_delete=django.db.models.deletion.CASCADE,
+                        related_name="metadata_entries",
+                        to="datasets.datasetversionfile",
+                    ),
+                ),
+            ],
+            options={
+                "abstract": False,
+            },
+        ),
+    ]

hexa/datasets/models.py

@@ -3,6 +3,8 @@
 from django.contrib.auth.models import AnonymousUser
 from django.core.exceptions import PermissionDenied
 from django.db import models
+from django.db.models import JSONField
+from django.utils.translation import gettext_lazy as _
 from dpq.models import BaseJob
 from slugify import slugify
 
@@ -251,10 +253,41 @@ class DatasetVersionFile(Base):
     def filename(self):
         return self.uri.split("/")[-1]
 
+    @property
+    def latest_metadata(self):
+        return self.metadata_entries.order_by("-created_at").first()
+
     class Meta:
         ordering = ["uri"]
 
 
+class DatasetFileMetadata(Base):
+    STATUS_PROCESSING = "PROCESSING"
+    STATUS_FAILED = "FAILED"
+    STATUS_FINISHED = "FINISHED"
+
+    STATUS_CHOICES = [
+        (STATUS_PROCESSING, _("Processing")),
+        (STATUS_FAILED, _("Failed")),
+        (STATUS_FINISHED, _("Finished")),
+    ]
+
+    sample = JSONField(blank=True, default=list, null=True)
+    status = models.CharField(
+        max_length=10,
+        choices=STATUS_CHOICES,
+        default=STATUS_PROCESSING,
+    )
+    status_reason = models.TextField(blank=True, null=True)
+    dataset_version_file = models.ForeignKey(
+        DatasetVersionFile,
+        null=False,
+        blank=False,
+        on_delete=models.CASCADE,
+        related_name="metadata_entries",
+    )
+
+
 class DatasetLinkQuerySet(BaseQuerySet):
     def filter_for_user(self, user: AnonymousUser | User):
         # FIXME: Use a generic permission system instead of differencing between User and PipelineRunUser

hexa/datasets/permissions.py

@@ -1,4 +1,8 @@
-from hexa.datasets.models import Dataset, DatasetLink, DatasetVersion
+from hexa.datasets.models import (
+    Dataset,
+    DatasetLink,
+    DatasetVersion,
+)
 from hexa.user_management.models import User
 from hexa.workspaces.models import (
     Workspace,

hexa/datasets/queue.py

@@ -1,18 +1,108 @@
+import json
 from logging import getLogger
 
+import pandas as pd
+from django.conf import settings
+from django.core.exceptions import ObjectDoesNotExist, ValidationError
+from django.db import DatabaseError, IntegrityError
 from dpq.queue import AtLeastOnceQueue
 
-from hexa.datasets.models import DatasetFileMetadataJob
+from hexa.core import mimetypes
+from hexa.datasets.api import generate_download_url
+from hexa.datasets.models import (
+    DatasetFileMetadata,
+    DatasetFileMetadataJob,
+    DatasetVersionFile,
+)
 
 logger = getLogger(__name__)
 
 
+def is_supported_mimetype(filename: str) -> bool:
+    supported_mimetypes = [
+        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
+        "application/vnd.ms-excel",
+        "application/vnd.apache.parquet",
+        "text/csv",
+    ]
+    supported_extensions = ["parquet"]
+    suffix = filename.split(".")[-1]
+    mime_type, encoding = mimetypes.guess_type(filename, strict=False)
+    return mime_type in supported_mimetypes or suffix in supported_extensions
+
+
+def download_file_as_dataframe(
+    dataset_version_file: DatasetVersionFile,
+) -> pd.DataFrame | None:
+    mime_type, encoding = mimetypes.guess_type(
+        dataset_version_file.filename, strict=False
+    )
+    download_url = generate_download_url(dataset_version_file)
+    if mime_type == "text/csv":
+        return pd.read_csv(download_url)
+    elif (
+        mime_type == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
+        or mime_type == "application/vnd.ms-excel"
+    ):
+        return pd.read_excel(download_url)
+    elif (
+        mime_type == "application/vnd.apache.parquet"
+        or dataset_version_file.filename.split(".")[-1] == "parquet"
+    ):
+        return pd.read_parquet(download_url)
+
+
 def generate_dataset_file_sample_task(
     queue: AtLeastOnceQueue, job: DatasetFileMetadataJob
 ):
-    # TODO: imlpement ticket PATHWAYS-98  - extract data in background task
-    dataset_version_file_id = job.args["fileId"]
-    logger.info(f"Creating dataset version file {dataset_version_file_id}")
+    dataset_version_file_id = job.args["file_id"]
+    try:
+        dataset_version_file = DatasetVersionFile.objects.get(
+            id=dataset_version_file_id
+        )
+    except ObjectDoesNotExist as e:
+        logger.error(
+            f"DatasetVersionFile with id {dataset_version_file_id} does not exist: {e}"
+        )
+        return
+
+    if not is_supported_mimetype(dataset_version_file.filename):
+        logger.info(f"Unsupported file format: {dataset_version_file.filename}")
+        return
+
+    logger.info(f"Creating dataset sample for version file {dataset_version_file.id}")
+    try:
+        dataset_file_metadata = DatasetFileMetadata.objects.create(
+            dataset_version_file=dataset_version_file,
+            status=DatasetFileMetadata.STATUS_PROCESSING,
+        )
+    except (IntegrityError, DatabaseError, ValidationError) as e:
+        logger.error(f"Error creating DatasetFileMetadata: {e}")
+        return
+
+    try:
+        file_content = download_file_as_dataframe(dataset_version_file)
+        if not file_content.empty:
+            random_seed = 22
+            file_sample = file_content.sample(
+                settings.WORKSPACE_DATASETS_FILE_SNAPSHOT_SIZE,
+                random_state=random_seed,
+                replace=True,
+            )
+            dataset_file_metadata.sample = file_sample.to_json(orient="records")
+        else:
+            dataset_file_metadata.sample = json.dumps([])
+        logger.info(f"Dataset sample saved for file {dataset_version_file_id}")
+        dataset_file_metadata.status = DatasetFileMetadata.STATUS_FINISHED
+        dataset_file_metadata.save()
+        logger.info(f"Dataset sample created for file {dataset_version_file_id}")
+    except Exception as e:
+        dataset_file_metadata.status = DatasetFileMetadata.STATUS_FAILED
+        dataset_file_metadata.status_reason = str(e)
+        dataset_file_metadata.save()
+        logger.exception(
+            f"Dataset file sample creation failed for file {dataset_version_file_id}: {e}"
+        )
 
 
 class DatasetsFileMetadataQueue(AtLeastOnceQueue):
@@ -21,18 +111,16 @@ class DatasetsFileMetadataQueue(AtLeastOnceQueue):
 
 dataset_file_metadata_queue = DatasetsFileMetadataQueue(
     tasks={
-        "generate_file_sample": generate_dataset_file_sample_task,
+        "generate_file_metadata": generate_dataset_file_sample_task,
     },
     notify_channel="dataset_file_metadata_queue",
 )
 
 
 def load_file_metadata(file_id: str):
     dataset_file_metadata_queue.enqueue(
+        "generate_file_metadata",
         {
-            "generate_file_metadata",
-            {
-                "file_id": str(file_id),
-            },
-        }
+            "file_id": str(file_id),
+        },
     )

hexa/datasets/schema/queries.py

@@ -2,7 +2,12 @@
 
 from hexa.core.graphql import result_page
 
-from ..models import Dataset, DatasetLink, DatasetVersion
+from ..models import (
+    Dataset,
+    DatasetLink,
+    DatasetVersion,
+    DatasetVersionFile,
+)
 
 datasets_queries = QueryType()
 
@@ -37,6 +42,17 @@ def resolve_dataset_version(_, info, **kwargs):
         return None
 
 
+@datasets_queries.field("datasetVersionFile")
+def resolve_dataset_version_file(_, info, **kwargs):
+    request = info.context["request"]
+    try:
+        return DatasetVersionFile.objects.filter_for_user(request.user).get(
+            id=kwargs["id"]
+        )
+    except DatasetVersionFile.DoesNotExist:
+        return None
+
+
 @datasets_queries.field("datasetLink")
 def resolve_dataset_link(_, info, **kwargs):
     request = info.context["request"]