<commit_before># Copyright 2011 Omniscale (http://omniscale.com)
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#     http://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import time

import psycopg2
import psycopg2.extensions

import logging
log = logging.getLogger(__name__)

from imposm import config
from imposm.mapping import UnionView, GeneralizedTable, Mapping

class PostGISDB(object):
    def __init__(self, db_conf):
        self.db_conf = db_conf
        self.srid = int(db_conf['proj'].split(':')[1])
        self._insert_stmts = {}
        self._connection = None
        self._cur = None

    @property
    def table_prefix(self):
        return self.db_conf.prefix.rstrip('_') + '_'

    def to_tablename(self, name):
        return self.table_prefix + name.lower()

    @property
    def connection(self):
        if not self._connection:
            kw = {}
            if self.db_conf.port:
                kw['port'] = int(self.db_conf.port)
            self._connection = psycopg2.connect(
                database=self.db_conf.db,
                host=self.db_conf.host,
                user=self.db_conf.user,
                password=self.db_conf.password,
                sslmode=self.db_conf.get('sslmode', 'allow'),
                **kw
            )
            self._connection.set_isolation_level(
                psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED)
        return self._connection
    
    def commit(self):
        self.connection.commit()
    
    @property
    def cur(self):
        if self._cur is None:
            self._cur = self.connection.cursor()
        return self._cur
    
    def insert(self, mapping, insert_data, tries=0):
        insert_stmt = self.insert_stmt(mapping)
        try:
            if tries:
                self.reconnect()
            self.cur.executemany(insert_stmt, insert_data)
        except psycopg2.OperationalError, ex:
            if tries >= 8:
                log.warn('%s, giving up', ex)
                raise
            seconds = 2 ** (tries + 1)
            log.warn('%s, retry in %d', ex, seconds)
            time.sleep(seconds)
            self.insert(mapping, insert_data, tries=tries + 1)
        except psycopg2.Error, ex:
            self.connection.rollback()
            for data in insert_data:
                try:
                    self.cur.execute(insert_stmt, data)
                except psycopg2.Error, ex:
                    log.warn('error while importing "%r": %s', data, ex)
                    self.connection.rollback()
                else:
                    self.connection.commit()

        self.connection.commit()
    
    def geom_wrapper(self, geom):
        return psycopg2.Binary(geom.wkb)
    
    def reconnect(self):
        if self._connection:
            try:
                self._connection.close()
            except psycopg2.InterfaceError:
                pass
        self._connection = None
        self._cur = None

    def insert_stmt(self, mapping):
        if mapping.name not in self._insert_stmts:
            self._insert_stmts[mapping.name] = self._insert_stmt(mapping)

        return self._insert_stmts[mapping.name]

    def _insert_stmt(self, mapping):
        extra_arg_names = extra_args = ''
        if mapping.fields:
            extra_arg_names = [n for n, t in mapping.fields]
            extra_args = ', %s' * len(extra_arg_names)
            extra_arg_names = ', ' + ', '.join('"' + name + '"' for name in extra_arg_names)
        return """INSERT INTO "%(tablename)s"
            (osm_id, geometry %(extra_arg_names)s)
            VALUES (%%s, ST_Transform(ST_GeomFromWKB(%%s, 4326), %(srid)s)
                %(extra_args)s)
        """.strip() % dict(tablename=self.table_prefix + mapping.name, srid=self.srid,
            extra_arg_names=extra_arg_names, extra_args=extra_args)


    def create_tables(self, mappings):
        for mapping in mappings:
            self.create_table(mapping)

    def create_table(self, mapping):
        tablename = self.table_prefix + mapping.name
        cur = self.connection.cursor()
        cur.execute('SAVEPOINT pre_drop_tables')
        try:
            cur.execute('DROP TABLE "' + tablename + '" CASCADE')
        except psycopg2.ProgrammingError:
            cur.execute('ROLLBACK TO SAVEPOINT pre_drop_tables')

        extra_fields = ''
        for n, t in mapping.fields:
            extra_fields += ', "%s" %s ' % (n, t.column_type)

        if config.imposm_pg_serial_id:
            serial_column = "id SERIAL PRIMARY KEY,"
        else:
            serial_column = ""
            
        cur.execute("""
            CREATE TABLE "%s" (
                %s
                osm_id INT4
                %s
            );
        """ % (tablename, serial_column, extra_fields))
        cur.execute("""
            SELECT AddGeometryColumn ('', '%(tablename)s', 'geometry',
                                      %(srid)s, '%(pg_geometry_type)s', 2)
        """ % dict(tablename=tablename, srid=self.srid,
                   pg_geometry_type=mapping.geom_type))

        for n, t in mapping.fields:
            if isinstance(t, TrigramIndex):
                cur.execute("""
                    CREATE INDEX "%(tablename)s_trgm_idx_%(column)s" ON "%(tablename)s" USING GIST (%(column)s gist_trgm_ops)
                """ % dict(tablename=tablename, column=n))

        cur.execute("""
            CREATE INDEX "%(tablename)s_geom" ON "%(tablename)s" USING GIST (geometry)
        """ % dict(tablename=tablename))
    
    def swap_tables(self, new_prefix, existing_prefix, backup_prefix):
        cur = self.connection.cursor()

        self.remove_tables(backup_prefix)
        
        cur.execute('SELECT tablename FROM pg_tables WHERE tablename like %s', (existing_prefix + '%', ))
        existing_tables = []
        for row in cur:
            table_name = row[0]
            if table_name.startswith(existing_prefix) and not table_name.startswith((new_prefix, backup_prefix)):
                # check for overlapping prefixes: osm_ but not osm_new_ or osm_backup_
                existing_tables.append(table_name)

        cur.execute('SELECT indexname FROM pg_indexes WHERE indexname like %s', (existing_prefix + '%', ))
        existing_indexes = set()
        for row in cur:
            index_name = row[0]
            if index_name.startswith(existing_prefix) and not index_name.startswith((new_prefix, backup_prefix)):
                # check for overlapping prefixes: osm_ but not osm_new_ or osm_backup_
                existing_indexes.add(index_name)
        
        cur.execute('SELECT relname FROM pg_class WHERE relname like %s', (existing_prefix + '%_id_seq', ))
        existing_seq = set()
        for row in cur:<commit_msg>put quotes around column name when creating trigam index<commit_after># Copyright 2011 Omniscale (http://omniscale.com)
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#     http://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import time

import psycopg2
import psycopg2.extensions

import logging
log = logging.getLogger(__name__)

from imposm import config
from imposm.mapping import UnionView, GeneralizedTable, Mapping

class PostGISDB(object):
    def __init__(self, db_conf):
        self.db_conf = db_conf
        self.srid = int(db_conf['proj'].split(':')[1])
        self._insert_stmts = {}
        self._connection = None
        self._cur = None

    @property
    def table_prefix(self):
        return self.db_conf.prefix.rstrip('_') + '_'

    def to_tablename(self, name):
        return self.table_prefix + name.lower()

    @property
    def connection(self):
        if not self._connection:
            kw = {}
            if self.db_conf.port:
                kw['port'] = int(self.db_conf.port)
            self._connection = psycopg2.connect(
                database=self.db_conf.db,
                host=self.db_conf.host,
                user=self.db_conf.user,
                password=<PASSWORD>,
                sslmode=self.db_conf.get('sslmode', 'allow'),
                **kw
            )
            self._connection.set_isolation_level(
                psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED)
        return self._connection
    
    def commit(self):
        self.connection.commit()
    
    @property
    def cur(self):
        if self._cur is None:
            self._cur = self.connection.cursor()
        return self._cur
    
    def insert(self, mapping, insert_data, tries=0):
        insert_stmt = self.insert_stmt(mapping)
        try:
            if tries:
                self.reconnect()
            self.cur.executemany(insert_stmt, insert_data)
        except psycopg2.OperationalError, ex:
            if tries >= 8:
                log.warn('%s, giving up', ex)
                raise
            seconds = 2 ** (tries + 1)
            log.warn('%s, retry in %d', ex, seconds)
            time.sleep(seconds)
            self.insert(mapping, insert_data, tries=tries + 1)
        except psycopg2.Error, ex:
            self.connection.rollback()
            for data in insert_data:
                try:
                    self.cur.execute(insert_stmt, data)
                except psycopg2.Error, ex:
                    log.warn('error while importing "%r": %s', data, ex)
                    self.connection.rollback()
                else:
                    self.connection.commit()

        self.connection.commit()
    
    def geom_wrapper(self, geom):
        return psycopg2.Binary(geom.wkb)
    
    def reconnect(self):
        if self._connection:
            try:
                self._connection.close()
            except psycopg2.InterfaceError:
                pass
        self._connection = None
        self._cur = None

    def insert_stmt(self, mapping):
        if mapping.name not in self._insert_stmts:
            self._insert_stmts[mapping.name] = self._insert_stmt(mapping)

        return self._insert_stmts[mapping.name]

    def _insert_stmt(self, mapping):
        extra_arg_names = extra_args = ''
        if mapping.fields:
            extra_arg_names = [n for n, t in mapping.fields]
            extra_args = ', %s' * len(extra_arg_names)
            extra_arg_names = ', ' + ', '.join('"' + name + '"' for name in extra_arg_names)
        return """INSERT INTO "%(tablename)s"
            (osm_id, geometry %(extra_arg_names)s)
            VALUES (%%s, ST_Transform(ST_GeomFromWKB(%%s, 4326), %(srid)s)
                %(extra_args)s)
        """.strip() % dict(tablename=self.table_prefix + mapping.name, srid=self.srid,
            extra_arg_names=extra_arg_names, extra_args=extra_args)


    def create_tables(self, mappings):
        for mapping in mappings:
            self.create_table(mapping)

    def create_table(self, mapping):
        tablename = self.table_prefix + mapping.name
        cur = self.connection.cursor()
        cur.execute('SAVEPOINT pre_drop_tables')
        try:
            cur.execute('DROP TABLE "' + tablename + '" CASCADE')
        except psycopg2.ProgrammingError:
            cur.execute('ROLLBACK TO SAVEPOINT pre_drop_tables')

        extra_fields = ''
        for n, t in mapping.fields:
            extra_fields += ', "%s" %s ' % (n, t.column_type)

        if config.imposm_pg_serial_id:
            serial_column = "id SERIAL PRIMARY KEY,"
        else:
            serial_column = ""
            
        cur.execute("""
            CREATE TABLE "%s" (
                %s
                osm_id INT4
                %s
            );
        """ % (tablename, serial_column, extra_fields))
        cur.execute("""
            SELECT AddGeometryColumn ('', '%(tablename)s', 'geometry',
                                      %(srid)s, '%(pg_geometry_type)s', 2)
        """ % dict(tablename=tablename, srid=self.srid,
                   pg_geometry_type=mapping.geom_type))

        for n, t in mapping.fields:
            if isinstance(t, TrigramIndex):
                cur.execute("""
                    CREATE INDEX "%(tablename)s_trgm_idx_%(column)s" ON "%(tablename)s" USING GIST ("%(column)s" gist_trgm_ops)
                """ % dict(tablename=tablename, column=n))

        cur.execute("""
            CREATE INDEX "%(tablename)s_geom" ON "%(tablename)s" USING GIST (geometry)
        """ % dict(tablename=tablename))
    
    def swap_tables(self, new_prefix, existing_prefix, backup_prefix):
        cur = self.connection.cursor()

        self.remove_tables(backup_prefix)
        
        cur.execute('SELECT tablename FROM pg_tables WHERE tablename like %s', (existing_prefix + '%', ))
        existing_tables = []
        for row in cur:
            table_name = row[0]
            if table_name.startswith(existing_prefix) and not table_name.startswith((new_prefix, backup_prefix)):
                # check for overlapping prefixes: osm_ but not osm_new_ or osm_backup_
                existing_tables.append(table_name)

        cur.execute('SELECT indexname FROM pg_indexes WHERE indexname like %s', (existing_prefix + '%', ))
        existing_indexes = set()
        for row in cur:
            index_name = row[0]
            if index_name.startswith(existing_prefix) and not index_name.startswith((new_prefix, backup_prefix)):
                # check for overlapping prefixes: osm_ but not osm_new_ or osm_backup_
                existing_indexes.add(index_name)
        
        cur.execute('SELECT relname FROM pg_class WHERE relname like %s', (existing_prefix + '%_id_seq', ))
        existing_seq = set()
        for row in cur:
{'fraction_non_alphanumeric': 0.10282313866592858, 'fraction_numerical': 0.004428452809299751, 'mean_word_length': 2.053665751109233, 'pattern_counts': {'":': 2, '<': 4, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, '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': '7017221', 'n_tokens_mistral': 4363, 'n_tokens_neox': 4166, 'n_words': 997}