# postgresql/base.py
# Copyright (C) 2005-2011 the SQLAlchemy authors and contributors <see AUTHORS file>
#
# This module is part of SQLAlchemy and is released under
# the MIT License: http://www.opensource.org/licenses/mit-license.php

"""Support for the PostgreSQL database.

For information on connecting using specific drivers, see the documentation
section regarding that driver.

Sequences/SERIAL
----------------

PostgreSQL supports sequences, and SQLAlchemy uses these as the default means
of creating new primary key values for integer-based primary key columns. When
creating tables, SQLAlchemy will issue the ``SERIAL`` datatype for
integer-based primary key columns, which generates a sequence and server side
default corresponding to the column.

To specify a specific named sequence to be used for primary key generation,
use the :func:`~sqlalchemy.schema.Sequence` construct::

    Table('sometable', metadata, 
            Column('id', Integer, Sequence('some_id_seq'), primary_key=True)
        )

When SQLAlchemy issues a single INSERT statement, to fulfill the contract of
having the "last insert identifier" available, a RETURNING clause is added to
the INSERT statement which specifies the primary key columns should be
returned after the statement completes. The RETURNING functionality only takes
place if Postgresql 8.2 or later is in use. As a fallback approach, the
sequence, whether specified explicitly or implicitly via ``SERIAL``, is
executed independently beforehand, the returned value to be used in the
subsequent insert. Note that when an
:func:`~sqlalchemy.sql.expression.insert()` construct is executed using
"executemany" semantics, the "last inserted identifier" functionality does not
apply; no RETURNING clause is emitted nor is the sequence pre-executed in this
case.

To force the usage of RETURNING by default off, specify the flag
``implicit_returning=False`` to :func:`.create_engine`.

Transaction Isolation Level
---------------------------

:func:`.create_engine` accepts an ``isolation_level`` parameter which results
in the command ``SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL
<level>`` being invoked for every new connection. Valid values for this
parameter are ``READ_COMMITTED``, ``READ_UNCOMMITTED``, ``REPEATABLE_READ``,
and ``SERIALIZABLE``.  Note that the psycopg2 dialect does *not* use this
technique and uses psycopg2-specific APIs (see that dialect for details).

INSERT/UPDATE...RETURNING
-------------------------

The dialect supports PG 8.2's ``INSERT..RETURNING``, ``UPDATE..RETURNING`` and
``DELETE..RETURNING`` syntaxes.   ``INSERT..RETURNING`` is used by default
for single-row INSERT statements in order to fetch newly generated
primary key identifiers.   To specify an explicit ``RETURNING`` clause,
use the :meth:`._UpdateBase.returning` method on a per-statement basis::

    # INSERT..RETURNING
    result = table.insert().returning(table.c.col1, table.c.col2).\\
        values(name='foo')
    print result.fetchall()

    # UPDATE..RETURNING
    result = table.update().returning(table.c.col1, table.c.col2).\\
        where(table.c.name=='foo').values(name='bar')
    print result.fetchall()

    # DELETE..RETURNING
    result = table.delete().returning(table.c.col1, table.c.col2).\\
        where(table.c.name=='foo')
    print result.fetchall()

Indexes
-------

PostgreSQL supports partial indexes. To create them pass a postgresql_where
option to the Index constructor::

  Index('my_index', my_table.c.id, postgresql_where=tbl.c.value > 10)

"""

import re

from sqlalchemy import sql, schema, exc, util
from sqlalchemy.engine import default, reflection
from sqlalchemy.sql import compiler, expression, util as sql_util
from sqlalchemy import types as sqltypes

try:
    from uuid import UUID as _python_UUID
except ImportError:
    _python_UUID = None

from sqlalchemy.types import INTEGER, BIGINT, SMALLINT, VARCHAR, \
        CHAR, TEXT, FLOAT, NUMERIC, \
        DATE, BOOLEAN, REAL

RESERVED_WORDS = set(
    ["all", "analyse", "analyze", "and", "any", "array", "as", "asc",
    "asymmetric", "both", "case", "cast", "check", "collate", "column",
    "constraint", "create", "current_catalog", "current_date",
    "current_role", "current_time", "current_timestamp", "current_user",
    "default", "deferrable", "desc", "distinct", "do", "else", "end",
    "except", "false", "fetch", "for", "foreign", "from", "grant", "group",
    "having", "in", "initially", "intersect", "into", "leading", "limit",
    "localtime", "localtimestamp", "new", "not", "null", "off", "offset",
    "old", "on", "only", "or", "order", "placing", "primary", "references",
    "returning", "select", "session_user", "some", "symmetric", "table",
    "then", "to", "trailing", "true", "union", "unique", "user", "using",
    "variadic", "when", "where", "window", "with", "authorization",
    "between", "binary", "cross", "current_schema", "freeze", "full",
    "ilike", "inner", "is", "isnull", "join", "left", "like", "natural",
    "notnull", "outer", "over", "overlaps", "right", "similar", "verbose"
    ])

_DECIMAL_TYPES = (1231, 1700)
_FLOAT_TYPES = (700, 701, 1021, 1022)
_INT_TYPES = (20, 21, 23, 26, 1005, 1007, 1016)

class BYTEA(sqltypes.LargeBinary):
    __visit_name__ = 'BYTEA'

class DOUBLE_PRECISION(sqltypes.Float):
    __visit_name__ = 'DOUBLE_PRECISION'

class INET(sqltypes.TypeEngine):
    __visit_name__ = "INET"
PGInet = INET

class CIDR(sqltypes.TypeEngine):
    __visit_name__ = "CIDR"
PGCidr = CIDR

class MACADDR(sqltypes.TypeEngine):
    __visit_name__ = "MACADDR"
PGMacAddr = MACADDR

class TIMESTAMP(sqltypes.TIMESTAMP):
    def __init__(self, timezone=False, precision=None):
        super(TIMESTAMP, self).__init__(timezone=timezone)
        self.precision = precision


class TIME(sqltypes.TIME):
    def __init__(self, timezone=False, precision=None):
        super(TIME, self).__init__(timezone=timezone)
        self.precision = precision

class INTERVAL(sqltypes.TypeEngine):
    """Postgresql INTERVAL type.

    The INTERVAL type may not be supported on all DBAPIs.
    It is known to work on psycopg2 and not pg8000 or zxjdbc.

    """
    __visit_name__ = 'INTERVAL'
    def __init__(self, precision=None):
        self.precision = precision

    @classmethod
    def _adapt_from_generic_interval(cls, interval):
        return INTERVAL(precision=interval.second_precision)

    @property
    def _type_affinity(self):
        return sqltypes.Interval

PGInterval = INTERVAL

class BIT(sqltypes.TypeEngine):
    __visit_name__ = 'BIT'
    def __init__(self, length=None, varying=False):
        if not varying:
            # BIT without VARYING defaults to length 1
            self.length = length or 1
        else:
            # but BIT VARYING can be unlimited-length, so no default
            self.length = length
        self.varying = varying

PGBit = BIT

class UUID(sqltypes.TypeEngine):
    """Postgresql UUID type.

    Represents the UUID column type, interpreting
    data either as natively returned by the DBAPI
    or as Python uuid objects.

    The UUID type may not be supported on all DBAPIs.
    It is known to work on psycopg2 and not pg8000.

    """
    __visit_name__ = 'UUID'

    def __init__(self, as_uuid=False):
        """Construct a UUID type.


        :param as_uuid=False: if True, values will be interpreted
         as Python uuid objects, converting to/from string via the
         DBAPI.

         """
        if as_uuid and _python_UUID is None:
            raise NotImplementedError(
                    "This version of Python does not support the native UUID type."
                )
        self.as_uuid = as_uuid

    def bind_processor(self, dialect):
        if self.as_uuid:
            def process(value):
                if value is not None:
                    value = str(value)
                return value
            return process
        else:
            return None

    def result_processor(self, dialect, coltype):
        if self.as_uuid:
            def process(value):
                if value is not None:
                    value = _python_UUID(value)
                return value
            return process
        else:
            return None

PGUuid = UUID

class ARRAY(sqltypes.MutableType, sqltypes.Concatenable, sqltypes.TypeEngine):
    """Postgresql ARRAY type.

    Represents values as Python lists.

    The ARRAY type may not be supported on all DBAPIs.
    It is known to work on psycopg2 and not pg8000.


    """
    __visit_name__ = 'ARRAY'

    def __init__(self, item_type, mutable=False, as_tuple=False):
        """Construct an ARRAY.

        E.g.::

          Column('myarray', ARRAY(Integer))

        Arguments are:

        :param item_type: The data type of items of this array. Note that
          dimensionality is irrelevant here, so multi-dimensional arrays like
          ``INTEGER[][]``, are constructed as ``ARRAY(Integer)``, not as
          ``ARRAY(ARRAY(Integer))`` or such. The type mapping figures out on
          the fly

        :param mutable=False: Specify whether lists passed to this
          class should be considered mutable - this enables 
          "mutable types" mode in the ORM.  Be sure to read the 
          notes for :class:`.MutableType` regarding ORM 
          performance implications (default changed from ``True`` in 
          0.7.0).

          .. note:: This functionality is now superseded by the
             ``sqlalchemy.ext.mutable`` extension described in 
             :ref:`mutable_toplevel`.

        :param as_tuple=False: Specify whether return results
          should be converted to tuples from lists. DBAPIs such
          as psycopg2 return lists by default. When tuples are
          returned, the results are hashable. This flag can only
          be set to ``True`` when ``mutable`` is set to
          ``False``. (new in 0.6.5)

        """
        if isinstance(item_type, ARRAY):
            raise ValueError("Do not nest ARRAY types; ARRAY(basetype) "
                            "handles multi-dimensional arrays of basetype")
        if isinstance(item_type, type):
            item_type = item_type()
        self.item_type = item_type
        self.mutable = mutable
        if mutable and as_tuple:
            raise exc.ArgumentError(
                "mutable must be set to False if as_tuple is True."
            )
        self.as_tuple = as_tuple

    def copy_value(self, value):
        if value is None:
            return None
        elif self.mutable:
            return list(value)
        else:
            return value

    def compare_values(self, x, y):
        return x == y

    def is_mutable(self):
        return self.mutable

    def bind_processor(self, dialect):
        item_proc = self.item_type.dialect_impl(dialect).bind_processor(dialect)
        if item_proc:
            def convert_item(item):
                if isinstance(item, (list, tuple)):
                    return [convert_item(child) for child in item]
                else:
                    return item_proc(item)
        else:
            def convert_item(item):
                if isinstance(item, (list, tuple)):
                    return [convert_item(child) for child in item]
                else:
                    return item
        def process(value):
            if value is None:
                return value
            return [convert_item(item) for item in value]
        return process

    def result_processor(self, dialect, coltype):
        item_proc = self.item_type.dialect_impl(dialect).result_processor(dialect, coltype)
        if item_proc:
            def convert_item(item):
                if isinstance(item, list):
                    r = [convert_item(child) for child in item]
                    if self.as_tuple:
                        r = tuple(r)
                    return r
                else:
                    return item_proc(item)
        else:
            def convert_item(item):
                if isinstance(item, list):
                    r = [convert_item(child) for child in item]
                    if self.as_tuple:
                        r = tuple(r)
                    return r
                else:
                    return item
        def process(value):
            if value is None:
                return value
            r = [convert_item(item) for item in value]
            if self.as_tuple:
                r = tuple(r)
            return r
        return process
PGArray = ARRAY

class ENUM(sqltypes.Enum):

    def create(self, bind=None, checkfirst=True):
        if not bind.dialect.supports_native_enum:
            return

        if not checkfirst or \
            not bind.dialect.has_type(bind, self.name, schema=self.schema):
            bind.execute(CreateEnumType(self))

    def drop(self, bind=None, checkfirst=True):
        if not bind.dialect.supports_native_enum:
            return

        if not checkfirst or \
            bind.dialect.has_type(bind, self.name, schema=self.schema):
            bind.execute(DropEnumType(self))

    def _on_table_create(self, event, target, bind, **kw):
        self.create(bind=bind, checkfirst=True)

    def _on_metadata_create(self, event, target, bind, **kw):
        if self.metadata is not None:
            self.create(bind=bind, checkfirst=True)

    def _on_metadata_drop(self, event, target, bind, **kw):
        self.drop(bind=bind, checkfirst=True)

colspecs = {
    sqltypes.Interval:INTERVAL,
    sqltypes.Enum:ENUM,
}

ischema_names = {
    'integer' : INTEGER,
    'bigint' : BIGINT,
    'smallint' : SMALLINT,
    'character varying' : VARCHAR,
    'character' : CHAR,
    '"char"' : sqltypes.String,
    'name' : sqltypes.String,
    'text' : TEXT,
    'numeric' : NUMERIC,
    'float' : FLOAT,
    'real' : REAL,
    'inet': INET,
    'cidr': CIDR,
    'uuid': UUID,
    'bit': BIT,
    'bit varying': BIT,
    'macaddr': MACADDR,
    'double precision' : DOUBLE_PRECISION,
    'timestamp' : TIMESTAMP,
    'timestamp with time zone' : TIMESTAMP,
    'timestamp without time zone' : TIMESTAMP,
    'time with time zone' : TIME,
    'time without time zone' : TIME,
    'date' : DATE,
    'time': TIME,
    'bytea' : BYTEA,
    'boolean' : BOOLEAN,
    'interval':INTERVAL,
    'interval year to month':INTERVAL,
    'interval day to second':INTERVAL,
}



class PGCompiler(compiler.SQLCompiler):

    def visit_match_op(self, binary, **kw):
        return "%s @@ to_tsquery(%s)" % (
                        self.process(binary.left), 
                        self.process(binary.right))

    def visit_ilike_op(self, binary, **kw):
        escape = binary.modifiers.get("escape", None)
        return '%s ILIKE %s' % \
                (self.process(binary.left), self.process(binary.right)) \
                + (escape and 
                        (' ESCAPE ' + self.render_literal_value(escape, None))
                        or '')

    def visit_notilike_op(self, binary, **kw):
        escape = binary.modifiers.get("escape", None)
        return '%s NOT ILIKE %s' % \
                (self.process(binary.left), self.process(binary.right)) \
                + (escape and 
                        (' ESCAPE ' + self.render_literal_value(escape, None))
                        or '')

    def render_literal_value(self, value, type_):
        value = super(PGCompiler, self).render_literal_value(value, type_)
        # TODO: need to inspect "standard_conforming_strings"
        if self.dialect._backslash_escapes:
            value = value.replace('\\', '\\\\')
        return value

    def visit_sequence(self, seq):
        return "nextval('%s')" % self.preparer.format_sequence(seq)

    def limit_clause(self, select):
        text = ""
        if select._limit is not None:
            text +=  " \n LIMIT " + self.process(sql.literal(select._limit))
        if select._offset is not None:
            if select._limit is None:
                text += " \n LIMIT ALL"
            text += " OFFSET " + self.process(sql.literal(select._offset))
        return text

    def get_select_precolumns(self, select):
        if select._distinct is not False:
            if select._distinct is True:
                return "DISTINCT "
            elif isinstance(select._distinct, (list, tuple)):
                return "DISTINCT ON (" + ', '.join(
                    [self.process(col) for col in select._distinct]
                )+ ") "
            else:
                return "DISTINCT ON (" + self.process(select._distinct) + ") "
        else:
            return ""

    def for_update_clause(self, select):
        if select.for_update == 'nowait':
            return " FOR UPDATE NOWAIT"
        else:
            return super(PGCompiler, self).for_update_clause(select)

    def returning_clause(self, stmt, returning_cols):

        columns = [
                self.process(
                    self.label_select_column(None, c, asfrom=False), 
                    within_columns_clause=True, 
                    result_map=self.result_map) 
                for c in expression._select_iterables(returning_cols)
            ]

        return 'RETURNING ' + ', '.join(columns)

    def visit_extract(self, extract, **kwargs):
        field = self.extract_map.get(extract.field, extract.field)
        if extract.expr.type:
            affinity = extract.expr.type._type_affinity
        else:
            affinity = None

        casts = {
                    sqltypes.Date:'date', 
                    sqltypes.DateTime:'timestamp', 
                    sqltypes.Interval:'interval', sqltypes.Time:'time'
                }
        cast = casts.get(affinity, None)
        if isinstance(extract.expr, sql.ColumnElement) and cast is not None:
            expr = extract.expr.op('::')(sql.literal_column(cast))
        else:
            expr = extract.expr
        return "EXTRACT(%s FROM %s)" % (
            field, self.process(expr))

class PGDDLCompiler(compiler.DDLCompiler):
    def get_column_specification(self, column, **kwargs):
        colspec = self.preparer.format_column(column)
        impl_type = column.type.dialect_impl(self.dialect)
        if column.primary_key and \
            column is column.table._autoincrement_column and \
            not isinstance(impl_type, sqltypes.SmallInteger) and \
            (
                column.default is None or 
                (
                    isinstance(column.default, schema.Sequence) and
                    column.default.optional
                )
            ):
            if isinstance(impl_type, sqltypes.BigInteger):
                colspec += " BIGSERIAL"
            else:
                colspec += " SERIAL"
        else:
            colspec += " " + self.dialect.type_compiler.process(column.type)
            default = self.get_column_default_string(column)
            if default is not None:
                colspec += " DEFAULT " + default

        if not column.nullable:
            colspec += " NOT NULL"
        return colspec

    def visit_create_enum_type(self, create):
        type_ = create.element

        return "CREATE TYPE %s AS ENUM (%s)" % (
            self.preparer.format_type(type_),
            ",".join("'%s'" % e for e in type_.enums)
        )

    def visit_drop_enum_type(self, drop):
        type_ = drop.element

        return "DROP TYPE %s" % (
            self.preparer.format_type(type_)
        )

    def visit_create_index(self, create):
        preparer = self.preparer
        index = create.element
        text = "CREATE "
        if index.unique:
            text += "UNIQUE "
        text += "INDEX %s ON %s (%s)" \
                % (preparer.quote(
                    self._index_identifier(index.name), index.quote),
                   preparer.format_table(index.table),