Update page_files/Categorized_Search.py

page_files/Categorized_Search.py

@@ -12,6 +12,40 @@ from data_loader import get_all_sections, load_material_data
 import streamlit.components.v1 as components
 import time
 t0 = time.time()
+import os
+from urllib.parse import urlparse
+
+S3_BUCKET = os.getenv("S3_BUCKET", "imagesstorageaims")
+S3_REGION = os.getenv("AWS_REGION", "us-east-2")
+
+_s3_client = None
+def _s3():
+    global _s3_client
+    if _s3_client is None:
+        import boto3
+        _s3_client = boto3.client(
+            "s3",
+            aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID"),
+            aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY"),
+            region_name=S3_REGION,
+        )
+    return _s3_client
+
+def _s3_bytes(ref: str):
+    """ref = S3 key, or full s3:// / https URL -> object bytes, or None."""
+    if not ref:
+        return None
+    try:
+        if ref.startswith("s3://"):
+            p = urlparse(ref); bucket, key = p.netloc, p.path.lstrip("/")
+        elif ref.startswith("http"):
+            p = urlparse(ref); bucket = p.netloc.split(".s3")[0]; key = p.path.lstrip("/")
+        else:
+            bucket, key = S3_BUCKET, ref
+        return _s3().get_object(Bucket=bucket, Key=key)["Body"].read()
+    except Exception as e:
+        st.caption(f"(S3 fetch failed: {e})")
+        return None
 
 # _class -> RDS category table
 _CLASS_TABLE = {"Composites": "Composites_materials", "Fibers": "Fibers", "Polymers": "Polymers"}
@@ -84,12 +118,24 @@ def _load_plot_bytes(selected_abbr, property_choice, result_df):
         for where, params in attempts:
             try:
                 with engine.connect() as conn:
-                    img = conn.execute(text(
-                        f'SELECT image FROM "{tbl}" '
-                        f'WHERE {where} AND image IS NOT NULL LIMIT 1'),
-                        params).scalar()
-                if img:
-                    return bytes(img)   # memoryview -> bytes for st.image
+                    try:
+                        row = conn.execute(text(
+                            f'SELECT image_url, image FROM "{tbl}" '
+                            f'WHERE {where} LIMIT 1'), params).fetchone()
+                        url, blob = (row[0], row[1]) if row else (None, None)
+                    except Exception:
+                        # image_url column not present yet -> blob only
+                        blob = conn.execute(text(
+                            f'SELECT image FROM "{tbl}" '
+                            f'WHERE {where} AND image IS NOT NULL LIMIT 1'),
+                            params).scalar()
+                        url = None
+                if url:
+                    b = _s3_bytes(url)
+                    if b:
+                        return b
+                if blob:
+                    return bytes(blob)   # memoryview -> bytes for st.image
             except Exception:
                 continue
     return None