diff options
| -rw-r--r-- | data/exclude_channel.txt | 1 | ||||
| -rw-r--r-- | nijitrack.py | 66 |
2 files changed, 58 insertions, 9 deletions
diff --git a/data/exclude_channel.txt b/data/exclude_channel.txt index 18663d6..2e75ce7 100644 --- a/data/exclude_channel.txt +++ b/data/exclude_channel.txt @@ -1,3 +1,2 @@ DELETE_THIS and write channel id to ignore on each line i.e UC-1234567890 -UC-0987654321
\ No newline at end of file diff --git a/nijitrack.py b/nijitrack.py index fd2c76b..4b2faea 100644 --- a/nijitrack.py +++ b/nijitrack.py @@ -4,15 +4,19 @@ from datetime import datetime import dotenv import pytz -from b2sdk.v2 import * from logger import * from sql.pg_handler import PostgresHandler from webapi.holodex import HolodexAPI from webapi.youtube import YouTubeAPI +from enum import Enum import fileutil as fs import graph +class BucketType(Enum): + B2 = 1 + R2 = 2 + dotenv.load_dotenv() DATA_SETTING = fs.load_json_file("sql_table_config.json") @@ -136,27 +140,58 @@ def combine_excluded_channel_ids(inactive_channel_data: list, excluded_channels: channel_ids.append(inactive_channel) return channel_ids -def uploadFileToBucket(filepath: str) -> bool: +def uploadFileToBucketB2(filepath: str) -> bool: + from b2sdk.v2 import InMemoryAccountInfo, B2Api try: info = InMemoryAccountInfo() b2_api = B2Api(info) application_key_id = os.environ.get("B2_APP_ID") application_key = os.environ.get("B2_APP_KEY") + bucket_name = os.environ.get("B2_BUCKET_NAME") file_info = {'how': 'good-file'} b2_api.authorize_account("production", application_key_id, application_key) b2_file_name = "graph.html" - bucket = b2_api.get_bucket_by_name("vtuber-rabbit-hole-archive") + bucket = b2_api.get_bucket_by_name(bucket_name) bucket.upload_local_file(local_file=filepath, file_name=b2_file_name, file_info=file_info) return True except Exception as e: print("An error occured while attempting to upload to B2") print(e) - return False; + return False + +def uploadFileToBucketR2(filepath: str) -> bool: + import boto3 + account_id = os.environ.get("R2_ACCOUNT_ID") + access_key = os.environ.get("R2_ACCESS_KEY") + secret_key = os.environ.get("R2_SECRET_KEY") + bucket_name = os.environ.get("R2_BUCKET_NAME") + s3 = boto3.resource( + 's3', + endpoint_url=f'https://{account_id}.r2.cloudflarestorage.com', + aws_access_key_id=access_key, + aws_secret_access_key=secret_key + ) + try: + with open(filepath, "rb") as f: + s3.Bucket(bucket_name).upload_fileobj(f, filepath) + return True + except Exception as e: + print("An error occurred while attempting to upload to R2") + print(e) + return False + +def generate_api_routes(bucket_type: BucketType): + import app + api_subscribers = app.api_subscribers() + print(api_subscribers) if __name__ == "__main__": parser = argparse.ArgumentParser(description="NijiTrack - A Subscriber Tracker") parser.add_argument('--mode', choices=['yt', 'holodex'], help='Specify the data source to use (yt or holodex)') - parser.add_argument('--b2', action='store_true', help="Upload graph html to Backblaze B2") + parser.add_argument('--b2', action='store_true', help="Use Backblaze B2 as the bucket upload source") + parser.add_argument('--r2', action='store_true', help="Use Cloudflare R2 as the bucket upload source") + parser.add_argument('--uploadGraph', action='store_true', help="Upload graph html to Backblaze B2") + parser.add_argument('--uploadRoutes', action='store_true', help="Pre-generate every API route and upload it") parser.add_argument('--ff', action='store_true', help="Force a full refresh of all data (override daily refresh)") args = parser.parse_args() server = create_database_connection() @@ -175,7 +210,22 @@ if __name__ == "__main__": graph_html = graph.plot_subscriber_count_over_time(server, DATA_SETTING["TABLE_HISTORICAL"], exclude_channels=combine_excluded_channel_ids(inactive_channels, fs.get_excluded_channels())) with open("index.html", "w", encoding="utf-8") as file: file.write(graph_html) + upstream_bucket = None if args.b2: - uploadFileToBucket("index.html") - else: - print("Skipping B2 Upload") + upstream_bucket = BucketType.B2 + elif args.r2: + upstream_bucket = BucketType.R2 + + if args.uploadGraph: + if upstream_bucket is None: + print("Tried to upload graph but no remote source has been specified. Skipping....") + match upstream_bucket: + case BucketType.B2: + uploadFileToBucketB2("index.html") + case BucketType.R2: + uploadFileToBucketR2("index.html") + + if args.uploadRoutes: + if upstream_bucket is None: + print("Tried to upload routes but no remote source has been specified. Skipping....") + generate_api_routes(upstream_bucket) |