ai-chatbot/scripts/storage/configure-multipart-lifecycle.ts

import { config as loadEnv } from "dotenv";

import {
  GetBucketLifecycleConfigurationCommand,
  PutBucketLifecycleConfigurationCommand,
  S3Client,
  type LifecycleRule,
} from "@aws-sdk/client-s3";
import { and, eq } from "drizzle-orm";

loadEnv({ path: ".env.local" });
loadEnv();

type S3UploadConfig = {
  endpoint: string;
  region: string;
  bucket: string;
  accessKeyId: string;
  secretAccessKey: string;
  provider?: "s3" | "qiniu-kodo";
};

const DEFAULT_RULE_ID = "abort-incomplete-direct-upload-parts";
const DEFAULT_PREFIX = "uploads/direct/";
const DEFAULT_ABORT_DAYS = 2;

function normalizeEndpoint(endpoint: string): string {
  const trimmed = endpoint.trim().replace(/\/+$/, "");
  if (!trimmed) return "";

  if (/aliyuncs\.com/i.test(trimmed)) {
    const withoutProtocol = trimmed.replace(/^https?:\/\//i, "");
    return `https://${withoutProtocol}`;
  }

  return trimmed;
}

function createS3Client(config: S3UploadConfig) {
  const normalizedEndpoint = normalizeEndpoint(config.endpoint);
  const lowerEndpoint = normalizedEndpoint.toLowerCase();
  const isCloudS3 =
    lowerEndpoint.includes("aliyuncs.com") ||
    lowerEndpoint.includes("amazonaws.com");

  return new S3Client({
    endpoint: normalizedEndpoint,
    region: config.region || "auto",
    credentials: {
      accessKeyId: config.accessKeyId,
      secretAccessKey: config.secretAccessKey,
    },
    forcePathStyle: config.provider === "qiniu-kodo" ? true : !isCloudS3,
  });
}

function toS3Config(value: unknown): S3UploadConfig | null {
  const raw = value as Record<string, unknown> | null | undefined;
  const endpoint = raw?.endpoint ? normalizeEndpoint(String(raw.endpoint)) : "";
  const bucket = raw?.bucket ? String(raw.bucket).trim() : "";
  const accessKeyId = raw?.accessKeyId ? String(raw.accessKeyId).trim() : "";
  const secretAccessKey = raw?.secretAccessKey
    ? String(raw.secretAccessKey).trim()
    : "";

  if (
    !raw ||
    raw.type !== "s3" ||
    raw.isEnabled !== true ||
    !endpoint ||
    !bucket ||
    !accessKeyId ||
    !secretAccessKey
  ) {
    return null;
  }

  return {
    endpoint,
    region: String(raw.region || "auto").trim() || "auto",
    bucket,
    accessKeyId,
    secretAccessKey,
    provider: raw.provider === "qiniu-kodo" ? "qiniu-kodo" : "s3",
  };
}

function readEnvS3Config(): S3UploadConfig | null {
  const endpoint = process.env.S3_ENDPOINT
    ? normalizeEndpoint(process.env.S3_ENDPOINT)
    : "";
  const bucket = process.env.S3_BUCKET?.trim() || "";
  const accessKeyId = process.env.S3_ACCESS_KEY_ID?.trim() || "";
  const secretAccessKey = process.env.S3_SECRET_ACCESS_KEY?.trim() || "";

  if (!endpoint || !bucket || !accessKeyId || !secretAccessKey) {
    return null;
  }

  return {
    endpoint,
    region: process.env.S3_REGION?.trim() || "auto",
    bucket,
    accessKeyId,
    secretAccessKey,
  };
}

async function resolveS3Config() {
  const [{ db }, { storageConfig }] = await Promise.all([
    import("@/lib/db"),
    import("@/lib/db/schema"),
  ]);
  const [defaultConfig] = await db
    .select()
    .from(storageConfig)
    .where(
      and(
        eq(storageConfig.type, "s3"),
        eq(storageConfig.isDefault, true),
        eq(storageConfig.isEnabled, true)
      )
    )
    .limit(1);

  return toS3Config(defaultConfig) || readEnvS3Config();
}

function readAbortDays() {
  const value = Number(process.env.DIRECT_UPLOAD_ABORT_INCOMPLETE_DAYS || "");
  if (Number.isFinite(value) && value >= 1 && value <= 30) {
    return Math.floor(value);
  }
  return DEFAULT_ABORT_DAYS;
}

function readRulePrefix() {
  const raw = process.env.DIRECT_UPLOAD_LIFECYCLE_PREFIX || DEFAULT_PREFIX;
  const normalized = raw.trim().replace(/^\/+/, "");
  return normalized || DEFAULT_PREFIX;
}

function normalizeRule(rule: LifecycleRule): LifecycleRule {
  return {
    ...rule,
    ID: rule.ID || DEFAULT_RULE_ID,
    Status: rule.Status || "Enabled",
  };
}

async function main() {
  const config = await resolveS3Config();
  if (!config) {
    console.log("[multipart-lifecycle] skip: no S3-compatible upload target");
    return;
  }

  const client = createS3Client(config);
  const prefix = readRulePrefix();
  const abortDays = readAbortDays();
  const ruleId = process.env.DIRECT_UPLOAD_LIFECYCLE_RULE_ID || DEFAULT_RULE_ID;

  let existingRules: LifecycleRule[] = [];
  try {
    const current = await client.send(
      new GetBucketLifecycleConfigurationCommand({ Bucket: config.bucket })
    );
    existingRules = current.Rules || [];
  } catch (error) {
    console.warn(
      "[multipart-lifecycle] existing lifecycle not found or unreadable, replacing with direct upload cleanup rule:",
      error instanceof Error ? error.message : String(error)
    );
  }

  const cleanupRule: LifecycleRule = {
    ID: ruleId,
    Status: "Enabled",
    Filter: { Prefix: prefix },
    AbortIncompleteMultipartUpload: {
      DaysAfterInitiation: abortDays,
    },
  };

  const mergedRules = [
    cleanupRule,
    ...existingRules.filter((rule) => rule.ID !== ruleId),
  ]
    .slice(0, 1000)
    .map(normalizeRule);

  await client.send(
    new PutBucketLifecycleConfigurationCommand({
      Bucket: config.bucket,
      LifecycleConfiguration: {
        Rules: mergedRules,
      },
    })
  );

  console.log("[multipart-lifecycle] configured", {
    bucket: config.bucket,
    prefix,
    abortDays,
    ruleCount: mergedRules.length,
  });
}

main()
  .then(() => {
    process.exit(0);
  })
  .catch((error) => {
    console.error("[multipart-lifecycle] failed:", error);
    process.exit(1);
  });