# -*- coding: utf-8 -*-
#
# Copyright (C) 2005 Edgewall Software
# Copyright (C) 2005 Christopher Lenz <cmlenz@gmx.de>
# All rights reserved.
#
# This software is licensed as described in the file COPYING, which
# you should have received as part of this distribution. The terms
# are also available at http://trac.edgewall.com/license.html.
#
# This software consists of voluntary contributions made by many
# individuals. For the exact contribution history, see the revision
# history and logs, available at http://projects.edgewall.com/trac/.
#
# Author: Christopher Lenz <cmlenz@gmx.de>

from trac.core import *
from trac.db.api import IDatabaseConnector
from trac.db.util import ConnectionWrapper

psycopg = None
PgSQL = None


class PostgreSQLConnector(Component):
    """PostgreSQL database support."""

    implements(IDatabaseConnector)

    def get_supported_schemes(self):
        return [('postgres', 1)]

    def get_connection(self, path, user=None, password=None, host=None,
                       port=None, params={}):
        return PostgreSQLConnection(path, user, password, host, port, params)

    def init_db(self, path, user=None, password=None, host=None, port=None,
                params={}):
        cnx = self.get_connection(path, user, password, host, port, params)
        cursor = cnx.cursor()
        from trac.db_default import schema
        for table in schema:
            for stmt in self.to_sql(table):
                cursor.execute(stmt)
        cnx.commit()

    def to_sql(self, table):
        sql = ["CREATE TABLE %s (" % table.name]
        coldefs = []
        for column in table.columns:
            ctype = column.type
            if column.auto_increment:
                ctype = "SERIAL"
            if len(table.key) == 1 and column.name in table.key:
                ctype += " PRIMARY KEY"
            coldefs.append("    %s %s" % (column.name, ctype))
        if len(table.key) > 1:
            coldefs.append("    CONSTRAINT %s_pk PRIMARY KEY (%s)"
                           % (table.name, ','.join(table.key)))
        sql.append(',\n'.join(coldefs) + '\n)')
        yield '\n'.join(sql)
        for index in table.indices:
            yield "CREATE INDEX %s_%s_idx ON %s (%s)" % (table.name, 
                   '_'.join(index.columns), table.name, ','.join(index.columns))


class PostgreSQLConnection(ConnectionWrapper):
    """Connection wrapper for PostgreSQL."""

    poolable = True

    def __init__(self, path, user=None, password=None, host=None, port=None,
                 params={}):
        if path.startswith('/'):
            path = path[1:]
        # We support both psycopg and PgSQL but prefer psycopg
        global psycopg
        global PgSQL
        global have_psycopg2
        
        if not psycopg and not PgSQL:
            try:
                try:
                    import psycopg2 as psycopg
                    import psycopg2.extensions
                    psycopg2.extensions.register_type(psycopg2.extensions.UNICODE)
                    have_psycopg2 = True
                except ImportError:
                    have_psycopg2 = False
                    import psycopg
            except ImportError:
                from pyPgSQL import PgSQL
        if psycopg:
            dsn = []
            if path:
                dsn.append('dbname=' + path)
            if user:
                dsn.append('user=' + user)
            if password:
                dsn.append('password=' + password)
            if host:
                dsn.append('host=' + host)
            if port:
                dsn.append('port=' + str(port))
            cnx = psycopg.connect(' '.join(dsn))
            if have_psycopg2:
                cnx.set_client_encoding('UNICODE')
        else:
            cnx = PgSQL.connect('', user, password, host, path, port, 
                                client_encoding='utf-8', unicode_results=True)
        ConnectionWrapper.__init__(self, cnx)

    def cast(self, column, type):
        # Temporary hack needed for the union of selects in the search module
        return 'CAST(%s AS %s)' % (column, type)

    def like(self):
        # Temporary hack needed for the case-insensitive string matching in the
        # search module
        return 'ILIKE'

    def get_last_id(self, cursor, table, column='id'):
        cursor.execute("SELECT CURRVAL('%s_%s_seq')" % (table, column))
        return cursor.fetchone()[0]

@@ -724,16 +724,37 @@
  = None%0A
+PGSchemaError = None%0A
 %0A%0Aclass 
@@ -1339,16 +1339,174 @@
 ursor()%0A
+        if cnx.schema:%0A            cursor.execute('CREATE SCHEMA %25s' %25 cnx.schema)%0A            cursor.execute('SET search_path TO %25s, public', (cnx.schema,))%0A
         
@@ -2891,24 +2891,53 @@
 lobal PgSQL%0A
+        global PGSchemaError%0A
         glob
@@ -3132,16 +3132,91 @@
 ensions%0A
+                    from psycopg2 import ProgrammingError as PGSchemaError%0A
         
@@ -3460,69 +3460,219 @@
     
-except ImportError:%0A                from pyPgSQL import PgSQL
+        from psycopg import ProgrammingError as PGSchemaError%0A            except ImportError:%0A                from pyPgSQL import PgSQL%0A                from pyPgSQL.libpq import OperationalError as PGSchemaError
 %0A   
@@ -4342,16 +4342,314 @@
 s=True)%0A
+        try:%0A            self.schema = None%0A            if 'schema' in params:%0A                self.schema = params%5B'schema'%5D%0A            cnx.cursor().execute('SET search_path TO %25s, public', %0A                                (self.schema,))%0A        except PGSchemaError:%0A            cnx.rollback()%0A
         
@@ -5160,8 +5160,305 @@
 ne()%5B0%5D%0A
+%0A    def rollback(self):%0A        self.cnx.rollback()%0A        if self.schema:%0A            try:%0A                self.cnx.cursor().execute(%22SET search_path TO %25s, public%22, %0A                                         (self.schema,))%0A            except PGSchemaError:%0A                self.cnx.rollback()%0A

Apply patch by pacopablo that adds support for using PostgreSQL schemas with Trac. Closes #2711.