Spaces:
Running
Running
File size: 3,052 Bytes
2f1e7ff | 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 | """Track committed content writes across all shared PostgreSQL client pools.
Only a boolean is retained; query text and parameters are never stored/logged.
Periodic semantic reconciliation also covers administrative connections and DDL.
"""
import re
from psycopg2.extensions import connection, cursor
from .state import mark_dirty
_FACTORIES = {}
_WRITES = {"INSERT", "UPDATE", "DELETE", "TRUNCATE", "COPY", "CREATE", "ALTER", "DROP"}
class _TrackingCursor:
def _note(self, query, *, copy=False):
command = (self.statusmessage or "").split(" ", 1)[0]
if not copy and command not in _WRITES:
return
conn = self.connection
if conn.info.dbname == conn.backup_system_database:
if hasattr(query, "as_string"):
query = query.as_string(conn)
if isinstance(query, bytes):
query = query.decode("utf-8", errors="replace")
if not re.search(r"\b(?:data_brains|data_workspaces)\b", str(query), re.I):
return
if conn.autocommit:
mark_dirty()
else:
conn._backup_dirty = True
def execute(self, query, vars=None):
result = super().execute(query, vars)
self._note(query)
return result
def executemany(self, query, vars_list):
result = super().executemany(query, vars_list)
self._note(query)
return result
def copy_expert(self, sql, file, size=8192):
result = super().copy_expert(sql, file, size)
if re.search(r"\bFROM\s+STDIN\b", str(sql), re.I):
self._note(sql, copy=True)
return result
def copy_from(self, file, table, sep="\t", null="\\N", size=8192, columns=None):
result = super().copy_from(file, table, sep, null, size, columns)
self._note(table, copy=True)
return result
class DirtyConnection(connection):
def cursor(self, name=None, cursor_factory=None, scrollable=None, withhold=False):
base = cursor_factory or self.cursor_factory or cursor
factory = _FACTORIES.get(base)
if factory is None:
factory = _FACTORIES.setdefault(base, type(f"Content{base.__name__}", (_TrackingCursor, base), {}))
return super().cursor(name=name, cursor_factory=factory, scrollable=scrollable, withhold=withhold)
def commit(self):
changed = getattr(self, "_backup_dirty", False)
super().commit()
self._backup_dirty = False
if changed:
mark_dirty()
def rollback(self):
super().rollback()
self._backup_dirty = False
def __exit__(self, exc_type, exc, tb):
result = super().__exit__(exc_type, exc, tb)
changed = getattr(self, "_backup_dirty", False)
self._backup_dirty = False
if exc_type is None and changed:
mark_dirty()
return result
def connection_factory(system_database):
class ContentConnection(DirtyConnection):
backup_system_database = system_database
return ContentConnection
|