def _make_columns_clause(columns: List[str] = None) -> psql.Composable:
"""Create the list of columns for a select statement in psql syntax"""
if isinstance(columns, str):
raise TypeError(
f"'columns' must be a list of column names. Got '{columns}'")
return (psql.SQL(",").join([psql.Identifier(c) for c in columns])
if columns else psql.SQL("*"))
def test_init(self):
assert isinstance(sql.SQL("foo"), sql.SQL)
assert isinstance(sql.SQL("foo"), sql.SQL)
with pytest.raises(TypeError):
sql.SQL(10)
with pytest.raises(TypeError):
sql.SQL(dt.date(2016, 12, 31))
def work():
with psycopg.connect(dsn) as conn:
with conn.cursor(binary=fmt) as cur:
cur.execute(faker.drop_stmt)
cur.execute(faker.create_stmt)
with faker.find_insert_problem(conn):
cur.executemany(faker.insert_stmt, faker.records)
stmt = sql.SQL(
"copy (select {} from {} order by id) to stdout (format {})"
).format(
sql.SQL(", ").join(faker.fields_names),
faker.table_name,
sql.SQL(fmt.name),
)
with cur.copy(stmt) as copy:
if set_types:
copy.set_types(faker.types_names)
if method == "read":
while 1:
tmp = copy.read()
if not tmp:
break
elif method == "iter":
list(copy)
elif method == "row":
while 1:
tmp = copy.read_row() # type: ignore[assignment]
if tmp is None:
break
elif method == "rows":
list(copy.rows())
def main() -> None:
cnn = psycopg.connect()
cnn.execute(
sql.SQL("create table testdec ({})").format(
sql.SQL(", ").join([
sql.SQL("{} numeric(10,2)").format(sql.Identifier(f"t{i}"))
for i in range(ncols)
])))
cur = cnn.cursor()
if test == "copy":
with cur.copy(
f"copy testdec from stdin (format {format.name})") as copy:
for j in range(nrows):
copy.write_row([
Decimal(randrange(10000000000)) / 100 for i in range(ncols)
])
elif test == "insert":
ph = ["%t", "%b"][format]
cur.executemany(
"insert into testdec values (%s)" % ", ".join([ph] * ncols),
([Decimal(randrange(10000000000)) / 100 for i in range(ncols)]
for j in range(nrows)),
)
else:
raise Exception(f"bad test: {test}")
def test_compose_badnargs_auto(self):
with pytest.raises(IndexError):
sql.SQL("select {};").format()
with pytest.raises(ValueError):
sql.SQL("select {} {1};").format(10, 20)
with pytest.raises(ValueError):
sql.SQL("select {0} {};").format(10, 20)
def test_as_bytes(self, conn):
assert sql.SQL("foo").as_bytes(conn) == b"foo"
conn.execute("set client_encoding to utf8")
assert sql.SQL(eur).as_bytes(conn) == eur.encode()
conn.execute("set client_encoding to latin9")
assert sql.SQL(eur).as_bytes(conn) == eur.encode("latin9")
def test_json_load_copy(conn, val, jtype, fmt_out):
cur = conn.cursor()
stmt = sql.SQL("copy (select {}::{}) to stdout (format {})").format(
val, sql.Identifier(jtype), sql.SQL(fmt_out.name))
with cur.copy(stmt) as copy:
copy.set_types([jtype])
(got, ) = copy.read_row()
assert got == json.loads(val)
def test_iter(self):
obj = sql.Composed([sql.SQL("foo"), sql.SQL("bar")])
it = iter(obj)
i = next(it)
assert i == sql.SQL("foo")
i = next(it)
assert i == sql.SQL("bar")
with pytest.raises(StopIteration):
next(it)
def create_stmt(self):
field_values = []
for name, type in zip(self.fields_names, self.types_names_sql):
field_values.append(sql.SQL("{} {}").format(name, type))
fields = sql.SQL(", ").join(field_values)
return sql.SQL(
"create table {table} (id serial primary key, {fields})"
).format(table=self.table_name, fields=fields)
def insert_stmt(self):
phs = [
sql.Placeholder(format=self.format)
for i in range(len(self.schema))
]
return sql.SQL("insert into {} ({}) values ({})").format(
self.table_name,
sql.SQL(", ").join(self.fields_names),
sql.SQL(", ").join(phs),
)
def test_sum_inplace(self, conn):
obj = sql.Composed([sql.SQL("foo ")])
obj += sql.Literal("bar")
assert isinstance(obj, sql.Composed)
assert noe(obj.as_string(conn)) == "foo 'bar'"
obj = sql.Composed([sql.SQL("foo ")])
obj += sql.Composed([sql.Literal("bar")])
assert isinstance(obj, sql.Composed)
assert noe(obj.as_string(conn)) == "foo 'bar'"
def test_quote_percent(conn):
cur = conn.cursor()
cur.execute(sql.SQL("select {ch}").format(ch=sql.Literal("%")))
assert cur.fetchone()[0] == "%"
cur.execute(
sql.SQL("select {ch} = chr(%s)").format(ch=sql.Literal("%")),
(ord("%"), ),
)
assert cur.fetchone()[0] is True
def test_as_bytes(self, conn):
obj = sql.Composed([sql.SQL("foo"), sql.SQL("bar")])
assert obj.as_bytes(conn) == b"foobar"
obj = sql.Composed([sql.SQL("foo"), sql.SQL(eur)])
conn.execute("set client_encoding to utf8")
assert obj.as_bytes(conn) == ("foo" + eur).encode()
conn.execute("set client_encoding to latin9")
assert obj.as_bytes(conn) == ("foo" + eur).encode("latin9")
def func(self, function_name: str, *args: Any) -> Generator:
"""Call a database function."""
placeholders = sql.SQL(", ").join(sql.Placeholder() * len(args))
func = sql.Identifier(function_name)
cleaned_args = []
for arg in args:
if isinstance(arg, dict):
cleaned_args.append(psycopg.types.json.Jsonb(arg))
else:
cleaned_args.append(arg)
base_query = sql.SQL("SELECT * FROM {}({});").format(func, placeholders)
return self.query(base_query, cleaned_args)
def test_pos_spec(self, conn):
s = sql.SQL("select {0} from {1}").format(sql.Identifier("field"),
sql.Identifier("table"))
s1 = s.as_string(conn)
assert isinstance(s1, str)
assert s1 == 'select "field" from "table"'
s = sql.SQL("select {1} from {0}").format(sql.Identifier("table"),
sql.Identifier("field"))
s1 = s.as_string(conn)
assert isinstance(s1, str)
assert s1 == 'select "field" from "table"'
def test_load_list_int(conn, obj, want, fmt_out):
cur = conn.cursor(binary=fmt_out)
cur.execute("select %s::int[]", (obj, ))
assert cur.fetchone()[0] == want
stmt = sql.SQL("copy (select {}::int[]) to stdout (format {})").format(
obj, sql.SQL(fmt_out.name))
with cur.copy(stmt) as copy:
copy.set_types(["int4[]"])
(got, ) = copy.read_row()
assert got == want
def _get_type_name(self, tx, schema, value):
# Special case it as it is passed as unknown so is returned as text
if schema == (list, str):
return sql.SQL("text[]")
registry = self.conn.adapters.types
dumper = tx.get_dumper(value, self.format)
dumper.dump(value) # load the oid if it's dynamic (e.g. array)
info = registry.get(dumper.oid) or registry.get("text")
if dumper.oid == info.array_oid:
return sql.SQL("{}[]").format(sql.Identifier(info.name))
else:
return sql.Identifier(info.name)
def test_uuid_load(conn, fmt_out):
cur = conn.cursor(binary=fmt_out)
val = "12345678123456781234567812345679"
cur.execute("select %s::uuid", (val, ))
assert cur.fetchone()[0] == UUID(val)
stmt = sql.SQL("copy (select {}::uuid) to stdout (format {})").format(
val, sql.SQL(fmt_out.name))
with cur.copy(stmt) as copy:
copy.set_types(["uuid"])
(res, ) = copy.read_row()
assert res == UUID(val)
def test_load_badenc(conn, typename, fmt_out):
conn.autocommit = True
cur = conn.cursor(binary=fmt_out)
conn.execute("set client_encoding to latin1")
with pytest.raises(psycopg.DataError):
cur.execute(f"select chr(%s)::{typename}", [ord(eur)])
stmt = sql.SQL("copy (select chr({})) to stdout (format {})").format(
ord(eur), sql.SQL(fmt_out.name))
with cur.copy(stmt) as copy:
copy.set_types([typename])
with pytest.raises(psycopg.DataError):
copy.read_row()
def test_load_ascii(conn, typename, fmt_out):
cur = conn.cursor(binary=fmt_out)
conn.execute("set client_encoding to sql_ascii")
cur.execute(f"select chr(%s)::{typename}", [ord(eur)])
assert cur.fetchone()[0] == eur.encode()
stmt = sql.SQL("copy (select chr({})) to stdout (format {})").format(
ord(eur), sql.SQL(fmt_out.name))
with cur.copy(stmt) as copy:
copy.set_types([typename])
(res, ) = copy.read_row()
assert res == eur.encode()
def test_load_enc(conn, typename, encoding, fmt_out):
cur = conn.cursor(binary=fmt_out)
conn.client_encoding = encoding
(res, ) = cur.execute(f"select chr(%s::int)::{typename}",
(ord(eur), )).fetchone()
assert res == eur
stmt = sql.SQL("copy (select chr({}::int)) to stdout (format {})").format(
ord(eur), sql.SQL(fmt_out.name))
with cur.copy(stmt) as copy:
copy.set_types([typename])
(res, ) = copy.read_row()
assert res == eur
def test_must_be_adaptable(self, conn):
class Foo:
pass
s = sql.SQL("select {0};").format(sql.Literal(Foo()))
with pytest.raises(ProgrammingError):
s.as_string(conn)
def test_quote_1byte(conn):
cur = conn.cursor()
query = sql.SQL("select {ch} = %s::bytea")
for i in range(0, 256):
cur.execute(query.format(ch=sql.Literal(bytes([i]))),
(fr"\x{i:02x}", ))
assert cur.fetchone()[0] is True, i
def test_dict(self, conn):
s = sql.SQL("select {f} from {t}").format(
f=sql.Identifier("field"), t=sql.Identifier("table")
)
s1 = s.as_string(conn)
assert isinstance(s1, str)
assert s1 == 'select "field" from "table"'
def test_quote_int(conn, val, expr):
tx = Transformer()
assert tx.get_dumper(val, Format.TEXT).quote(val) == expr
cur = conn.cursor()
cur.execute(sql.SQL("select {v}, -{v}").format(v=sql.Literal(val)))
assert cur.fetchone() == (val, -val)
def test_inet_load_network(conn, fmt_out, val):
pyval = ipaddress.ip_interface(val)
cur = conn.cursor(binary=fmt_out)
cur.execute("select %s::inet", (val, ))
assert cur.fetchone()[0] == pyval
cur.execute("select array[null, %s::inet]", (val, ))
assert cur.fetchone()[0] == [None, pyval]
stmt = sql.SQL("copy (select {}::inet) to stdout (format {})").format(
val, sql.SQL(fmt_out.name))
with cur.copy(stmt) as copy:
copy.set_types(["inet"])
(got, ) = copy.read_row()
assert got == pyval
def test_inet_load_address(conn, fmt_out, val):
addr = ipaddress.ip_address(val.split("/", 1)[0])
cur = conn.cursor(binary=fmt_out)
cur.execute("select %s::inet", (val, ))
assert cur.fetchone()[0] == addr
cur.execute("select array[null, %s::inet]", (val, ))
assert cur.fetchone()[0] == [None, addr]
stmt = sql.SQL("copy (select {}::inet) to stdout (format {})").format(
val, sql.SQL(fmt_out.name))
with cur.copy(stmt) as copy:
copy.set_types(["inet"])
(got, ) = copy.read_row()
assert got == addr
def test_quote_none(conn):
tx = Transformer()
assert tx.get_dumper(None, Format.TEXT).quote(None) == b"NULL"
cur = conn.cursor()
cur.execute(sql.SQL("select {v}").format(v=sql.Literal(None)))
assert cur.fetchone()[0] is None
def test_quote_1char(conn):
cur = conn.cursor()
query = sql.SQL("select {ch} = chr(%s::int)")
for i in range(1, 256):
if chr(i) == "%":
continue
cur.execute(query.format(ch=sql.Literal(chr(i))), (i, ))
assert cur.fetchone()[0] is True, chr(i)
def _make_limit_clause(limit: int) -> psql.Composed:
"""Create a limit clause in psql syntax if limit != -1, otherwise return empty clause"""
if limit != -1:
if not isinstance(limit, int):
raise TypeError(
f"'limit' must be a positive integer. Got {limit}")
return psql.SQL(" LIMIT {limit}").format(limit=limit)
return psql.Composed([])