Merge pull request #116 from DanCardin/dc/quote-audit-columns

DanCardin · web-flow · commit b7c0d8c209b0 · 2025-06-23T15:51:37.000-04:00
fix: Quote audit column names.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## 0.15
 
+### 0.15.11
+- fix: Exclude extension-created objects from comparisons.
+- fix: Quote audit column names.
+
 ### 0.15.10
 
 - feat: Support alembic check.
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "sqlalchemy-declarative-extensions"
-version = "0.15.10"
+version = "0.15.11"
 authors = [
     {name = "Dan Cardin", email = "ddcardin@gmail.com"},
 ]
diff --git a/src/sqlalchemy_declarative_extensions/audit.py b/src/sqlalchemy_declarative_extensions/audit.py
@@ -10,6 +10,7 @@
     register_trigger,
 )
 from sqlalchemy_declarative_extensions.dialects.postgresql.trigger import Trigger
+from sqlalchemy_declarative_extensions.sql import quote_name
 
 default_primary_key = Column(
     "audit_pk", types.Integer(), primary_key=True, autoincrement=True
@@ -192,7 +193,7 @@ def create_audit_functions(
         if column.name == AUDIT_PK:
             continue
 
-        audit_columns.append(column.name)
+        audit_columns.append(quote_name(column.name))
 
         if column.name in {
             AUDIT_PK,
@@ -211,8 +212,8 @@ def create_audit_functions(
             old_elements.append(value)
             new_elements.append(value)
         else:
-            old_elements.append(f"OLD.{column.name}")
-            new_elements.append(f"NEW.{column.name}")
+            old_elements.append(f'OLD."{column.name}"')
+            new_elements.append(f'NEW."{column.name}"')
 
     audit_columns_str = ", ".join(audit_columns)
     old_elements_str = ", ".join(old_elements)
@@ -234,7 +235,7 @@ def create_audit_functions(
             "_".join([function_name, op.lower()]),
             f"""
             BEGIN
-            INSERT INTO {audit_table.fullname} ({audit_columns_str})
+            INSERT INTO {quote_name(audit_table.fullname)} ({audit_columns_str})
             SELECT '{op_key}', now(), current_user, {elements};
             RETURN NULL;
             END
diff --git a/src/sqlalchemy_declarative_extensions/dialects/postgresql/function.py b/src/sqlalchemy_declarative_extensions/dialects/postgresql/function.py
@@ -5,6 +5,7 @@
 from dataclasses import dataclass, replace
 
 from sqlalchemy_declarative_extensions.function import base
+from sqlalchemy_declarative_extensions.sql import quote_name
 
 
 @enum.unique
@@ -31,7 +32,7 @@ def to_sql_create(self, replace=False) -> list[str]:
             components.append("OR REPLACE")
 
         components.append("FUNCTION")
-        components.append(self.qualified_name + "()")
+        components.append(quote_name(self.qualified_name) + "()")
         components.append(f"RETURNS {self.returns}")
 
         if self.security == FunctionSecurity.definer:
diff --git a/src/sqlalchemy_declarative_extensions/dialects/postgresql/procedure.py b/src/sqlalchemy_declarative_extensions/dialects/postgresql/procedure.py
@@ -7,6 +7,7 @@
 from typing_extensions import Self
 
 from sqlalchemy_declarative_extensions.procedure import base
+from sqlalchemy_declarative_extensions.sql import quote_name
 
 
 @enum.unique
@@ -33,7 +34,7 @@ def to_sql_create(self, replace=False) -> list[str]:
             components.append("OR REPLACE")
 
         components.append("PROCEDURE")
-        components.append(self.qualified_name + "()")
+        components.append(quote_name(self.qualified_name) + "()")
 
         if self.security == ProcedureSecurity.definer:
             components.append("SECURITY DEFINER")
diff --git a/src/sqlalchemy_declarative_extensions/dialects/postgresql/trigger.py b/src/sqlalchemy_declarative_extensions/dialects/postgresql/trigger.py
@@ -5,6 +5,7 @@
 from sqlalchemy.engine import Connection
 
 from sqlalchemy_declarative_extensions.dialects.from_string import FromStrings
+from sqlalchemy_declarative_extensions.sql import quote_name
 from sqlalchemy_declarative_extensions.trigger import base
 
 
@@ -155,14 +156,12 @@ def to_sql_create(self, replace=False):
             components.append("OR REPLACE")
 
         components.append("TRIGGER")
-        components.append(f'"{self.name}"')
+        components.append(quote_name(self.name))
         components.append(self.time.value)
         components.append(" OR ".join([e.value for e in self.events]))
         components.append("ON")
 
-        on_components = [f'"{component}"' for component in self.on.split(".")]
-        on = ".".join(on_components)
-        components.append(on)
+        components.append(quote_name(self.on))
 
         components.append("FOR EACH")
         components.append(self.for_each.value)
@@ -175,7 +174,7 @@ def to_sql_create(self, replace=False):
         args_quoted = (
             tuple(f"'{arg}'" for arg in self.arguments) if self.arguments else ()
         )
-        components.append(self.execute + f"({','.join(args_quoted)})")
+        components.append(quote_name(self.execute) + f"({','.join(args_quoted)})")
         return " ".join(components) + ";"
 
     def to_sql_update(self, connection: Connection | None = None):
diff --git a/src/sqlalchemy_declarative_extensions/function/base.py b/src/sqlalchemy_declarative_extensions/function/base.py
@@ -6,7 +6,7 @@
 from sqlalchemy import MetaData
 from typing_extensions import Self
 
-from sqlalchemy_declarative_extensions.sql import qualify_name
+from sqlalchemy_declarative_extensions.sql import qualify_name, quote_name
 from sqlalchemy_declarative_extensions.sqlalchemy import HasMetaData
 
 
@@ -54,7 +54,7 @@ def to_sql_update(self) -> list[str]:
         ]
 
     def to_sql_drop(self) -> list[str]:
-        return [f"DROP FUNCTION {self.qualified_name}();"]
+        return [f"DROP FUNCTION {quote_name(self.qualified_name)}();"]
 
     def with_name(self, name: str):
         return replace(self, name=name)
diff --git a/src/sqlalchemy_declarative_extensions/procedure/base.py b/src/sqlalchemy_declarative_extensions/procedure/base.py
@@ -6,7 +6,7 @@
 from sqlalchemy import MetaData
 from typing_extensions import Self
 
-from sqlalchemy_declarative_extensions.sql import qualify_name
+from sqlalchemy_declarative_extensions.sql import qualify_name, quote_name
 from sqlalchemy_declarative_extensions.sqlalchemy import HasMetaData
 
 
@@ -52,7 +52,7 @@ def to_sql_update(self) -> list[str]:
         ]
 
     def to_sql_drop(self) -> list[str]:
-        return [f"DROP PROCEDURE {self.qualified_name}();"]
+        return [f"DROP PROCEDURE {quote_name(self.qualified_name)}();"]
 
     def with_name(self, name: str):
         return replace(self, name=name)
diff --git a/src/sqlalchemy_declarative_extensions/sql.py b/src/sqlalchemy_declarative_extensions/sql.py
@@ -12,9 +12,10 @@ def qualify_name(schema: str | None, name: str, quote=False) -> str:
             return f'"{name}"'
         return name
 
+    result = f"{schema}.{name}"
     if quote:
-        return f'"{schema}"."{name}"'
-    return f"{schema}.{name}"
+        return quote_name(result)
+    return result
 
 
 def split_schema(
@@ -46,3 +47,8 @@ def coerce_name(name: str | HasName):
     if isinstance(name, HasName):
         return name.name
     return name
+
+
+def quote_name(name: str) -> str:
+    components = [f'"{component}"' for component in name.split(".")]
+    return ".".join(components)
diff --git a/tests/audit/test_quoting.py b/tests/audit/test_quoting.py
@@ -21,6 +21,7 @@ class Something(Base):
     __tablename__ = "Something"
 
     id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
+    end: Mapped[int] = mapped_column()
 
 
 register_sqlalchemy_events(Base.metadata, functions=True, triggers=True)
@@ -31,3 +32,6 @@ class Something(Base):
 def test_quoted_name(pg):
     Base.metadata.create_all(bind=pg.connection())
     pg.commit()
+
+    pg.add(Something(end=2))
+    pg.commit()
diff --git a/tests/function/test_quote.py b/tests/function/test_quote.py
@@ -0,0 +1,52 @@
+from pytest_mock_resources import create_postgres_fixture
+from sqlalchemy import Column, Integer, text
+
+from sqlalchemy_declarative_extensions import (
+    Function,
+    Functions,
+    declarative_database,
+    register_sqlalchemy_events,
+)
+from sqlalchemy_declarative_extensions.function.compare import compare_functions
+from sqlalchemy_declarative_extensions.sqlalchemy import declarative_base
+
+_Base = declarative_base()
+
+
+@declarative_database
+class Base(_Base):  # type: ignore
+    __abstract__ = True
+
+    functions = Functions().are(
+        Function(
+            "table",
+            "INSERT INTO foo (id) VALUES (DEFAULT); SELECT count(*) FROM foo;",
+            returns="INTEGER",
+        )
+    )
+
+
+class Foo(Base):
+    __tablename__ = "foo"
+
+    id = Column(Integer, primary_key=True)
+
+
+register_sqlalchemy_events(Base.metadata, functions=True)
+
+pg = create_postgres_fixture(engine_kwargs={"echo": True}, session=True)
+
+
+def test_create(pg):
+    Base.metadata.create_all(bind=pg.connection())
+    pg.commit()
+
+    result = pg.execute(text('SELECT "table"()')).scalar()
+    assert result == 1
+
+    result = pg.execute(text('SELECT "table"()')).scalar()
+    assert result == 2
+
+    connection = pg.connection()
+    diff = compare_functions(connection, Base.metadata.info["functions"])
+    assert diff == []
diff --git a/tests/procedure/test_quote.py b/tests/procedure/test_quote.py
@@ -0,0 +1,51 @@
+from pytest_mock_resources import create_postgres_fixture
+from sqlalchemy import Column, Integer, text
+
+from sqlalchemy_declarative_extensions import (
+    Procedure,
+    Procedures,
+    declarative_database,
+    register_sqlalchemy_events,
+)
+from sqlalchemy_declarative_extensions.procedure.compare import compare_procedures
+from sqlalchemy_declarative_extensions.sqlalchemy import declarative_base
+
+_Base = declarative_base()
+
+
+@declarative_database
+class Base(_Base):  # type: ignore
+    __abstract__ = True
+
+    procedures = Procedures().are(
+        Procedure(
+            "table",
+            "INSERT INTO foo (id) VALUES (DEFAULT);",
+        )
+    )
+
+
+class Foo(Base):
+    __tablename__ = "foo"
+
+    id = Column(Integer, primary_key=True)
+
+
+register_sqlalchemy_events(Base.metadata, procedures=True)
+
+pg = create_postgres_fixture(engine_kwargs={"echo": True}, session=True)
+
+
+def test_create(pg):
+    Base.metadata.create_all(bind=pg.connection())
+    pg.commit()
+
+    pg.execute(text('CALL "table"()'))
+    pg.execute(text('CALL "table"()'))
+
+    result = pg.execute(text("SELECT count(*) FROM foo")).scalar()
+    assert result == 2
+
+    connection = pg.connection()
+    diff = compare_procedures(connection, Base.metadata.info["procedures"])
+    assert diff == []
diff --git a/uv.lock b/uv.lock

Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,6 @@`
`1`	`1`	`[project]`
`2`	`2`	`name = "sqlalchemy-declarative-extensions"`
`3`		`-version = "0.15.10"`
	`3`	`+version = "0.15.11"`
`4`	`4`	`authors = [`
`5`	`5`	`{name = "Dan Cardin", email = "ddcardin@gmail.com"},`
`6`	`6`	`]`