cosmos-explorer/utils/cleanupDBs.js

const { AzureCliCredential } = require("@azure/identity");
const { CosmosDBManagementClient } = require("@azure/arm-cosmosdb");
const ms = require("ms");

const subscriptionId = process.env["AZURE_SUBSCRIPTION_ID"];
const resourceGroupName = process.env["E2ETESTS_RESOURCEGROUP_NAME"];

function parseCleanupMinimumAge(configuredAge) {
  let parsedAge;
  try {
    parsedAge = ms(configuredAge === undefined ? "6h" : configuredAge);
  } catch {
    parsedAge = undefined;
  }

  if (!Number.isFinite(parsedAge) || parsedAge <= 0) {
    throw new Error("E2E_CLEANUP_MINIMUM_AGE must be a positive duration");
  }

  return parsedAge;
}
const cleanupMinimumAge = parseCleanupMinimumAge(process.env["E2E_CLEANUP_MINIMUM_AGE"]);
const cleanupThreshold = Date.now() - cleanupMinimumAge;

function shouldDeleteResource(name, timestamp, threshold = cleanupThreshold) {
  return Boolean(name?.startsWith("t_") && timestamp && timestamp < threshold);
}

function friendlyTime(date) {
  try {
    return ms(date);
  } catch (error) {
    return "Unknown";
  }
}

async function main() {
  const credentials = new AzureCliCredential();
  const client = new CosmosDBManagementClient(credentials, subscriptionId);
  const accounts = client.databaseAccounts.listByResourceGroup(resourceGroupName);
  for await (const account of accounts) {
    if (account.name.endsWith("-readonly")) {
      console.log(`SKIPPED: ${account.name}`);
      continue;
    }
    if (account.kind === "MongoDB") {
      const mongoDatabases = client.mongoDBResources.listMongoDBDatabases(resourceGroupName, account.name);
      for await (const database of mongoDatabases) {
        // Unfortunately Mongo does not provide a timestamp in ARM. There is no way to tell how old the DB is other thn encoding it in the ID :(
        const timestamp = Number(database.name.split("_").pop());
        if (shouldDeleteResource(database.name, timestamp)) {
          await client.mongoDBResources.beginDeleteMongoDBDatabaseAndWait(
            resourceGroupName,
            account.name,
            database.name,
          );
          console.log(`DELETED: ${account.name} | ${database.name} | Age: ${friendlyTime(Date.now() - timestamp)}`);
        } else {
          console.log(`SKIPPED: ${account.name} | ${database.name} | Age: ${friendlyTime(Date.now() - timestamp)}`);
        }
      }
    } else if (account.capabilities.find((c) => c.name === "EnableCassandra")) {
      const cassandraDatabases = client.cassandraResources.listCassandraKeyspaces(resourceGroupName, account.name);
      for await (const database of cassandraDatabases) {
        const timestamp = Number(database.resource.ts) * 1000;
        if (shouldDeleteResource(database.name, timestamp)) {
          await client.cassandraResources.beginDeleteCassandraKeyspaceAndWait(
            resourceGroupName,
            account.name,
            database.name,
          );
          console.log(`DELETED: ${account.name} | ${database.name} | Age: ${friendlyTime(Date.now() - timestamp)}`);
        } else {
          console.log(`SKIPPED: ${account.name} | ${database.name} | Age: ${friendlyTime(Date.now() - timestamp)}`);
        }
      }
    } else if (account.capabilities.find((c) => c.name === "EnableTable")) {
      const tablesDatabase = client.tableResources.listTables(resourceGroupName, account.name);
      for await (const database of tablesDatabase) {
        const timestamp = Number(database.resource.ts) * 1000;
        if (shouldDeleteResource(database.name, timestamp)) {
          await client.tableResources.beginDeleteTableAndWait(resourceGroupName, account.name, database.name);
          console.log(`DELETED: ${account.name} | ${database.name} | Age: ${friendlyTime(Date.now() - timestamp)}`);
        } else {
          console.log(`SKIPPED: ${account.name} | ${database.name} | Age: ${friendlyTime(Date.now() - timestamp)}`);
        }
      }
    } else if (account.capabilities.find((c) => c.name === "EnableGremlin")) {
      const graphDatabases = client.gremlinResources.listGremlinDatabases(resourceGroupName, account.name);
      for await (const database of graphDatabases) {
        const timestamp = Number(database.resource.ts) * 1000;
        if (shouldDeleteResource(database.name, timestamp)) {
          await client.gremlinResources.beginDeleteGremlinDatabaseAndWait(
            resourceGroupName,
            account.name,
            database.name,
          );
          console.log(`DELETED: ${account.name} | ${database.name} | Age: ${friendlyTime(Date.now() - timestamp)}`);
        } else {
          console.log(`SKIPPED: ${account.name} | ${database.name} | Age: ${friendlyTime(Date.now() - timestamp)}`);
        }
      }
    } else if (account.kind === "GlobalDocumentDB") {
      const sqlDatabases = client.sqlResources.listSqlDatabases(resourceGroupName, account.name);
      const sqlDatabasesToDelete = [];
      for await (const database of sqlDatabases) {
        sqlDatabasesToDelete.push(deleteWithRetry(client, database, account.name));
      }
      await Promise.all(sqlDatabasesToDelete);
    }
  }
}

// Retry logic for handling throttling
async function deleteWithRetry(client, database, accountName) {
  const maxRetries = 5;
  let attempt = 0;
  let backoffTime = 1000; // Start with 1 second

  while (attempt < maxRetries) {
    try {
      const timestamp = Number(database.resource.ts) * 1000;
      if (shouldDeleteResource(database.name, timestamp)) {
        await client.sqlResources.beginDeleteSqlDatabaseAndWait(resourceGroupName, accountName, database.name);
        console.log(`DELETED: ${accountName} | ${database.name} | Age: ${friendlyTime(Date.now() - timestamp)}`);
      } else {
        console.log(`SKIPPED: ${accountName} | ${database.name} | Age: ${friendlyTime(Date.now() - timestamp)}`);
      }
      return;
    } catch (error) {
      if (error.statusCode === 429) {
        // Throttling error (HTTP 429), apply exponential backoff
        console.log(`Throttling detected, retrying ${database.name}... (Attempt ${attempt + 1})`);
        await delay(backoffTime);
        attempt++;
        backoffTime *= 2; // Exponential backoff
      } else {
        // For other errors, log and break
        console.error(`Error deleting ${database.name}:`, error);
        break;
      }
    }
  }
  console.log(`Failed to delete ${database.name} after ${maxRetries} attempts.`);
}

// Helper function to delay the retry attempts
function delay(ms) {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

if (require.main === module) {
  main()
    .then(() => {
      console.log("Completed");
      process.exit(0);
    })
    .catch((err) => {
      console.error(err);
      process.exit(1);
    });
}

module.exports = { parseCleanupMinimumAge, shouldDeleteResource };