diff --git a/python/semantic_kernel/connectors/sql_server.py b/python/semantic_kernel/connectors/sql_server.py index 3d0d35abe0d2..c8aa904e8d7a 100644 --- a/python/semantic_kernel/connectors/sql_server.py +++ b/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 diff --git a/python/tests/unit/connectors/memory/test_sql_server.py b/python/tests/unit/connectors/memory/test_sql_server.py index 40ef5539d0db..24d235de524f 100644 --- a/python/tests/unit/connectors/memory/test_sql_server.py +++ b/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" ), (),