Skip to content

Commit

Permalink
Address review comments
Browse files Browse the repository at this point in the history
  • Loading branch information
@Eden-D-Zhang
Eden-D-Zhang committed Jan 16, 2025
1 parent c210992 commit 812ce45
Show file tree
Hide file tree
Showing 2 changed files with 68 additions and 54 deletions.
64 changes: 36 additions & 28 deletions components/clp-package-utils/clp_package_utils/scripts/archive_manager.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,19 @@
validate_and_load_db_credentials_file,
)

# Command/Argument Constants
FIND_COMMAND = "find"
DEL_COMMAND = "del"
BY_IDS_COMMAND = "by-ids"
BY_FILTER_COMMAND = "by-filter"
BEGIN_TS_ARG = "--begin-ts"
END_TS_ARG = "--end-ts"
DRY_RUN_ARG = "--dry-run"

logger = logging.getLogger(__file__)


def validate_timestamps(begin_ts, end_ts):
def _validate_timestamps(begin_ts, end_ts):
if begin_ts > end_ts:
logger.error("begin-ts must be <= end-ts")
return False
Expand Down Expand Up @@ -49,32 +58,32 @@ def main(argv):
required=True,
)
find_parser = subparsers.add_parser(
"find",
FIND_COMMAND,
help="List IDs of archives.",
)
del_parser = subparsers.add_parser(
"del",
DEL_COMMAND,
help="Delete archives.",
)

# Find options
find_parser.add_argument(
"--begin-ts",
BEGIN_TS_ARG,
dest="begin_ts",
type=int,
default=0,
help="Time-range lower-bound (inclusive) as milliseconds from the UNIX epoch.",
)
find_parser.add_argument(
"--end-ts",
END_TS_ARG,
dest="end_ts",
type=int,
help="Time-range upper-bound (include) as milliseconds from the UNIX epoch.",
)

# Delete options
del_parser.add_argument(
"--dry-run",
DRY_RUN_ARG,
dest="dry_run",
action="store_true",
help="Preview delete without making changes. Lists errors and files to be deleted.",
Expand All @@ -86,11 +95,11 @@ def main(argv):
required=True,
)
del_id_parser = del_subparsers.add_parser(
"by-ids",
BY_IDS_COMMAND,
help="Delete archives by ID.",
)
del_filter_parser = del_subparsers.add_parser(
"by-filter",
BY_FILTER_COMMAND,
help="Delete archives within time frame.",
)

Expand All @@ -103,13 +112,13 @@ def main(argv):

# Delete by filter arguments
del_filter_parser.add_argument(
"--begin-ts",
BEGIN_TS_ARG,
type=int,
default=0,
help="Time-range lower-bound (inclusive) as milliseconds from the UNIX epoch.",
)
del_filter_parser.add_argument(
"--end-ts",
END_TS_ARG,
type=int,
required=True,
help="Time-range upper-bound (include) as milliseconds from the UNIX epoch.",
Expand All @@ -135,18 +144,18 @@ def main(argv):
return -1

# Validate input depending on subcommands
if "del" == parsed_args.subcommand:
if "by-filter" == parsed_args.del_subcommand:
if DEL_COMMAND == parsed_args.subcommand:
if BY_FILTER_COMMAND == parsed_args.del_subcommand:

# Validate the input timestamp
if not validate_timestamps(parsed_args.begin_ts, parsed_args.end_ts):
if not _validate_timestamps(parsed_args.begin_ts, parsed_args.end_ts):
return -1

elif "find" == parsed_args.subcommand:
if hasattr(parsed_args, "end_ts") and parsed_args.end_ts is not None:
elif FIND_COMMAND == parsed_args.subcommand:
if parsed_args.end_ts is not None:

# Validate the input timestamp
if not validate_timestamps(parsed_args.begin_ts, parsed_args.end_ts):
if not _validate_timestamps(parsed_args.begin_ts, parsed_args.end_ts):
return -1

container_name = generate_container_name("archive-manager")
Expand All @@ -168,21 +177,20 @@ def main(argv):
"--config", str(generated_config_path_on_container),
str(parsed_args.subcommand),
]

# fmt : on
# Add subcommand-specific arguments
if "del" == parsed_args.subcommand:
if DEL_COMMAND == parsed_args.subcommand:
if True == parsed_args.dry_run:
archive_manager_cmd.extend(["--dry-run"])
if "by-ids" == parsed_args.del_subcommand:
archive_manager_cmd.extend(["by-ids"])
archive_manager_cmd.append(DRY_RUN_ARG)
if BY_IDS_COMMAND == parsed_args.del_subcommand:
archive_manager_cmd.append(BY_IDS_COMMAND)
archive_manager_cmd.extend(parsed_args.ids)
elif "by-filter" == parsed_args.del_subcommand:
archive_manager_cmd.extend(["by-filter", str(parsed_args.begin_ts), str(parsed_args.end_ts)])
elif "find" == parsed_args.subcommand:
if hasattr(parsed_args, "end_ts") and parsed_args.end_ts is not None:
archive_manager_cmd.extend(["--begin-ts", str(parsed_args.begin_ts), "--end-ts", str(parsed_args.end_ts)])

# fmt: on
elif BY_FILTER_COMMAND == parsed_args.del_subcommand:
archive_manager_cmd.extend([BY_FILTER_COMMAND, str(parsed_args.begin_ts), str(parsed_args.end_ts)])
elif FIND_COMMAND == parsed_args.subcommand:
archive_manager_cmd.extend([BEGIN_TS_ARG, str(parsed_args.begin_ts)])
if parsed_args.end_ts is not None:
archive_manager_cmd.extend([END_TS_ARG, str(parsed_args.end_ts)])

cmd = container_start_cmd + archive_manager_cmd

Expand Down
58 changes: 32 additions & 26 deletions components/clp-package-utils/clp_package_utils/scripts/native/archive_manager.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,15 @@
load_config_file,
)

# Command/Argument Constants
FIND_COMMAND = "find"
DEL_COMMAND = "del"
BY_IDS_COMMAND = "by-ids"
BY_FILTER_COMMAND = "by-filter"
BEGIN_TS_ARG = "--begin-ts"
END_TS_ARG = "--end-ts"
DRY_RUN_ARG = "--dry-run"

logger = logging.getLogger(__file__)


Expand All @@ -38,31 +47,31 @@ def main(argv):
required=True,
)
find_parser = subparsers.add_parser(
"find",
FIND_COMMAND,
help="List IDs of archives.",
)
del_parser = subparsers.add_parser(
"del",
DEL_COMMAND,
help="Delete archives.",
)

# Find options
find_parser.add_argument(
"--begin-ts",
BEGIN_TS_ARG,
dest="begin_ts",
type=int,
help="Time-range lower-bound (inclusive) as milliseconds from the UNIX epoch.",
)
find_parser.add_argument(
"--end-ts",
END_TS_ARG,
dest="end_ts",
type=int,
help="Time-range upper-bound (include) as milliseconds from the UNIX epoch.",
)

# Delete options
del_parser.add_argument(
"--dry-run",
DRY_RUN_ARG,
dest="dry_run",
action="store_true",
help="Preview delete without making changes. Lists errors and files to be deleted.",
Expand All @@ -74,11 +83,11 @@ def main(argv):
required=True,
)
del_id_parser = del_subparsers.add_parser(
"by-ids",
BY_IDS_COMMAND,
help="Delete archives by ID.",
)
del_filter_parser = del_subparsers.add_parser(
"by-filter",
BY_FILTER_COMMAND,
help="delte archives within time frame.",
)

Expand Down Expand Up @@ -118,25 +127,22 @@ def main(argv):
logger.error("`archive_output.directory` doesn't exist.")
return -1

if "find" == parsed_args.subcommand:
if parsed_args.begin_ts is None and parsed_args.end_ts is None:
return _find_archives(archives_dir, database_config)
else:
return _find_archives(
if FIND_COMMAND == parsed_args.subcommand:
return _find_archives(
archives_dir,
database_config,
parsed_args.begin_ts,
parsed_args.end_ts,
)
elif "del" == parsed_args.subcommand:
if "by-ids" == parsed_args.del_subcommand:
elif DEL_COMMAND == parsed_args.subcommand:
if BY_IDS_COMMAND == parsed_args.del_subcommand:
return _delete_archives_by_ids(
archives_dir,
database_config,
parsed_args.ids,
parsed_args.dry_run,
)
elif "by-filter" == parsed_args.del_subcommand:
elif BY_FILTER_COMMAND == parsed_args.del_subcommand:
return _delete_archives_by_filter(
archives_dir,
database_config,
Expand All @@ -149,7 +155,7 @@ def main(argv):
def _find_archives(
archives_dir: Path,
database_config: Database,
begin_ts: int = None,
begin_ts: int,
end_ts: int = None,
) -> int:
"""
Expand All @@ -171,10 +177,10 @@ def _find_archives(
db_conn.cursor(dictionary=True)
) as db_cursor:

params = ()
query = f"SELECT id FROM `{table_prefix}archives`"
if begin_ts is not None and end_ts is not None:
query += " WHERE begin_timestamp >= %s AND end_timestamp <= %s"
params = (begin_ts,)
query = f"SELECT id FROM `{table_prefix}archives` WHERE begin_timestamp >= %s"
if end_ts is not None:
query += " AND end_timestamp <= %s"
params = (begin_ts, end_ts)

db_cursor.execute(query, params)
Expand Down Expand Up @@ -255,20 +261,20 @@ def _delete_archives(
"""
)

id_string = ", ".join(f"'{archive_id}'" for archive_id in archive_ids)

db_cursor.execute(
f"""
DELETE FROM `{table_prefix}files`
WHERE archive_id in ({', '.join(['%s'] * len(archive_ids))})
""",
archive_ids,
WHERE archive_id in ({id_string})
"""
)

db_cursor.execute(
f"""
DELETE FROM `{table_prefix}archive_tags`
WHERE archive_id in ({', '.join(['%s'] * len(archive_ids))})
""",
archive_ids,
WHERE archive_id in ({id_string})
"""
)

if dry_run:
Expand Down

0 comments on commit 812ce45

Please sign in to comment.