559 lines
No EOL
21 KiB
Python
559 lines
No EOL
21 KiB
Python
# /management/commands/import_data.py
|
|
|
|
import os
|
|
import json
|
|
import requests
|
|
from blake3 import blake3
|
|
from tqdm.auto import tqdm
|
|
from PIL import Image
|
|
|
|
from datetime import datetime
|
|
|
|
from django.core.management.base import BaseCommand
|
|
from django.core.files.base import ContentFile
|
|
from django.utils.text import slugify
|
|
from django.utils import timezone
|
|
from django.core.exceptions import ObjectDoesNotExist
|
|
|
|
from django.contrib.contenttypes.models import ContentType
|
|
|
|
|
|
from apps.files.models import User_Profile_Images, User_Banner_Images, Submission_File, Metadata_Files
|
|
|
|
from apps.sites.models import Category, Submissions, Users, Tags
|
|
|
|
from apps.sites.furaffinity.models import FA_Submission, FA_Tags, FA_User, FA_Species, FA_Gender, FA_Mature
|
|
from apps.sites.twitter.models import Twitter_Submissions, Twitter_Users, Twitter_Tags
|
|
|
|
from utils.files import get_mime_type
|
|
from utils.strings import get_urls
|
|
|
|
|
|
class Command(BaseCommand):
|
|
help = 'Import data from JSON files in a folder or a single JSON file to the Twitter archive'
|
|
|
|
def add_arguments(self, parser):
|
|
parser.add_argument('path', type=str, help='Path to the folder containing JSON files or a single JSON file')
|
|
parser.add_argument('--delete', action='store_true', help='Delete imported files')
|
|
|
|
|
|
def handle(self, *args, **kwargs):
|
|
path = kwargs['path']
|
|
delete = kwargs['delete']
|
|
|
|
if os.path.isfile(path):
|
|
self.process_json_file(path, delete)
|
|
elif os.path.isdir(path):
|
|
self.process_json_folder(path, delete)
|
|
else:
|
|
self.stdout.write(self.style.ERROR(f"The path '{path}' is not a valid file or folder."))
|
|
return
|
|
|
|
|
|
def process_json_file(self, file_path, delete):
|
|
#self.stdout.write(self.style.NOTICE(f"Importing data from: {file_path}"))
|
|
tqdm.write(f"Importing data from: {file_path}")
|
|
|
|
with open(file_path) as f:
|
|
data = json.load(f)
|
|
|
|
self.import_data(data, file_path, delete)
|
|
|
|
tqdm.write(self.style.SUCCESS('Data imported successfully.'))
|
|
|
|
|
|
def process_json_folder(self, folder_path, delete):
|
|
if not os.path.exists(folder_path):
|
|
#self.stdout.write(self.style.ERROR(f"The folder '{folder_path}' does not exist."))
|
|
tqdm.write(self.style.ERROR(f"The folder '{folder_path}' does not exist."))
|
|
return
|
|
|
|
for root, dirs, files in tqdm(os.walk(folder_path), dynamic_ncols=True):
|
|
for file_name in files:
|
|
if file_name.endswith('.json'):
|
|
file_path = os.path.join(root, file_name)
|
|
self.process_json_file(file_path, delete)
|
|
|
|
|
|
def compute_file_hash(self, file_path):
|
|
""" Compute BLAKE3 hash of the file """
|
|
try:
|
|
hasher = blake3()
|
|
with open(file_path, 'rb') as f:
|
|
while chunk := f.read(65536):
|
|
hasher.update(chunk)
|
|
return hasher.hexdigest()
|
|
except Exception as e:
|
|
tqdm.write(self.style.WARNING(f"Error computing file hash: {e}"))
|
|
return None
|
|
def compute_string_hash(self, string):
|
|
""" Compute BLAKE3 hash of the string """
|
|
try:
|
|
hasher = blake3()
|
|
hasher.update(string.encode())
|
|
return hasher.hexdigest()
|
|
except Exception as e:
|
|
tqdm.write(self.style.WARNING(f"Error computing string hash: {e}"))
|
|
return None
|
|
def import_file(self, file_path, model, delete=False):
|
|
"""
|
|
Imports a file if it doesn't already exist in the database and returns the instance.
|
|
|
|
:param file_path: The path to the file to import.
|
|
:param model: The model class to which the file instance should be linked.
|
|
:param delete: Whether to delete the imported file after processing.
|
|
:return: The file instance.
|
|
"""
|
|
|
|
file_instance = None # Initialize file_instance to None
|
|
|
|
if os.path.exists(file_path):
|
|
file_hash = self.compute_file_hash(file_path)
|
|
|
|
file_name = os.path.basename(file_path)
|
|
Null, file_ext = os.path.splitext(file_name)
|
|
hash_file_name = file_hash + file_ext
|
|
|
|
try:
|
|
file_instance = model.objects.get(file_hash=file_hash)
|
|
|
|
file_instance.file_ext = file_ext
|
|
file_instance.size = os.path.getsize(file_path)
|
|
file_instance.file_mime = get_mime_type(file_path)
|
|
|
|
if file_instance.file_mime.startswith("image/"):
|
|
im = Image.open(file_instance.file)
|
|
file_instance.image_height, file_instance.image_width = im.size
|
|
else:
|
|
file_instance.image_height = None
|
|
file_instance.image_width = None
|
|
|
|
file_instance.save()
|
|
|
|
tqdm.write(self.style.NOTICE(f"Skipping: {file_path} file, already imported"))
|
|
|
|
except model.DoesNotExist:
|
|
# If the file doesn't exist, create a new file instance
|
|
with open(file_path, 'rb') as file:
|
|
file_instance = model()
|
|
file_instance.file_hash = file_hash
|
|
|
|
file_instance.file.save(hash_file_name, file)
|
|
|
|
file_instance.file_ext = file_ext
|
|
file_instance.file_mime = get_mime_type(file_path)
|
|
file_instance.size = os.path.getsize(file_path)
|
|
|
|
if file_instance.file_mime.startswith("image/"):
|
|
im = Image.open(file_instance.file)
|
|
file_instance.image_height, file_instance.image_width = im.size
|
|
else:
|
|
file_instance.image_height = None
|
|
file_instance.image_width = None
|
|
|
|
file_instance.file_name = file_name
|
|
file_instance.save()
|
|
|
|
tqdm.write(self.style.NOTICE(f"Import file: {file_path}"))
|
|
|
|
if delete:
|
|
self.delete_imported_file(file_path)
|
|
|
|
return file_instance
|
|
|
|
|
|
def delete_imported_file(self, file_path, delete=False):
|
|
"""
|
|
Delete the file if the --delete flag is used
|
|
|
|
:param delete: Whether to delete the imported file after processing.
|
|
"""
|
|
if delete:
|
|
if os.path.exists(file_path):
|
|
os.remove(file_path)
|
|
tqdm.write(self.style.SUCCESS(f"Deleted: {file_path}"))
|
|
else:
|
|
tqdm.write(self.style.WARNING(f"File not found: {file_path}"))
|
|
|
|
|
|
def import_data(self, data, json_file_path, delete):
|
|
|
|
category = data['category']
|
|
|
|
if category == "twitter":
|
|
self.import_from_twitter(data, json_file_path, delete)
|
|
|
|
elif category == "furaffinity":
|
|
self.import_from_furaffinity(data, json_file_path, delete)
|
|
|
|
else:
|
|
tqdm.write(f"Skipping '{category}' not implemented")
|
|
|
|
def import_twitter_user(self, data, file_path, category, delete=False):
|
|
"""
|
|
Import a Twitter user from the provided data into the database.
|
|
|
|
Parameters:
|
|
data (dict): The data containing information about the Twitter user.
|
|
file_path (str): The file path for importing user images.
|
|
delete (bool): Flag indicating whether to delete user images after importing it.
|
|
|
|
Returns:
|
|
Twitter_Users: The Twitter user object imported or retrieved from the database.
|
|
"""
|
|
|
|
content_type = ContentType.objects.get_for_model(Twitter_Users)
|
|
|
|
author, created = Twitter_Users.objects.get_or_create(artist_id=data['author']['id'])
|
|
|
|
author.artist = data['author']['nick']
|
|
author.artist_url = data['author']['name']
|
|
author.date = timezone.make_aware(datetime.strptime(data['author']["date"], "%Y-%m-%d %H:%M:%S"))
|
|
author.description = data['author']['description']
|
|
if 'url' in data['author'].keys():
|
|
author.extra_url = data['author']['url']
|
|
author.location = data['author']['location']
|
|
author.verified = data['author']['verified']
|
|
|
|
if author.favourites_count == None or data['author']["favourites_count"] > author.favourites_count:
|
|
author.favourites_count = data['author']["favourites_count"]
|
|
if author.followers_count == None or data['author']["followers_count"] > author.followers_count:
|
|
author.followers_count = data['author']["followers_count"]
|
|
if author.friends_count == None or data['author']["friends_count"] > author.friends_count:
|
|
author.friends_count = data['author']["friends_count"]
|
|
if author.media_count == None or data['author']["media_count"] > author.media_count:
|
|
author.media_count = data['author']["media_count"]
|
|
if author.listed_count == None or data['author']["listed_count"] > author.listed_count:
|
|
author.listed_count = data['author']["listed_count"]
|
|
if author.statuses_count == None or data['author']["statuses_count"] > author.statuses_count:
|
|
author.statuses_count = data['author']["statuses_count"]
|
|
|
|
if data['subcategory'] == "avatar":
|
|
author.profile_image = data['author']['profile_image']
|
|
|
|
author.icon = self.import_file(file_path, User_Profile_Images, delete)
|
|
|
|
elif data['subcategory'] == "background":
|
|
|
|
author.profile_banner = data['author']['profile_banner']
|
|
|
|
author.banner = self.import_file(file_path, User_Banner_Images, delete)
|
|
|
|
|
|
author_hash = self.compute_string_hash(data['author']['name'] + data['category'])
|
|
|
|
site_user, created = Users.objects.get_or_create(user_hash=author_hash)
|
|
|
|
site_user.category = category
|
|
|
|
|
|
# Get the primary key of the twitter_submission instance
|
|
site_user_id = author.pk
|
|
|
|
# Create the SubmissionsLink instance
|
|
site_user.content_type=content_type
|
|
site_user.object_id=site_user_id
|
|
|
|
site_user.save()
|
|
|
|
author.save()
|
|
|
|
return author, site_user
|
|
|
|
|
|
def import_twitter_tags(self, data: dict, category: str) -> list[Twitter_Tags]:
|
|
"""
|
|
Import a Twitter tag from the provided data into the database.
|
|
|
|
Parameters:
|
|
data (dict): The data containing information about the Twitter tag.
|
|
|
|
Returns:
|
|
list[Twitter_Tags]: A list of imported or retrieved Twitter tag objects.
|
|
"""
|
|
|
|
content_type = ContentType.objects.get_for_model(Twitter_Tags)
|
|
|
|
tags: list[Twitter_Tags] = []
|
|
|
|
if "hashtags" in data:
|
|
for t_tag_name in data["hashtags"]:
|
|
t_tag_slug = slugify(t_tag_name)
|
|
try:
|
|
# Check if the tag already exists in the database by name
|
|
tag: Twitter_Tags = Twitter_Tags.objects.get(tag_slug=t_tag_slug)
|
|
|
|
tag_id = tag.pk
|
|
|
|
except ObjectDoesNotExist:
|
|
# If the tag does not exist, create a new tag and generate the slug
|
|
tag = Twitter_Tags(tag=t_tag_name)
|
|
tag.tag_slug = t_tag_slug
|
|
|
|
tag_id = tag.pk
|
|
|
|
site_tags, created = Tags.objects.get_or_create(tag_slug=t_tag_slug)
|
|
|
|
site_tags.category.add(category)
|
|
|
|
site_tags.content_type=content_type
|
|
site_tags.object_id=tag_id
|
|
|
|
site_tags.save()
|
|
|
|
tag.save() # Save the tag (either new or existing)
|
|
|
|
tags.append(tag)
|
|
|
|
return tags
|
|
|
|
|
|
def import_from_twitter(self, data, json_file_path, delete):
|
|
|
|
category, created = Category.objects.get_or_create(name=data['category'])
|
|
|
|
category.save()
|
|
|
|
twitter_submission, created = Twitter_Submissions.objects.get_or_create(submission_id=data["tweet_id"])
|
|
|
|
file_path = json_file_path.removesuffix(".json")
|
|
|
|
# Handle author import
|
|
author, site_user = self.import_twitter_user(data, file_path, category, delete)
|
|
|
|
twitter_submission.author = author
|
|
|
|
# Handle tag import
|
|
tags = self.import_twitter_tags(data, category)
|
|
|
|
for tag in tags:
|
|
twitter_submission.tags.add(tag) # Add the tag to the submission
|
|
|
|
twitter_submission.gallery_type = data['subcategory']
|
|
|
|
# Handle file import
|
|
twitter_submission.files.add(self.import_file(file_path, Submission_File, delete))
|
|
|
|
# Handle metadata file import
|
|
twitter_submission.metadata.add(self.import_file(json_file_path, Metadata_Files, delete))
|
|
|
|
twitter_submission.description = data['content']
|
|
|
|
twitter_submission.date = timezone.make_aware(datetime.strptime(data['date'], "%Y-%m-%d %H:%M:%S"))
|
|
|
|
twitter_submission.origin_site = data['category']
|
|
|
|
twitter_submission.file_extension = data['extension']
|
|
|
|
twitter_submission.origin_filename = data['filename']
|
|
|
|
if twitter_submission.media_num is None or data['num'] > twitter_submission.media_num:
|
|
twitter_submission.media_num = data['num']
|
|
|
|
if "height" in data.keys():
|
|
twitter_submission.image_height = data['height']
|
|
if "width" in data.keys():
|
|
twitter_submission.image_width = data['width']
|
|
|
|
if "sensitive" in data.keys():
|
|
twitter_submission.sensitive = data['sensitive']
|
|
|
|
if "favorite_count" in data.keys():
|
|
twitter_submission.favorites_count = data['favorite_count']
|
|
|
|
if "quote_count" in data.keys():
|
|
twitter_submission.quote_count = data['quote_count']
|
|
|
|
if "reply_count" in data.keys():
|
|
twitter_submission.reply_count = data['reply_count']
|
|
|
|
if "retweet_count" in data.keys():
|
|
twitter_submission.retweet_count = data['retweet_count']
|
|
|
|
twitter_submission.lang = data['lang']
|
|
|
|
twitter_submission.save()
|
|
|
|
submission_hash = self.compute_string_hash(category.name + data['author']['name'] + str(data["tweet_id"]))
|
|
|
|
submission, created = Submissions.objects.get_or_create(submission_hash=submission_hash)
|
|
|
|
submission.category = category
|
|
|
|
submission.author = site_user
|
|
|
|
if twitter_submission.sensitive is not None:
|
|
submission.mature = twitter_submission.sensitive
|
|
else:
|
|
submission.mature = False
|
|
|
|
submission.date = timezone.make_aware(datetime.strptime(data['date'], "%Y-%m-%d %H:%M:%S"))
|
|
|
|
content_type = ContentType.objects.get_for_model(Twitter_Submissions)
|
|
|
|
# Get the primary key of the twitter_submission instance
|
|
twitter_submission_id = twitter_submission.pk
|
|
|
|
# Create the SubmissionsLink instance
|
|
submission.content_type=content_type
|
|
submission.object_id=twitter_submission_id
|
|
|
|
submission.save()
|
|
|
|
self.delete_imported_file(json_file_path, delete)
|
|
self.delete_imported_file(file_path, delete)
|
|
|
|
|
|
def import_furaffinity_user(self, data, json_file_path, category, delete):
|
|
content_type = ContentType.objects.get_for_model(FA_User)
|
|
|
|
artist, created = FA_User.objects.get_or_create(artist_url=data["artist_url"], artist=data["artist"])
|
|
|
|
author_hash = self.compute_string_hash(data["artist_url"] + data['category'])
|
|
|
|
site_user, created = Users.objects.get_or_create(user_hash=author_hash)
|
|
|
|
site_user.category = category
|
|
|
|
|
|
# Get the primary key of the furaffinity_submission instance
|
|
site_user_id = artist.pk
|
|
|
|
# Create the SubmissionsLink instance
|
|
site_user.content_type=content_type
|
|
site_user.object_id=site_user_id
|
|
|
|
site_user.save()
|
|
|
|
return artist, site_user
|
|
|
|
|
|
def import_furaffinity_tags(self, data, category):
|
|
|
|
content_type = ContentType.objects.get_for_model(FA_Tags)
|
|
|
|
tags: list[FA_Tags] = []
|
|
|
|
site_tags: list[Tags] = []
|
|
|
|
if "tags" in data:
|
|
for t_tag_name in data["tags"]:
|
|
t_tag_slug = slugify(t_tag_name)
|
|
try:
|
|
# Check if the tag already exists in the database by name
|
|
tag: FA_Tags = FA_Tags.objects.get(tag_slug=t_tag_slug)
|
|
|
|
tag_id = tag.pk
|
|
|
|
except ObjectDoesNotExist:
|
|
# If the tag does not exist, create a new tag and generate the slug
|
|
tag = FA_Tags(tag=t_tag_name)
|
|
tag.tag_slug = t_tag_slug
|
|
|
|
tag_id = tag.pk
|
|
|
|
site_tag, created = Tags.objects.get_or_create(tag_slug=t_tag_slug)
|
|
|
|
site_tag.category.add(category)
|
|
|
|
site_tag.content_type=content_type
|
|
site_tag.object_id=tag_id
|
|
|
|
site_tag.save()
|
|
|
|
tag.save() # Save the tag (either new or existing)
|
|
|
|
tags.append(tag)
|
|
|
|
site_tags.append(site_tag)
|
|
|
|
return tags, site_tags
|
|
|
|
|
|
def import_from_furaffinity(self, data, json_file_path, delete):
|
|
|
|
category, created = Category.objects.get_or_create(name=data['category'])
|
|
|
|
category.save()
|
|
|
|
furaffinity_submission, created = FA_Submission.objects.get_or_create(submission_id=data["id"])
|
|
|
|
furaffinity_submission.media_url = data["url"]
|
|
furaffinity_submission.title = data["title"]
|
|
furaffinity_submission.description = data["description"]
|
|
|
|
furaffinity_submission.date = timezone.make_aware(datetime.strptime(data["date"], "%Y-%m-%d %H:%M:%S"))
|
|
|
|
file_path = json_file_path.removesuffix(".json")
|
|
|
|
# Handle author import
|
|
author, site_user = self.import_furaffinity_user(data, file_path, category, delete)
|
|
|
|
furaffinity_submission.artist = author
|
|
|
|
# Handle tag import
|
|
tags, site_tags = self.import_furaffinity_tags(data, category)
|
|
|
|
for tag in tags:
|
|
furaffinity_submission.tags.add(tag) # Add the tag to the submission
|
|
|
|
species, created = FA_Species.objects.get_or_create(species=data["species"])
|
|
furaffinity_submission.species = species
|
|
|
|
# Handle mature rating import
|
|
mature, created = FA_Mature.objects.get_or_create(mature=data["rating"])
|
|
furaffinity_submission.mature_rating = mature
|
|
|
|
furaffinity_submission.number_of_comments = data["comments"]
|
|
furaffinity_submission.views = data["views"]
|
|
|
|
gender, created = FA_Gender.objects.get_or_create(gender=data["gender"])
|
|
furaffinity_submission.gender = gender
|
|
|
|
furaffinity_submission.fa_theme = data["theme"]
|
|
furaffinity_submission.fa_category = data["fa_category"]
|
|
furaffinity_submission.gallery_type = data["subcategory"]
|
|
furaffinity_submission.file_extension = data["extension"]
|
|
furaffinity_submission.image_height = data["height"]
|
|
furaffinity_submission.image_width = data["width"]
|
|
|
|
# Handle file import
|
|
furaffinity_submission.files.add(self.import_file(file_path, Submission_File, delete))
|
|
|
|
# Handle metadata file import
|
|
furaffinity_submission.metadata.add(self.import_file(json_file_path, Metadata_Files, delete))
|
|
|
|
furaffinity_submission.save()
|
|
|
|
submission_hash = self.compute_string_hash(category.name + data["artist_url"] + str(data["id"]))
|
|
|
|
submission, created = Submissions.objects.get_or_create(submission_hash=submission_hash)
|
|
|
|
submission.category = category
|
|
|
|
submission.tags.add(*site_tags)
|
|
|
|
submission.author = site_user
|
|
|
|
if furaffinity_submission.mature_rating.mature != "General" and not None:
|
|
print("Mature")
|
|
submission.mature = True
|
|
else:
|
|
submission.mature = False
|
|
|
|
submission.date = timezone.make_aware(datetime.strptime(data['date'], "%Y-%m-%d %H:%M:%S"))
|
|
|
|
content_type = ContentType.objects.get_for_model(FA_Submission)
|
|
|
|
# Get the primary key of the twitter_submission instance
|
|
furaffinity_submission_id = furaffinity_submission.pk
|
|
|
|
# Create the SubmissionsLink instance
|
|
submission.content_type=content_type
|
|
submission.object_id=furaffinity_submission_id
|
|
|
|
submission.save()
|
|
|
|
self.delete_imported_file(json_file_path, delete)
|
|
self.delete_imported_file(file_path, delete)
|
|
|