SnowflakeのS3外部ステージにStreamlitアプリからオブジェクトをアップロード・削除する

import io
import uuid
from pathlib import PurePosixPath

import streamlit as st
from snowflake.snowpark.context import get_active_session

INTERNAL_STAGE_FQN = "データベース名.スキーマ名.内部ステージ名"
INTERNAL_STAGE_REF = f"@{INTERNAL_STAGE_FQN}"

EXTERNAL_STAGE_FQN = "データベース名.スキーマ名.外部ステージ名"
EXTERNAL_STAGE_REF = f"@{EXTERNAL_STAGE_FQN}"

DELETE_CANDIDATES_KEY = "delete_candidate_names"

def stage_location_for_key(stage_ref: str, relative_key: str) -> str:
    file_name = PurePosixPath(relative_key).name
    return f"{stage_ref}/{file_name}"

def describe_stages(session) -> bool:
    can_desc_stage = True
    for label, fqn in (
        ("内部", INTERNAL_STAGE_FQN),
        ("外部", EXTERNAL_STAGE_FQN),
    ):
        try:
            session.sql(f"DESC STAGE {fqn}").collect()
        except Exception as exc:
            st.error(
                f"{label}ステージ `{fqn}` を参照できません。名前・権限（USAGE 等）を確認してください。"
                f"\n\n詳細: {exc}"
            )
            can_desc_stage = False
    return can_desc_stage

def list_stage_files(session):
    try:
        return session.sql(f"LIST {EXTERNAL_STAGE_REF}").collect()
    except Exception as exc:
        st.warning(f"ファイル一覧の取得に失敗しました: {exc}")
        return []

def format_list_row(row) -> str:
    try:
        name = row["name"]
    except (KeyError, TypeError, IndexError):
        name = row[0] if row is not None else ""
    if isinstance(name, str) and "/" in name:
        return name.split("/", 1)[1]
    return str(name)

def main():
    st.set_page_config(page_title="ファイルアップローダー", layout="centered")
    st.title("ファイルアップローダー")

    session = get_active_session()

    if not describe_stages(session):
        st.stop()

    uploaded_file = st.file_uploader(
        "アップロードするファイルを選択してください",
        help=(
            "ブラウザから SiS に読み込んだあと、内部ステージへ put_stream し、"
            "copy_files で外部ステージ（S3）へコピーします。内部ステージの一時オブジェクトは "
            "成功・失敗にかかわらず finally で削除します。"
        ),
    )

    if uploaded_file is not None:
        file_name = PurePosixPath(uploaded_file.name).name
        tmp_prefix = f"streamlit_upload_tmp/{uuid.uuid4()}/"
        internal_dir = f"{INTERNAL_STAGE_REF}/{tmp_prefix}"
        internal_file = f"{internal_dir}{file_name}"
        try:
            file_stream = io.BytesIO(uploaded_file.getvalue())
            session.file.put_stream(
                file_stream,
                internal_file,
                auto_compress=False,
                overwrite=True,
            )
            session.file.copy_files(
                source=internal_dir,
                target_stage_location=f"{EXTERNAL_STAGE_REF}/",
                files=[file_name],
                detailed_output=True,
            )
            st.success(f"S3（外部ステージ）へアップロードしました: `{file_name}`")
        except Exception as exc:
            st.error(f"アップロードに失敗しました: {exc}")
        finally:
            try:
                session.file.remove(internal_dir)
            except Exception as cleanup_exc:
                st.warning(
                    f"内部ステージ上の一時ファイルの削除に失敗しました: {cleanup_exc}"
                )

    with st.expander("アップロード済みファイル一覧", expanded=False):
        if st.button("一覧を取得", key="list_stage_files"):
            rows = list_stage_files(session)
            if rows:
                names = [format_list_row(r) for r in rows]
                st.dataframe({"name": names}, use_container_width=True, hide_index=True)
            else:
                st.info(
                    "ファイルが見つかりません（空、または LIST 権限がありません）。"
                )

    with st.expander("アップロード済みファイル削除", expanded=False):
        if st.button("候補一覧を読み込み", key="load_delete_candidates"):
            rows = list_stage_files(session)
            st.session_state[DELETE_CANDIDATES_KEY] = (
                [format_list_row(r) for r in rows] if rows else []
            )

        candidates = st.session_state.get(DELETE_CANDIDATES_KEY)
        if candidates is None:
            st.info("「候補一覧を読み込み」で LIST 結果を読み込んでください。")
        elif not candidates:
            st.info("削除対象がありません（空、または LIST 権限がありません）。")
        else:
            selected = st.selectbox(
                "削除するオブジェクトキー（ステージ直下の相対パス）",
                options=candidates,
                key="delete_select",
            )
            confirm = st.checkbox(
                "選択したファイルを削除する（取り消しできません）",
                value=False,
                key="delete_confirm",
            )
            if st.button("削除を実行", type="primary", disabled=not confirm):
                if selected not in candidates:
                    st.error("選択が無効です。候補一覧を再読み込みしてください。")
                else:
                    try:
                        loc = stage_location_for_key(EXTERNAL_STAGE_REF, selected)
                        session.file.remove(loc)
                        st.success(f"削除しました: `{selected}`")
                        st.session_state[DELETE_CANDIDATES_KEY] = [
                            n for n in candidates if n != selected
                        ]
                    except Exception as exc:
                        st.error(f"削除に失敗しました: {exc}")

if __name__ == "__main__":
    main()

import io
from pathlib import PurePosixPath

import streamlit as st
from snowflake.snowpark.context import get_active_session

STAGE_FQN = "データベース名.スキーマ名.外部ステージ名"
STAGE_REF = f"@{STAGE_FQN}"

DELETE_CANDIDATES_KEY = "delete_candidate_names"

def stage_location_for_key(stage_ref: str, relative_key: str) -> str:
    file_name = PurePosixPath(relative_key).name
    return f"{stage_ref}/{file_name}"

def describe_stage(session) -> bool:
    try:
        session.sql(f"DESC STAGE {STAGE_FQN}").collect()
        return True
    except Exception as exc:
        st.error(
            "ステージを参照できません。名前・権限（USAGE 等）を確認してください。"
            f"\n\n詳細: {exc}"
        )
        return False

def list_stage_files(session):
    try:
        return session.sql(f"LIST {STAGE_REF}").collect()
    except Exception as exc:
        st.warning(f"ファイル一覧の取得に失敗しました: {exc}")
        return []

def format_list_row(row) -> str:
    try:
        name = row["name"]
    except (KeyError, TypeError, IndexError):
        name = row[0] if row is not None else ""
    if isinstance(name, str) and "/" in name:
        return name.split("/", 1)[1]
    return str(name)

def main():
    st.set_page_config(page_title="ファイルアップローダー", layout="centered")
    st.title("ファイルアップローダー")

    session = get_active_session()

    if not describe_stage(session):
        st.stop()

    uploaded_file = st.file_uploader(
        "アップロードするファイルを選択してください",
        help="ブラウザから SiS に読み込んだあと、外部ステージ経由で S3 に PUT します。",
    )

    if uploaded_file is not None:
        try:
            file_stream = io.BytesIO(uploaded_file.getvalue())
            session.file.put_stream(
                file_stream,
                f"{STAGE_REF}/{uploaded_file.name}",
                auto_compress=False,
                overwrite=True,
            )
            st.success(
                f"S3（外部ステージ）へアップロードしました: `{uploaded_file.name}`"
            )
        except Exception as exc:
            st.error(f"アップロードに失敗しました: {exc}")

    with st.expander("アップロード済みファイル一覧", expanded=False):
        if st.button("一覧を取得", key="list_stage_files"):
            rows = list_stage_files(session)
            if rows:
                names = [format_list_row(r) for r in rows]
                st.dataframe({"name": names}, use_container_width=True, hide_index=True)
            else:
                st.info(
                    "ファイルが見つかりません（空、または LIST 権限がありません）。"
                )

    with st.expander("アップロード済みファイル削除", expanded=False):
        if st.button("候補一覧を読み込み", key="load_delete_candidates"):
            rows = list_stage_files(session)
            st.session_state[DELETE_CANDIDATES_KEY] = (
                [format_list_row(r) for r in rows] if rows else []
            )

        candidates = st.session_state.get(DELETE_CANDIDATES_KEY)
        if candidates is None:
            st.info("「候補一覧を読み込み」で LIST 結果を読み込んでください。")
        elif not candidates:
            st.info("削除対象がありません（空、または LIST 権限がありません）。")
        else:
            selected = st.selectbox(
                "削除するオブジェクトキー（ステージ直下の相対パス）",
                options=candidates,
                key="delete_select",
            )
            confirm = st.checkbox(
                "選択したファイルを削除する（取り消しできません）",
                value=False,
                key="delete_confirm",
            )
            if st.button("削除を実行", type="primary", disabled=not confirm):
                if selected not in candidates:
                    st.error("選択が無効です。候補一覧を再読み込みしてください。")
                else:
                    try:
                        loc = stage_location_for_key(STAGE_REF, selected)
                        session.file.remove(loc)
                        st.success(f"削除しました: `{selected}`")
                        st.session_state[DELETE_CANDIDATES_KEY] = [
                            n for n in candidates if n != selected
                        ]
                    except Exception as exc:
                        st.error(f"削除に失敗しました: {exc}")

if __name__ == "__main__":
    main()