Python psycopg2.extensions.adapt() Examples

def as_string(self, context):
        # is it a connection or cursor?
        if isinstance(context, ext.connection):
            conn = context
        elif isinstance(context, ext.cursor):
            conn = context.connection
        else:
            raise TypeError("context must be a connection or a cursor")

        a = ext.adapt(self._wrapped)
        if hasattr(a, 'prepare'):
            a.prepare(conn)

        rv = a.getquoted()
        if sys.version_info[0] >= 3 and isinstance(rv, bytes):
            rv = rv.decode(ext.encodings[conn.encoding])

        return rv 

def getquoted(self):
        r = self.adapted
        if r.isempty:
            return b("'empty'")

        if not r.lower_inf:
            # not exactly: we are relying that none of these object is really
            # quoted (they are numbers). Also, I'm lazy and not preparing the
            # adapter because I assume encoding doesn't matter for these
            # objects.
            lower = adapt(r.lower).getquoted().decode('ascii')
        else:
            lower = ''

        if not r.upper_inf:
            upper = adapt(r.upper).getquoted().decode('ascii')
        else:
            upper = ''

        return ("'%s%s,%s%s'" % (
            r._bounds[0], lower, upper, r._bounds[1])).encode('ascii')

# TODO: probably won't work with infs, nans and other tricky cases. 

def getquoted(self):
        r = self.adapted
        if r.isempty:
            return b"'empty'"

        if not r.lower_inf:
            # not exactly: we are relying that none of these object is really
            # quoted (they are numbers). Also, I'm lazy and not preparing the
            # adapter because I assume encoding doesn't matter for these
            # objects.
            lower = adapt(r.lower).getquoted().decode('ascii')
        else:
            lower = ''

        if not r.upper_inf:
            upper = adapt(r.upper).getquoted().decode('ascii')
        else:
            upper = ''

        return ("'%s%s,%s%s'" % (
            r._bounds[0], lower, upper, r._bounds[1])).encode('ascii')

# TODO: probably won't work with infs, nans and other tricky cases. 

def test_none_fast_path(self):
        # the None adapter is not actually invoked in regular adaptation

        class WonkyAdapter(object):
            def __init__(self, obj): pass
            def getquoted(self): return "NOPE!"

        curs = self.conn.cursor()

        orig_adapter = ext.adapters[type(None), ext.ISQLQuote]
        try:
            ext.register_adapter(type(None), WonkyAdapter)
            self.assertEqual(ext.adapt(None).getquoted(), "NOPE!")

            s = curs.mogrify("SELECT %s;", (None,))
            self.assertEqual(b("SELECT NULL;"), s)

        finally:
            ext.register_adapter(type(None), orig_adapter) 

def test_inet_conform(self):
        from psycopg2.extras import Inet
        i = Inet("192.168.1.0/24")
        a = psycopg2.extensions.adapt(i)
        a.prepare(self.conn)
        self.assertEqual(
            filter_scs(self.conn, b("E'192.168.1.0/24'::inet")),
            a.getquoted())

        # adapts ok with unicode too
        i = Inet(u"192.168.1.0/24")
        a = psycopg2.extensions.adapt(i)
        a.prepare(self.conn)
        self.assertEqual(
            filter_scs(self.conn, b("E'192.168.1.0/24'::inet")),
            a.getquoted()) 

def test_adapt_most_specific(self):
        from psycopg2.extensions import adapt, register_adapter, AsIs

        class A(object): pass
        class B(A): pass
        class C(B): pass

        register_adapter(A, lambda a: AsIs("a"))
        register_adapter(B, lambda b: AsIs("b"))
        try:
            self.assertEqual(b('b'), adapt(C()).getquoted())
        finally:
           del psycopg2.extensions.adapters[A, psycopg2.extensions.ISQLQuote]
           del psycopg2.extensions.adapters[B, psycopg2.extensions.ISQLQuote] 

def getquoted(self):
        if self.name is None:
            raise NotImplementedError(
                'RangeAdapter must be subclassed overriding its name '
                'or the getquoted() method')

        r = self.adapted
        if r.isempty:
            return b("'empty'::" + self.name)

        if r.lower is not None:
            a = adapt(r.lower)
            if hasattr(a, 'prepare'):
                a.prepare(self._conn)
            lower = a.getquoted()
        else:
            lower = b('NULL')

        if r.upper is not None:
            a = adapt(r.upper)
            if hasattr(a, 'prepare'):
                a.prepare(self._conn)
            upper = a.getquoted()
        else:
            upper = b('NULL')

        return b(self.name + '(') + lower + b(', ') + upper \
                + b(", '%s')" % r._bounds) 

def getquoted(self):
        obj = _A(self.addr)
        if hasattr(obj, 'prepare'):
            obj.prepare(self._conn)
        return obj.getquoted() + b("::inet") 

def test_str(self):
        snowman = u"\u2603"
        obj = {'a': [1, 2, snowman]}
        j = psycopg2.extensions.adapt(psycopg2.extras.Json(obj))
        s = str(j)
        self.assert_(isinstance(s, str))
        # no pesky b's
        self.assert_(s.startswith("'"))
        self.assert_(s.endswith("'")) 

def test_adapt_subtype_3(self):
        from psycopg2.extensions import adapt, register_adapter, AsIs

        class A: pass
        class B(A): pass

        register_adapter(A, lambda a: AsIs("a"))
        try:
            self.assertEqual(b("a"), adapt(B()).getquoted())
        finally:
           del psycopg2.extensions.adapters[A, psycopg2.extensions.ISQLQuote] 

def test_no_mro_no_joy(self):
        from psycopg2.extensions import adapt, register_adapter, AsIs

        class A: pass
        class B(A): pass

        register_adapter(A, lambda a: AsIs("a"))
        try:
            self.assertRaises(psycopg2.ProgrammingError, adapt, B())
        finally:
           del psycopg2.extensions.adapters[A, psycopg2.extensions.ISQLQuote] 

def adapt_array(arr):
    conn = arr.field.model_class._meta.database.get_conn()
    items = adapt(arr.items)
    items.prepare(conn)
    return AsIs('%s::%s%s' % (
        items,
        arr.field.get_column_type(),
        '[]' * arr.field.dimensions)) 

def test_adapt_subtype(self):
        from psycopg2.extensions import adapt
        class Sub(str): pass
        s1 = "hel'lo"
        s2 = Sub(s1)
        self.assertEqual(adapt(s1).getquoted(), adapt(s2).getquoted()) 

def quote_value(self, value):
        if value ==None:
            return SQL_NULL
        if is_list(value):
            json = value2json(value)
            return self.quote_value(json)

        if is_text(value) and len(value) > 256:
            value = value[:256]
        return SQL(adapt(value)) 

def sanitize_for_db(field):
    if field == None:
        return 'NULL'
    elif type(field) == int:
        return repr(int(field))
    elif type(field) == str and "'" in field:
        from psycopg2.extensions import adapt
        return str(adapt(field))
    else:
        return repr(field) 

def getquoted(self):
        if self.name is None:
            raise NotImplementedError(
                'RangeAdapter must be subclassed overriding its name '
                'or the getquoted() method')

        r = self.adapted
        if r.isempty:
            return b"'empty'::" + self.name.encode('utf8')

        if r.lower is not None:
            a = adapt(r.lower)
            if hasattr(a, 'prepare'):
                a.prepare(self._conn)
            lower = a.getquoted()
        else:
            lower = b'NULL'

        if r.upper is not None:
            a = adapt(r.upper)
            if hasattr(a, 'prepare'):
                a.prepare(self._conn)
            upper = a.getquoted()
        else:
            upper = b'NULL'

        return self.name.encode('utf8') + b'(' + lower + b', ' + upper \
            + b", '" + r._bounds.encode('utf8') + b"')" 

def getquoted(self):
        obj = _A(self.addr)
        if hasattr(obj, 'prepare'):
            obj.prepare(self._conn)
        return obj.getquoted() + b"::inet" 

def modify_recovery_conf_to_point_at_new_master(self, new_master_instance):
        with open(os.path.join(self.config.get("pg_data_directory"), "PG_VERSION"), "r") as fp:
            pg_version = fp.read().strip()

        if LooseVersion(pg_version) >= "12":
            recovery_conf_filename = "postgresql.auto.conf"
        else:
            recovery_conf_filename = "recovery.conf"

        path_to_recovery_conf = os.path.join(self.config.get("pg_data_directory"), recovery_conf_filename)
        with open(path_to_recovery_conf, "r") as fp:
            old_conf = fp.read().splitlines()
        has_recovery_target_timeline = False
        new_conf = []
        old_conn_info = None
        for line in old_conf:
            if line.startswith("recovery_target_timeline"):
                has_recovery_target_timeline = True
            if line.startswith("primary_conninfo"):
                # grab previous entry: strip surrounding quotes and replace two quotes with one
                try:
                    old_conn_info = get_connection_info_from_config_line(line)
                except ValueError:
                    self.log.exception("failed to parse previous %r, ignoring", line)
                continue  # skip this line
            new_conf.append(line)

        # If has_recovery_target_timeline is set and old_conn_info matches
        # new info we don't have to do anything
        new_conn_info = get_connection_info(self.primary_conninfo_template)
        master_instance_conn_info = get_connection_info(self.config["remote_conns"][new_master_instance])
        assert "host" in master_instance_conn_info
        new_conn_info["host"] = master_instance_conn_info["host"]
        if "port" in master_instance_conn_info:
            new_conn_info["port"] = master_instance_conn_info["port"]
        if new_conn_info == old_conn_info:
            self.log.debug("recovery.conf already contains conninfo matching %r, not updating", new_master_instance)
            return False
        # Otherwise we append the new primary_conninfo
        new_conf.append("primary_conninfo = {0}".format(adapt(create_connection_string(new_conn_info))))
        # The timeline of the recovery.conf will require a higher timeline target
        if not has_recovery_target_timeline:
            new_conf.append("recovery_target_timeline = 'latest'")
        # prepend our tag
        new_conf.insert(0,
                        "# pglookout updated primary_conninfo for instance {0} at {1}"
                        .format(new_master_instance, get_iso_timestamp()))
        # Replace old recovery.conf with a fresh copy
        with open(path_to_recovery_conf + "_temp", "w") as fp:
            fp.write("\n".join(new_conf) + "\n")

        os.rename(path_to_recovery_conf + "_temp", path_to_recovery_conf)
        return True 

def start_replication(self, slot_name=None, slot_type=None, start_lsn=0,
                          timeline=0, options=None, decode=False):
        """Start replication stream."""

        command = "START_REPLICATION "

        if slot_type is None:
            slot_type = self.connection.replication_type

        if slot_type == REPLICATION_LOGICAL:
            if slot_name:
                command += "SLOT %s " % quote_ident(slot_name, self)
            else:
                raise psycopg2.ProgrammingError(
                    "slot name is required for logical replication")

            command += "LOGICAL "

        elif slot_type == REPLICATION_PHYSICAL:
            if slot_name:
                command += "SLOT %s " % quote_ident(slot_name, self)
            # don't add "PHYSICAL", before 9.4 it was just START_REPLICATION XXX/XXX

        else:
            raise psycopg2.ProgrammingError(
                "unrecognized replication type: %s" % repr(slot_type))

        if type(start_lsn) is str:
            lsn = start_lsn.split('/')
            lsn = "%X/%08X" % (int(lsn[0], 16), int(lsn[1], 16))
        else:
            lsn = "%X/%08X" % ((start_lsn >> 32) & 0xFFFFFFFF,
                               start_lsn & 0xFFFFFFFF)

        command += lsn

        if timeline != 0:
            if slot_type == REPLICATION_LOGICAL:
                raise psycopg2.ProgrammingError(
                    "cannot specify timeline for logical replication")

            command += " TIMELINE %d" % timeline

        if options:
            if slot_type == REPLICATION_PHYSICAL:
                raise psycopg2.ProgrammingError(
                    "cannot specify output plugin options for physical replication")

            command += " ("
            for k, v in options.items():
                if not command.endswith('('):
                    command += ", "
                command += "%s %s" % (quote_ident(k, self), _A(str(v)))
            command += ")"

        self.start_replication_expert(command, decode=decode)

    # allows replication cursors to be used in select.select() directly 

def execute_sql(self, csv_path_or_obj=None):
        """
        Run the COPY TO query.
        """
        logger.debug("Copying data to {}".format(csv_path_or_obj))

        # adapt SELECT query parameters to SQL syntax
        params = self.as_sql()[1]
        adapted_params = tuple(adapt(p) for p in params)

        # use stdout to avoid file permission issues
        with connections[self.using].cursor() as c:
            # compile the SELECT query
            select_sql = self.as_sql()[0] % adapted_params
            # then the COPY TO query
            copy_to_sql = "COPY ({}) TO STDOUT {} CSV"
            copy_to_sql = copy_to_sql.format(select_sql, self.query.copy_to_delimiter)
            # Optional extras
            options_list = [
                self.query.copy_to_header,
                self.query.copy_to_null_string,
                self.query.copy_to_quote_char,
                self.query.copy_to_force_quote,
                self.query.copy_to_encoding,
                self.query.copy_to_escape
            ]
            options_sql = " ".join([o for o in options_list if o]).strip()
            if options_sql:
                copy_to_sql = copy_to_sql + " " + options_sql
            # then execute
            logger.debug(copy_to_sql)

            # If a file-like object was provided, write it out there.
            if hasattr(csv_path_or_obj, 'write'):
                c.cursor.copy_expert(copy_to_sql, csv_path_or_obj)
                return
            # If a file path was provided, write it out there.
            elif csv_path_or_obj:
                with open(csv_path_or_obj, 'wb') as stdout:
                    c.cursor.copy_expert(copy_to_sql, stdout)
                    return
            # If there's no csv_path, return the output as a string.
            else:
                stdout = BytesIO()
                c.cursor.copy_expert(copy_to_sql, stdout)
                return stdout.getvalue()