<commit_before>import sqlite3
import sys
import os
import datetime
from django.conf import settings
from django.core.management import call_command
from django.core.management.base import BaseCommand
from django.core.files.uploadedfile import SimpleUploadedFile
from django.db import transaction
from django.core.files import File as DJFile
from le_utils.constants import content_kinds,file_formats, format_presets, licenses, exercises
from contentcuration import models
from contentcuration.api import write_file_to_storage
import logging as logmodule
logging = logmodule.getLogger(__name__)

CHANNEL_TABLE = 'content_channelmetadata'
NODE_TABLE = 'content_contentnode'
FILE_TABLE = 'content_file'
TAG_TABLE = 'content_contenttag'
NODE_TAG_TABLE = 'content_contentnode_tags'
LICENSE_TABLE = 'content_license'
NODE_COUNT = 0
FILE_COUNT = 0
TAG_COUNT = 0


class EarlyExit(BaseException):
    def __init__(self, message, db_path):
        self.message = message
        self.db_path = db_path

class Command(BaseCommand):

    help = 'Restores a channel based on its exported database (Usage: restore_channel [source-channel-id] [target-channel-id]'

    def add_arguments(self, parser):

        # ID of channel to read data from
        parser.add_argument('source_id', type=str)

        # ID of channel to write data to (can be same as source channel)
        parser.add_argument('target_id', type=str)

    def handle(self, *args, **options):
        try:
          # Set up variables for restoration process
          logging.info("\n\n********** STARTING CHANNEL RESTORATION **********")
          start = datetime.datetime.now()
          source_id = options['source_id']
          target_id = options.get('target_id') or source_id

          # Test connection to database
          logging.info("Connecting to database for channel {}...".format(source_id))
          conn = sqlite3.connect(os.path.join(settings.DB_ROOT,'{}.sqlite3'.format(source_id)))
          cursor=conn.cursor()

          # Start by creating channel
          logging.info("Creating channel...")
          channel, root_pk = create_channel(conn, target_id)

          # Create nodes mapping to channel
          logging.info("   Creating nodes...")
          root = None
          with transaction.atomic():
            root = create_nodes(cursor, target_id)

          # Save tree to target tree
          channel.main_tree = root
          channel.save()

          # Print stats
          logging.info("\n\nChannel has been restored (time: {ms}, {node_count} nodes, {file_count} files, {tag_count} tags)\n".format(
            ms = datetime.datetime.now() - start,
            node_count = NODE_COUNT,
            file_count = FILE_COUNT,
            tag_count = TAG_COUNT)
          )
          logging.info("\n\n********** RESTORATION COMPLETE **********\n\n")

        except EarlyExit as e:
            logging.warning("Exited early due to {message}.".format(
                message=e.message))
            self.stdout.write("You can find your database in {path}".format(
                path=e.db_path))


def create_channel(cursor, target_id):
    """ create_channel: Create channel at target id
        Args:
            cursor (sqlite3.Connection): connection to export database
            target_id (str): channel_id to write to
        Returns: channel model created and id of root node
    """
    id, name, description, author, thumbnail, root_pk, version = cursor.execute('SELECT * FROM {table}'.format(table=CHANNEL_TABLE)).fetchone()
    channel, is_new = models.Channel.objects.get_or_create(pk=target_id)
    channel.name = name
    channel.description = description
    channel.thumbnail = write_to_thumbnail_file(thumbnail)
    channel.version = version
    channel.save()
    logging.info("\tCreated channel {} with name {}".format(target_id, name))
    return channel, root_pk

def write_to_thumbnail_file(raw_thumbnail):
    """ write_to_thumbnail_file: Convert base64 thumbnail to file
        Args:
            raw_thumbnail (str): base64 encoded thumbnail
        Returns: thumbnail filename
    """
    if raw_thumbnail and isinstance(raw_thumbnail, str) and raw_thumbnail != "":
        with SimpleUploadedFile('temp.png',raw_thumbnail.replace('data:image/png;base64,', '').decode('base64')) as tempf:
            filename = write_file_to_storage(tempf, check_valid=False)
            logging.info("\tCreated thumbnail {}".format(filename))
            return filename

def create_nodes(cursor, target_id, parent=None, indent=1):
    """ create_channel: Create channel at target id
        Args:
            cursor (sqlite3.Connection): connection to export database
            target_id (str): channel_id to write to
            parent (models.ContentNode): node's parent
            indent (int): How far to indent print statements
        Returns: newly created node
    """
    # Read database rows that match parent
    parent_query = 'parent_id IS NULL'
    if parent:
        parent_query = "parent_id=\'{}\'".format(parent.node_id)

    sql_command = 'SELECT id, title, content_id, description, sort_order, '\
        'license_owner, author, license_id, kind FROM {table} WHERE {query};'\<commit_msg>Update restore_channel.py<commit_after>import sqlite3
import sys
import os
import datetime
from django.conf import settings
from django.core.management import call_command
from django.core.management.base import BaseCommand
from django.core.files.uploadedfile import SimpleUploadedFile
from django.db import transaction
from django.core.files import File as DJFile
from le_utils.constants import content_kinds,file_formats, format_presets, licenses, exercises
from contentcuration import models
from contentcuration.api import write_file_to_storage
import logging as logmodule
logging = logmodule.getLogger(__name__)

CHANNEL_TABLE = 'content_channelmetadata'
NODE_TABLE = 'content_contentnode'
FILE_TABLE = 'content_file'
TAG_TABLE = 'content_contenttag'
NODE_TAG_TABLE = 'content_contentnode_tags'
LICENSE_TABLE = 'content_license'
NODE_COUNT = 0
FILE_COUNT = 0
TAG_COUNT = 0


class EarlyExit(BaseException):
    def __init__(self, message, db_path):
        self.message = message
        self.db_path = db_path

class Command(BaseCommand):

    help = 'Restores a channel based on its exported database (Usage: restore_channel [source-channel-id] [target-channel-id]'

    def add_arguments(self, parser):

        # ID of channel to read data from
        parser.add_argument('source_id', type=str)

        # ID of channel to write data to (can be same as source channel)
        parser.add_argument('target_id', type=str)

    def handle(self, *args, **options):
        try:
          # Set up variables for restoration process
          logging.info("\n\n********** STARTING CHANNEL RESTORATION **********")
          start = datetime.datetime.now()
          source_id = options['source_id']
          target_id = options.get('target_id') or source_id

          # Test connection to database
          logging.info("Connecting to database for channel {}...".format(source_id))
          conn = sqlite3.connect(os.path.join(settings.DB_ROOT,'{}.sqlite3'.format(source_id)))
          cursor=conn.cursor()

          # Start by creating channel
          logging.info("Creating channel...")
          channel, root_pk = create_channel(conn, target_id)

          # Create nodes mapping to channel
          logging.info("   Creating nodes...")
          root = None
          with transaction.atomic():
            root = create_nodes(cursor, target_id)

          # Save tree to target tree
          channel.main_tree = root
          channel.save()

          # Print stats
          logging.info("\n\nChannel has been restored (time: {ms}, {node_count} nodes, {file_count} files, {tag_count} tags)\n".format(
            ms = datetime.datetime.now() - start,
            node_count = NODE_COUNT,
            file_count = FILE_COUNT,
            tag_count = TAG_COUNT)
          )
          logging.info("\n\n********** RESTORATION COMPLETE **********\n\n")

        except EarlyExit as e:
            logging.warning("Exited early due to {message}.".format(
                message=e.message))
            self.stdout.write("You can find your database in {path}".format(
                path=e.db_path))


def create_channel(cursor, target_id):
    """ create_channel: Create channel at target id
        Args:
            cursor (sqlite3.Connection): connection to export database
            target_id (str): channel_id to write to
        Returns: channel model created and id of root node
    """
    id, name, description, author, thumbnail, root_pk, version = cursor.execute('SELECT * FROM {table}'.format(table=CHANNEL_TABLE)).fetchone()
    channel, is_new = models.Channel.objects.get_or_create(pk=target_id)
    channel.name = name
    channel.description = description
    channel.thumbnail = write_to_thumbnail_file(thumbnail)
    channel.version = version
    channel.save()
    logging.info("\tCreated channel {} with name {}".format(target_id, name))
    return channel, root_pk

def write_to_thumbnail_file(raw_thumbnail):
    """ write_to_thumbnail_file: Convert base64 thumbnail to file
        Args:
            raw_thumbnail (str): base64 encoded thumbnail
        Returns: thumbnail filename
    """
    if raw_thumbnail and isinstance(raw_thumbnail, str) and raw_thumbnail != "" and 'static' not in raw_thumbnail:
        with SimpleUploadedFile('temp.png',raw_thumbnail.replace('data:image/png;base64,', '').decode('base64')) as tempf:
            filename = write_file_to_storage(tempf, check_valid=False)
            logging.info("\tCreated thumbnail {}".format(filename))
            return filename

def create_nodes(cursor, target_id, parent=None, indent=1):
    """ create_channel: Create channel at target id
        Args:
            cursor (sqlite3.Connection): connection to export database
            target_id (str): channel_id to write to
            parent (models.ContentNode): node's parent
            indent (int): How far to indent print statements
        Returns: newly created node
    """
    # Read database rows that match parent
    parent_query = 'parent_id IS NULL'
    if parent:
        parent_query = "parent_id=\'{}\'".format(parent.node_id)

    sql_command = 'SELECT id, title, content_id, description, sort_order, '\
        'license_owner, author, license_id, kind FROM {table} WHERE {query};'\
{'fraction_non_alphanumeric': 0.1011439466158246, 'fraction_numerical': 0.0032411820781696856, 'mean_word_length': 3.063129357087529, 'pattern_counts': {'":': 1, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2277777', 'n_tokens_mistral': 2918, 'n_tokens_neox': 2847, 'n_words': 819}