<commit_before>"""Search API.

It's assumed that the API is used with an object stored in a
RelStorage with a Postgres back end.
"""

import re
from ZODB.utils import p64

def _try_to_close_cursor(cursor):
    try:
        cursor.close()
    except Exception:
        pass

def search(conn, query, *args, **kw):
    """Search for newt objects using an SQL query.

    Query parameters may be provided as either positional
    arguments or keyword arguments.  They are inserted into the
    query where there are placeholders of the form ``%s`` for
    positional arguments or ``%(NAME)s`` for keyword arguments.

    The query results must contain the columns ``zoid`` and
    ``ghost_pickle``.  It's simplest and costs nothing to simply
    select all columns (using ``*``) from the ``newt`` table.

    A sequence of newt objects is returned.
    """
    if kw:
        if args:
            raise TypeError("Only positional or keyword arguments can be used,"
                            " not both")
        args = kw
    get = conn.ex_get
    cursor = read_only_cursor(conn)
    try:
        cursor.execute((b"select zoid, ghost_pickle from (%s)_"
                        if isinstance(query, bytes) else
                        "select zoid, ghost_pickle from (%s)_"
                        ) % query,
                       args or kw)
        return [get(p64(zoid), ghost_pickle) for (zoid, ghost_pickle) in cursor]
    finally:
        _try_to_close_cursor(cursor)

def search_batch(conn, query, args, batch_start, batch_size=None):
    """Query for a batch of newt objects.

    Query parameters are provided using the ``args``
    argument, which may be a tuple or a dictionary.  They are
    inserted into the query where there are placeholders of the
    form ``%s`` for an arguments tuple or ``%(NAME)s`` for an
    arguments dict.

    The ``batch_size`` and ``batch_size`` arguments are used to
    specify the result batch.  An ``ORDER BY`` clause should be
    used to order results.

    The total result count and sequence of batch result objects
    are returned.

    The query parameters, ``args``, may be omitted. (In this case,
    ``batch_size`` will be None and the other arguments will be
    re-arranged appropriately. ``batch_size`` *is required*.)  You
    might use this feature if you pre-inserted data using a database
    cursor `mogrify
    <http://initd.org/psycopg/docs/cursor.html#cursor.mogrify>`_
    method.
    """
    if not batch_size:
        if isinstance(args, int):
            batch_size = batch_start
            batch_start = args
            args = ()
        else:
            raise AssertionError("Invalid batch size %r" % batch_size)

    query = (
        b"""select zoid, ghost_pickle, count(*) over()
        from (%s) _
        offset %d limit %d
        """
        if isinstance(query, bytes) else
        """select zoid, ghost_pickle, count(*) over()
        from (%s) _
        offset %d limit %d
        """
        ) % (query, batch_start, batch_size)
    get = conn.ex_get
    cursor = read_only_cursor(conn)
    try:
        cursor.execute(query, args)
        count = 0
        result = []
        for zoid, ghost_pickle, count in cursor:
            result.append(get(p64(zoid), ghost_pickle))
        return count, result
    finally:
        _try_to_close_cursor(cursor)


text_extraction_template = """\
create or replace function %s(state jsonb) returns tsvector as $$<commit_msg>workaround: 3.4 doesn't support byte format strings<commit_after>"""Search API.

It's assumed that the API is used with an object stored in a
RelStorage with a Postgres back end.
"""

import re
from ZODB.utils import p64

def _try_to_close_cursor(cursor):
    try:
        cursor.close()
    except Exception:
        pass

def search(conn, query, *args, **kw):
    """Search for newt objects using an SQL query.

    Query parameters may be provided as either positional
    arguments or keyword arguments.  They are inserted into the
    query where there are placeholders of the form ``%s`` for
    positional arguments or ``%(NAME)s`` for keyword arguments.

    The query results must contain the columns ``zoid`` and
    ``ghost_pickle``.  It's simplest and costs nothing to simply
    select all columns (using ``*``) from the ``newt`` table.

    A sequence of newt objects is returned.
    """
    if kw:
        if args:
            raise TypeError("Only positional or keyword arguments can be used,"
                            " not both")
        args = kw
    get = conn.ex_get
    cursor = read_only_cursor(conn)
    try:
        cursor.execute(b"select zoid, ghost_pickle from (" + query + b")_"
                       if isinstance(query, bytes) else
                       "select zoid, ghost_pickle from (" + query + ")_",
                       args or kw)
        return [get(p64(zoid), ghost_pickle) for (zoid, ghost_pickle) in cursor]
    finally:
        _try_to_close_cursor(cursor)

def search_batch(conn, query, args, batch_start, batch_size=None):
    """Query for a batch of newt objects.

    Query parameters are provided using the ``args``
    argument, which may be a tuple or a dictionary.  They are
    inserted into the query where there are placeholders of the
    form ``%s`` for an arguments tuple or ``%(NAME)s`` for an
    arguments dict.

    The ``batch_size`` and ``batch_size`` arguments are used to
    specify the result batch.  An ``ORDER BY`` clause should be
    used to order results.

    The total result count and sequence of batch result objects
    are returned.

    The query parameters, ``args``, may be omitted. (In this case,
    ``batch_size`` will be None and the other arguments will be
    re-arranged appropriately. ``batch_size`` *is required*.)  You
    might use this feature if you pre-inserted data using a database
    cursor `mogrify
    <http://initd.org/psycopg/docs/cursor.html#cursor.mogrify>`_
    method.
    """
    if not batch_size:
        if isinstance(args, int):
            batch_size = batch_start
            batch_start = args
            args = ()
        else:
            raise AssertionError("Invalid batch size %r" % batch_size)

    if isinstance(query, str):
        query = """select zoid, ghost_pickle, count(*) over()
        from (%s) _
        offset %d limit %d
        """ % (query, batch_start, batch_size)
    else:
        # Python 3.4, whimper
        query = (
            b"select zoid, ghost_pickle, count(*) over()\nfrom (" +
            query +
            (") _\noffset %s limit %d" % (batch_start, batch_size)
             ).encode('ascii')
            )

    get = conn.ex_get
    cursor = read_only_cursor(conn)
    try:
        cursor.execute(query, args)
        count = 0
        result = []
        for zoid, ghost_pickle, count in cursor:
            result.append(get(p64(zoid), ghost_pickle))
        return count, result
    finally:
        _try_to_close_cursor(cursor)


text_extraction_template = """\
create or replace function %s(state jsonb) returns tsvector as $$
{'fraction_non_alphanumeric': 0.0930892578402157, 'fraction_numerical': 0.002554278416347382, 'mean_word_length': 2.39172281039461, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3910580', 'n_tokens_mistral': 2041, 'n_tokens_neox': 1985, 'n_words': 756}