Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions src/psycopack/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,10 @@
InvalidIndexes,
InvalidPrimaryKeyTypeForConversion,
InvalidStageForReset,
NoCreateAndUsagePrivilegeOnSchema,
NoReferencesPrivilege,
NoReferringTableOwnership,
NotTableOwner,
PrimaryKeyNotFound,
ReferringForeignKeyInDifferentSchema,
Repack,
Expand All @@ -30,6 +34,10 @@
"InvalidIndexes",
"InvalidPrimaryKeyTypeForConversion",
"InvalidStageForReset",
"NoCreateAndUsagePrivilegeOnSchema",
"NoReferencesPrivilege",
"NoReferringTableOwnership",
"NotTableOwner",
"PrimaryKeyNotFound",
"ReferringForeignKeyInDifferentSchema",
"Repack",
Expand Down
137 changes: 134 additions & 3 deletions src/psycopack/_introspect.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,14 @@ class ReferringForeignKey:
is_validated: bool
referring_table: str
schema: str
is_owned_by_user: bool


@dataclasses.dataclass
class ReferredForeignKey:
name: str
schema: str
privileges: list[str]


@dataclasses.dataclass
Expand Down Expand Up @@ -239,7 +247,16 @@ def get_referring_fks(self, *, table: str) -> list[ReferringForeignKey]:
pg_get_constraintdef(pg_constraint.oid) AS definition,
pg_constraint.convalidated AS is_validated,
referring_pg_class.relname AS referring_table,
pg_namespace.nspname AS schema
pg_namespace.nspname AS schema,
(
SELECT
tableowner = current_user
FROM
pg_tables
WHERE
schemaname = {schema}
AND tablename = referring_pg_class.relname
) AS is_owned_by_user
FROM
pg_catalog.pg_constraint
INNER JOIN
Expand All @@ -260,7 +277,10 @@ def get_referring_fks(self, *, table: str) -> list[ReferringForeignKey]:
definition;
""")
)
.format(table=psycopg.sql.Literal(table))
.format(
table=psycopg.sql.Literal(table),
schema=psycopg.sql.Literal(self.schema),
)
.as_string(self.conn)
)
results = self.cur.fetchall()
Expand All @@ -272,8 +292,74 @@ def get_referring_fks(self, *, table: str) -> list[ReferringForeignKey]:
is_validated=is_validated,
referring_table=referring_table,
schema=schema,
is_owned_by_user=is_owned_by_user,
)
for name, definition, is_validated, referring_table, schema in results
for name, definition, is_validated, referring_table, schema, is_owned_by_user in results
]

def get_referred_fks(self, *, table: str, schema: str) -> list[ReferredForeignKey]:
self.cur.execute(
psycopg.sql.SQL(
dedent("""
SELECT
referred_pg_class.relname AS referred_table,
referred_pg_namespace.nspname AS referred_schema,
array_agg(
table_privileges.privilege_type
ORDER BY table_privileges.privilege_type
) AS privileges
FROM
pg_catalog.pg_constraint
INNER JOIN
pg_catalog.pg_class AS referring_pg_class
ON (pg_constraint.conrelid = referring_pg_class.oid)
INNER JOIN
pg_catalog.pg_class AS referred_pg_class
ON (pg_constraint.confrelid = referred_pg_class.oid)
INNER JOIN
pg_catalog.pg_namespace AS referring_pg_namespace
ON (referring_pg_namespace.oid = referring_pg_class.relnamespace)
INNER JOIN
pg_catalog.pg_namespace AS referred_pg_namespace
ON (referred_pg_namespace.oid = referred_pg_class.relnamespace)
LEFT OUTER JOIN
information_schema.table_privileges
ON (
table_privileges.table_name = referred_pg_class.relname
AND table_privileges.table_schema = referred_pg_namespace.nspname
)
WHERE
pg_constraint.conrelid = referring_pg_class.oid
AND referring_pg_class.relname = {table}
AND referring_pg_namespace.nspname = {schema}
AND pg_constraint.contype = 'f'
GROUP BY
referred_pg_namespace.nspname,
referred_pg_class.relname
ORDER BY
referred_pg_namespace.nspname,
referred_pg_class.relname;
""")
)
.format(
table=psycopg.sql.Literal(table),
schema=psycopg.sql.Literal(schema),
)
.as_string(self.conn)
)
results = self.cur.fetchall()
assert results is not None
return [
ReferredForeignKey(
name=name,
schema=schema,
privileges=(
privileges.strip("{}").split(",")
if "NULL" not in privileges
else []
),
)
for name, schema, privileges in results
]

def table_is_empty(self, *, table: str) -> int:
Expand Down Expand Up @@ -529,3 +615,48 @@ def get_backfill_batch(self, *, table: str) -> BackfillBatch | None:
if not (result := self.cur.fetchone()):
return None
return BackfillBatch(id=result[0], start=result[1], end=result[2])

def get_user(self) -> str:
self.cur.execute(psycopg.sql.SQL("SELECT current_user;").as_string(self.conn))
result = self.cur.fetchone()
assert result is not None
return str(result[0])

def has_create_and_usage_privilege_on_schema(self) -> bool:
self.cur.execute(
psycopg.sql.SQL(
dedent("""
SELECT
has_schema_privilege({schema}, 'CREATE'),
has_schema_privilege({schema}, 'USAGE');
""")
)
.format(schema=psycopg.sql.Literal(self.schema))
.as_string(self.conn)
)
result = self.cur.fetchone()
assert result is not None
return bool(result[0] and result[1])

def is_table_owner(self, *, table: str, schema: str) -> bool:
self.cur.execute(
psycopg.sql.SQL(
dedent("""
SELECT
tableowner = current_user
FROM
pg_tables
WHERE
schemaname = {schema}
AND tablename = {table}
""")
)
.format(
schema=psycopg.sql.Literal(schema),
table=psycopg.sql.Literal(table),
)
.as_string(self.conn)
)
result = self.cur.fetchone()
assert result is not None
return bool(result[0])
74 changes: 74 additions & 0 deletions src/psycopack/_repack.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,22 @@ class ReferringForeignKeyInDifferentSchema(BaseRepackError):
pass


class NoCreateAndUsagePrivilegeOnSchema(BaseRepackError):
pass


class NotTableOwner(BaseRepackError):
pass


class NoReferringTableOwnership(BaseRepackError):
pass


class NoReferencesPrivilege(BaseRepackError):
pass


class PostBackfillBatchCallback(typing.Protocol):
def __call__(
self, batch: _introspect.BackfillBatch, /
Expand Down Expand Up @@ -175,6 +191,7 @@ def __init__(
)
self._pk_column = ""
self.schema = schema
self._check_user_permissions()

@property
def pk_column(self) -> str:
Expand Down Expand Up @@ -776,3 +793,60 @@ def _create_referring_fks(self) -> None:
self.command.validate_constraint(
table=fk.referring_table, constraint=constraint_name
)

def _check_user_permissions(self) -> None:
if not self.introspector.has_create_and_usage_privilege_on_schema():
user = self.introspector.get_user()
raise NoCreateAndUsagePrivilegeOnSchema(
f"Psycopack requires the database user to have CREATE and USAGE "
f"privilege on the {self.schema} schema to create auxiliary "
f"objects. You can grant it to your user via:\n"
f"GRANT CREATE, USAGE ON SCHEMA {self.schema} TO {user};"
)

if not self.introspector.is_table_owner(table=self.table, schema=self.schema):
user = self.introspector.get_user()
raise NotTableOwner(
f"Psycopack requires the database user to have ownership of "
f"the table {self.schema}.{self.table}. You can grant it to "
f"your user via:\n"
f"ALTER TABLE {self.schema}.{self.table} OWNER TO {user};"
)

referring_tables_without_ownership = [
f"{fk.schema}.{fk.referring_table}"
for fk in self.introspector.get_referring_fks(table=self.table)
if not fk.is_owned_by_user
]
if referring_tables_without_ownership:
user = self.introspector.get_user()
message = (
f"Psycopack requires the user to have ownership of tables that "
f"have foreign keys pointing to the table {self.schema}.{self.table}. "
f"You can grant ownership to your user via the following commands:\n"
)
for table in referring_tables_without_ownership:
message += f"ALTER TABLE {table} OWNER TO {user};\n"
raise NoReferringTableOwnership(message)

referred_fks_without_references_privilege = [
fk
for fk in self.introspector.get_referred_fks(
table=self.table, schema=self.schema
)
if "REFERENCES" not in fk.privileges
]
if referred_fks_without_references_privilege:
user = self.introspector.get_user()
tables_missing_ownership = [
f"{fk.schema}.{fk.name}"
for fk in referred_fks_without_references_privilege
]
message = (
f"Psycopack requires the user to REFERENCES privilege on tables that "
f"{self.schema}.{self.table} has foreign keys to. "
f"You can grant this privilege to your user via the following commands:\n"
)
for table in tables_missing_ownership:
message += f"GRANT REFERENCES ON TABLE {table} TO {user};\n"
raise NoReferencesPrivilege(message)
Loading