README.md

---
library_name: transformers
tags:
- text-to-SQL
- SQL
- code-generation
- NLQ-to-SQL
- text2SQL
- Security
- Vulnerability detection
datasets:
- salmane11/SQLShield
language:
- en
base_model:
- microsoft/codebert-base
---

# SQLQueryShield

## Model Description

SQLQueryShield is a vulnerable SQL query detection model. It classifies SQL queries as either vulnerable (e.g., prone to SQL injection or unsafe execution) or benign (safe to execute).

The checkpoint included in this repository is based on [microsoft/codebert-base](https://huggingface.co/microsoft/codebert-base) and further finetuned on [SQLShield](https://huggingface.co/datasets/salmane11/SQLShield), a dataset dedicated to text-to-SQL vulnerability detection composed of vulnerable and safe NLQs and their related SQL queries.


## Finetuning Procedure
The model was fine-tuned using the Hugging Face Transformers library. The following steps were used:

1. Dataset: SSQLShield, only the SQL queries from the (NLQ, SQL) pairs were used for training.

2. Preprocessing:

    - Input Format: Raw SQL query strings.

    - Tokenization: Tokenized using microsoft/codebert-base.

    - Max Length: 128 tokens.

    - Padding and truncation applied.

## Intended Use and Limitations

SQLQueryShield is intended for use as a post-generation filter or analysis tool in any system that executes or generates SQL queries. Its main role is to detect whether a SQL query is potentially harmful due to vulnerability patterns such as SQL injection, improper string concatenation, or unsafe expressions.

Ideal use cases:

    - Filtering SQL queries in Text-to-SQL applications

    - Post-processing or validating user-generated SQL before execution


## How to Use

Example 1: Malicious

```python
from transformers import pipeline

sql_query_shield = pipeline("text-classification", model="salmane11/SQLQueryShield")

# For the following Table schema
# CREATE TABLE campuses
#   (
#      campus   VARCHAR,
#      location VARCHAR
#   )

query = "SELECT campus FROM campuses WHERE location = '' UNION SELECT database() --"

prediction = sql_query_shield(query)
print(prediction)
#[{'label': 'MALICIOUS', 'score': 0.9995294809341431}]
```


Example 2: Safe

```python
from transformers import pipeline

sql_query_shield = pipeline("text-classification", model="salmane11/SQLQueryShield")

# For the following Table schema
# CREATE TABLE tv_channel
#   (
#      package_option VARCHAR,
#      series_name    VARCHAR
#   ) 

query = "SELECT package_option FROM tv_channel WHERE series_name = 'Sky Radio'"


prediction = sql_query_shield(query)
print(prediction)
#[{'label': 'SAFE', 'score': 0.999503493309021}]
```


## Cite our work

Citation