<commit_before>import numpy

from smqtk.representation import DescriptorElement


__author__ = '<EMAIL>'


# Try to import required modules
try:
    import psycopg2
except ImportError:
    psycopg2 = None


class PostgresDescriptorElement (DescriptorElement):
    """
    Descriptor element whose vector is stored in a Postgres database.

    We assume we will work with a Postgres version of at least 9.4 (due to
    versions tested).

    Efficient connection pooling may be achieved via external utilities like
    PGBounder.

    """

    ARRAY_DTYPE = numpy.float64

    SELECT_TMPL = ' '.join("""
        SELECT {binary_col:s}
          FROM {table_name:s}
          WHERE {type_col:s} = %(type_val)s
            AND {uuid_col:s} = %(uuid_val)s
        ;
    """.split())

    UPSERT_TMPL = ' '.join("""
        WITH upsert AS (
          UPDATE {table_name:s}
            SET {binary_col:s} = %(binary_val)s
            WHERE {type_col:s} = %(type_val)s
              AND {uuid_col:s} = %(uuid_val)s
            RETURNING *
          )
        INSERT INTO {table_name:s} ({type_col:s}, {uuid_col:s}, {binary_col:s})
          SELECT %(type_val)s, %(uuid_val)s, %(binary_val)s
            WHERE NOT EXISTS (SELECT * FROM upsert);
    """.split())

    @classmethod
    def is_usable(cls):
        return psycopg2 is not None

    def __init__(self, type_str, uuid,
                 table_name='descriptors',
                 uuid_col='uid', type_col='type_str', binary_col='vector',
                 db_name='postgres', db_host=None, db_port=None, db_user=None,
                 db_pass=None):
        """
        Initialize new PgsqlDescriptorElement attached to some database
        credentials.

        We require that storage tables treat uuid AND type string columns as
        primary keys. The type and uuid columns should be of the 'text' type.
        The binary column should be of the 'bytea' type.

        Default argument values assume a local PostgreSQL database with a table
        created via the
        ``smqtk/representation/descriptor_element/postgres_element/example_table_init.sql``
        file.

        NOTES:
            - Not all uuid types used here are necessarily of the ``uuid.UUID``<commit_msg>Added warning message when postgres element not usable<commit_after>import numpy

from smqtk.representation import DescriptorElement


__author__ = '<EMAIL>'


# Try to import required modules
try:
    import psycopg2
except ImportError:
    psycopg2 = None


class PostgresDescriptorElement (DescriptorElement):
    """
    Descriptor element whose vector is stored in a Postgres database.

    We assume we will work with a Postgres version of at least 9.4 (due to
    versions tested).

    Efficient connection pooling may be achieved via external utilities like
    PGBounder.

    """

    ARRAY_DTYPE = numpy.float64

    SELECT_TMPL = ' '.join("""
        SELECT {binary_col:s}
          FROM {table_name:s}
          WHERE {type_col:s} = %(type_val)s
            AND {uuid_col:s} = %(uuid_val)s
        ;
    """.split())

    UPSERT_TMPL = ' '.join("""
        WITH upsert AS (
          UPDATE {table_name:s}
            SET {binary_col:s} = %(binary_val)s
            WHERE {type_col:s} = %(type_val)s
              AND {uuid_col:s} = %(uuid_val)s
            RETURNING *
          )
        INSERT INTO {table_name:s} ({type_col:s}, {uuid_col:s}, {binary_col:s})
          SELECT %(type_val)s, %(uuid_val)s, %(binary_val)s
            WHERE NOT EXISTS (SELECT * FROM upsert);
    """.split())

    @classmethod
    def is_usable(cls):
        if psycopg2 is None:
            cls.logger().warning("Not usable. Requires psycopg2 module")
            return False
        return True

    def __init__(self, type_str, uuid,
                 table_name='descriptors',
                 uuid_col='uid', type_col='type_str', binary_col='vector',
                 db_name='postgres', db_host=None, db_port=None, db_user=None,
                 db_pass=None):
        """
        Initialize new PgsqlDescriptorElement attached to some database
        credentials.

        We require that storage tables treat uuid AND type string columns as
        primary keys. The type and uuid columns should be of the 'text' type.
        The binary column should be of the 'bytea' type.

        Default argument values assume a local PostgreSQL database with a table
        created via the
        ``smqtk/representation/descriptor_element/postgres_element/example_table_init.sql``
        file.

        NOTES:
            - Not all uuid types used here are necessarily of the ``uuid.UUID``
{'fraction_non_alphanumeric': 0.09039792387543252, 'fraction_numerical': 0.0032439446366782005, 'mean_word_length': 2.3636363636363638, '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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1465077', 'n_tokens_mistral': 1390, 'n_tokens_neox': 1333, 'n_words': 416}