468 lines
16 KiB
Python
468 lines
16 KiB
Python
|
"""
|
||
|
|
SQL composition utility module
|
||
|
|
"""
|
||
|
|
|
||
|
|
# Copyright (C) 2020 The Psycopg Team
|
||
|
|
|
||
|
|
import codecs
|
||
|
|
import string
|
||
|
|
from abc import ABC, abstractmethod
|
||
|
|
from typing import Any, Iterator, Iterable, List, Optional, Sequence, Union
|
||
|
|
|
||
|
|
from .pq import Escaping
|
||
|
|
from .abc import AdaptContext
|
||
|
|
from .adapt import Transformer, PyFormat
|
||
|
|
from ._compat import LiteralString
|
||
|
|
from ._encodings import conn_encoding
|
||
|
|
|
||
|
|
|
||
|
|
def quote(obj: Any, context: Optional[AdaptContext] = None) -> str:
|
||
|
|
"""
|
||
|
|
Adapt a Python object to a quoted SQL string.
|
||
|
|
|
||
|
|
Use this function only if you absolutely want to convert a Python string to
|
||
|
|
an SQL quoted literal to use e.g. to generate batch SQL and you won't have
|
||
|
|
a connection available when you will need to use it.
|
||
|
|
|
||
|
|
This function is relatively inefficient, because it doesn't cache the
|
||
|
|
adaptation rules. If you pass a `!context` you can adapt the adaptation
|
||
|
|
rules used, otherwise only global rules are used.
|
||
|
|
|
||
|
|
"""
|
||
|
|
return Literal(obj).as_string(context)
|
||
|
|
|
||
|
|
|
||
|
|
class Composable(ABC):
|
||
|
|
"""
|
||
|
|
Abstract base class for objects that can be used to compose an SQL string.
|
||
|
|
|
||
|
|
`!Composable` objects can be passed directly to
|
||
|
|
`~psycopg.Cursor.execute()`, `~psycopg.Cursor.executemany()`,
|
||
|
|
`~psycopg.Cursor.copy()` in place of the query string.
|
||
|
|
|
||
|
|
`!Composable` objects can be joined using the ``+`` operator: the result
|
||
|
|
will be a `Composed` instance containing the objects joined. The operator
|
||
|
|
``*`` is also supported with an integer argument: the result is a
|
||
|
|
`!Composed` instance containing the left argument repeated as many times as
|
||
|
|
requested.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, obj: Any):
|
||
|
|
self._obj = obj
|
||
|
|
|
||
|
|
def __repr__(self) -> str:
|
||
|
|
return f"{self.__class__.__name__}({self._obj!r})"
|
||
|
|
|
||
|
|
@abstractmethod
|
||
|
|
def as_bytes(self, context: Optional[AdaptContext]) -> bytes:
|
||
|
|
"""
|
||
|
|
Return the value of the object as bytes.
|
||
|
|
|
||
|
|
:param context: the context to evaluate the object into.
|
||
|
|
:type context: `connection` or `cursor`
|
||
|
|
|
||
|
|
The method is automatically invoked by `~psycopg.Cursor.execute()`,
|
||
|
|
`~psycopg.Cursor.executemany()`, `~psycopg.Cursor.copy()` if a
|
||
|
|
`!Composable` is passed instead of the query string.
|
||
|
|
|
||
|
|
"""
|
||
|
|
raise NotImplementedError
|
||
|
|
|
||
|
|
def as_string(self, context: Optional[AdaptContext]) -> str:
|
||
|
|
"""
|
||
|
|
Return the value of the object as string.
|
||
|
|
|
||
|
|
:param context: the context to evaluate the string into.
|
||
|
|
:type context: `connection` or `cursor`
|
||
|
|
|
||
|
|
"""
|
||
|
|
conn = context.connection if context else None
|
||
|
|
enc = conn_encoding(conn)
|
||
|
|
b = self.as_bytes(context)
|
||
|
|
if isinstance(b, bytes):
|
||
|
|
return b.decode(enc)
|
||
|
|
else:
|
||
|
|
# buffer object
|
||
|
|
return codecs.lookup(enc).decode(b)[0]
|
||
|
|
|
||
|
|
def __add__(self, other: "Composable") -> "Composed":
|
||
|
|
if isinstance(other, Composed):
|
||
|
|
return Composed([self]) + other
|
||
|
|
if isinstance(other, Composable):
|
||
|
|
return Composed([self]) + Composed([other])
|
||
|
|
else:
|
||
|
|
return NotImplemented
|
||
|
|
|
||
|
|
def __mul__(self, n: int) -> "Composed":
|
||
|
|
return Composed([self] * n)
|
||
|
|
|
||
|
|
def __eq__(self, other: Any) -> bool:
|
||
|
|
return type(self) is type(other) and self._obj == other._obj
|
||
|
|
|
||
|
|
def __ne__(self, other: Any) -> bool:
|
||
|
|
return not self.__eq__(other)
|
||
|
|
|
||
|
|
|
||
|
|
class Composed(Composable):
|
||
|
|
"""
|
||
|
|
A `Composable` object made of a sequence of `!Composable`.
|
||
|
|
|
||
|
|
The object is usually created using `!Composable` operators and methods.
|
||
|
|
However it is possible to create a `!Composed` directly specifying a
|
||
|
|
sequence of objects as arguments: if they are not `!Composable` they will
|
||
|
|
be wrapped in a `Literal`.
|
||
|
|
|
||
|
|
Example::
|
||
|
|
|
||
|
|
>>> comp = sql.Composed(
|
||
|
|
... [sql.SQL("INSERT INTO "), sql.Identifier("table")])
|
||
|
|
>>> print(comp.as_string(conn))
|
||
|
|
INSERT INTO "table"
|
||
|
|
|
||
|
|
`!Composed` objects are iterable (so they can be used in `SQL.join` for
|
||
|
|
instance).
|
||
|
|
"""
|
||
|
|
|
||
|
|
_obj: List[Composable]
|
||
|
|
|
||
|
|
def __init__(self, seq: Sequence[Any]):
|
||
|
|
seq = [obj if isinstance(obj, Composable) else Literal(obj) for obj in seq]
|
||
|
|
super().__init__(seq)
|
||
|
|
|
||
|
|
def as_bytes(self, context: Optional[AdaptContext]) -> bytes:
|
||
|
|
return b"".join(obj.as_bytes(context) for obj in self._obj)
|
||
|
|
|
||
|
|
def __iter__(self) -> Iterator[Composable]:
|
||
|
|
return iter(self._obj)
|
||
|
|
|
||
|
|
def __add__(self, other: Composable) -> "Composed":
|
||
|
|
if isinstance(other, Composed):
|
||
|
|
return Composed(self._obj + other._obj)
|
||
|
|
if isinstance(other, Composable):
|
||
|
|
return Composed(self._obj + [other])
|
||
|
|
else:
|
||
|
|
return NotImplemented
|
||
|
|
|
||
|
|
def join(self, joiner: Union["SQL", LiteralString]) -> "Composed":
|
||
|
|
"""
|
||
|
|
Return a new `!Composed` interposing the `!joiner` with the `!Composed` items.
|
||
|
|
|
||
|
|
The `!joiner` must be a `SQL` or a string which will be interpreted as
|
||
|
|
an `SQL`.
|
||
|
|
|
||
|
|
Example::
|
||
|
|
|
||
|
|
>>> fields = sql.Identifier('foo') + sql.Identifier('bar') # a Composed
|
||
|
|
>>> print(fields.join(', ').as_string(conn))
|
||
|
|
"foo", "bar"
|
||
|
|
|
||
|
|
"""
|
||
|
|
if isinstance(joiner, str):
|
||
|
|
joiner = SQL(joiner)
|
||
|
|
elif not isinstance(joiner, SQL):
|
||
|
|
raise TypeError(
|
||
|
|
"Composed.join() argument must be strings or SQL,"
|
||
|
|
f" got {joiner!r} instead"
|
||
|
|
)
|
||
|
|
|
||
|
|
return joiner.join(self._obj)
|
||
|
|
|
||
|
|
|
||
|
|
class SQL(Composable):
|
||
|
|
"""
|
||
|
|
A `Composable` representing a snippet of SQL statement.
|
||
|
|
|
||
|
|
`!SQL` exposes `join()` and `format()` methods useful to create a template
|
||
|
|
where to merge variable parts of a query (for instance field or table
|
||
|
|
names).
|
||
|
|
|
||
|
|
The `!obj` string doesn't undergo any form of escaping, so it is not
|
||
|
|
suitable to represent variable identifiers or values: you should only use
|
||
|
|
it to pass constant strings representing templates or snippets of SQL
|
||
|
|
statements; use other objects such as `Identifier` or `Literal` to
|
||
|
|
represent variable parts.
|
||
|
|
|
||
|
|
Example::
|
||
|
|
|
||
|
|
>>> query = sql.SQL("SELECT {0} FROM {1}").format(
|
||
|
|
... sql.SQL(', ').join([sql.Identifier('foo'), sql.Identifier('bar')]),
|
||
|
|
... sql.Identifier('table'))
|
||
|
|
>>> print(query.as_string(conn))
|
||
|
|
SELECT "foo", "bar" FROM "table"
|
||
|
|
"""
|
||
|
|
|
||
|
|
_obj: LiteralString
|
||
|
|
_formatter = string.Formatter()
|
||
|
|
|
||
|
|
def __init__(self, obj: LiteralString):
|
||
|
|
super().__init__(obj)
|
||
|
|
if not isinstance(obj, str):
|
||
|
|
raise TypeError(f"SQL values must be strings, got {obj!r} instead")
|
||
|
|
|
||
|
|
def as_string(self, context: Optional[AdaptContext]) -> str:
|
||
|
|
return self._obj
|
||
|
|
|
||
|
|
def as_bytes(self, context: Optional[AdaptContext]) -> bytes:
|
||
|
|
enc = "utf-8"
|
||
|
|
if context:
|
||
|
|
enc = conn_encoding(context.connection)
|
||
|
|
return self._obj.encode(enc)
|
||
|
|
|
||
|
|
def format(self, *args: Any, **kwargs: Any) -> Composed:
|
||
|
|
"""
|
||
|
|
Merge `Composable` objects into a template.
|
||
|
|
|
||
|
|
:param args: parameters to replace to numbered (``{0}``, ``{1}``) or
|
||
|
|
auto-numbered (``{}``) placeholders
|
||
|
|
:param kwargs: parameters to replace to named (``{name}``) placeholders
|
||
|
|
:return: the union of the `!SQL` string with placeholders replaced
|
||
|
|
:rtype: `Composed`
|
||
|
|
|
||
|
|
The method is similar to the Python `str.format()` method: the string
|
||
|
|
template supports auto-numbered (``{}``), numbered (``{0}``,
|
||
|
|
``{1}``...), and named placeholders (``{name}``), with positional
|
||
|
|
arguments replacing the numbered placeholders and keywords replacing
|
||
|
|
the named ones. However placeholder modifiers (``{0!r}``, ``{0:<10}``)
|
||
|
|
are not supported.
|
||
|
|
|
||
|
|
If a `!Composable` objects is passed to the template it will be merged
|
||
|
|
according to its `as_string()` method. If any other Python object is
|
||
|
|
passed, it will be wrapped in a `Literal` object and so escaped
|
||
|
|
according to SQL rules.
|
||
|
|
|
||
|
|
Example::
|
||
|
|
|
||
|
|
>>> print(sql.SQL("SELECT * FROM {} WHERE {} = %s")
|
||
|
|
... .format(sql.Identifier('people'), sql.Identifier('id'))
|
||
|
|
... .as_string(conn))
|
||
|
|
SELECT * FROM "people" WHERE "id" = %s
|
||
|
|
|
||
|
|
>>> print(sql.SQL("SELECT * FROM {tbl} WHERE name = {name}")
|
||
|
|
... .format(tbl=sql.Identifier('people'), name="O'Rourke"))
|
||
|
|
... .as_string(conn))
|
||
|
|
SELECT * FROM "people" WHERE name = 'O''Rourke'
|
||
|
|
|
||
|
|
"""
|
||
|
|
rv: List[Composable] = []
|
||
|
|
autonum: Optional[int] = 0
|
||
|
|
# TODO: this is probably not the right way to whitelist pre
|
||
|
|
# pyre complains. Will wait for mypy to complain too to fix.
|
||
|
|
pre: LiteralString
|
||
|
|
for pre, name, spec, conv in self._formatter.parse(self._obj):
|
||
|
|
if spec:
|
||
|
|
raise ValueError("no format specification supported by SQL")
|
||
|
|
if conv:
|
||
|
|
raise ValueError("no format conversion supported by SQL")
|
||
|
|
if pre:
|
||
|
|
rv.append(SQL(pre))
|
||
|
|
|
||
|
|
if name is None:
|
||
|
|
continue
|
||
|
|
|
||
|
|
if name.isdigit():
|
||
|
|
if autonum:
|
||
|
|
raise ValueError(
|
||
|
|
"cannot switch from automatic field numbering to manual"
|
||
|
|
)
|
||
|
|
rv.append(args[int(name)])
|
||
|
|
autonum = None
|
||
|
|
|
||
|
|
elif not name:
|
||
|
|
if autonum is None:
|
||
|
|
raise ValueError(
|
||
|
|
"cannot switch from manual field numbering to automatic"
|
||
|
|
)
|
||
|
|
rv.append(args[autonum])
|
||
|
|
autonum += 1
|
||
|
|
|
||
|
|
else:
|
||
|
|
rv.append(kwargs[name])
|
||
|
|
|
||
|
|
return Composed(rv)
|
||
|
|
|
||
|
|
def join(self, seq: Iterable[Composable]) -> Composed:
|
||
|
|
"""
|
||
|
|
Join a sequence of `Composable`.
|
||
|
|
|
||
|
|
:param seq: the elements to join.
|
||
|
|
:type seq: iterable of `!Composable`
|
||
|
|
|
||
|
|
Use the `!SQL` object's string to separate the elements in `!seq`.
|
||
|
|
Note that `Composed` objects are iterable too, so they can be used as
|
||
|
|
argument for this method.
|
||
|
|
|
||
|
|
Example::
|
||
|
|
|
||
|
|
>>> snip = sql.SQL(', ').join(
|
||
|
|
... sql.Identifier(n) for n in ['foo', 'bar', 'baz'])
|
||
|
|
>>> print(snip.as_string(conn))
|
||
|
|
"foo", "bar", "baz"
|
||
|
|
"""
|
||
|
|
rv = []
|
||
|
|
it = iter(seq)
|
||
|
|
try:
|
||
|
|
rv.append(next(it))
|
||
|
|
except StopIteration:
|
||
|
|
pass
|
||
|
|
else:
|
||
|
|
for i in it:
|
||
|
|
rv.append(self)
|
||
|
|
rv.append(i)
|
||
|
|
|
||
|
|
return Composed(rv)
|
||
|
|
|
||
|
|
|
||
|
|
class Identifier(Composable):
|
||
|
|
"""
|
||
|
|
A `Composable` representing an SQL identifier or a dot-separated sequence.
|
||
|
|
|
||
|
|
Identifiers usually represent names of database objects, such as tables or
|
||
|
|
fields. PostgreSQL identifiers follow `different rules`__ than SQL string
|
||
|
|
literals for escaping (e.g. they use double quotes instead of single).
|
||
|
|
|
||
|
|
.. __: https://www.postgresql.org/docs/current/sql-syntax-lexical.html# \
|
||
|
|
SQL-SYNTAX-IDENTIFIERS
|
||
|
|
|
||
|
|
Example::
|
||
|
|
|
||
|
|
>>> t1 = sql.Identifier("foo")
|
||
|
|
>>> t2 = sql.Identifier("ba'r")
|
||
|
|
>>> t3 = sql.Identifier('ba"z')
|
||
|
|
>>> print(sql.SQL(', ').join([t1, t2, t3]).as_string(conn))
|
||
|
|
"foo", "ba'r", "ba""z"
|
||
|
|
|
||
|
|
Multiple strings can be passed to the object to represent a qualified name,
|
||
|
|
i.e. a dot-separated sequence of identifiers.
|
||
|
|
|
||
|
|
Example::
|
||
|
|
|
||
|
|
>>> query = sql.SQL("SELECT {} FROM {}").format(
|
||
|
|
... sql.Identifier("table", "field"),
|
||
|
|
... sql.Identifier("schema", "table"))
|
||
|
|
>>> print(query.as_string(conn))
|
||
|
|
SELECT "table"."field" FROM "schema"."table"
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
_obj: Sequence[str]
|
||
|
|
|
||
|
|
def __init__(self, *strings: str):
|
||
|
|
# init super() now to make the __repr__ not explode in case of error
|
||
|
|
super().__init__(strings)
|
||
|
|
|
||
|
|
if not strings:
|
||
|
|
raise TypeError("Identifier cannot be empty")
|
||
|
|
|
||
|
|
for s in strings:
|
||
|
|
if not isinstance(s, str):
|
||
|
|
raise TypeError(
|
||
|
|
f"SQL identifier parts must be strings, got {s!r} instead"
|
||
|
|
)
|
||
|
|
|
||
|
|
def __repr__(self) -> str:
|
||
|
|
return f"{self.__class__.__name__}({', '.join(map(repr, self._obj))})"
|
||
|
|
|
||
|
|
def as_bytes(self, context: Optional[AdaptContext]) -> bytes:
|
||
|
|
conn = context.connection if context else None
|
||
|
|
if not conn:
|
||
|
|
raise ValueError("a connection is necessary for Identifier")
|
||
|
|
esc = Escaping(conn.pgconn)
|
||
|
|
enc = conn_encoding(conn)
|
||
|
|
escs = [esc.escape_identifier(s.encode(enc)) for s in self._obj]
|
||
|
|
return b".".join(escs)
|
||
|
|
|
||
|
|
|
||
|
|
class Literal(Composable):
|
||
|
|
"""
|
||
|
|
A `Composable` representing an SQL value to include in a query.
|
||
|
|
|
||
|
|
Usually you will want to include placeholders in the query and pass values
|
||
|
|
as `~cursor.execute()` arguments. If however you really really need to
|
||
|
|
include a literal value in the query you can use this object.
|
||
|
|
|
||
|
|
The string returned by `!as_string()` follows the normal :ref:`adaptation
|
||
|
|
rules <types-adaptation>` for Python objects.
|
||
|
|
|
||
|
|
Example::
|
||
|
|
|
||
|
|
>>> s1 = sql.Literal("fo'o")
|
||
|
|
>>> s2 = sql.Literal(42)
|
||
|
|
>>> s3 = sql.Literal(date(2000, 1, 1))
|
||
|
|
>>> print(sql.SQL(', ').join([s1, s2, s3]).as_string(conn))
|
||
|
|
'fo''o', 42, '2000-01-01'::date
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
def as_bytes(self, context: Optional[AdaptContext]) -> bytes:
|
||
|
|
tx = Transformer.from_context(context)
|
||
|
|
return tx.as_literal(self._obj)
|
||
|
|
|
||
|
|
|
||
|
|
class Placeholder(Composable):
|
||
|
|
"""A `Composable` representing a placeholder for query parameters.
|
||
|
|
|
||
|
|
If the name is specified, generate a named placeholder (e.g. ``%(name)s``,
|
||
|
|
``%(name)b``), otherwise generate a positional placeholder (e.g. ``%s``,
|
||
|
|
``%b``).
|
||
|
|
|
||
|
|
The object is useful to generate SQL queries with a variable number of
|
||
|
|
arguments.
|
||
|
|
|
||
|
|
Examples::
|
||
|
|
|
||
|
|
>>> names = ['foo', 'bar', 'baz']
|
||
|
|
|
||
|
|
>>> q1 = sql.SQL("INSERT INTO my_table ({}) VALUES ({})").format(
|
||
|
|
... sql.SQL(', ').join(map(sql.Identifier, names)),
|
||
|
|
... sql.SQL(', ').join(sql.Placeholder() * len(names)))
|
||
|
|
>>> print(q1.as_string(conn))
|
||
|
|
INSERT INTO my_table ("foo", "bar", "baz") VALUES (%s, %s, %s)
|
||
|
|
|
||
|
|
>>> q2 = sql.SQL("INSERT INTO my_table ({}) VALUES ({})").format(
|
||
|
|
... sql.SQL(', ').join(map(sql.Identifier, names)),
|
||
|
|
... sql.SQL(', ').join(map(sql.Placeholder, names)))
|
||
|
|
>>> print(q2.as_string(conn))
|
||
|
|
INSERT INTO my_table ("foo", "bar", "baz") VALUES (%(foo)s, %(bar)s, %(baz)s)
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, name: str = "", format: Union[str, PyFormat] = PyFormat.AUTO):
|
||
|
|
super().__init__(name)
|
||
|
|
if not isinstance(name, str):
|
||
|
|
raise TypeError(f"expected string as name, got {name!r}")
|
||
|
|
|
||
|
|
if ")" in name:
|
||
|
|
raise ValueError(f"invalid name: {name!r}")
|
||
|
|
|
||
|
|
if type(format) is str:
|
||
|
|
format = PyFormat(format)
|
||
|
|
if not isinstance(format, PyFormat):
|
||
|
|
raise TypeError(
|
||
|
|
f"expected PyFormat as format, got {type(format).__name__!r}"
|
||
|
|
)
|
||
|
|
|
||
|
|
self._format: PyFormat = format
|
||
|
|
|
||
|
|
def __repr__(self) -> str:
|
||
|
|
parts = []
|
||
|
|
if self._obj:
|
||
|
|
parts.append(repr(self._obj))
|
||
|
|
if self._format is not PyFormat.AUTO:
|
||
|
|
parts.append(f"format={self._format.name}")
|
||
|
|
|
||
|
|
return f"{self.__class__.__name__}({', '.join(parts)})"
|
||
|
|
|
||
|
|
def as_string(self, context: Optional[AdaptContext]) -> str:
|
||
|
|
code = self._format.value
|
||
|
|
return f"%({self._obj}){code}" if self._obj else f"%{code}"
|
||
|
|
|
||
|
|
def as_bytes(self, context: Optional[AdaptContext]) -> bytes:
|
||
|
|
conn = context.connection if context else None
|
||
|
|
enc = conn_encoding(conn)
|
||
|
|
return self.as_string(context).encode(enc)
|
||
|
|
|
||
|
|
|
||
|
|
# Literals
|
||
|
|
NULL = SQL("NULL")
|
||
|
|
DEFAULT = SQL("DEFAULT")
|