Example #1
0
    def dump_to_dest(self, src_conn_id: str, t_schema: str, t_name: str):
        """
            Carrega dump da origem para o destino
        """
        # avoid parse to python dictionary (keeps postgres json)
        register_adapter(dict, Json)
        register_json(oid=3802, array_oid=3807, globally=True)

        src_hook = PostgresHook(postgres_conn_id=src_conn_id)
        src_conn = src_hook.get_conn()
        src_cursor = src_conn.cursor()
        src_cursor.execute(f'select count(0) from {t_name};')
        qtd = src_cursor.fetchone()[0]

        dest_cursor = self.conn.cursor()
        dest_cursor.execute(f'TRUNCATE TABLE {t_schema}.{t_name};')
        self.conn.commit()

        if qtd > 0:
            with tempfile.NamedTemporaryFile() as temp_file:
                print('Gerando dump tabela:', t_name, 'linhas:', qtd)
                src_hook.bulk_dump(t_name, temp_file.name)
                print('Carregando dump tabela:', f'{t_schema}.{t_name}',
                      'linhas:', qtd)
                self.hook.bulk_load(f'{t_schema}.{t_name}', temp_file.name)
        else:
            print('Não foi gerado dump tabela:', t_name,
                  'pois possui 0 registros')
Example #2
0
def setup_connection(conn):
    # We want to use unicode everywhere
    register_type(UNICODE, conn)
    register_type(UNICODEARRAY, conn)
    conn.set_client_encoding("UTF8")
    cur = conn.cursor()
    cur.execute("SELECT NULL::numeric")
    oid = cur.description[0][1]
    NUMERIC = new_type((oid,), "NUMERIC", numeric_converter)
    cur.execute("SELECT NULL::numeric[]")
    oid = cur.description[0][1]
    NUMERICL = new_array_type((oid,), "NUMERIC[]", NUMERIC)
    register_type(NUMERIC, conn)
    register_type(NUMERICL, conn)
    register_adapter(dict, Json)
    register_json(conn, loads=Json.loads)
    try:
        from sage.all import Integer, RealNumber
    except ImportError:
        pass
    else:
        register_adapter(Integer, AsIs)
        from .encoding import RealEncoder, LmfdbRealLiteral
        register_adapter(RealNumber, RealEncoder)
        register_adapter(LmfdbRealLiteral, RealEncoder)
Example #3
0
    def _make_conn(self, conn_info):
        # if multiple hosts are provided, select one at random as a kind of
        # simple load balancing.
        host = conn_info.get('host')
        if host and isinstance(host, list):
            host = random.choice(host)
            conn_info = conn_info.copy()
            conn_info['host'] = host

        conn = psycopg2.connect(**conn_info)
        conn.set_session(readonly=self.readonly, autocommit=True)
        register_hstore(conn)
        register_json(conn, loads=ujson.loads)
        return conn
Example #4
0
def setup_connection(conn):
    # We want to use unicode everywhere
    register_type(UNICODE, conn)
    register_type(UNICODEARRAY, conn)
    cur = conn.cursor()
    cur.execute("SELECT NULL::numeric")
    oid = cur.description[0][1]
    NUMERIC = new_type((oid,), "NUMERIC", numeric_converter)
    register_type(NUMERIC, conn)
    register_adapter(Integer, AsIs)
    register_adapter(RealNumber, RealEncoder)
    register_adapter(list, Json)
    register_adapter(tuple, Json)
    register_adapter(dict, Json)
    register_json(conn, loads=Json.loads)
Example #5
0
def dial(uri):
    """
        Parses URI's that loook like this and returns a psycopg connection:
            db://test_user:test_pw@localhost/test_db
    """
    import urlparse
    uri = urlparse.urlparse(uri.strip())
    username = uri.username
    password = uri.password
    database = uri.path[1:]
    hostname = uri.hostname
    connection = psycopg2.connect(
        database=database,
        user=username,
        password=password,
        host=hostname
    )
    register_json(connection, loads=loads)
    return connection
    def select(self, query, parameter=[]):
        """
        Connects to database and extracts
        raw tweets and any other columns we
        need
        Parameters:
        ----------------
        arg1: string: SQL query
        Returns: pandas dataframe
        ----------------
        """

        cursor = self.conn.cursor()
        try:
            register_json(oid=3802, array_oid=3807)
            cursor.execute(query, parameter)
        except:
            print('cant execute query')
            raise
        rows = cursor.fetchall()
        return rows
Example #7
0
    def exec_insert(self, data: list, table: str, truncate: bool = False):
        """
            Insere lista em tabela
        """

        register_adapter(
            dict,
            Json)  # avoid parse to python dictionary (keeps postgres json)
        register_json(
            oid=3802, array_oid=3807, globally=True
        )  # avoid parse to python dictionary (keeps postgres json)

        print(f'Inserting in: {table}')

        self.conn.autocommit = False
        dest_cursor = self.conn.cursor()

        if truncate:
            dest_cursor.execute(f'TRUNCATE TABLE {table};')

        inserted = 0
        while True:
            lines = data[0:1000]
            del data[0:1000]
            inserted += len(lines)
            if not lines:
                break
            try:
                execute_values(
                    dest_cursor,
                    f'INSERT INTO {table} VALUES %s;'.format(
                        table=sql.Identifier(table)),
                    lines,
                )
            except Exception as error:
                print(f'Line - {lines}')
                raise Exception(error) from error

            print(f'Inserted: {inserted}')
        self.conn.commit()
Example #8
0
 def _make_conn(self, conn_info):
     conn = psycopg2.connect(**conn_info)
     conn.set_session(readonly=True, autocommit=True)
     register_hstore(conn)
     register_json(conn)
     return conn
Example #9
0
from __future__ import unicode_literals

import json
import decimal

from django.core.serializers.json import DjangoJSONEncoder
from django.db import models
from django.db.models.lookups import BuiltinLookup, Transform
from django.utils import six
from django.conf import settings
from psycopg2.extras import register_json

# We can register jsonb to be loaded as json!
# http://schinckel.net/2014/05/24/python%2C-postgres-and-jsonb/
register_json(oid=3802, array_oid=3807)
# However, it may be that we want to use specific decoding on
# the json object... which if we wanted to do it on a per-field
# basis, we'd need to not have run that line.


class JSONField(models.Field):
    description = 'JSON Field'

    def __init__(self, *args, **kwargs):
        self.decode_kwargs = kwargs.pop('decode_kwargs', {
            'parse_float': decimal.Decimal
        })
        self.encode_kwargs = kwargs.pop('encode_kwargs', {
            'cls': DjangoJSONEncoder,
        })
        super(JSONField, self).__init__(*args, **kwargs)
Example #10
0
from __future__ import unicode_literals

import json
import decimal

from django.core.serializers.json import DjangoJSONEncoder
from django.db import models
from django.db.models.lookups import BuiltinLookup, Transform
from django.utils import six
from django.conf import settings
from psycopg2.extras import register_json

# We can register jsonb to be loaded as json!
# http://schinckel.net/2014/05/24/python%2C-postgres-and-jsonb/
register_json(oid=3802, array_oid=3807)
# However, it may be that we want to use specific decoding on
# the json object... which if we wanted to do it on a per-field
# basis, we'd need to not have run that line.


class JSONField(models.Field):
    description = 'JSON Field'

    def __init__(self, *args, **kwargs):
        self.decode_kwargs = kwargs.pop('decode_kwargs',
                                        {'parse_float': decimal.Decimal})
        self.encode_kwargs = kwargs.pop('encode_kwargs', {
            'cls': DjangoJSONEncoder,
        })
        super(JSONField, self).__init__(*args, **kwargs)
Example #11
0
 def _make_conn(self, conn_info):
     conn = psycopg2.connect(**conn_info)
     conn.set_session(readonly=self.readonly, autocommit=True)
     register_hstore(conn)
     register_json(conn, loads=ujson.loads)
     return conn
Example #12
0
    def _execute(self,
                 native,
                 command,
                 data=None,
                 returning=True,
                 mapper=dict):
        """
        Executes the inputted command into the current \
        connection cursor.
        
        :param      command    | <str>
                    data       | <dict> || None
                    autoCommit | <bool> | commit database changes immediately
                    autoClose  | <bool> | closes connections immediately
        
        :return     [{<str> key: <variant>, ..}, ..], <int> count
        """
        if data is None:
            data = {}

        cursor = native.cursor(cursor_factory=DictCursor)

        # register the hstore option
        try:
            register_hstore(cursor, unicode=True)
        except pg.ProgrammingError:
            log.warning('HSTORE is not supported in this version of Postgres!')

        # register the json option
        try:
            register_json(cursor)
        except pg.ProgrammingError:
            log.warning('JSON is not supported in this version of Postgres!')

        start = datetime.datetime.now()

        log.debug('***********************')
        log.debug(command % data)
        log.debug('***********************')

        try:
            cursor.execute(command, data)
            rowcount = cursor.rowcount

        # look for a cancelled query
        except QueryCanceledError as cancelled:
            try:
                native.rollback()
            except StandardError as err:
                log.error('Rollback error: {0}'.format(err))
            log.critical(command)
            if data:
                log.critical(str(data))

            # raise more useful errors
            if 'statement timeout' in str(cancelled):
                raise orb.errors.QueryTimeout(
                    command, (datetime.datetime.now() - start).total_seconds())
            else:
                raise orb.errors.Interruption()

        # look for a disconnection error
        except pg.InterfaceError:
            raise orb.errors.ConnectionLost()

        # look for integrity errors
        except (pg.IntegrityError, pg.OperationalError) as err:
            try:
                native.rollback()
            except StandardError:
                pass

            # look for a duplicate error
            duplicate_error = re.search('Key (.*) already exists.', nstr(err))
            if duplicate_error:
                key = duplicate_error.group(1)
                result = re.match(
                    '^\(lower\((?P<column>[^\)]+)::text\)\)=\((?P<value>[^\)]+)\)$',
                    key)
                if not result:
                    result = re.match('^(?P<column>\w+)=(?P<value>\w+)', key)

                if result:
                    msg = '{value} is already being used.'.format(
                        **result.groupdict())
                    raise orb.errors.DuplicateEntryFound(msg)
                else:
                    raise orb.errors.DuplicateEntryFound(
                        duplicate_error.group())

            # look for a reference error
            reference_error = re.search(
                'Key .* is still referenced from table ".*"', nstr(err))
            if reference_error:
                msg = 'Cannot remove this record, it is still being referenced.'
                raise orb.errors.CannotDelete(msg)

            # unknown error
            log.debug(traceback.print_exc())
            raise orb.errors.QueryFailed(command, data, nstr(err))

        # connection has closed underneath the hood
        except (pg.Error, pg.ProgrammingError) as err:
            try:
                native.rollback()
            except StandardError:
                pass

            log.error(traceback.print_exc())
            raise orb.errors.QueryFailed(command, data, nstr(err))

        try:
            results = [mapper(record) for record in cursor.fetchall()]
        except pg.ProgrammingError:
            results = []

        return results, rowcount
Example #13
0
    def _execute(self,
                 native,
                 command,
                 data=None,
                 returning=True,
                 mapper=dict):
        """
        Executes the inputted command into the current \
        connection cursor.
        
        :param      command    | <str>
                    data       | <dict> || None
                    autoCommit | <bool> | commit database changes immediately
                    autoClose  | <bool> | closes connections immediately
        
        :return     [{<str> key: <variant>, ..}, ..], <int> count
        """
        if data is None:
            data = {}

        cursor = native.cursor(cursor_factory=DictCursor)

        # register the hstore option
        try:
            register_hstore(cursor, unicode=True)
        except pg.ProgrammingError:
            log.warning('HSTORE is not supported in this version of Postgres!')

        # register the json option
        try:
            register_json(cursor)
        except pg.ProgrammingError:
            log.warning('JSON is not supported in this version of Postgres!')

        start = datetime.datetime.now()

        log.debug('***********************')
        log.debug(command % data)
        log.debug('***********************')

        try:
            cursor.execute(command, data)
            rowcount = cursor.rowcount

        # look for a cancelled query
        except pg_ext.QueryCanceledError as cancelled:
            try:
                native.rollback()
            except StandardError as err:
                log.error('Rollback error: {0}'.format(err))
            log.critical(command)
            if data:
                log.critical(str(data))

            # raise more useful errors
            if 'statement timeout' in str(cancelled):
                raise orb.errors.QueryTimeout(command, (datetime.datetime.now() - start).total_seconds())
            else:
                raise orb.errors.Interruption()

        # look for a disconnection error
        except pg.InterfaceError:
            raise orb.errors.ConnectionLost()

        # look for integrity errors
        except (pg.IntegrityError, pg.OperationalError) as err:
            try:
                native.rollback()
            except StandardError:
                pass

            # look for a duplicate error
            duplicate_error = re.search('Key (.*) already exists.', nstr(err))
            if duplicate_error:
                key = duplicate_error.group(1)
                result = re.match('^\(lower\((?P<column>[^\)]+)::text\)\)=\((?P<value>[^\)]+)\)$', key)
                if not result:
                    result = re.match('^(?P<column>\w+)=(?P<value>\w+)', key)

                if result:
                    msg = '{value} is already being used.'.format(**result.groupdict())
                    raise orb.errors.DuplicateEntryFound(msg)
                else:
                    raise orb.errors.DuplicateEntryFound(duplicate_error.group())

            # look for a reference error
            reference_error = re.search('Key .* is still referenced from table ".*"', nstr(err))
            if reference_error:
                msg = 'Cannot remove this record, it is still being referenced.'
                raise orb.errors.CannotDelete(msg)

            # unknown error
            log.debug(traceback.print_exc())
            raise orb.errors.QueryFailed(command, data, nstr(err))

        # connection has closed underneath the hood
        except (pg.Error, pg.ProgrammingError) as err:
            try:
                native.rollback()
            except StandardError:
                pass

            log.error(traceback.print_exc())
            raise orb.errors.QueryFailed(command, data, nstr(err))

        try:
            results = [mapper(record) for record in cursor.fetchall()]
        except pg.ProgrammingError:
            results = []

        return results, rowcount
Example #14
0
def register_flexmap(conn, types):
   """
   Create flexmap type from
   
    1) 'schema.typename' str or
    2) (typename, oid, attrs, array_oid, schema) tuple where attrs = [(attname, atttypid)]
   
   For 1), type info will be looked up in database catalog. For 2), type info is
   taken as specified.
   """
   
   register_hstore(conn)
   register_json(conn)

   casters = {}

   class DictComposite(CompositeCaster):
      """
      A type caster returning composite types as Python dicts
      enriched with a type field containing 'schema.typename'.
      """
      def make(self, attrs):
         o = {}
         for i in xrange(len(self.attnames)):
            if attrs[i] is not None:
               o[self.attnames[i]] = attrs[i]
         o['type'] = self.schema + '.' + self.name
         return o

   ## create type casters for whole list
   ##
   for t in types:
      if type(t) == str:
         caster = DictComposite._from_db(t, conn)
      elif type(t) in [tuple, list]:
         caster = CompositeCaster(*t)
      else:
         raise Exception("invalid type %s in flexmap type list" % type(t))
         
      ## register item and array casters
      ##
      register_type(caster.typecaster, conn)     
      if caster.array_typecaster is not None:
         register_type(caster.array_typecaster, conn)
         
      ## remember caster under 'schema.typename'
      ##
      casters['%s.%s' % (caster.schema, caster.name)] = caster


   class DictAdapter(object):
      """
      A dictionary adapter converting Python dicts to PostgreSQL
      JSON, Hstore or Composite Types depending on the dict field 'type'.      
      """
      def __init__(self, adapted):
         ## remember value to be adaptated - a Python dict
         self.adapted = adapted
         
         ## remember type field of dict-value to be adapted
         if adapted.has_key('type'):
            self._type = adapted['type']
            #del(adapted['type'])
         else:
            self._type = None
         
         ## create adapter to hstore if requested
         if self._type == 'hstore':
            self._hstoreAdapter = HstoreAdapter(adapted)
         
      def prepare(self, conn):
         self._conn = conn
         if self._type == 'hstore':
            self._hstoreAdapter.prepare(conn)
         
      def getquoted(self):
         if self._type is not None:
            if self._type == 'json':
               return adapt(Json(self.adapted)).getquoted() + '::json'
            elif self._type == 'hstore':
               return self._hstoreAdapter.getquoted()
            elif casters.has_key(self._type):
               c = casters[self._type]
               v = []
               for n in c.attnames:
                  v.append(self.adapted.get(n, None))
               a = adapt(tuple(v))
               a.prepare(self._conn)
               return a.getquoted() + '::' + self._type
            else:
               raise psycopg2.ProgrammingError("unknown type %s in dictionary type hint" % self._type)
         else:
            raise psycopg2.ProgrammingError("dictionary is missing type hint")
   
   register_adapter(dict, DictAdapter)
Example #15
0
def register_flexmap(conn, types):
    """
   Create flexmap type from
   
    1) 'schema.typename' str or
    2) (typename, oid, attrs, array_oid, schema) tuple where attrs = [(attname, atttypid)]
   
   For 1), type info will be looked up in database catalog. For 2), type info is
   taken as specified.
   """

    register_hstore(conn)
    register_json(conn)

    casters = {}

    class DictComposite(CompositeCaster):
        """
      A type caster returning composite types as Python dicts
      enriched with a type field containing 'schema.typename'.
      """
        def make(self, attrs):
            o = {}
            for i in xrange(len(self.attnames)):
                if attrs[i] is not None:
                    o[self.attnames[i]] = attrs[i]
            o['type'] = self.schema + '.' + self.name
            return o

    ## create type casters for whole list
    ##
    for t in types:
        if type(t) == str:
            caster = DictComposite._from_db(t, conn)
        elif type(t) in [tuple, list]:
            caster = CompositeCaster(*t)
        else:
            raise Exception("invalid type %s in flexmap type list" % type(t))

        ## register item and array casters
        ##
        register_type(caster.typecaster, conn)
        if caster.array_typecaster is not None:
            register_type(caster.array_typecaster, conn)

        ## remember caster under 'schema.typename'
        ##
        casters['%s.%s' % (caster.schema, caster.name)] = caster

    class DictAdapter(object):
        """
      A dictionary adapter converting Python dicts to PostgreSQL
      JSON, Hstore or Composite Types depending on the dict field 'type'.      
      """
        def __init__(self, adapted):
            ## remember value to be adaptated - a Python dict
            self.adapted = adapted

            ## remember type field of dict-value to be adapted
            if adapted.has_key('type'):
                self._type = adapted['type']
                #del(adapted['type'])
            else:
                self._type = None

            ## create adapter to hstore if requested
            if self._type == 'hstore':
                self._hstoreAdapter = HstoreAdapter(adapted)

        def prepare(self, conn):
            self._conn = conn
            if self._type == 'hstore':
                self._hstoreAdapter.prepare(conn)

        def getquoted(self):
            if self._type is not None:
                if self._type == 'json':
                    return adapt(Json(self.adapted)).getquoted() + '::json'
                elif self._type == 'hstore':
                    return self._hstoreAdapter.getquoted()
                elif casters.has_key(self._type):
                    c = casters[self._type]
                    v = []
                    for n in c.attnames:
                        v.append(self.adapted.get(n, None))
                    a = adapt(tuple(v))
                    a.prepare(self._conn)
                    return a.getquoted() + '::' + self._type
                else:
                    raise psycopg2.ProgrammingError(
                        "unknown type %s in dictionary type hint" % self._type)
            else:
                raise psycopg2.ProgrammingError(
                    "dictionary is missing type hint")

    register_adapter(dict, DictAdapter)