<commit_before># -*- coding: utf-8 -*-
import logging

from sqlalchemy import Table, MetaData
from sqlalchemy.orm import mapper
from sqlalchemy.sql.expression import text

from ambry.util import get_logger

logger = get_logger(__name__, level=logging.DEBUG, propagate=False)


def add_partition(connection, partition):
    FOREIGN_SERVER_NAME = 'partition_server'
    _create_if_not_exists(connection, FOREIGN_SERVER_NAME)

    columns = ['{} {}'.format(c.name, c.type.compile()) for c in partition.table.columns]
    query = """
        CREATE FOREIGN TABLE {table_name} (
            {columns}
        ) server {server_name} options (
            filename '{file_name}'
        );
    """.format(table_name=_table_name(partition), columns=',\n'.join(columns),
               server_name=FOREIGN_SERVER_NAME, file_name=partition.datafile.syspath)
    logging.debug('Create foreign table for {} partition. Query:\n{}.'.format(partition.vid, query))
    connection.execute(query)


def _as_orm(connection, partition):
    """ Returns sqlalchemy table for partition rows.

    Example:
        PartitionRow = _as_orm(connection, partition)
        print session.query(PartitionRow).all()

    Returns:
        FIXME:
    """

    # FIXME: That solution is not documented by multicorn. Try documented solution again.

    class PartitionRow(object):
        pass

    table_name = _table_name(partition)
    metadata = MetaData(bind=connection.engine)
    t = Table(table_name, metadata, *partition.table.columns)
    mapper(PartitionRow, t)
    return PartitionRow


def _server_exists(connection, server_name):
    """ Returns True is foreign server with given name exists. Otherwise returns False. """
    query = text("""
        SELECT 1 FROM pg_foreign_server WHERE srvname=:server_name;
    """)
    return connection.execute(query, server_name=server_name).fetchall() == [(1,)]


def _create_if_not_exists(connection, server_name):
    """ Creates foreign server if it does not exist. """
    if not _server_exists(connection, server_name):
        logging.info('Create {} foreign server because it does not exist.'.format(server_name))
        query = """
            CREATE SERVER {} FOREIGN DATA WRAPPER multicorn
            options (
                wrapper 'ambryfdw.PartitionMsgpackForeignDataWrapper'<commit_msg>Partition row orm simplified. #99.<commit_after># -*- coding: utf-8 -*-
import logging

from sqlalchemy import Table, MetaData
from sqlalchemy.orm import mapper
from sqlalchemy.sql.expression import text

from ambry.util import get_logger

logger = get_logger(__name__, level=logging.DEBUG, propagate=False)


def add_partition(connection, partition):
    FOREIGN_SERVER_NAME = 'partition_server'
    _create_if_not_exists(connection, FOREIGN_SERVER_NAME)

    columns = ['{} {}'.format(c.name, c.type.compile()) for c in partition.table.columns]
    query = """
        CREATE FOREIGN TABLE {table_name} (
            {columns}
        ) server {server_name} options (
            filename '{file_name}'
        );
    """.format(table_name=_table_name(partition), columns=',\n'.join(columns),
               server_name=FOREIGN_SERVER_NAME, file_name=partition.datafile.syspath)
    logging.debug('Create foreign table for {} partition. Query:\n{}.'.format(partition.vid, query))
    connection.execute(query)


def _as_orm(connection, partition):
    """ Returns sqlalchemy table for partition rows.

    Example:
        PartitionRow = _as_orm(connection, partition)
        print session.query(PartitionRow).all()

    Returns:
        FIXME:
    """

    # FIXME: That solution is not documented by multicorn. Try documented solution again.

    table_name = _table_name(partition)
    metadata = MetaData(bind=connection.engine)
    PartitionRow = Table(table_name, metadata, *partition.table.columns)
    return PartitionRow


def _server_exists(connection, server_name):
    """ Returns True is foreign server with given name exists. Otherwise returns False. """
    query = text("""
        SELECT 1 FROM pg_foreign_server WHERE srvname=:server_name;
    """)
    return connection.execute(query, server_name=server_name).fetchall() == [(1,)]


def _create_if_not_exists(connection, server_name):
    """ Creates foreign server if it does not exist. """
    if not _server_exists(connection, server_name):
        logging.info('Create {} foreign server because it does not exist.'.format(server_name))
        query = """
            CREATE SERVER {} FOREIGN DATA WRAPPER multicorn
            options (
                wrapper 'ambryfdw.PartitionMsgpackForeignDataWrapper'
{'fraction_non_alphanumeric': 0.1052517361111111, 'fraction_numerical': 0.001736111111111111, 'mean_word_length': 3.5408866995073893, 'pattern_counts': {'":': 0, '<': 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': '4296845', 'n_tokens_mistral': 1329, 'n_tokens_neox': 1274, 'n_words': 332}