<commit_before>from __future__ import absolute_import
from __future__ import print_function
import datetime
from boto.s3.key import Key
from boto.s3.connection import S3Connection
from django.conf import settings
from django.db import connection
from django.forms.models import model_to_dict
from django.utils import timezone
import glob
import logging
import os
import ujson
import shutil
import subprocess
import tempfile
from zerver.lib.avatar import user_avatar_hash
from zerver.lib.create_user import random_api_key
from zerver.models import UserProfile, Realm, Client, Huddle, Stream, \
    UserMessage, Subscription, Message, RealmEmoji, RealmFilter, \
    RealmAlias, Recipient, DefaultStream, get_user_profile_by_id, \
    UserPresence, UserActivity, UserActivityInterval, get_user_profile_by_email, \
    get_display_recipient, Attachment
from zerver.lib.parallel import run_parallel
from zerver.lib.utils import mkdir_p
from six import text_type
from six.moves import range
from typing import Any, Dict, List, Tuple

# Custom mypy types follow:
Record = Dict[str, Any]
TableName = str
TableData = Dict[TableName, List[Record]]
Field = str
Path = text_type

# The keys of our MessageOutput variables are normally
# List[Record], but when we write partials, we can get
# lists of integers or a single integer.
# TODO: tighten this up with a union.
MessageOutput = Dict[str, Any]

realm_tables = [("zerver_defaultstream", DefaultStream),
                ("zerver_realmemoji", RealmEmoji),
                ("zerver_realmalias", RealmAlias),
                ("zerver_realmfilter", RealmFilter)] # List[Tuple[TableName, Any]]


DATE_FIELDS = {
    'zerver_attachment': ['create_time'],
    'zerver_message': ['last_edit_time', 'pub_date'],
    'zerver_realm': ['date_created'],
    'zerver_stream': ['date_created'],
    'zerver_useractivity': ['last_visit'],
    'zerver_useractivityinterval': ['start', 'end'],
    'zerver_userpresence': ['timestamp'],
    'zerver_userprofile': ['date_joined', 'last_login', 'last_reminder'],
} # type: Dict[TableName, List[Field]]


def make_raw(query):
    # type: (Any) -> List[Record]
    '''
    Takes a Django query and returns a JSONable list
    of dictionaries corresponding to the database rows.
    '''
    return [model_to_dict(x) for x in query]

def floatify_datetime_fields(data, table):
    # type: (TableData, TableName) -> None
    for item in data[table]:
        for field in DATE_FIELDS[table]:
            orig_dt = item[field]
            if orig_dt is None:
                continue
            if timezone.is_naive(orig_dt):
                logging.warning("Naive datetime:", item)
                dt = timezone.make_aware(orig_dt)
            else:
                dt = orig_dt
            utc_naive  = dt.replace(tzinfo=None) - dt.utcoffset()
            item[field] = (utc_naive - datetime.datetime(1970, 1, 1)).total_seconds()

# Export common, public information about the realm that we can share
# with all realm users
def export_realm_data(realm, response):
    # type: (Realm, TableData) -> None
    response['zerver_realm'] = make_raw(Realm.objects.filter(id=realm.id))
    floatify_datetime_fields(response, 'zerver_realm')

    for (table, model) in realm_tables:
        # mypy does not know that model is a Django model that
        # supports "objects"
        table_query =  model.objects.filter(realm_id=realm.id) # type: ignore
        response[table] =  make_raw(table_query)
    response["zerver_client"] = make_raw(Client.objects.select_related())

# To export only some users, you can tweak the below UserProfile query
# to give the target users, but then you should create any users not
# being exported in a separate
# response['zerver_userprofile_mirrordummy'] export so that
# conversations with those users can still be exported.
def export_with_admin_auth(realm, response, include_invite_only=True, include_private=True):
    # type: (Realm, TableData, bool, bool) -> None
    response['zerver_userprofile'] = [model_to_dict(x, exclude=["password", "api_key"])
                                      for x in UserProfile.objects.filter(realm=realm)]
    if realm.domain == "zulip.com":
        response['zerver_userprofile_crossrealm'] = []
    else:
        response['zerver_userprofile_crossrealm'] = [dict(email=x.email, id=x.id) for x in [
        get_user_profile_by_email(settings.NOTIFICATION_BOT),
        get_user_profile_by_email(settings.EMAIL_GATEWAY_BOT),
        get_user_profile_by_email(settings.WELCOME_BOT),
    ]]
    floatify_datetime_fields(response, 'zerver_userprofile')
    user_profile_ids = set(userprofile["id"] for userprofile in response['zerver_userprofile'])
    user_recipient_query = Recipient.objects.filter(type=Recipient.PERSONAL,
                                                    type_id__in=user_profile_ids)
    user_recipients = make_raw(user_recipient_query)
    user_recipient_ids = set(x["id"] for x in user_recipients)
    user_subscription_query = Subscription.objects.filter(user_profile__in=user_profile_ids,
                                                          recipient_id__in=user_recipient_ids)
    user_subscription_dicts = make_raw(user_subscription_query)

    user_presence_query = UserPresence.objects.filter(user_profile__in=user_profile_ids)
    response["zerver_userpresence"] = make_raw(user_presence_query)
    floatify_datetime_fields(response, 'zerver_userpresence')

    user_activity_query = UserActivity.objects.filter(user_profile__in=user_profile_ids)
    response["zerver_useractivity"] = make_raw(user_activity_query)
    floatify_datetime_fields(response, 'zerver_useractivity')

    user_activity_interval_query = UserActivityInterval.objects.filter(user_profile__in=user_profile_ids)
    response["zerver_useractivityinterval"] = make_raw(user_activity_interval_query)
    floatify_datetime_fields(response, 'zerver_useractivityinterval')

    stream_query = Stream.objects.filter(realm=realm)
    if not include_invite_only:
        stream_query = stream_query.filter(invite_only=False)
    response['zerver_stream'] = [model_to_dict(x, exclude=["email_token"]) for x in stream_query]
    floatify_datetime_fields(response, 'zerver_stream')
    stream_ids = set(x["id"] for x in response['zerver_stream'])

    stream_recipient_query = Recipient.objects.filter(type=Recipient.STREAM,
                                                      type_id__in=stream_ids)
    stream_recipients = make_raw(stream_recipient_query)
    stream_recipient_ids = set(x["id"] for x in stream_recipients)

    stream_subscription_query = Subscription.objects.filter(user_profile__in=user_profile_ids,
                                                            recipient_id__in=stream_recipient_ids)
    stream_subscription_dicts = make_raw(stream_subscription_query)

    if include_private:
        # First we get all huddles involving someone in the realm.
        realm_huddle_subs = Subscription.objects.select_related("recipient").filter(recipient__type=Recipient.HUDDLE,
                                                                                    user_profile__in=user_profile_ids)
        realm_huddle_recipient_ids = set(sub.recipient_id for sub in realm_huddle_subs)

        # Mark all Huddles whose recipient ID contains a cross-realm user.
        unsafe_huddle_recipient_ids = set()
        for sub in Subscription.objects.select_related().filter(recipient__in=realm_huddle_recipient_ids):
            if sub.user_profile.realm != realm:
                # In almost every case the other realm will be zulip.com
                unsafe_huddle_recipient_ids.add(sub.recipient_id)

        # Now filter down to just those huddles that are entirely within the realm.
        #
        # This is important for ensuring that the User objects needed
        # to import it on the other end exist (since we're only
        # exporting the users from this realm), at the cost of losing
        # some of these cross-realm messages.
        huddle_subs = [sub for sub in realm_huddle_subs if sub.recipient_id not in unsafe_huddle_recipient_ids]
        huddle_recipient_ids = set(sub.recipient_id for sub in huddle_subs)
        huddle_ids = set(sub.recipient.type_id for sub in huddle_subs)

        huddle_subscription_dicts = make_raw(huddle_subs)<commit_msg>export: Extract filter_by_users().<commit_after>from __future__ import absolute_import
from __future__ import print_function
import datetime
from boto.s3.key import Key
from boto.s3.connection import S3Connection
from django.conf import settings
from django.db import connection
from django.forms.models import model_to_dict
from django.utils import timezone
import glob
import logging
import os
import ujson
import shutil
import subprocess
import tempfile
from zerver.lib.avatar import user_avatar_hash
from zerver.lib.create_user import random_api_key
from zerver.models import UserProfile, Realm, Client, Huddle, Stream, \
    UserMessage, Subscription, Message, RealmEmoji, RealmFilter, \
    RealmAlias, Recipient, DefaultStream, get_user_profile_by_id, \
    UserPresence, UserActivity, UserActivityInterval, get_user_profile_by_email, \
    get_display_recipient, Attachment
from zerver.lib.parallel import run_parallel
from zerver.lib.utils import mkdir_p
from six import text_type
from six.moves import range
from typing import Any, Dict, List, Tuple

# Custom mypy types follow:
Record = Dict[str, Any]
TableName = str
TableData = Dict[TableName, List[Record]]
Field = str
Path = text_type

# The keys of our MessageOutput variables are normally
# List[Record], but when we write partials, we can get
# lists of integers or a single integer.
# TODO: tighten this up with a union.
MessageOutput = Dict[str, Any]

realm_tables = [("zerver_defaultstream", DefaultStream),
                ("zerver_realmemoji", RealmEmoji),
                ("zerver_realmalias", RealmAlias),
                ("zerver_realmfilter", RealmFilter)] # List[Tuple[TableName, Any]]


DATE_FIELDS = {
    'zerver_attachment': ['create_time'],
    'zerver_message': ['last_edit_time', 'pub_date'],
    'zerver_realm': ['date_created'],
    'zerver_stream': ['date_created'],
    'zerver_useractivity': ['last_visit'],
    'zerver_useractivityinterval': ['start', 'end'],
    'zerver_userpresence': ['timestamp'],
    'zerver_userprofile': ['date_joined', 'last_login', 'last_reminder'],
} # type: Dict[TableName, List[Field]]


def make_raw(query):
    # type: (Any) -> List[Record]
    '''
    Takes a Django query and returns a JSONable list
    of dictionaries corresponding to the database rows.
    '''
    return [model_to_dict(x) for x in query]

def floatify_datetime_fields(data, table):
    # type: (TableData, TableName) -> None
    for item in data[table]:
        for field in DATE_FIELDS[table]:
            orig_dt = item[field]
            if orig_dt is None:
                continue
            if timezone.is_naive(orig_dt):
                logging.warning("Naive datetime:", item)
                dt = timezone.make_aware(orig_dt)
            else:
                dt = orig_dt
            utc_naive  = dt.replace(tzinfo=None) - dt.utcoffset()
            item[field] = (utc_naive - datetime.datetime(1970, 1, 1)).total_seconds()

# Export common, public information about the realm that we can share
# with all realm users
def export_realm_data(realm, response):
    # type: (Realm, TableData) -> None
    response['zerver_realm'] = make_raw(Realm.objects.filter(id=realm.id))
    floatify_datetime_fields(response, 'zerver_realm')

    for (table, model) in realm_tables:
        # mypy does not know that model is a Django model that
        # supports "objects"
        table_query =  model.objects.filter(realm_id=realm.id) # type: ignore
        response[table] =  make_raw(table_query)
    response["zerver_client"] = make_raw(Client.objects.select_related())

# To export only some users, you can tweak the below UserProfile query
# to give the target users, but then you should create any users not
# being exported in a separate
# response['zerver_userprofile_mirrordummy'] export so that
# conversations with those users can still be exported.
def export_with_admin_auth(realm, response, include_invite_only=True, include_private=True):
    # type: (Realm, TableData, bool, bool) -> None
    response['zerver_userprofile'] = [model_to_dict(x, exclude=["password", "api_key"])
                                      for x in UserProfile.objects.filter(realm=realm)]
    if realm.domain == "zulip.com":
        response['zerver_userprofile_crossrealm'] = []
    else:
        response['zerver_userprofile_crossrealm'] = [dict(email=x.email, id=x.id) for x in [
        get_user_profile_by_email(settings.NOTIFICATION_BOT),
        get_user_profile_by_email(settings.EMAIL_GATEWAY_BOT),
        get_user_profile_by_email(settings.WELCOME_BOT),
    ]]
    floatify_datetime_fields(response, 'zerver_userprofile')
    user_profile_ids = set(userprofile["id"] for userprofile in response['zerver_userprofile'])


    user_recipient_query = Recipient.objects.filter(type=Recipient.PERSONAL,
                                                    type_id__in=user_profile_ids)
    user_recipients = make_raw(user_recipient_query)
    user_recipient_ids = set(x["id"] for x in user_recipients)


    def filter_by_users(model, **kwargs):
        # type: (Any, **Any) -> Any
        return model.objects.filter(user_profile__in=user_profile_ids, **kwargs)


    user_subscription_query = filter_by_users(Subscription,
                                              recipient_id__in=user_recipient_ids)
    user_subscription_dicts = make_raw(user_subscription_query)

    user_presence_query = filter_by_users(UserPresence)
    response["zerver_userpresence"] = make_raw(user_presence_query)
    floatify_datetime_fields(response, 'zerver_userpresence')

    user_activity_query = filter_by_users(UserActivity)
    response["zerver_useractivity"] = make_raw(user_activity_query)
    floatify_datetime_fields(response, 'zerver_useractivity')

    user_activity_interval_query = filter_by_users(UserActivityInterval)
    response["zerver_useractivityinterval"] = make_raw(user_activity_interval_query)
    floatify_datetime_fields(response, 'zerver_useractivityinterval')

    stream_query = Stream.objects.filter(realm=realm)
    if not include_invite_only:
        stream_query = stream_query.filter(invite_only=False)
    response['zerver_stream'] = [model_to_dict(x, exclude=["email_token"]) for x in stream_query]
    floatify_datetime_fields(response, 'zerver_stream')
    stream_ids = set(x["id"] for x in response['zerver_stream'])

    stream_recipient_query = Recipient.objects.filter(type=Recipient.STREAM,
                                                      type_id__in=stream_ids)
    stream_recipients = make_raw(stream_recipient_query)
    stream_recipient_ids = set(x["id"] for x in stream_recipients)

    stream_subscription_query = filter_by_users(Subscription,
                                                recipient_id__in=stream_recipient_ids)
    stream_subscription_dicts = make_raw(stream_subscription_query)

    if include_private:
        # First we get all huddles involving someone in the realm.
        realm_huddle_subs = Subscription.objects.select_related("recipient").filter(recipient__type=Recipient.HUDDLE,
                                                                                    user_profile__in=user_profile_ids)
        realm_huddle_recipient_ids = set(sub.recipient_id for sub in realm_huddle_subs)

        # Mark all Huddles whose recipient ID contains a cross-realm user.
        unsafe_huddle_recipient_ids = set()
        for sub in Subscription.objects.select_related().filter(recipient__in=realm_huddle_recipient_ids):
            if sub.user_profile.realm != realm:
                # In almost every case the other realm will be zulip.com
                unsafe_huddle_recipient_ids.add(sub.recipient_id)

        # Now filter down to just those huddles that are entirely within the realm.
        #
        # This is important for ensuring that the User objects needed
        # to import it on the other end exist (since we're only
        # exporting the users from this realm), at the cost of losing
        # some of these cross-realm messages.
        huddle_subs = [sub for sub in realm_huddle_subs if sub.recipient_id not in unsafe_huddle_recipient_ids]
        huddle_recipient_ids = set(sub.recipient_id for sub in huddle_subs)
        huddle_ids = set(sub.recipient.type_id for sub in huddle_subs)

        huddle_subscription_dicts = make_raw(huddle_subs)
{'fraction_non_alphanumeric': 0.08335835685544411, 'fraction_numerical': 0.0010810161551858747, 'mean_word_length': 3.5509702104400107, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 12, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 26, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4585948', 'n_tokens_mistral': 5077, 'n_tokens_neox': 4918, 'n_words': 1185}