Skip to content

Commit

Permalink
feat: improvements to media/s3 pulling
Browse files Browse the repository at this point in the history
  • Loading branch information
dbirman committed Nov 25, 2024
1 parent 5e86f6b commit 4c99382
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 17 deletions.
38 changes: 23 additions & 15 deletions src/aind_qc_portal/panel/media.py
Original file line number Diff line number Diff line change
@@ -1,15 +1,16 @@
import io
import tempfile
import panel as pn
from io import BytesIO
from urllib.parse import urlparse
import param
import boto3
from pathlib import Path
from panel.reactive import ReactiveHTML
import requests
import time
import os

s3_client = boto3.client("s3")

CSS = """
:not(:root):fullscreen::backdrop {
background: white;
Expand Down Expand Up @@ -147,13 +148,16 @@ def parse_reference(self, reference: str):
elif "s3" in reference:
bucket = reference.split("/")[2]
key = "/".join(reference.split("/")[3:])
reference_data = _get_s3_data(self.parent.s3_client, bucket, key)
reference_data = _get_s3_data(bucket, key)
elif "sha" in reference:
reference_data = _get_kachery_cloud_url(reference)
else:
# assume local data asset_get_s3_asset

# if a user appends extra things up to results/, strip that
if "results/" in reference:
reference = reference.split("results/")[1]
reference_data = _get_s3_data(
self.parent.s3_client,
self.parent.s3_bucket,
str(Path(self.parent.s3_prefix) / reference),
)
Expand All @@ -179,13 +183,17 @@ def _is_image(reference):


def _get_s3_file(url, ext):
response = requests.get(url)
if response.status_code == 200:
with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as temp_file:
temp_file.write(response.content)
return temp_file.name
else:
print(f"[ERROR] Failed to fetch asset {url}: {response.status_code} / {response.text}")
try:
response = requests.get(url)
if response.status_code == 200:
with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as temp_file:
temp_file.write(response.content)
return temp_file.name
else:
print(f"[ERROR] Failed to fetch asset {url}: {response.status_code} / {response.text}")
return None
except Exception as e:
print(f"[ERROR] Failed to fetch asset {url}, error: {e}")
return None


Expand All @@ -199,7 +207,7 @@ def _parse_type(reference, data):
data : _type_
_description_
"""
if "s3" in data:
if "s3://" in data:
data = _get_s3_file(data, os.path.splitext(reference)[1])

if not data:
Expand Down Expand Up @@ -228,18 +236,18 @@ def _parse_type(reference, data):


@pn.cache()
def _get_s3_data(s3_client, bucket, key):
def _get_s3_data(bucket, key):
"""Get an S3 asset from the given bucket and key
Parameters
----------
s3_client : boto3.client
S3 client object
bucket : str
S3 bucket name
key : str
S3 key name
"""

print((f"Getting S3 data for {bucket}/{key}"))
try:
response = s3_client.get_object(Bucket=bucket, Key=key)
data = BytesIO(response["Body"].read())
Expand Down
2 changes: 0 additions & 2 deletions src/aind_qc_portal/panel/quality_control.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
"""Build the Quality Control Panel object"""

import json
import boto3

import panel as pn
import pandas as pd
Expand All @@ -28,7 +27,6 @@ def __init__(self, id, **params):
super().__init__(**params)

self.id = id
self.s3_client = boto3.client("s3")

# Set up the submission area
self.submit_button = pn.widgets.Button(
Expand Down

0 comments on commit 4c99382

Please sign in to comment.