1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961 962 963 964 965 966 967 968 969 970 971 972 973 974 975 976 977 978 979 980 981 982 983 984 985 986 987 988 989 990 991 992 993 994 995 996 997 998 999 1000 1001 1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015 1016 1017 1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036 1037 1038 1039 1040 1041 1042 1043 1044 1045 1046 1047 1048 1049 1050 1051 1052 1053 1054 1055 1056 1057 1058 1059 1060 1061 1062 1063 1064 1065 1066 1067 1068 1069 1070 1071 1072 1073 1074 1075 1076 1077 1078 1079 1080 1081 1082 1083 1084 1085 1086 1087 1088 1089 1090 1091 1092 1093 1094 1095 1096 1097 1098 1099 1100 1101 1102 1103 1104 1105 1106 1107 1108 1109 1110 1111 1112 1113 1114 1115 1116 1117 1118 1119 1120 1121 1122 1123 1124 1125 1126 1127 1128 1129 1130 1131 1132 1133 1134 1135 1136 1137 1138 1139 1140 1141 1142 1143 1144 1145 1146 1147 1148 1149 1150 1151 1152 1153 1154 1155 1156 1157 1158 1159 1160 1161 1162 1163 1164 1165 1166 1167 1168 1169 1170 1171 1172 1173 1174 1175 1176 1177 1178 1179 1180 1181 1182 1183 1184 1185 1186 1187 1188 1189 1190 1191 1192 1193 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206 1207 1208 1209 1210 1211 1212 1213 1214 1215 1216 1217 1218 1219 1220 1221 1222 1223 1224 1225 1226 1227 1228 1229 1230 1231 1232 1233 1234 1235 1236 1237 1238 1239 1240 1241 1242 1243 1244 1245 1246 1247 1248 1249 1250 1251 1252 1253 1254 1255 1256 1257 1258 1259 1260 1261 1262 1263 1264 1265 1266 1267 1268 1269 1270 1271 1272 1273 1274 1275 1276 1277 1278 1279 1280 1281 1282 1283 1284 1285 1286 1287 1288 1289 1290 1291 1292 1293 1294 1295 1296 1297 1298 1299 1300 1301 1302 1303 1304 1305 1306 1307 1308 1309 1310 1311 1312 1313 1314 1315 1316 1317 1318 1319 1320 1321 1322 1323 1324 1325 1326 1327 1328 1329 1330 1331 1332 1333 1334 1335 1336 1337 1338 1339 1340 1341 1342 1343 1344 1345 1346 1347 1348 1349 1350 1351 1352 1353 1354 1355 1356 1357 1358 1359 1360 1361 1362 1363 1364 1365 1366 1367 1368 1369 1370 1371 1372 1373 1374 1375 1376 1377 1378 1379 1380 1381 1382 1383 1384 1385 1386 1387 1388 1389 1390 1391 1392 1393 1394 1395 1396 1397 1398 1399 1400 1401 1402 1403 1404 1405 1406 1407 1408 1409 1410 1411 1412 1413 1414 1415 1416 1417 1418 1419 1420 1421 1422 1423 1424 1425 1426 1427 1428 1429 1430 1431 1432 1433 1434 1435 1436 1437 1438 1439 1440 1441 1442 1443 1444 1445 1446 1447 1448 1449 1450 1451 1452 1453 1454 1455 1456 1457 1458 1459 1460 1461 1462 1463 1464 1465 1466 1467 1468 1469 1470 1471 1472 1473 1474 1475 1476 1477 1478 1479 1480 1481 1482 1483 1484 1485 1486 1487 1488 1489 1490 1491 1492 1493 1494 1495 1496 1497 1498 1499 1500 1501 1502
|
"""Connection wrapper."""
from __future__ import annotations
from contextlib import suppress
from json import dumps as jsonencode
from json import loads as jsondecode
from operator import itemgetter
from typing import (
TYPE_CHECKING,
Any,
Callable,
Iterator,
Sequence,
TypeVar,
overload,
)
from . import Connection, connect
from .adapt import Adapter, DbTypes
from .attrs import AttrDict
from .core import (
InternalError,
LargeObject,
ProgrammingError,
Query,
get_bool,
get_jsondecode,
unescape_bytea,
)
from .error import db_error, int_error, prg_error
from .helpers import namediter, oid_key, quote_if_unqualified
from .notify import NotificationHandler
if TYPE_CHECKING:
from pgdb.connection import Connection as DbApi2Connection
try:
AnyStr = TypeVar('AnyStr', str, bytes, str | bytes)
except TypeError: # Python < 3.10
AnyStr = Any # type: ignore
__all__ = ['DB']
# The actual PostgreSQL database connection interface:
class DB:
"""Wrapper class for the core connection type."""
dbname: str
host: str
port: int
options: str
error: str
status: int
user : str
protocol_version: int
server_version: int
socket: int
backend_pid: int
ssl_in_use: bool
ssl_attributes: dict[str, str | None]
db: Connection | None = None # invalid fallback for underlying connection
_db_args: Any # either the connect args or the underlying connection
@overload
def __init__(self, dbname: str | None = None,
host: str | None = None, port: int = -1,
opt: str | None = None,
user: str | None = None, passwd: str | None = None,
nowait: bool = False) -> None:
... # create a new connection using the specified parameters
@overload
def __init__(self, db: Connection | DB | DbApi2Connection) -> None:
... # create a connection wrapper based on an existing connection
def __init__(self, *args: Any, **kw: Any) -> None:
"""Create a new connection.
You can pass either the connection parameters or an existing
pg or pgdb Connection. This allows you to use the methods
of the classic pg interface with a DB-API 2 pgdb Connection.
"""
if kw:
db = kw.get('db')
if db is not None and (args or len(kw) > 1):
raise TypeError("Conflicting connection parameters")
elif len(args) == 1 and not isinstance(args[0], str):
db = args[0]
else:
db = None
if db:
if isinstance(db, DB):
db = db.db # allow db to be a wrapped Connection
else:
with suppress(AttributeError):
db = db._cnx # allow db to be a pgdb Connection
if not isinstance(db, Connection):
raise TypeError(
"The 'db' argument must be a valid database connection.")
self._db_args = db
self._closeable = False
else:
db = connect(*args, **kw)
self._db_args = args, kw
self._closeable = True
self.db = db
self.dbname = db.db
self._regtypes = False
self._attnames: dict[str, AttrDict] = {}
self._generated: dict[str, frozenset[str]] = {}
self._pkeys: dict[str, str | tuple[str, ...]] = {}
self._privileges: dict[tuple[str, str], bool] = {}
self.adapter = Adapter(self)
self.dbtypes = DbTypes(self)
self._query_attnames = (
"SELECT a.attname,"
" t.oid, t.typname, t.oid::pg_catalog.regtype,"
" t.typlen, t.typtype, t.typcategory, t.typdelim, t.typrelid"
" FROM pg_catalog.pg_attribute a"
" JOIN pg_catalog.pg_type t"
" ON t.oid OPERATOR(pg_catalog.=) a.atttypid"
" WHERE a.attrelid OPERATOR(pg_catalog.=)"
" {}::pg_catalog.regclass"
" AND {} AND NOT a.attisdropped ORDER BY a.attnum")
if db.server_version < 120000:
self._query_generated = (
"a.attidentity OPERATOR(pg_catalog.=) 'a'"
)
else:
self._query_generated = (
"(a.attidentity OPERATOR(pg_catalog.=) 'a' OR"
" a.attgenerated OPERATOR(pg_catalog.!=) '')"
)
db.set_cast_hook(self.dbtypes.typecast)
# For debugging scripts, self.debug can be set
# * to a string format specification (e.g. in CGI set to "%s<BR>"),
# * to a file object to write debug statements or
# * to a callable object which takes a string argument
# * to any other true value to just print debug statements
self.debug: Any = None
def __getattr__(self, name: str) -> Any:
"""Get the specified attribute of the connection."""
# All undefined members are same as in underlying connection:
if self.db:
return getattr(self.db, name)
else:
raise int_error('Connection is not valid')
def __dir__(self) -> list[str]:
"""List all attributes of the connection."""
# Custom dir function including the attributes of the connection:
attrs = set(self.__class__.__dict__)
attrs.update(self.__dict__)
attrs.update(dir(self.db))
return sorted(attrs)
# Context manager methods
def __enter__(self) -> DB:
"""Enter the runtime context. This will start a transaction."""
self.begin()
return self
def __exit__(self, et: type[BaseException] | None,
ev: BaseException | None, tb: Any) -> None:
"""Exit the runtime context. This will end the transaction."""
if et is None and ev is None and tb is None:
self.commit()
else:
self.rollback()
def __del__(self) -> None:
"""Delete the connection."""
try:
db = self.db
except AttributeError:
db = None
if db:
with suppress(TypeError): # when already closed
db.set_cast_hook(None)
if self._closeable:
with suppress(InternalError): # when already closed
db.close()
# Auxiliary methods
def _do_debug(self, *args: Any) -> None:
"""Print a debug message."""
if self.debug:
s = '\n'.join(str(arg) for arg in args)
if isinstance(self.debug, str):
print(self.debug % s)
elif hasattr(self.debug, 'write'):
# noinspection PyCallingNonCallable
self.debug.write(s + '\n')
elif callable(self.debug):
self.debug(s)
else:
print(s)
def _escape_qualified_name(self, s: str) -> str:
"""Escape a qualified name.
Escapes the name for use as an SQL identifier, unless the
name contains a dot, in which case the name is ambiguous
(could be a qualified name or just a name with a dot in it)
and must be quoted manually by the caller.
"""
if '.' not in s:
s = self.escape_identifier(s)
return s
@staticmethod
def _make_bool(d: Any) -> bool | str:
"""Get boolean value corresponding to d."""
return bool(d) if get_bool() else ('t' if d else 'f')
@staticmethod
def _list_params(params: Sequence) -> str:
"""Create a human readable parameter list."""
return ', '.join(f'${n}={v!r}' for n, v in enumerate(params, 1))
@property
def _valid_db(self) -> Connection:
"""Get underlying connection and make sure it is not closed."""
db = self.db
if not db:
raise int_error('Connection already closed')
return db
# Public methods
# escape_string and escape_bytea exist as methods,
# so we define unescape_bytea as a method as well
unescape_bytea = staticmethod(unescape_bytea)
@staticmethod
def decode_json(s: str) -> Any:
"""Decode a JSON string coming from the database."""
return (get_jsondecode() or jsondecode)(s)
@staticmethod
def encode_json(d: Any) -> str:
"""Encode a JSON string for use within SQL."""
return jsonencode(d)
def close(self) -> None:
"""Close the database connection."""
# Wraps shared library function so we can track state.
db = self._valid_db
with suppress(TypeError): # when already closed
db.set_cast_hook(None)
if self._closeable:
db.close()
self.db = None
def reset(self) -> None:
"""Reset connection with current parameters.
All derived queries and large objects derived from this connection
will not be usable after this call.
"""
self._valid_db.reset()
def reopen(self) -> None:
"""Reopen connection to the database.
Used in case we need another connection to the same database.
Note that we can still reopen a database that we have closed.
"""
# There is no such shared library function.
if self._closeable:
args, kw = self._db_args
db = connect(*args, **kw)
if self.db:
self.db.set_cast_hook(None)
self.db.close()
db.set_cast_hook(self.dbtypes.typecast)
self.db = db
else:
self.db = self._db_args
def begin(self, mode: str | None = None) -> Query:
"""Begin a transaction."""
qstr = 'BEGIN'
if mode:
qstr += ' ' + mode
return self.query(qstr)
start = begin
def commit(self) -> Query:
"""Commit the current transaction."""
return self.query('COMMIT')
end = commit
def rollback(self, name: str | None = None) -> Query:
"""Roll back the current transaction."""
qstr = 'ROLLBACK'
if name:
qstr += ' TO ' + name
return self.query(qstr)
abort = rollback
def savepoint(self, name: str) -> Query:
"""Define a new savepoint within the current transaction."""
return self.query('SAVEPOINT ' + name)
def release(self, name: str) -> Query:
"""Destroy a previously defined savepoint."""
return self.query('RELEASE ' + name)
def get_parameter(self,
parameter: str | list[str] | tuple[str, ...] |
set[str] | frozenset[str] | dict[str, Any]
) -> str | list[str] | dict[str, str]:
"""Get the value of a run-time parameter.
If the parameter is a string, the return value will also be a string
that is the current setting of the run-time parameter with that name.
You can get several parameters at once by passing a list, set or dict.
When passing a list of parameter names, the return value will be a
corresponding list of parameter settings. When passing a set of
parameter names, a new dict will be returned, mapping these parameter
names to their settings. Finally, if you pass a dict as parameter,
its values will be set to the current parameter settings corresponding
to its keys.
By passing the special name 'all' as the parameter, you can get a dict
of all existing configuration parameters.
"""
values: Any
if isinstance(parameter, str):
parameter = [parameter]
values = None
elif isinstance(parameter, (list, tuple)):
values = []
elif isinstance(parameter, (set, frozenset)):
values = {}
elif isinstance(parameter, dict):
values = parameter
else:
raise TypeError(
'The parameter must be a string, list, set or dict')
if not parameter:
raise TypeError('No parameter has been specified')
query = self._valid_db.query
params: Any = {} if isinstance(values, dict) else []
for param_key in parameter:
param = param_key.strip().lower() if isinstance(
param_key, (bytes, str)) else None
if not param:
raise TypeError('Invalid parameter')
if param == 'all':
cmd = 'SHOW ALL'
values = query(cmd).getresult()
values = {value[0]: value[1] for value in values}
break
if isinstance(params, dict):
params[param] = param_key
else:
params.append(param)
else:
for param in params:
cmd = f'SHOW {param}'
value = query(cmd).singlescalar()
if values is None:
values = value
elif isinstance(values, list):
values.append(value)
else:
values[params[param]] = value
return values
def set_parameter(self,
parameter: str | list[str] | tuple[str, ...] |
set[str] | frozenset[str] | dict[str, Any],
value: str | list[str] | tuple[str, ...] |
set[str] | frozenset[str]| None = None,
local: bool = False) -> None:
"""Set the value of a run-time parameter.
If the parameter and the value are strings, the run-time parameter
will be set to that value. If no value or None is passed as a value,
then the run-time parameter will be restored to its default value.
You can set several parameters at once by passing a list of parameter
names, together with a single value that all parameters should be
set to or with a corresponding list of values. You can also pass
the parameters as a set if you only provide a single value.
Finally, you can pass a dict with parameter names as keys. In this
case, you should not pass a value, since the values for the parameters
will be taken from the dict.
By passing the special name 'all' as the parameter, you can reset
all existing settable run-time parameters to their default values.
If you set local to True, then the command takes effect for only the
current transaction. After commit() or rollback(), the session-level
setting takes effect again. Setting local to True will appear to
have no effect if it is executed outside a transaction, since the
transaction will end immediately.
"""
if isinstance(parameter, str):
parameter = {parameter: value}
elif isinstance(parameter, (list, tuple)):
if isinstance(value, (list, tuple)):
parameter = dict(zip(parameter, value))
else:
parameter = dict.fromkeys(parameter, value)
elif isinstance(parameter, (set, frozenset)):
if isinstance(value, (list, tuple, set, frozenset)):
value = set(value)
if len(value) == 1:
value = next(iter(value))
if not (value is None or isinstance(value, str)):
raise ValueError(
'A single value must be specified'
' when parameter is a set')
parameter = dict.fromkeys(parameter, value)
elif isinstance(parameter, dict):
if value is not None:
raise ValueError(
'A value must not be specified'
' when parameter is a dictionary')
else:
raise TypeError(
'The parameter must be a string, list, set or dict')
if not parameter:
raise TypeError('No parameter has been specified')
params: dict[str, str | None] = {}
for param, param_value in parameter.items():
param = param.strip().lower() if isinstance(param, str) else None
if not param:
raise TypeError('Invalid parameter')
if param == 'all':
if param_value is not None:
raise ValueError(
'A value must not be specified'
" when parameter is 'all'")
params = {'all': None}
break
params[param] = param_value
local_clause = ' LOCAL' if local else ''
for param, param_value in params.items():
cmd = (f'RESET{local_clause} {param}'
if param_value is None else
f'SET{local_clause} {param} TO {param_value}')
self._do_debug(cmd)
self._valid_db.query(cmd)
def query(self, command: str, *args: Any) -> Query:
"""Execute a SQL command string.
This method simply sends a SQL query to the database. If the query is
an insert statement that inserted exactly one row into a table that
has OIDs, the return value is the OID of the newly inserted row.
If the query is an update or delete statement, or an insert statement
that did not insert exactly one row in a table with OIDs, then the
number of rows affected is returned as a string. If it is a statement
that returns rows as a result (usually a select statement, but maybe
also an "insert/update ... returning" statement), this method returns
a Query object that can be accessed via getresult() or dictresult()
or simply printed. Otherwise, it returns `None`.
The query can contain numbered parameters of the form $1 in place
of any data constant. Arguments given after the query string will
be substituted for the corresponding numbered parameter. Parameter
values can also be given as a single list or tuple argument.
"""
# Wraps shared library function for debugging.
db = self._valid_db
if args:
self._do_debug(command, args)
return db.query(command, args)
self._do_debug(command)
return db.query(command)
def query_formatted(self, command: str,
parameters: tuple | list | dict | None = None,
types: tuple | list | dict | None = None,
inline: bool =False) -> Query:
"""Execute a formatted SQL command string.
Similar to query, but using Python format placeholders of the form
%s or %(names)s instead of PostgreSQL placeholders of the form $1.
The parameters must be passed as a tuple, list or dict. You can
also pass a corresponding tuple, list or dict of database types in
order to format the parameters properly in case there is ambiguity.
If you set inline to True, the parameters will be sent to the database
embedded in the SQL command, otherwise they will be sent separately.
"""
return self.query(*self.adapter.format_query(
command, parameters, types, inline))
def query_prepared(self, name: str, *args: Any) -> Query:
"""Execute a prepared SQL statement.
This works like the query() method, except that instead of passing
the SQL command, you pass the name of a prepared statement. If you
pass an empty name, the unnamed statement will be executed.
"""
if name is None:
name = ''
db = self._valid_db
if args:
self._do_debug('EXECUTE', name, args)
return db.query_prepared(name, args)
self._do_debug('EXECUTE', name)
return db.query_prepared(name)
def prepare(self, name: str, command: str) -> None:
"""Create a prepared SQL statement.
This creates a prepared statement for the given command with the
given name for later execution with the query_prepared() method.
The name can be empty to create an unnamed statement, in which case
any pre-existing unnamed statement is automatically replaced;
otherwise it is an error if the statement name is already
defined in the current database session. We recommend always using
named queries, since unnamed queries have a limited lifetime and
can be automatically replaced or destroyed by various operations.
"""
if name is None:
name = ''
self._do_debug('prepare', name, command)
self._valid_db.prepare(name, command)
def describe_prepared(self, name: str | None = None) -> Query:
"""Describe a prepared SQL statement.
This method returns a Query object describing the result columns of
the prepared statement with the given name. If you omit the name,
the unnamed statement will be described if you created one before.
"""
if name is None:
name = ''
return self._valid_db.describe_prepared(name)
def delete_prepared(self, name: str | None = None) -> Query:
"""Delete a prepared SQL statement.
This deallocates a previously prepared SQL statement with the given
name, or deallocates all prepared statements if you do not specify a
name. Note that prepared statements are also deallocated automatically
when the current session ends.
"""
if not name:
name = 'ALL'
cmd = f"DEALLOCATE {name}"
self._do_debug(cmd)
return self._valid_db.query(cmd)
def pkey(self, table: str, composite: bool = False, flush: bool = False
) -> str | tuple[str, ...]:
"""Get the primary key of a table.
Single primary keys are returned as strings unless you
set the composite flag. Composite primary keys are always
represented as tuples. Note that this raises a KeyError
if the table does not have a primary key.
If flush is set then the internal cache for primary keys will
be flushed. This may be necessary after the database schema or
the search path has been changed.
"""
pkeys = self._pkeys
if flush:
pkeys.clear()
self._do_debug('The pkey cache has been flushed')
try: # cache lookup
pkey = pkeys[table]
except KeyError as e: # cache miss, check the database
cmd = ("SELECT" # noqa: S608
" a.attname, a.attnum, i.indkey"
" FROM pg_catalog.pg_index i"
" JOIN pg_catalog.pg_attribute a"
" ON a.attrelid OPERATOR(pg_catalog.=) i.indrelid"
" AND a.attnum OPERATOR(pg_catalog.=) ANY(i.indkey)"
" AND NOT a.attisdropped"
" WHERE i.indrelid OPERATOR(pg_catalog.=)"
" {}::pg_catalog.regclass"
" AND i.indisprimary ORDER BY a.attnum").format(
quote_if_unqualified('$1', table))
res = self._valid_db.query(cmd, (table,)).getresult()
if not res:
raise KeyError(f'Table {table} has no primary key') from e
# we want to use the order defined in the primary key index here,
# not the order as defined by the columns in the table
if len(res) > 1:
indkey = res[0][2]
pkey = tuple(row[0] for row in sorted(
res, key=lambda row: indkey.index(row[1])))
else:
pkey = res[0][0]
pkeys[table] = pkey # cache it
if composite and not isinstance(pkey, tuple):
pkey = (pkey,)
return pkey
def pkeys(self, table: str) -> tuple[str, ...]:
"""Get the primary key of a table as a tuple.
Same as pkey() with 'composite' set to True.
"""
return self.pkey(table, True) # type: ignore
def get_databases(self) -> list[str]:
"""Get list of databases in the system."""
return [r[0] for r in self._valid_db.query(
'SELECT datname FROM pg_catalog.pg_database').getresult()]
def get_relations(self, kinds: str | Sequence[str] | None = None,
system: bool = False) -> list[str]:
"""Get list of relations in connected database of specified kinds.
If kinds is None or empty, all kinds of relations are returned.
Otherwise, kinds can be a string or sequence of type letters
specifying which kind of relations you want to list.
Set the system flag if you want to get the system relations as well.
"""
where_parts = []
if kinds:
where_parts.append(
"r.relkind IN ({})".format(','.join(f"'{k}'" for k in kinds)))
if not system:
where_parts.append("s.nspname NOT SIMILAR"
" TO 'pg/_%|information/_schema' ESCAPE '/'")
where = " WHERE " + ' AND '.join(where_parts) if where_parts else ''
cmd = ("SELECT" # noqa: S608
" pg_catalog.quote_ident(s.nspname) OPERATOR(pg_catalog.||)"
" '.' OPERATOR(pg_catalog.||) pg_catalog.quote_ident(r.relname)"
" FROM pg_catalog.pg_class r"
" JOIN pg_catalog.pg_namespace s"
f" ON s.oid OPERATOR(pg_catalog.=) r.relnamespace{where}"
" ORDER BY s.nspname, r.relname")
return [r[0] for r in self._valid_db.query(cmd).getresult()]
def get_tables(self, system: bool = False) -> list[str]:
"""Return list of tables in connected database.
Set the system flag if you want to get the system tables as well.
"""
return self.get_relations('r', system)
def get_attnames(self, table: str, with_oid: bool=True, flush: bool=False
) -> AttrDict:
"""Given the name of a table, dig out the set of attribute names.
Returns a read-only dictionary of attribute names (the names are
the keys, the values are the names of the attributes' types)
with the column names in the proper order if you iterate over it.
If flush is set, then the internal cache for attribute names will
be flushed. This may be necessary after the database schema or
the search path has been changed.
By default, only a limited number of simple types will be returned.
You can get the registered types after calling use_regtypes(True).
"""
attnames = self._attnames
if flush:
attnames.clear()
self._do_debug('The attnames cache has been flushed')
try: # cache lookup
names = attnames[table]
except KeyError: # cache miss, check the database
cmd = "a.attnum OPERATOR(pg_catalog.>) 0"
if with_oid:
cmd = f"({cmd} OR a.attname OPERATOR(pg_catalog.=) 'oid')"
cmd = self._query_attnames.format(
quote_if_unqualified('$1', table), cmd)
res = self._valid_db.query(cmd, (table,)).getresult()
types = self.dbtypes
names = AttrDict((name[0], types.add(*name[1:])) for name in res)
attnames[table] = names # cache it
return names
def get_generated(self, table: str, flush: bool = False) -> frozenset[str]:
"""Given the name of a table, dig out the set of generated columns.
Returns a set of column names that are generated and unalterable.
If flush is set, then the internal cache for generated columns will
be flushed. This may be necessary after the database schema or
the search path has been changed.
"""
generated = self._generated
if flush:
generated.clear()
self._do_debug('The generated cache has been flushed')
try: # cache lookup
names = generated[table]
except KeyError: # cache miss, check the database
cmd = "a.attnum OPERATOR(pg_catalog.>) 0"
cmd = f"{cmd} AND {self._query_generated}"
cmd = self._query_attnames.format(
quote_if_unqualified('$1', table), cmd)
res = self._valid_db.query(cmd, (table,)).getresult()
names = frozenset(name[0] for name in res)
generated[table] = names # cache it
return names
def use_regtypes(self, regtypes: bool | None = None) -> bool:
"""Use registered type names instead of simplified type names."""
if regtypes is None:
return self.dbtypes._regtypes
regtypes = bool(regtypes)
if regtypes != self.dbtypes._regtypes:
self.dbtypes._regtypes = regtypes
self._attnames.clear()
self.dbtypes.clear()
return regtypes
def has_table_privilege(self, table: str, privilege: str = 'select',
flush: bool = False) -> bool:
"""Check whether current user has specified table privilege.
If flush is set, then the internal cache for table privileges will
be flushed. This may be necessary after privileges have been changed.
"""
privileges = self._privileges
if flush:
privileges.clear()
self._do_debug('The privileges cache has been flushed')
privilege = privilege.lower()
try: # ask cache
ret = privileges[table, privilege]
except KeyError: # cache miss, ask the database
cmd = "SELECT pg_catalog.has_table_privilege({}, $2)".format(
quote_if_unqualified('$1', table))
query = self._valid_db.query(cmd, (table, privilege))
ret = query.singlescalar() == self._make_bool(True)
privileges[table, privilege] = ret # cache it
return ret
def get(self, table: str, row: Any,
keyname: str | tuple[str, ...] | None = None) -> dict[str, Any]:
"""Get a row from a database table or view.
This method is the basic mechanism to get a single row. It assumes
that the keyname specifies a unique row. It must be the name of a
single column or a tuple of column names. If the keyname is not
specified, then the primary key for the table is used.
If row is a dictionary, then the value for the key is taken from it.
Otherwise, the row must be a single value or a tuple of values
corresponding to the passed keyname or primary key. The fetched row
from the table will be returned as a new dictionary or used to replace
the existing values when row was passed as a dictionary.
The OID is also put into the dictionary if the table has one, but
in order to allow the caller to work with multiple tables, it is
munged as "oid(table)" using the actual name of the table.
"""
if table.endswith('*'): # hint for descendant tables can be ignored
table = table[:-1].rstrip()
attnames = self.get_attnames(table)
qoid = oid_key(table) if 'oid' in attnames else None
if keyname and isinstance(keyname, str):
keyname = (keyname,)
if qoid and isinstance(row, dict) and qoid in row and 'oid' not in row:
row['oid'] = row[qoid]
if not keyname:
try: # if keyname is not specified, try using the primary key
keyname = self.pkeys(table)
except KeyError as e: # the table has no primary key
# try using the oid instead
if qoid and isinstance(row, dict) and 'oid' in row:
keyname = ('oid',)
else:
raise prg_error(
f'Table {table} has no primary key') from e
else: # the table has a primary key
# check whether all key columns have values
if isinstance(row, dict) and not set(keyname).issubset(row):
# try using the oid instead
if qoid and 'oid' in row:
keyname = ('oid',)
else:
raise KeyError(
'Missing value in row for specified keyname')
if not isinstance(row, dict):
if not isinstance(row, (tuple, list)):
row = [row]
if len(keyname) != len(row):
raise KeyError(
'Differing number of items in keyname and row')
row = dict(zip(keyname, row))
params = self.adapter.parameter_list()
adapt = params.add
col = self.escape_identifier
what = 'oid, *' if qoid else '*'
where = ' AND '.join(
f'{col(k)} OPERATOR(pg_catalog.=) {adapt(row[k], attnames[k])}'
for k in keyname)
if 'oid' in row:
if qoid:
row[qoid] = row['oid']
del row['oid']
t = self._escape_qualified_name(table)
cmd = f'SELECT {what} FROM {t} WHERE {where} LIMIT 1' # noqa: S608s
self._do_debug(cmd, params)
query = self._valid_db.query(cmd, params)
res = query.dictresult()
if not res:
# make where clause in error message better readable
where = where.replace('OPERATOR(pg_catalog.=)', '=')
raise db_error(
f'No such record in {table}\nwhere {where}\nwith '
+ self._list_params(params))
for n, value in res[0].items():
if qoid and n == 'oid':
n = qoid
row[n] = value
return row
def insert(self, table: str, row: dict[str, Any] | None = None, **kw: Any
) -> dict[str, Any]:
"""Insert a row into a database table.
This method inserts a row into a table. The name of the table must
be passed as the first parameter. The other parameters are used for
providing the data of the row that shall be inserted into the table.
If a dictionary is supplied as the second parameter, it starts with
that. Otherwise, it uses a blank dictionary.
Either way the dictionary is updated from the keywords.
The dictionary is then reloaded with the values actually inserted in
order to pick up values modified by rules, triggers, etc.
"""
if table.endswith('*'): # hint for descendant tables can be ignored
table = table[:-1].rstrip()
if row is None:
row = {}
row.update(kw)
if 'oid' in row:
del row['oid'] # do not insert oid
attnames = self.get_attnames(table)
generated = self.get_generated(table)
qoid = oid_key(table) if 'oid' in attnames else None
params = self.adapter.parameter_list()
adapt = params.add
col = self.escape_identifier
name_list, value_list = [], []
for n in attnames:
if n in row and n not in generated:
name_list.append(col(n))
value_list.append(adapt(row[n], attnames[n]))
if not name_list:
raise prg_error('No column found that can be inserted')
names, values = ', '.join(name_list), ', '.join(value_list)
ret = 'oid, *' if qoid else '*'
t = self._escape_qualified_name(table)
cmd = (f'INSERT INTO {t} ({names})' # noqa: S608
f' VALUES ({values}) RETURNING {ret}')
self._do_debug(cmd, params)
query = self._valid_db.query(cmd, params)
res = query.dictresult()
if res: # this should always be true
for n, value in res[0].items():
if qoid and n == 'oid':
n = qoid
row[n] = value
return row
def update(self, table: str, row: dict[str, Any] | None = None, **kw : Any
) -> dict[str, Any]:
"""Update an existing row in a database table.
Similar to insert, but updates an existing row. The update is based
on the primary key of the table or the OID value as munged by get()
or passed as keyword. The OID will take precedence if provided, so
that it is possible to update the primary key itself.
The dictionary is then modified to reflect any changes caused by the
update due to triggers, rules, default values, etc.
"""
if table.endswith('*'):
table = table[:-1].rstrip() # need parent table name
attnames = self.get_attnames(table)
generated = self.get_generated(table)
qoid = oid_key(table) if 'oid' in attnames else None
if row is None:
row = {}
elif 'oid' in row:
del row['oid'] # only accept oid key from named args for safety
row.update(kw)
if qoid and qoid in row and 'oid' not in row:
row['oid'] = row[qoid]
if qoid and 'oid' in row: # try using the oid
keynames: tuple[str, ...] = ('oid',)
keyset = set(keynames)
else: # try using the primary key
try:
keynames = self.pkeys(table)
except KeyError as e: # the table has no primary key
raise prg_error(f'Table {table} has no primary key') from e
keyset = set(keynames)
# check whether all key columns have values
if not keyset.issubset(row):
raise KeyError('Missing value for primary key in row')
params = self.adapter.parameter_list()
adapt = params.add
col = self.escape_identifier
where = ' AND '.join(
f'{col(k)} OPERATOR(pg_catalog.=) {adapt(row[k], attnames[k])}'
for k in keynames)
if 'oid' in row:
if qoid:
row[qoid] = row['oid']
del row['oid']
values_list = []
for n in attnames:
if n in row and n not in keyset and n not in generated:
values_list.append(f'{col(n)} = {adapt(row[n], attnames[n])}')
if not values_list:
return row
values = ', '.join(values_list)
ret = 'oid, *' if qoid else '*'
t = self._escape_qualified_name(table)
cmd = (f'UPDATE {t} SET {values}' # noqa: S608
f' WHERE {where} RETURNING {ret}')
self._do_debug(cmd, params)
query = self._valid_db.query(cmd, params)
res = query.dictresult()
if res: # may be empty when row does not exist
for n, value in res[0].items():
if qoid and n == 'oid':
n = qoid
row[n] = value
return row
def upsert(self, table: str, row: dict[str, Any] | None = None, **kw: Any
) -> dict[str, Any]:
"""Insert a row into a database table with conflict resolution.
This method inserts a row into a table, but instead of raising a
ProgrammingError exception in case a row with the same primary key
already exists, an update will be executed instead. This will be
performed as a single atomic operation on the database, so race
conditions can be avoided.
Like the insert method, the first parameter is the name of the
table and the second parameter can be used to pass the values to
be inserted as a dictionary.
Unlike the insert und update statement, keyword parameters are not
used to modify the dictionary, but to specify which columns shall
be updated in case of a conflict, and in which way:
A value of False or None means the column shall not be updated,
a value of True means the column shall be updated with the value
that has been proposed for insertion, i.e. has been passed as value
in the dictionary. Columns that are not specified by keywords but
appear as keys in the dictionary are also updated like in the case
keywords had been passed with the value True.
So if in the case of a conflict you want to update every column
that has been passed in the dictionary row, you would call
upsert(table, row). If you don't want to do anything in case
of a conflict, i.e. leave the existing row as it is, call
upsert(table, row, **dict.fromkeys(row)).
If you need more fine-grained control of what gets updated, you can
also pass strings in the keyword parameters. These strings will
be used as SQL expressions for the update columns. In these
expressions you can refer to the value that already exists in
the table by prefixing the column name with "included.", and to
the value that has been proposed for insertion by prefixing the
column name with the "excluded."
The dictionary is modified in any case to reflect the values in
the database after the operation has completed.
Note: The method uses the PostgreSQL "upsert" feature which is
only available since PostgreSQL 9.5.
"""
if table.endswith('*'): # hint for descendant tables can be ignored
table = table[:-1].rstrip()
if row is None:
row = {}
if 'oid' in row:
del row['oid'] # do not insert oid
if 'oid' in kw:
del kw['oid'] # do not update oid
attnames = self.get_attnames(table)
generated = self.get_generated(table)
qoid = oid_key(table) if 'oid' in attnames else None
params = self.adapter.parameter_list()
adapt = params.add
col = self.escape_identifier
name_list, value_list = [], []
for n in attnames:
if n in row and n not in generated:
name_list.append(col(n))
value_list.append(adapt(row[n], attnames[n]))
names, values = ', '.join(name_list), ', '.join(value_list)
try:
keynames = self.pkeys(table)
except KeyError as e:
raise prg_error(f'Table {table} has no primary key') from e
target = ', '.join(col(k) for k in keynames)
update = []
keyset = set(keynames)
keyset.add('oid')
for n in attnames:
if n not in keyset and n not in generated:
value = kw.get(n, n in row)
if value:
if not isinstance(value, str):
value = f'excluded.{col(n)}'
update.append(f'{col(n)} = {value}')
if not values:
return row
do = 'update set ' + ', '.join(update) if update else 'nothing'
ret = 'oid, *' if qoid else '*'
t = self._escape_qualified_name(table)
cmd = (f'INSERT INTO {t} AS included ({names})' # noqa: S608
f' VALUES ({values})'
f' ON CONFLICT ({target}) DO {do} RETURNING {ret}')
self._do_debug(cmd, params)
query = self._valid_db.query(cmd, params)
res = query.dictresult()
if res: # may be empty with "do nothing"
for n, value in res[0].items():
if qoid and n == 'oid':
n = qoid
row[n] = value
else:
self.get(table, row)
return row
def clear(self, table: str, row: dict[str, Any] | None = None
) -> dict[str, Any]:
"""Clear all the attributes to values determined by the types.
Numeric types are set to 0, Booleans are set to false, and everything
else is set to the empty string. If the row argument is present,
it is used as the row dictionary and any entries matching attribute
names are cleared with everything else left unchanged.
"""
# At some point we will need a way to get defaults from a table.
if row is None:
row = {} # empty if argument is not present
attnames = self.get_attnames(table)
for n, t in attnames.items():
if n == 'oid':
continue
t = t.simple
if t in DbTypes._num_types:
row[n] = 0
elif t == 'bool':
row[n] = self._make_bool(False)
else:
row[n] = ''
return row
def delete(self, table: str, row: dict[str, Any] | None = None, **kw: Any
) -> int:
"""Delete an existing row in a database table.
This method deletes the row from a table. It deletes based on the
primary key of the table or the OID value as munged by get() or
passed as keyword. The OID will take precedence if provided.
The return value is the number of deleted rows (i.e. 0 if the row
did not exist and 1 if the row was deleted).
Note that if the row cannot be deleted because e.g. it is still
referenced by another table, this method raises a ProgrammingError.
"""
if table.endswith('*'): # hint for descendant tables can be ignored
table = table[:-1].rstrip()
attnames = self.get_attnames(table)
qoid = oid_key(table) if 'oid' in attnames else None
if row is None:
row = {}
elif 'oid' in row:
del row['oid'] # only accept oid key from named args for safety
row.update(kw)
if qoid and qoid in row and 'oid' not in row:
row['oid'] = row[qoid]
if qoid and 'oid' in row: # try using the oid
keynames: tuple[str, ...] = ('oid',)
else: # try using the primary key
try:
keynames = self.pkeys(table)
except KeyError as e: # the table has no primary key
raise prg_error(f'Table {table} has no primary key') from e
# check whether all key columns have values
if not set(keynames).issubset(row):
raise KeyError('Missing value for primary key in row')
params = self.adapter.parameter_list()
adapt = params.add
col = self.escape_identifier
where = ' AND '.join(
f'{col(k)} OPERATOR(pg_catalog.=) {adapt(row[k], attnames[k])}'
for k in keynames)
if 'oid' in row:
if qoid:
row[qoid] = row['oid']
del row['oid']
t = self._escape_qualified_name(table)
cmd = f'DELETE FROM {t} WHERE {where}' # noqa: S608
self._do_debug(cmd, params)
res = self._valid_db.query(cmd, params)
return int(res) # type: ignore
def truncate(self, table: str | list[str] | tuple[str, ...] |
set[str] | frozenset[str], restart: bool = False,
cascade: bool = False, only: bool = False) -> Query:
"""Empty a table or set of tables.
This method quickly removes all rows from the given table or set
of tables. It has the same effect as an unqualified DELETE on each
table, but since it does not actually scan the tables it is faster.
Furthermore, it reclaims disk space immediately, rather than requiring
a subsequent VACUUM operation. This is most useful on large tables.
If restart is set to True, sequences owned by columns of the truncated
table(s) are automatically restarted. If cascade is set to True, it
also truncates all tables that have foreign-key references to any of
the named tables. If the parameter 'only' is not set to True, all the
descendant tables (if any) will also be truncated. Optionally, a '*'
can be specified after the table name to explicitly indicate that
descendant tables are included.
"""
if isinstance(table, str):
table_only = {table: only}
table = [table]
elif isinstance(table, (list, tuple)):
if isinstance(only, (list, tuple)):
table_only = dict(zip(table, only))
else:
table_only = dict.fromkeys(table, only)
elif isinstance(table, (set, frozenset)):
table_only = dict.fromkeys(table, only)
else:
raise TypeError('The table must be a string, list or set')
if not (restart is None or isinstance(restart, (bool, int))):
raise TypeError('Invalid type for the restart option')
if not (cascade is None or isinstance(cascade, (bool, int))):
raise TypeError('Invalid type for the cascade option')
tables = []
for t in table:
u = table_only.get(t)
if not (u is None or isinstance(u, (bool, int))):
raise TypeError('Invalid type for the only option')
if t.endswith('*'):
if u:
raise ValueError(
'Contradictory table name and only options')
t = t[:-1].rstrip()
t = self._escape_qualified_name(t)
if u:
t = f'ONLY {t}'
tables.append(t)
cmd_parts = ['TRUNCATE', ', '.join(tables)]
if restart:
cmd_parts.append('RESTART IDENTITY')
if cascade:
cmd_parts.append('CASCADE')
cmd = ' '.join(cmd_parts)
self._do_debug(cmd)
return self._valid_db.query(cmd)
def get_as_list(
self, table: str,
what: str | list[str] | tuple[str, ...] | None = None,
where: str | list[str] | tuple[str, ...] | None = None,
order: str | list[str] | tuple[str, ...] | bool | None = None,
limit: int | None = None, offset: int | None = None,
scalar: bool = False) -> list:
"""Get a table as a list.
This gets a convenient representation of the table as a list
of named tuples in Python. You only need to pass the name of
the table (or any other SQL expression returning rows). Note that
by default this will return the full content of the table which
can be huge and overflow your memory. However, you can control
the amount of data returned using the other optional parameters.
The parameter 'what' can restrict the query to only return a
subset of the table columns. It can be a string, list or a tuple.
The parameter 'where' can restrict the query to only return a
subset of the table rows. It can be a string, list or a tuple
of SQL expressions that all need to be fulfilled.
The parameter 'order' specifies the ordering of the rows. It can
also be a string, list or a tuple. If no ordering is specified,
the result will be ordered by the primary key(s) or all columns if
no primary key exists. You can set 'order' to False if you don't
care about the ordering. The parameters 'limit' and 'offset' can be
integers specifying the maximum number of rows returned and a number
of rows skipped over.
If you set the 'scalar' option to True, then instead of the
named tuples you will get the first items of these tuples.
This is useful if the result has only one column anyway.
"""
if not table:
raise TypeError('The table name is missing')
if what:
if isinstance(what, (list, tuple)):
what = ', '.join(map(str, what))
if order is None:
order = what
else:
what = '*'
cmd_parts = ['SELECT', what, 'FROM', table]
if where:
if isinstance(where, (list, tuple)):
where = ' AND '.join(map(str, where))
cmd_parts.extend(['WHERE', where])
if order is None or order is True:
try:
order = self.pkeys(table)
except (KeyError, ProgrammingError):
with suppress(KeyError, ProgrammingError):
order = list(self.get_attnames(table))
if order and not isinstance(order, bool):
if isinstance(order, (list, tuple)):
order = ', '.join(map(str, order))
cmd_parts.extend(['ORDER BY', order])
if limit:
cmd_parts.append(f'LIMIT {limit}')
if offset:
cmd_parts.append(f'OFFSET {offset}')
cmd = ' '.join(cmd_parts)
self._do_debug(cmd)
query = self._valid_db.query(cmd)
res = query.namedresult()
if res and scalar:
res = [row[0] for row in res]
return res
def get_as_dict(
self, table: str,
keyname: str | list[str] | tuple[str, ...] | None = None,
what: str | list[str] | tuple[str, ...] | None = None,
where: str | list[str] | tuple[str, ...] | None = None,
order: str | list[str] | tuple[str, ...] | bool | None = None,
limit: int | None = None, offset: int | None = None,
scalar: bool = False) -> dict:
"""Get a table as a dictionary.
This method is similar to get_as_list(), but returns the table
as a Python dict instead of a Python list, which can be even
more convenient. The primary key column(s) of the table will
be used as the keys of the dictionary, while the other column(s)
will be the corresponding values. The keys will be named tuples
if the table has a composite primary key. The rows will be also
named tuples unless the 'scalar' option has been set to True.
With the optional parameter 'keyname' you can specify an alternative
set of columns to be used as the keys of the dictionary. It must
be set as a string, list or a tuple.
The dictionary will be ordered using the order specified with the
'order' parameter or the key column(s) if not specified. You can
set 'order' to False if you don't care about the ordering.
"""
if not table:
raise TypeError('The table name is missing')
if not keyname:
try:
keyname = self.pkeys(table)
except (KeyError, ProgrammingError) as e:
raise prg_error(f'Table {table} has no primary key') from e
if isinstance(keyname, str):
keynames: list[str] | tuple[str, ...] = (keyname,)
elif isinstance(keyname, (list, tuple)):
keynames = keyname
else:
raise KeyError('The keyname must be a string, list or tuple')
if what:
if isinstance(what, (list, tuple)):
what = ', '.join(map(str, what))
if order is None:
order = what
else:
what = '*'
cmd_parts = ['SELECT', what, 'FROM', table]
if where:
if isinstance(where, (list, tuple)):
where = ' AND '.join(map(str, where))
cmd_parts.extend(['WHERE', where])
if order is None or order is True:
order = keyname
if order and not isinstance(order, bool):
if isinstance(order, (list, tuple)):
order = ', '.join(map(str, order))
cmd_parts.extend(['ORDER BY', order])
if limit:
cmd_parts.append(f'LIMIT {limit}')
if offset:
cmd_parts.append(f'OFFSET {offset}')
cmd = ' '.join(cmd_parts)
self._do_debug(cmd)
query = self._valid_db.query(cmd)
res = query.getresult()
if not res:
return {}
keyset = set(keynames)
fields = query.listfields()
if not keyset.issubset(fields):
raise KeyError('Missing keyname in row')
key_index: list[int] = []
row_index: list[int] = []
for i, f in enumerate(fields):
(key_index if f in keyset else row_index).append(i)
key_tuple = len(key_index) > 1
get_key = itemgetter(*key_index)
keys = map(get_key, res)
if scalar:
row_index = row_index[:1]
row_is_tuple = False
else:
row_is_tuple = len(row_index) > 1
if scalar or row_is_tuple:
get_row: Callable[[tuple], tuple] = itemgetter( # pyright: ignore
*row_index)
else:
frst_index = row_index[0]
def get_row(row : tuple) -> tuple:
return row[frst_index], # tuple with one item
row_is_tuple = True
rows = map(get_row, res)
if key_tuple or row_is_tuple:
if key_tuple:
keys = namediter(_MemoryQuery(keys, keynames)) # type: ignore
if row_is_tuple:
fields = tuple(f for f in fields if f not in keyset)
rows = namediter(_MemoryQuery(rows, fields)) # type: ignore
# noinspection PyArgumentList
return dict(zip(keys, rows))
def notification_handler(self, event: str, callback: Callable,
arg_dict: dict | None = None,
timeout: int | float | None = None,
stop_event: str | None = None
) -> NotificationHandler:
"""Get notification handler that will run the given callback."""
return NotificationHandler(self, event, callback,
arg_dict, timeout, stop_event)
# immediately wrapped methods
def send_query(self, cmd: str, args: Sequence | None = None) -> Query:
"""Create a new asynchronous query object for this connection."""
if args is None:
return self._valid_db.send_query(cmd)
return self._valid_db.send_query(cmd, args)
def poll(self) -> int:
"""Complete an asynchronous connection and get its state."""
return self._valid_db.poll()
def cancel(self) -> None:
"""Abandon processing of current SQL command."""
self._valid_db.cancel()
def fileno(self) -> int:
"""Get the socket used to connect to the database."""
return self._valid_db.fileno()
def get_cast_hook(self) -> Callable | None:
"""Get the function that handles all external typecasting."""
return self._valid_db.get_cast_hook()
def set_cast_hook(self, hook: Callable | None) -> None:
"""Set a function that will handle all external typecasting."""
self._valid_db.set_cast_hook(hook)
def get_notice_receiver(self) -> Callable | None:
"""Get the current notice receiver."""
return self._valid_db.get_notice_receiver()
def set_notice_receiver(self, receiver: Callable | None) -> None:
"""Set a custom notice receiver."""
self._valid_db.set_notice_receiver(receiver)
def getnotify(self) -> tuple[str, int, str] | None:
"""Get the last notify from the server."""
return self._valid_db.getnotify()
def inserttable(self, table: str, values: Sequence[list|tuple],
columns: list[str] | tuple[str, ...] | None = None) -> int:
"""Insert a Python iterable into a database table."""
if columns is None:
return self._valid_db.inserttable(table, values)
return self._valid_db.inserttable(table, values, columns)
def transaction(self) -> int:
"""Get the current in-transaction status of the server.
The status returned by this method can be TRANS_IDLE (currently idle),
TRANS_ACTIVE (a command is in progress), TRANS_INTRANS (idle, in a
valid transaction block), or TRANS_INERROR (idle, in a failed
transaction block). TRANS_UNKNOWN is reported if the connection is
bad. The status TRANS_ACTIVE is reported only when a query has been
sent to the server and not yet completed.
"""
return self._valid_db.transaction()
def parameter(self, name: str) -> str | None:
"""Look up a current parameter setting of the server."""
return self._valid_db.parameter(name)
def date_format(self) -> str:
"""Look up the date format currently being used by the database."""
return self._valid_db.date_format()
def escape_literal(self, s: AnyStr) -> AnyStr:
"""Escape a literal constant for use within SQL."""
return self._valid_db.escape_literal(s)
def escape_identifier(self, s: AnyStr) -> AnyStr:
"""Escape an identifier for use within SQL."""
return self._valid_db.escape_identifier(s)
def escape_string(self, s: AnyStr) -> AnyStr:
"""Escape a string for use within SQL."""
return self._valid_db.escape_string(s)
def escape_bytea(self, s: AnyStr) -> AnyStr:
"""Escape binary data for use within SQL as type 'bytea'."""
return self._valid_db.escape_bytea(s)
def putline(self, line: str) -> None:
"""Write a line to the server socket."""
self._valid_db.putline(line)
def getline(self) -> str:
"""Get a line from server socket."""
return self._valid_db.getline()
def endcopy(self) -> None:
"""Synchronize client and server."""
self._valid_db.endcopy()
def set_non_blocking(self, nb: bool) -> None:
"""Set the non-blocking mode of the connection."""
self._valid_db.set_non_blocking(nb)
def is_non_blocking(self) -> bool:
"""Get the non-blocking mode of the connection."""
return self._valid_db.is_non_blocking()
def locreate(self, mode: int) -> LargeObject:
"""Create a large object in the database.
The valid values for 'mode' parameter are defined as the module level
constants INV_READ and INV_WRITE.
"""
return self._valid_db.locreate(mode)
def getlo(self, oid: int) -> LargeObject:
"""Build a large object from given oid."""
return self._valid_db.getlo(oid)
def loimport(self, filename: str) -> LargeObject:
"""Import a file to a large object."""
return self._valid_db.loimport(filename)
class _MemoryQuery:
"""Class that embodies a given query result."""
result: Any
fields: tuple[str, ...]
def __init__(self, result: Any, fields: Sequence[str]) -> None:
"""Create query from given result rows and field names."""
self.result = result
self.fields = tuple(fields)
def listfields(self) -> tuple[str, ...]:
"""Return the stored field names of this query."""
return self.fields
def getresult(self) -> Any:
"""Return the stored result of this query."""
return self.result
def __iter__(self) -> Iterator[Any]:
return iter(self.result)
|