Python: Add field and table name escaping for python SqlServer connector

python/semantic_kernel/connectors/sql_server.py

@@ -154,6 +154,20 @@ def append_list(self, strings: Sequence[str], sep: str = ", ", suffix: str | Non
             self.append(string, suffix=sep)
         self.append(strings[-1], suffix=suffix)
 
+    @staticmethod
+    def escape_identifier(identifier: str) -> str:
+        """Escape and bracket-quote a SQL Server identifier.
+
+        Escapes `]` to `]]` and wraps in square brackets, matching
+        the SQL Server standard for delimited identifiers.
+        """
+        return f"[{identifier.replace(']', ']]')}]"
+
+    @staticmethod
+    def format_table_name(schema: str, table_name: str) -> str:
+        """Return an escaped [schema].[table] string."""
+        return f"{QueryBuilder.escape_identifier(schema)}.{QueryBuilder.escape_identifier(table_name)}"
+
     def append_table_name(
         self, schema: str, table_name: str, prefix: str = "", suffix: str | None = None, newline: bool = False
     ) -> None:
@@ -169,7 +183,10 @@ def append_table_name(
             suffix: Optional suffix to add after the table name.
             newline: Whether to add a newline after the table name or suffix.
         """
-        self.append(f"{prefix} [{schema}].[{table_name}] {suffix or ''}", suffix="\n" if newline else "")
+        self.append(
+            f"{prefix} {self.format_table_name(schema, table_name)} {suffix or ''}",
+            suffix="\n" if newline else "",
+        )
 
     def remove_last(self, number_of_chars: int):
         """Remove the last number_of_chars from the StringBuilder."""
@@ -631,14 +648,14 @@ def parse(node: ast.AST) -> str:
                         raise VectorStoreOperationException(
                             f"Field '{node.attr}' not in data model (storage property names are used)."
                         )
-                    return f"[{node.attr}]"
+                    return QueryBuilder.escape_identifier(node.attr)
                 case ast.Name():
                     # Only allow names that are in the data model
                     if node.id not in self.definition.storage_names:
                         raise VectorStoreOperationException(
                             f"Field '{node.id}' not in data model (storage property names are used)."
                         )
-                    return f"[{node.id}]"
+                    return QueryBuilder.escape_identifier(node.id)
                 case ast.Constant():
                     value = node.value
                     if isinstance(value, (str, int, float, bool, bytes)) or value is None:
@@ -860,32 +877,34 @@ def _build_create_table_query(
     """Build the CREATE TABLE query based on the data model."""
     command = SqlCommand()
     if if_not_exists:
-        command.query.append_table_name(
-            schema, table, prefix="IF OBJECT_ID(N'", suffix="', N'U') IS NULL", newline=True
-        )
+        # OBJECT_ID takes a string literal — escape ' → '' in addition to identifier escaping
+        table_ref = QueryBuilder.format_table_name(schema, table).replace("'", "''")
+        command.query.append(f"IF OBJECT_ID(N'{table_ref}', N'U') IS NULL", suffix="\n")
     with command.query.in_logical_group():
         command.query.append_table_name(schema, table, prefix="CREATE TABLE", newline=True)
         with command.query.in_parenthesis(suffix=";"):
             # add the key field
             command.query.append(
-                f'"{key_field.storage_name or key_field.name}" '
+                f"{QueryBuilder.escape_identifier(key_field.storage_name or key_field.name)} "
                 f"{_python_type_to_sql(key_field.type_, is_key=True)} NOT NULL,\n"
             )
             # add the data fields
-            [
-                command.query.append(f'"{field.storage_name or field.name}" {_python_type_to_sql(field.type_)} NULL,\n')
-                for field in data_fields
-            ]
+            for field in data_fields:
+                command.query.append(
+                    f"{QueryBuilder.escape_identifier(field.storage_name or field.name)}"
+                    f" {_python_type_to_sql(field.type_)} NULL,\n"
+                )
             # add the vector fields
             for field in vector_fields:
                 if field.index_kind not in INDEX_KIND_MAP:
                     raise VectorStoreOperationException(
                         f"Index kind '{field.index_kind}' is not supported for field '{field.name}'"
                     )
-                command.query.append(f'"{field.storage_name or field.name}" VECTOR({field.dimensions}) NULL,\n')
+                escaped_name = QueryBuilder.escape_identifier(field.storage_name or field.name)
+                command.query.append(f"{escaped_name} VECTOR({field.dimensions}) NULL,\n")
             # set the primary key
             with command.query.in_parenthesis("PRIMARY KEY", "\n"):
-                command.query.append(key_field.name)
+                command.query.append(QueryBuilder.escape_identifier(key_field.storage_name or key_field.name))
     return command
 
 
@@ -954,9 +973,11 @@ def _add_field_names(
     """
     fields = chain([key_field], data_fields, vector_fields or [])
     if table_identifier:
-        strings = [f"{table_identifier}.{field.storage_name or field.name}" for field in fields]
+        strings = [
+            f"{table_identifier}.{QueryBuilder.escape_identifier(field.storage_name or field.name)}" for field in fields
+        ]
     else:
-        strings = [field.storage_name or field.name for field in fields]
+        strings = [QueryBuilder.escape_identifier(field.storage_name or field.name) for field in fields]
     command.query.append_list(strings)
 
 
@@ -997,17 +1018,17 @@ def _build_merge_query(
         _add_field_names(command, key_field, data_fields, vector_fields)
     # add the ON clause
     with command.query.in_parenthesis("ON", "\n"):
-        command.query.append(
-            f"t.{key_field.storage_name or key_field.name} = s.{key_field.storage_name or key_field.name}"
-        )
+        escaped_key = QueryBuilder.escape_identifier(key_field.storage_name or key_field.name)
+        command.query.append(f"t.{escaped_key} = s.{escaped_key}")
     # Set the Matched clause
     command.query.append("WHEN MATCHED THEN\n")
     command.query.append("UPDATE SET ")
+    update_parts = []
+    for field in chain(data_fields, vector_fields):
+        escaped = QueryBuilder.escape_identifier(field.storage_name or field.name)
+        update_parts.append(f"t.{escaped} = s.{escaped}")
     command.query.append_list(
-        [
-            f"t.{field.storage_name or field.name} = s.{field.storage_name or field.name}"
-            for field in chain(data_fields, vector_fields)
-        ],
+        update_parts,
         suffix="\n",
     )
     # Set the Not Matched clause
@@ -1018,7 +1039,8 @@ def _build_merge_query(
     with command.query.in_parenthesis("VALUES", " \n"):
         _add_field_names(command, key_field, data_fields, vector_fields, table_identifier="s")
     # add the closing parenthesis
-    command.query.append(f"OUTPUT inserted.{key_field.name} INTO @UpsertedKeys (KeyColumn);\n")
+    escaped_key_out = QueryBuilder.escape_identifier(key_field.storage_name or key_field.name)
+    command.query.append(f"OUTPUT inserted.{escaped_key_out} INTO @UpsertedKeys (KeyColumn);\n")
     command.query.append("SELECT KeyColumn FROM @UpsertedKeys;\n")
     return command
 
@@ -1041,7 +1063,7 @@ def _build_select_query(
     command.query.append_table_name(schema, table, prefix=" FROM", newline=True)
     # add the WHERE clause
     if keys:
-        command.query.append(f"WHERE {key_field.storage_name or key_field.name} IN\n")
+        command.query.append(f"WHERE {QueryBuilder.escape_identifier(key_field.storage_name or key_field.name)} IN\n")
         with command.query.in_parenthesis():
             # add the keys
             command.query.append_list(["?"] * len(keys))
@@ -1061,7 +1083,7 @@ def _build_delete_query(
     # start the DELETE statement
     command.query.append_table_name(schema, table)
     # add the WHERE clause
-    command.query.append(f"WHERE [{key_field.storage_name or key_field.name}] IN")
+    command.query.append(f"WHERE {QueryBuilder.escape_identifier(key_field.storage_name or key_field.name)} IN")
     with command.query.in_parenthesis():
         # add the keys
         command.query.append_list(["?"] * len(keys))
@@ -1109,7 +1131,7 @@ def _build_search_query(
     asc = DISTANCE_FUNCTION_DIRECTION_HELPER[vector_field.distance_function](0, 1)
 
     command.query.append(
-        f", VECTOR_DISTANCE('{distance_function}', {vector_field.storage_name or vector_field.name}, CAST(? AS VECTOR({vector_field.dimensions}))) as {SCORE_FIELD_NAME}\n",  # noqa: E501
+        f", VECTOR_DISTANCE('{distance_function}', {QueryBuilder.escape_identifier(vector_field.storage_name or vector_field.name)}, CAST(? AS VECTOR({vector_field.dimensions}))) as {SCORE_FIELD_NAME}\n",  # noqa: E501
     )
     command.add_parameter(_cast_value(vector))
     # add the FROM clause

python/tests/unit/connectors/memory/test_sql_server.py

@@ -42,12 +42,32 @@ def test_query_builder_append_list(self):
         result = str(qb).strip()
         assert result == "id, name, age;"
 
+    def test_query_builder_escape_identifier(self):
+        assert QueryBuilder.escape_identifier("simple") == "[simple]"
+        assert QueryBuilder.escape_identifier("has]bracket") == "[has]]bracket]"
+        assert QueryBuilder.escape_identifier("two]]brackets") == "[two]]]]brackets]"
+        assert QueryBuilder.escape_identifier("") == "[]"
+
     def test_query_builder_append_table_name(self):
         qb = QueryBuilder()
         qb.append_table_name("dbo", "Users", prefix="SELECT * FROM", suffix=";", newline=False)
         result = str(qb).strip()
         assert result == "SELECT * FROM [dbo].[Users] ;"
 
+    def test_query_builder_append_table_name_escapes_closing_bracket(self):
+        qb = QueryBuilder()
+        qb.append_table_name("my]schema", "my]table", prefix="SELECT * FROM", suffix=";")
+        result = str(qb).strip()
+        assert result == "SELECT * FROM [my]]schema].[my]]table] ;"
+
+    def test_query_builder_append_table_name_prevents_sql_injection(self):
+        qb = QueryBuilder()
+        qb.append("DROP TABLE IF EXISTS")
+        qb.append_table_name("dbo", "]; EXEC xp_cmdshell('whoami'); --", suffix=";")
+        result = str(qb)
+        assert "EXEC xp_cmdshell" not in result.split("].[")[0], "SQL injection should not escape bracket quoting"
+        assert "[dbo].[]]; EXEC xp_cmdshell('whoami'); --]" in result
+
     def test_query_builder_remove_last(self):
         qb = QueryBuilder("SELECT * FROM table;")
         qb.remove_last(1)  # remove trailing semicolon
@@ -121,10 +141,31 @@ def test_build_create_table_query(self):
         cmd_str = str(cmd.query)
         assert (
             cmd_str
-            == 'BEGIN\nCREATE TABLE [dbo].[Test] \n ("id" nvarchar(255) NOT NULL,\n"name" nvarchar(max) NULL,\n"age" '
-            'int NULL,\n"embedding" VECTOR(1536) NULL,\nPRIMARY KEY (id) \n) ;\nEND\n'
+            == "BEGIN\nCREATE TABLE [dbo].[Test] \n ([id] nvarchar(255) NOT NULL,\n[name] nvarchar(max) NULL,\n[age] "
+            "int NULL,\n[embedding] VECTOR(1536) NULL,\nPRIMARY KEY ([id]) \n) ;\nEND\n"
         )
 
+    def test_build_create_table_query_escapes_single_quote_in_object_id(self):
+        key_field = VectorStoreField("key", name="id", type="str")
+        cmd = _build_create_table_query("dbo", "test'table", key_field, [], [], if_not_exists=True)
+        cmd_str = str(cmd.query)
+        # Single quote must be escaped inside the OBJECT_ID N'...' string literal
+        assert "OBJECT_ID(N'[dbo].[test''table]'" in cmd_str
+
+    def test_build_create_table_query_uses_storage_name_for_primary_key(self):
+        key_field = VectorStoreField("key", name="id", type="str", storage_name="pk_id")
+        cmd = _build_create_table_query("dbo", "Test", key_field, [], [])
+        cmd_str = str(cmd.query)
+        assert "[pk_id] nvarchar" in cmd_str
+        assert "PRIMARY KEY ([pk_id])" in cmd_str
+
+    def test_build_merge_query_output_uses_storage_name(self):
+        key_field = VectorStoreField("key", name="id", type="str", storage_name="pk_id")
+        records = [{"pk_id": "1"}]
+        cmd = _build_merge_query("dbo", "Test", key_field, [], [], records)
+        cmd_str = str(cmd.query)
+        assert "OUTPUT inserted.[pk_id] INTO @UpsertedKeys" in cmd_str
+
     def test_delete_table_query(self):
         schema = "dbo"
         table = "Test"
@@ -170,11 +211,17 @@ def test_build_merge_query(self):
         assert cmd.parameters[3] == json.dumps(records[0]["embedding"])
         str_cmd = str(cmd)
         assert str_cmd == (
-            "DECLARE @UpsertedKeys TABLE (KeyColumn nvarchar(255));\nMERGE INTO [dbo].[Test] AS t\nUSING ( "
-            "VALUES  (?, ?, ?, ?) ) AS s (id, name, age, embedding)  ON (t.id = s.id) \nWHEN MATCHED THEN\nUPDATE "
-            "SET t.name = s.name, t.age = s.age, t.embedding = s.embedding\nWHEN NOT MATCHED THEN\nINSERT "
-            "(id, name, age, embedding)  VALUES (s.id, s.name, s.age, s.embedding)  \nOUTPUT inserted.id "
-            "INTO @UpsertedKeys (KeyColumn);\nSELECT KeyColumn FROM @UpsertedKeys;\n"
+            "DECLARE @UpsertedKeys TABLE (KeyColumn nvarchar(255));\n"
+            "MERGE INTO [dbo].[Test] AS t\n"
+            "USING ( VALUES  (?, ?, ?, ?) ) AS s ([id], [name], [age], [embedding])  "
+            "ON (t.[id] = s.[id]) \n"
+            "WHEN MATCHED THEN\n"
+            "UPDATE SET t.[name] = s.[name], t.[age] = s.[age], t.[embedding] = s.[embedding]\n"
+            "WHEN NOT MATCHED THEN\n"
+            "INSERT ([id], [name], [age], [embedding])  "
+            "VALUES (s.[id], s.[name], s.[age], s.[embedding])  \n"
+            "OUTPUT inserted.[id] INTO @UpsertedKeys (KeyColumn);\n"
+            "SELECT KeyColumn FROM @UpsertedKeys;\n"
         )
 
     def test_build_select_query(self):
@@ -192,7 +239,7 @@ def test_build_select_query(self):
         cmd = _build_select_query(schema, table, key_field, data_fields, vector_fields, keys)
         assert cmd.parameters == ["test"]
         str_cmd = str(cmd)
-        assert str_cmd == "SELECT\nid, name, age, embedding FROM [dbo].[Test] \nWHERE id IN\n (?) ;"
+        assert str_cmd == "SELECT\n[id], [name], [age], [embedding] FROM [dbo].[Test] \nWHERE [id] IN\n (?) ;"
 
     def test_build_delete_query(self):
         schema = "dbo"
@@ -230,7 +277,7 @@ def test_build_search_query(self):
         assert cmd.parameters[0] == json.dumps(vector)
         str_cmd = str(cmd)
         assert (
-            str_cmd == "SELECT id, name, age, VECTOR_DISTANCE('cosine', embedding, CAST(? AS VECTOR(5))) as "
+            str_cmd == "SELECT [id], [name], [age], VECTOR_DISTANCE('cosine', [embedding], CAST(? AS VECTOR(5))) as "
             "_vector_distance_value\n FROM [dbo].[Test] \nORDER BY "
             "_vector_distance_value ASC\nOFFSET 0 ROWS FETCH NEXT 3 ROWS ONLY;"
         )
@@ -378,11 +425,17 @@ async def test_upsert(
         await collection.upsert(record)
         mock_connection.cursor.return_value.__enter__.return_value.execute.assert_called_with(
             (
-                "DECLARE @UpsertedKeys TABLE (KeyColumn nvarchar(255));\nMERGE INTO [dbo].[test] AS t\nUSING ( VALUES"
-                "  (?, ?, ?) ) AS s (id, content, vector)  ON (t.id = s.id) \nWHEN MATCHED THEN\nUPDATE SET t.content"
-                " = s.content, t.vector = s.vector\nWHEN NOT MATCHED THEN\nINSERT (id, content, vector)  VALUES (s.id, "
-                "s.content, s.vector)  \nOUTPUT inserted.id INTO @UpsertedKeys (KeyColumn);\nSELECT KeyColumn "
-                "FROM @UpsertedKeys;\n"
+                "DECLARE @UpsertedKeys TABLE (KeyColumn nvarchar(255));\n"
+                "MERGE INTO [dbo].[test] AS t\n"
+                "USING ( VALUES  (?, ?, ?) ) AS s ([id], [content], [vector])  "
+                "ON (t.[id] = s.[id]) \n"
+                "WHEN MATCHED THEN\n"
+                "UPDATE SET t.[content] = s.[content], t.[vector] = s.[vector]\n"
+                "WHEN NOT MATCHED THEN\n"
+                "INSERT ([id], [content], [vector])  "
+                "VALUES (s.[id], s.[content], s.[vector])  \n"
+                "OUTPUT inserted.[id] INTO @UpsertedKeys (KeyColumn);\n"
+                "SELECT KeyColumn FROM @UpsertedKeys;\n"
             ),
             ("1", "test", json.dumps([0.1, 0.2, 0.3, 0.4, 0.5])),
         )
@@ -415,7 +468,7 @@ class MockRow(NamedTuple):
         mock_cursor.__iter__.return_value = [row]
         record = await collection.get(key, include_vectors=True)
         mock_cursor.execute.assert_called_with(
-            "SELECT\nid, content, vector FROM [dbo].[test] \nWHERE id IN\n (?) ;", ("1",)
+            "SELECT\n[id], [content], [vector] FROM [dbo].[test] \nWHERE [id] IN\n (?) ;", ("1",)
         )
         assert record["id"] == "1"
         assert record["content"] == "test"
@@ -478,7 +531,7 @@ class MockRow:
             assert record.score == 0.1
         mock_cursor.execute.assert_called_with(
             (
-                "SELECT id, content, VECTOR_DISTANCE('cosine', vector, CAST(? AS VECTOR(5))) as "
+                "SELECT [id], [content], VECTOR_DISTANCE('cosine', [vector], CAST(? AS VECTOR(5))) as "
                 "_vector_distance_value\n FROM [dbo].[test] \n WHERE [content] = ? \nORDER BY _vector_distance_value "
                 "ASC\nOFFSET 0 ROWS FETCH NEXT 3 ROWS ONLY;"
             ),
@@ -502,8 +555,10 @@ async def test_ensure_collection_exists(
         await collection.ensure_collection_exists()
         mock_connection.cursor.return_value.__enter__.return_value.execute.assert_called_with(
             (
-                "IF OBJECT_ID(N' [dbo].[test] ', N'U') IS NULL\nBEGIN\nCREATE TABLE [dbo].[test] \n (\"id\" nvarchar"
-                '(255) NOT NULL,\n"content" nvarchar(max) NULL,\n"vector" VECTOR(5) NULL,\nPRIMARY KEY (id) \n) ;'
+                "IF OBJECT_ID(N'[dbo].[test]', N'U') IS NULL\n"
+                "BEGIN\nCREATE TABLE [dbo].[test] \n ([id] nvarchar"
+                "(255) NOT NULL,\n[content] nvarchar(max) NULL,\n[vector] VECTOR(5) NULL,\n"
+                "PRIMARY KEY ([id]) \n) ;"
                 "\nEND\n"
             ),
             (),