Python psycopg2.extras.DictCursor() Examples

def get_replication_status(db):
    result = {"result": False, "status": None}
    db_conn = connect('Destination', db_name=db)
    result["result"] = False
    try:
        cur = db_conn.cursor(cursor_factory=extras.DictCursor)
        cur.execute("SELECT status FROM pglogical.show_subscription_status(subscription_name := 'subscription');")
        r = cur.fetchone()
        if r:
            result["result"] = True
            result["status"] = r['status']

    except (psycopg2.InternalError, psycopg2.OperationalError, psycopg2.ProgrammingError) as e:
        print(e)
        result["result"] = False

    return result 

def connect(self):
        host, _, port = self.config.postgresql_host.partition(':')

        ssl_kwargs = {}
        if self.config.postgresql_ssl_enabled:
            ssl_kwargs['sslmode'] = 'verify-full'
            ssl_kwargs['sslrootcert'] = self.config.postgresql_ca_cert_path

            # It only makes sense to check this if SSL is on
            if self.config.postgresql_ssl_client_verification:
                ssl_kwargs['sslcert'] = self.config.postgresql_ssl_cert_path
                ssl_kwargs['sslkey'] = self.config.postgresql_ssl_key_path

        return psycopg2.connect(
            dbname=self.config.postgresql_db_name,
            host=host,
            port=port or 5432,
            user=self.config.postgresql_username,
            password=self.config.postgresql_password,
            cursor_factory=DictCursor,
            **ssl_kwargs
        ) 

def _fetchMetrics(self, conn, query):
        cursor = False

        try:
            # Open a cursor
            cursor = conn.cursor(cursor_factory=DictCursor)

            # Fetch statistics
            cursor.execute(query)

            return cursor.fetchall()
        except Exception as error:
            logging.getLogger().error("Unable run query {query} on {dsn}".format(query=query, dsn=self.config.getDsnWithMaskedPassword()), extra={"exception": str(error)})

            return False
        finally:
            if cursor:
                cursor.close() 

def archive_timeout(self):
        """
        Retrieve the archive_timeout setting in PostgreSQL

        :return: The archive timeout (in seconds)
        """
        try:
            cur = self._cursor(cursor_factory=DictCursor)
            # We can't use the `get_setting` method here, because it
            # uses `SHOW`, returning an human readable value such as "5min",
            # while we prefer a raw value such as 300.
            cur.execute("SELECT setting "
                        "FROM pg_settings "
                        "WHERE name='archive_timeout'")
            result = cur.fetchone()
            archive_timeout = int(result[0])

            return archive_timeout
        except ValueError as e:
            _logger.error("Error retrieving archive_timeout: %s",
                          force_str(e).strip())
            return None 

def set_field(self, identifier, field, value, group=False):
        shaid = sha1(str(identifier)).hexdigest()
        cur = self.conn.cursor(cursor_factory=DictCursor)
        if group:
            sql = "SELECT count(shaid) as count FROM groups WHERE shaid = %s"
            cur.execute(sql, (shaid,))
        else:
            sql = "SELECT count(shaid) as count FROM users WHERE shaid = %s"
            cur.execute(sql, (shaid,))
        num = cur.fetchone()
        if num['count'] == 0:
            if group:
                sql = "INSERT INTO groups (shaid,{0}) VALUES (%s,%s)"
                cur.execute(sql.format(field), (shaid, value))
            else:
                sql = "INSERT INTO users (shaid,{0}) VALUES (%s,%s)"
                cur.execute(sql.format(field), (shaid, value))
        else:
            if group:
                sql = "UPDATE groups SET {0} = %s WHERE shaid = %s"
                cur.execute(sql.format(field), (value, shaid))
            else:
                sql = "UPDATE users SET {0} = %s WHERE shaid = %s"
                cur.execute(sql.format(field), (value, shaid))
        self.conn.commit()
        cur.close()
        return cur.rowcount != 0 

def get_user(self, identifier, group=False):
        """
        Retrives the information from a user

        :param identifier: 
        :param group:
        :return:
        """
        shaid = sha1(str(identifier)).hexdigest()
        cur = self.conn.cursor(cursor_factory=DictCursor)
        if group:
            sql = 'SELECT * FROM groups WHERE shaid = %s LIMIT 1'
            cur.execute(sql, (shaid,))
        else:
            sql = 'SELECT * FROM users WHERE shaid = %s LIMIT 1'
            cur.execute(sql, (shaid,))
        d = cur.fetchone()
        cur.close()
        data = dict()
        if d is None:
            return self.get_defaultconfig()
        else:
            data.update(d)
            if 'lang' not in d:
                data['lang'] = 'en'
                data['lang_set'] = False
            else:
                data['lang_set'] = True
            if data['lang'] is None:
                data['lang'] = 'en'
                data['lang_set'] = False
            return data 

def _generate_constraints(conn, table: str):
    """
    Using the existing table as a template, generate ALTER TABLE ADD CONSTRAINT
    statements pointing to the new table.

    :return: A list of SQL statements.
    """
    # List all active constraints across the database.
    get_all_constraints = '''
        SELECT conrelid::regclass AS table, conname, pg_get_constraintdef(c.oid)
        FROM pg_constraint c
        JOIN pg_namespace n ON n.oid = c.connamespace
        AND n.nspname = 'public'
        ORDER BY conrelid::regclass::text, contype DESC;
    '''
    with conn.cursor(cursor_factory=DictCursor) as cur:
        cur.execute(get_all_constraints)
        all_constraints = cur.fetchall()
    # Find all constraints that either exist inside of the table or
    # reference it from another table. Ignore PRIMARY KEY statements.
    remap_constraints = []
    drop_orphans = []
    for constraint in all_constraints:
        statement = constraint['pg_get_constraintdef']
        con_table = constraint['table']
        is_fk = _is_foreign_key(statement, table)
        if (con_table == table or is_fk) and 'PRIMARY KEY' not in statement:
            alter_stmnts = _remap_constraint(
                constraint['conname'], con_table, statement, table
            )
            remap_constraints.extend(alter_stmnts)
            if is_fk:
                del_orphans = _generate_delete_orphans(statement, con_table)
                drop_orphans.append(del_orphans)

    constraint_statements = []
    constraint_statements.extend(drop_orphans)
    constraint_statements.extend(remap_constraints)
    return constraint_statements 

def pgespresso_stop_backup(self, backup_label):
        """
        Execute a pgespresso_stop_backup

        This method returns a dictionary containing the following data:

         * end_wal
         * timestamp

        :param str backup_label: backup label as returned
            by pgespress_start_backup
        :rtype: psycopg2.extras.DictRow
        """
        try:
            conn = self.connect()
            # Issue a rollback to release any unneeded lock
            conn.rollback()
            cur = conn.cursor(cursor_factory=DictCursor)
            cur.execute("SELECT pgespresso_stop_backup(%s) AS end_wal, "
                        "now() AS timestamp",
                        (backup_label,))
            return cur.fetchone()
        except (PostgresConnectionError, psycopg2.Error) as e:
            msg = "Error issuing pgespresso_stop_backup() command: %s" % (
                force_str(e).strip())
            _logger.debug(msg)
            raise PostgresException(
                '%s\n'
                'HINT: You might have to manually execute '
                'pgespresso_abort_backup() on your PostgreSQL '
                'server' % msg) 

def pgespresso_start_backup(self, label):
        """
        Execute a pgespresso_start_backup

        This method returns a dictionary containing the following data:

         * backup_label
         * timestamp

        :param str label: descriptive string to identify the backup
        :rtype: psycopg2.extras.DictRow
        """
        try:
            conn = self.connect()

            # Rollback to release the transaction,
            # as the pgespresso_start_backup invocation can last
            # up to PostgreSQL's checkpoint_timeout
            conn.rollback()

            # Start the concurrent backup using pgespresso
            cur = conn.cursor(cursor_factory=DictCursor)
            cur.execute(
                'SELECT pgespresso_start_backup(%s,%s) AS backup_label, '
                'now() AS timestamp',
                (label, self.immediate_checkpoint))

            start_row = cur.fetchone()

            # Rollback to release the transaction, as the connection
            # is to be retained until the end of backup
            conn.rollback()

            return start_row
        except (PostgresConnectionError, psycopg2.Error) as e:
            msg = "pgespresso_start_backup(): %s" % force_str(e).strip()
            _logger.debug(msg)
            raise PostgresException(msg) 

def stop_concurrent_backup(self):
        """
        Calls pg_stop_backup on the PostgreSQL server using the
        API introduced with version 9.6

        This method returns a dictionary containing the following data:

         * location
         * timeline
         * backup_label
         * timestamp

        :rtype: psycopg2.extras.DictRow
        """
        try:
            conn = self.connect()

            # Rollback to release the transaction, as the pg_stop_backup
            # invocation could will wait until the current WAL file is shipped
            conn.rollback()

            # Stop the backup  using the api introduced with version 9.6
            cur = conn.cursor(cursor_factory=DictCursor)
            cur.execute(
                'SELECT end_row.lsn AS location, '
                '(SELECT CASE WHEN pg_is_in_recovery() '
                'THEN min_recovery_end_timeline ELSE timeline_id END '
                'FROM pg_control_checkpoint(), pg_control_recovery()'
                ') AS timeline, '
                'end_row.labelfile AS backup_label, '
                'now() AS timestamp FROM pg_stop_backup(FALSE) AS end_row')

            return cur.fetchone()
        except (PostgresConnectionError, psycopg2.Error) as e:
            msg = ("Error issuing pg_stop_backup command: %s" %
                   force_str(e).strip())
            _logger.debug(msg)
            raise PostgresException(msg) 

def stop_exclusive_backup(self):
        """
        Calls pg_stop_backup() on the PostgreSQL server

        This method returns a dictionary containing the following data:

         * location
         * file_name
         * file_offset
         * timestamp

        :rtype: psycopg2.extras.DictRow
        """
        try:
            conn = self.connect()

            # Rollback to release the transaction, as the pg_stop_backup
            # invocation could will wait until the current WAL file is shipped
            conn.rollback()

            # Stop the backup
            cur = conn.cursor(cursor_factory=DictCursor)
            cur.execute(
                "SELECT location, "
                "({pg_walfile_name_offset}(location)).*, "
                "now() AS timestamp "
                "FROM pg_stop_backup() AS location"
                .format(**self.name_map)
            )

            return cur.fetchone()
        except (PostgresConnectionError, psycopg2.Error) as e:
            msg = ("Error issuing pg_stop_backup command: %s" %
                   force_str(e).strip())
            _logger.debug(msg)
            raise PostgresException(
                'Cannot terminate exclusive backup. '
                'You might have to manually execute pg_stop_backup '
                'on your PostgreSQL server') 

def start_concurrent_backup(self, label):
        """
        Calls pg_start_backup on the PostgreSQL server using the
        API introduced with version 9.6

        This method returns a dictionary containing the following data:

         * location
         * timeline
         * timestamp

        :param str label: descriptive string to identify the backup
        :rtype: psycopg2.extras.DictRow
        """
        try:
            conn = self.connect()

            # Rollback to release the transaction, as the pg_start_backup
            # invocation can last up to PostgreSQL's checkpoint_timeout
            conn.rollback()

            # Start the backup using the api introduced in postgres 9.6
            cur = conn.cursor(cursor_factory=DictCursor)
            cur.execute(
                "SELECT location, "
                "(SELECT timeline_id "
                "FROM pg_control_checkpoint()) AS timeline, "
                "now() AS timestamp "
                "FROM pg_start_backup(%s, %s, FALSE) AS location",
                (label, self.immediate_checkpoint))
            start_row = cur.fetchone()

            # Rollback to release the transaction, as the connection
            # is to be retained until the end of backup
            conn.rollback()

            return start_row
        except (PostgresConnectionError, psycopg2.Error) as e:
            msg = "pg_start_backup command: %s" % (force_str(e).strip(),)
            _logger.debug(msg)
            raise PostgresException(msg) 

def get_archiver_stats(self):
        """
        This method gathers statistics from pg_stat_archiver.
        Only for Postgres 9.4+ or greater. If not available, returns None.

        :return dict|None: a dictionary containing Postgres statistics from
            pg_stat_archiver or None
        """
        try:
            # pg_stat_archiver is only available from Postgres 9.4+
            if self.server_version < 90400:
                return None
            cur = self._cursor(cursor_factory=DictCursor)
            # Select from pg_stat_archiver statistics view,
            # retrieving statistics about WAL archiver process activity,
            # also evaluating if the server is archiving without issues
            # and the archived WALs per second rate.
            #
            # We are using current_settings to check for archive_mode=always.
            # current_setting does normalise its output so we can just
            # check for 'always' settings using a direct string
            # comparison
            cur.execute(
                "SELECT *, "
                "current_setting('archive_mode') IN ('on', 'always') "
                "AND (last_failed_wal IS NULL "
                "OR last_failed_wal LIKE '%.history' "
                "AND substring(last_failed_wal from 1 for 8) "
                "<= substring(last_archived_wal from 1 for 8) "
                "OR last_failed_time <= last_archived_time) "
                "AS is_archiving, "
                "CAST (archived_count AS NUMERIC) "
                "/ EXTRACT (EPOCH FROM age(now(), stats_reset)) "
                "AS current_archived_wals_per_second "
                "FROM pg_stat_archiver")
            return cur.fetchone()
        except (PostgresConnectionError, psycopg2.Error) as e:
            _logger.debug("Error retrieving pg_stat_archive data: %s",
                          force_str(e).strip())
            return None 

def checkpoint_timeout(self):
        """
        Retrieve the checkpoint_timeout setting in PostgreSQL

        :return: The checkpoint timeout (in seconds)
        """
        try:
            cur = self._cursor(cursor_factory=DictCursor)
            # We can't use the `get_setting` method here, because it
            # uses `SHOW`, returning an human readable value such as "5min",
            # while we prefer a raw value such as 300.
            cur.execute("SELECT setting "
                        "FROM pg_settings "
                        "WHERE name='checkpoint_timeout'")
            result = cur.fetchone()
            checkpoint_timeout = int(result[0])

            return checkpoint_timeout
        except ValueError as e:
            _logger.error("Error retrieving checkpoint_timeout: %s",
                          force_str(e).strip())
            return None 

def current_xlog_info(self):
        """
        Get detailed information about the current WAL position in PostgreSQL.

        This method returns a dictionary containing the following data:

         * location
         * file_name
         * file_offset
         * timestamp

        When executed on a standby server file_name and file_offset are always
        None

        :rtype: psycopg2.extras.DictRow
        """
        try:
            cur = self._cursor(cursor_factory=DictCursor)
            if not self.is_in_recovery:
                cur.execute(
                    "SELECT location, "
                    "({pg_walfile_name_offset}(location)).*, "
                    "CURRENT_TIMESTAMP AS timestamp "
                    "FROM {pg_current_wal_lsn}() AS location"
                    .format(**self.name_map))
                return cur.fetchone()
            else:
                cur.execute(
                    "SELECT location, "
                    "NULL AS file_name, "
                    "NULL AS file_offset, "
                    "CURRENT_TIMESTAMP AS timestamp "
                    "FROM {pg_last_wal_replay_lsn}() AS location"
                    .format(**self.name_map))
                return cur.fetchone()
        except (PostgresConnectionError, psycopg2.Error) as e:
            _logger.debug("Error retrieving current xlog "
                          "detailed information: %s",
                          force_str(e).strip())
        return None 

def get_row(
    table_name: str,
    key_name: str,
    key_value,
    column_names: Optional[List[str]] = None,
    filter_formula: Optional[Mapping] = None,
    filter_pairs: Optional[Mapping] = None,
):
    """Get a single row in the DB with the key name/value pair.

    :param table_name: Name of the table
    :param key_name: Key name to uniquely identify the row
    :param key_value: Key value to uniquely identify the row
    :param column_names: Columns to access (all of them if empty)
    :param filter_formula: Optional filter formula
    :param filter_pairs: Optional dictionary to restrict the clause
    :return: Dictionary with the row
    """
    key_pair = {key_name: key_value}
    if filter_pairs:
        filter_pairs = dict(key_pair, **filter_pairs)
    else:
        filter_pairs = key_pair

    query, fields = get_select_query(
        table_name,
        column_names=column_names,
        filter_formula=filter_formula,
        filter_pairs=filter_pairs,
    )

    # Execute the query
    cursor = connection.connection.cursor(cursor_factory=DictCursor)
    cursor.execute(query, fields)

    if cursor.rowcount != 1:
        raise Exception('Query returned more than one row.')

    return cursor.fetchone() 

def get_rows(
    table_name: str,
    column_names: Optional[List[str]] = None,
    filter_formula: Optional[Mapping] = None,
    filter_pairs: Optional[Mapping] = None,
):
    """Get columns in a row selected by filter and/or pairs.

    Execute a select query in the database with an optional filter and
    pairs and return a subset of columns (or all of them if empty)

    :param table_name: Primary key of the workflow storing the data
    :param column_names: optional list of columns to select
    :param filter_formula: Optional JSON formula to use in the WHERE clause
    :param filter_pairs: Pairs key: value to filter in the WHERE clause
    :return: cursor resulting from the query
    """
    query, fields = get_select_query(
        table_name,
        column_names=column_names,
        filter_formula=filter_formula,
        filter_pairs=filter_pairs,
    )

    # Execute the query
    cursor = connection.connection.cursor(cursor_factory=DictCursor)
    cursor.execute(query, fields)
    return cursor 

def __init__(self, host, database, user, password):
        self.connection = pg.connect(host=host, database=database,
            user=user, password=password, cursor_factory=DictCursor)

        techs = self.query("SELECT * FROM lab_techs ORDER BY name")
        labs = self.query("SELECT id FROM labs ORDER BY id")
        plots = self.query("SELECT DISTINCT plot FROM plots ORDER BY plot")
        self.fields['Technician']['values'] = [x['name'] for x in techs]
        self.fields['Lab']['values'] = [x['id'] for x in labs]
        self.fields['Plot']['values'] = [str(x['plot']) for x in plots] 

def __init__(self, host, database, user, password):
        self.connection = pg.connect(host=host, database=database,
            user=user, password=password, cursor_factory=DictCursor)

        techs = self.query("SELECT * FROM lab_techs ORDER BY name")
        labs = self.query("SELECT id FROM labs ORDER BY id")
        plots = self.query("SELECT DISTINCT plot FROM plots ORDER BY plot")
        self.fields['Technician']['values'] = [x['name'] for x in techs]
        self.fields['Lab']['values'] = [x['id'] for x in labs]
        self.fields['Plot']['values'] = [str(x['plot']) for x in plots] 

def __init__(self, host, database, user, password):
        self.connection = pg.connect(host=host, database=database,
            user=user, password=password, cursor_factory=DictCursor)

        techs = self.query("SELECT * FROM lab_techs ORDER BY name")
        labs = self.query("SELECT id FROM labs ORDER BY id")
        plots = self.query("SELECT DISTINCT plot FROM plots ORDER BY plot")
        self.fields['Technician']['values'] = [x['name'] for x in techs]
        self.fields['Lab']['values'] = [x['id'] for x in labs]
        self.fields['Plot']['values'] = [str(x['plot']) for x in plots] 

def __init__(self, host, database, user, password):
        self.connection = pg.connect(host=host, database=database,
            user=user, password=password, cursor_factory=DictCursor)

        techs = self.query("SELECT * FROM lab_techs ORDER BY name")
        labs = self.query("SELECT id FROM labs ORDER BY id")
        plots = self.query("SELECT DISTINCT plot FROM plots ORDER BY plot")
        self.fields['Technician']['values'] = [x['name'] for x in techs]
        self.fields['Lab']['values'] = [x['id'] for x in labs]
        self.fields['Plot']['values'] = [str(x['plot']) for x in plots] 

def __init__(self, host, database, user, password):
        self.connection = pg.connect(host=host, database=database,
            user=user, password=password, cursor_factory=DictCursor)

        techs = self.query("SELECT * FROM lab_techs ORDER BY name")
        labs = self.query("SELECT id FROM labs ORDER BY id")
        plots = self.query("SELECT DISTINCT plot FROM plots ORDER BY plot")
        self.fields['Technician']['values'] = [x['name'] for x in techs]
        self.fields['Lab']['values'] = [x['id'] for x in labs]
        self.fields['Plot']['values'] = [str(x['plot']) for x in plots] 

def xlog_segment_size(self):
        """
        Retrieve the size of one WAL file.

        In PostgreSQL 11, users will be able to change the WAL size
        at runtime. Up to PostgreSQL 10, included, the WAL size can be changed
        at compile time

        :return: The wal size (In bytes)
        """

        # Prior to PostgreSQL 8.4, the wal segment size was not configurable,
        # even in compilation
        if self.server_version < 80400:
            return DEFAULT_XLOG_SEG_SIZE

        try:
            cur = self._cursor(cursor_factory=DictCursor)
            # We can't use the `get_setting` method here, because it
            # use `SHOW`, returning an human readable value such as "16MB",
            # while we prefer a raw value such as 16777216.
            cur.execute("SELECT setting "
                        "FROM pg_settings "
                        "WHERE name='wal_segment_size'")
            result = cur.fetchone()
            wal_segment_size = int(result[0])

            # Prior to PostgreSQL 11, the wal segment size is returned in
            # blocks
            if self.server_version < 110000:
                cur.execute("SELECT setting "
                            "FROM pg_settings "
                            "WHERE name='wal_block_size'")
                result = cur.fetchone()
                wal_block_size = int(result[0])

                wal_segment_size *= wal_block_size

            return wal_segment_size
        except ValueError as e:
            _logger.error("Error retrieving current xlog "
                          "segment size: %s",
                          force_str(e).strip())
            return None 

def _clean_data_worker(rows, temp_table, sources_config):
    log.info('Starting data cleaning worker')
    global_field_to_func = sources_config['*']['fields']
    worker_conn = database_connect()
    log.info('Data cleaning worker connected to database')
    write_cur = worker_conn.cursor(cursor_factory=DictCursor)
    log.info('Cleaning {} rows'.format(len(rows)))
    tls_cache = {}
    start_time = time.time()
    for row in rows:
        # Map fields that need updating to their cleaning functions
        source = row['source']
        _id = row['id']
        if source in sources_config:
            source_field_to_func = sources_config[source]['fields']
            # Merge source-local and global function field mappings
            fields_to_update = \
                {**global_field_to_func, **source_field_to_func}
        else:
            fields_to_update = global_field_to_func
        # Map fields to their cleaned data
        cleaned_data = {}
        for update_field in fields_to_update:
            dirty_value = row[update_field]
            if not dirty_value:
                continue
            cleaning_func = fields_to_update[update_field]
            if cleaning_func == CleanupFunctions.cleanup_url:
                clean = cleaning_func(url=dirty_value, tls_support=tls_cache)
            else:
                clean = cleaning_func(dirty_value)
            if clean:
                cleaned_data[update_field] = clean
        # Generate SQL update for all the fields we just cleaned
        update_field_expressions = []
        for field in cleaned_data:
            update_field_expressions.append(
                '{field} = {cleaned}'.format(
                    field=field,
                    cleaned=cleaned_data[field]
                )
            )
        if len(update_field_expressions) > 0:
            update_query = '''
                UPDATE {temp_table} SET {field_expressions} WHERE id = {_id}
            '''.format(
                temp_table=temp_table,
                field_expressions=', '.join(update_field_expressions),
                _id=_id
            )
            write_cur.execute(update_query)
    log.info('TLS cache: {}'.format(tls_cache))
    log.info('Worker committing changes...')
    worker_conn.commit()
    write_cur.close()
    worker_conn.close()
    end_time = time.time()
    total_time = end_time - start_time
    log.info('Worker finished batch in {}'.format(total_time))
    return True