File Upload in Machine Learning

import boto3
import hashlib
import requests
from concurrent.futures import ThreadPoolExecutor, as_completed

# === Backend: Generate presigned URLs for multipart upload ===

def initiate_multipart_upload(bucket: str, key: str) -> dict:
    """Initiate a multipart upload and return upload_id + presigned URLs."""
    s3 = boto3.client('s3')
    
    # Start the multipart upload
    response = s3.create_multipart_upload(
        Bucket=bucket,
        Key=key,
        ContentType='application/octet-stream',
        Metadata={'uploaded-by': 'ml-pipeline', 'status': 'uploading'}
    )
    upload_id = response['UploadId']
    return {'upload_id': upload_id, 'bucket': bucket, 'key': key}


def generate_part_urls(bucket: str, key: str, upload_id: str, num_parts: int,
                       expiry: int = 3600) -> list[dict]:
    """Generate presigned URLs for each part."""
    s3 = boto3.client('s3')
    part_urls = []
    for part_number in range(1, num_parts + 1):
        url = s3.generate_presigned_url(
            'upload_part',
            Params={
                'Bucket': bucket,
                'Key': key,
                'UploadId': upload_id,
                'PartNumber': part_number,
            },
            ExpiresIn=expiry,
        )
        part_urls.append({'part_number': part_number, 'url': url})
    return part_urls


def complete_multipart_upload(bucket: str, key: str, upload_id: str,
                              parts: list[dict]) -> dict:
    """Complete the multipart upload after all parts are uploaded."""
    s3 = boto3.client('s3')
    response = s3.complete_multipart_upload(
        Bucket=bucket,
        Key=key,
        UploadId=upload_id,
        MultipartUpload={'Parts': parts},  # [{'PartNumber': 1, 'ETag': '...'}]
    )
    return {'location': response['Location'], 'etag': response['ETag']}


# === Client: Upload file chunks in parallel ===

def upload_file_multipart(file_path: str, part_urls: list[dict],
                          chunk_size: int = 16 * 1024 * 1024,
                          max_workers: int = 4) -> list[dict]:
    """Upload a file in parallel chunks using presigned URLs."""
    parts = []
    
    def upload_part(part_info: dict, data: bytes) -> dict:
        md5 = hashlib.md5(data).digest()
        response = requests.put(
            part_info['url'],
            data=data,
            headers={'Content-MD5': __import__('base64').b64encode(md5).decode()},
        )
        response.raise_for_status()
        return {
            'PartNumber': part_info['part_number'],
            'ETag': response.headers['ETag'],
        }
    
    with open(file_path, 'rb') as f:
        chunks = []
        for part_info in part_urls:
            chunk = f.read(chunk_size)
            if not chunk:
                break
            chunks.append((part_info, chunk))
    
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {
            executor.submit(upload_part, part_info, data): part_info
            for part_info, data in chunks
        }
        for future in as_completed(futures):
            parts.append(future.result())
    
    return sorted(parts, key=lambda p: p['PartNumber'])

import os
import hashlib
from tusclient import client as tus_client

def resumable_upload(file_path: str, upload_url: str,
                     chunk_size: int = 32 * 1024 * 1024) -> str:
    """Upload a file using the tus resumable upload protocol.
    
    If interrupted, calling this again with the same file will
    resume from where it left off.
    """
    file_size = os.path.getsize(file_path)
    file_hash = hashlib.sha256(open(file_path, 'rb').read(8192)).hexdigest()
    
    # Create tus client
    my_client = tus_client.TusClient(upload_url)
    
    # Create uploader with fingerprint for resume capability
    uploader = my_client.uploader(
        file_path,
        chunk_size=chunk_size,
        metadata={
            'filename': os.path.basename(file_path),
            'filetype': _detect_mime_type(file_path),
            'sha256_prefix': file_hash,
        },
        # Store upload URL locally for resumption
        store_url=True,
        url_storage=tus_client.fingerprint.Fingerprint(),
    )
    
    # Upload with automatic resume
    uploader.upload()
    return uploader.url


def _detect_mime_type(file_path: str) -> str:
    """Detect MIME type from file magic bytes."""
    import magic
    mime = magic.Magic(mime=True)
    return mime.from_file(file_path)


# Usage
upload_location = resumable_upload(
    file_path='/data/training_images.tar.gz',
    upload_url='https://upload.example.com/files/',
    chunk_size=64 * 1024 * 1024,  # 64 MB chunks
)

import pyarrow as pa
import pyarrow.parquet as pq
import pyarrow.csv as csv
import json
from pathlib import Path


def infer_schema(file_path: str, sample_rows: int = 10000) -> dict:
    """Infer schema from uploaded file (CSV, Parquet, or JSON).
    
    Returns a schema dict with column names, types, nullability,
    and basic statistics.
    """
    path = Path(file_path)
    ext = path.suffix.lower()
    
    if ext == '.parquet':
        return _infer_parquet_schema(file_path)
    elif ext == '.csv':
        return _infer_csv_schema(file_path, sample_rows)
    elif ext in ('.json', '.jsonl'):
        return _infer_json_schema(file_path, sample_rows)
    else:
        return {'format': ext, 'schema': None, 'error': 'Unsupported format'}


def _infer_parquet_schema(file_path: str) -> dict:
    """Parquet stores schema in footer -- O(1) inference."""
    pf = pq.ParquetFile(file_path)
    schema = pf.schema_arrow
    metadata = pf.metadata
    
    columns = []
    for i in range(len(schema)):
        field = schema.field(i)
        columns.append({
            'name': field.name,
            'type': str(field.type),
            'nullable': field.nullable,
        })
    
    return {
        'format': 'parquet',
        'num_rows': metadata.num_rows,
        'num_columns': len(schema),
        'num_row_groups': metadata.num_row_groups,
        'serialized_size_bytes': metadata.serialized_size,
        'columns': columns,
        'created_by': metadata.created_by,
    }


def _infer_csv_schema(file_path: str, sample_rows: int) -> dict:
    """CSV requires sampling rows for type inference."""
    read_options = csv.ReadOptions(block_size=1024 * 1024)  # 1 MB blocks
    convert_options = csv.ConvertOptions(
        auto_dict_encode=True,
        auto_dict_max_cardinality=50,
    )
    
    # Read a sample for inference
    table = csv.read_csv(
        file_path,
        read_options=read_options,
        convert_options=convert_options,
    )
    
    # Sample if the file is large
    if table.num_rows > sample_rows:
        table = table.slice(0, sample_rows)
    
    columns = []
    for i, name in enumerate(table.column_names):
        col = table.column(i)
        columns.append({
            'name': name,
            'type': str(col.type),
            'nullable': col.null_count > 0,
            'null_count': col.null_count,
            'null_fraction': round(col.null_count / len(col), 4),
        })
    
    return {
        'format': 'csv',
        'num_rows_sampled': table.num_rows,
        'num_columns': len(table.column_names),
        'columns': columns,
    }


def _infer_json_schema(file_path: str, sample_rows: int) -> dict:
    """Infer schema from JSON/JSONL by sampling records."""
    table = pa.json_.read_json(file_path)
    
    if table.num_rows > sample_rows:
        table = table.slice(0, sample_rows)
    
    columns = []
    for i, name in enumerate(table.column_names):
        col = table.column(i)
        columns.append({
            'name': name,
            'type': str(col.type),
            'nullable': col.null_count > 0,
        })
    
    return {
        'format': 'json',
        'num_rows_sampled': table.num_rows,
        'num_columns': len(table.column_names),
        'columns': columns,
    }


# Usage
schema_info = infer_schema('/uploads/training_data.parquet')
print(json.dumps(schema_info, indent=2))

import clamd
import logging
from enum import Enum
from dataclasses import dataclass

logger = logging.getLogger(__name__)


class ScanResult(Enum):
    CLEAN = 'clean'
    INFECTED = 'infected'
    ERROR = 'error'


@dataclass
class FileScanReport:
    file_path: str
    result: ScanResult
    threat_name: str | None = None
    error_message: str | None = None


class FileScanner:
    """Scan uploaded files for malware using ClamAV daemon."""
    
    def __init__(self, host: str = 'localhost', port: int = 3310):
        self.scanner = clamd.ClamdNetworkSocket(host=host, port=port)
        # Verify ClamAV daemon is running
        try:
            self.scanner.ping()
        except clamd.ConnectionError:
            raise RuntimeError(
                'ClamAV daemon is not reachable. '
                'Ensure clamd is running on {host}:{port}'
            )
    
    def scan_file(self, file_path: str) -> FileScanReport:
        """Scan a single file and return the result."""
        try:
            result = self.scanner.scan(file_path)
            if result is None:
                return FileScanReport(
                    file_path=file_path,
                    result=ScanResult.CLEAN
                )
            
            # ClamAV returns: {'/path/to/file': ('FOUND', 'Malware.Name')}
            status, threat = result[file_path]
            if status == 'FOUND':
                logger.warning(f'Malware detected in {file_path}: {threat}')
                return FileScanReport(
                    file_path=file_path,
                    result=ScanResult.INFECTED,
                    threat_name=threat,
                )
            return FileScanReport(
                file_path=file_path,
                result=ScanResult.CLEAN
            )
        except Exception as e:
            logger.error(f'Scan error for {file_path}: {e}')
            return FileScanReport(
                file_path=file_path,
                result=ScanResult.ERROR,
                error_message=str(e),
            )
    
    def scan_stream(self, file_stream) -> FileScanReport:
        """Scan a file stream without writing to disk."""
        try:
            result = self.scanner.instream(file_stream)
            status, threat = result['stream']
            if status == 'FOUND':
                return FileScanReport(
                    file_path='<stream>',
                    result=ScanResult.INFECTED,
                    threat_name=threat,
                )
            return FileScanReport(
                file_path='<stream>',
                result=ScanResult.CLEAN
            )
        except Exception as e:
            return FileScanReport(
                file_path='<stream>',
                result=ScanResult.ERROR,
                error_message=str(e),
            )


# Usage
scanner = FileScanner(host='clamav-service', port=3310)
report = scanner.scan_file('/tmp/uploads/user_dataset.csv')
if report.result == ScanResult.INFECTED:
    # Quarantine the file
    move_to_quarantine(report.file_path)
    notify_security_team(report)

# File Upload Service Configuration (YAML)
upload:
  max_file_size_gb: 100
  max_files_per_upload: 1000
  allowed_formats:
    - csv
    - json
    - jsonl
    - parquet
    - avro
    - png
    - jpg
    - jpeg
    - tiff
    - wav
    - mp3
  chunk_size_mb: 16
  max_parallel_chunks: 8
  presigned_url_expiry_seconds: 3600
  checksum_algorithm: sha256

storage:
  provider: s3  # s3 | gcs | azure_blob
  bucket: ml-data-uploads
  region: ap-south-1  # Mumbai for India-based deployments
  quarantine_bucket: ml-data-quarantine
  lifecycle:
    incomplete_upload_cleanup_days: 7
    quarantine_retention_days: 30

scanning:
  enabled: true
  engine: clamav
  clamav_host: clamav-service
  clamav_port: 3310
  max_scan_size_gb: 10
  timeout_seconds: 300

schema_inference:
  enabled: true
  sample_rows: 10000
  max_columns: 5000
  type_coercion: true
  register_schema: true

quotas:
  per_user_daily_gb: 500
  per_upload_max_gb: 100
  concurrent_uploads_per_user: 5