Skip to content

Commit 3215a39

Browse files
vishalshenoyvishalshenoy
authored
SQLAlchemy Soft Delete Example with Graph Disabled (#43)
* init * . * . * . * Automated pre-commit update
1 parent 1438f80 commit 3215a39

File tree

2 files changed

+258
-0
lines changed

2 files changed

+258
-0
lines changed

‎examples/sqlalchemy_soft_delete/README.md

+153
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
# SQLAlchemy Soft Delete Codemod
2+
3+
This codemod automatically adds soft delete conditions to SQLAlchemy join queries in your codebase. It ensures that joins only include non-deleted records by adding appropriate `deleted_at` checks.
4+
5+
## Overview
6+
7+
The codemod analyzes your codebase and automatically adds soft delete conditions to SQLAlchemy join methods (`join`, `outerjoin`, `innerjoin`) for specified models. This helps prevent accidentally including soft-deleted records in query results.
8+
9+
## How It Works
10+
11+
The codemod processes your codebase in several steps:
12+
13+
1. **Join Detection**
14+
```python
15+
def should_process_join_call(call, soft_delete_models, join_methods):
16+
if str(call.name) not in join_methods:
17+
return False
18+
19+
call_args = list(call.args)
20+
if not call_args:
21+
return False
22+
23+
model_name = str(call_args[0].value)
24+
return model_name in soft_delete_models
25+
```
26+
- Scans for SQLAlchemy join method calls (`join`, `outerjoin`, `innerjoin`)
27+
- Identifies joins involving soft-deletable models
28+
- Analyzes existing join conditions
29+
30+
2. **Condition Addition**
31+
```python
32+
def add_deleted_at_check(file, call, model_name):
33+
call_args = list(call.args)
34+
deleted_at_check = f"{model_name}.deleted_at.is_(None)"
35+
36+
if len(call_args) == 1:
37+
call_args.append(deleted_at_check)
38+
return
39+
40+
second_arg = call_args[1].value
41+
if isinstance(second_arg, FunctionCall) and second_arg.name == "and_":
42+
second_arg.args.append(deleted_at_check)
43+
else:
44+
call_args[1].edit(f"and_({second_arg.source}, {deleted_at_check})")
45+
```
46+
- Adds `deleted_at.is_(None)` checks to qualifying joins
47+
- Handles different join condition patterns:
48+
- Simple joins with no conditions
49+
- Joins with existing conditions (combines using `and_`)
50+
- Preserves existing conditions while adding soft delete checks
51+
52+
3. **Import Management**
53+
```python
54+
def ensure_and_import(file):
55+
if not any("and_" in imp.name for imp in file.imports):
56+
file.add_import_from_import_string("from sqlalchemy import and_")
57+
```
58+
- Automatically adds required SQLAlchemy imports (`and_`)
59+
- Prevents duplicate imports
60+
61+
## Configuration
62+
63+
### Soft Delete Models
64+
65+
The codemod processes joins for the following models:
66+
```python
67+
soft_delete_models = {
68+
"User",
69+
"Update",
70+
"Proposal",
71+
"Comment",
72+
"Project",
73+
"Team",
74+
"SavedSession"
75+
}
76+
```
77+
78+
### Join Methods
79+
80+
The codemod handles these SQLAlchemy join methods:
81+
```python
82+
join_methods = {"join", "outerjoin", "innerjoin"}
83+
```
84+
85+
## Code Transformations
86+
87+
### Simple Join with Model Reference
88+
```python
89+
# Before
90+
query.join(Project, Session.project)
91+
92+
# After
93+
from sqlalchemy import and_
94+
query.join(Project, and_(Session.project, Project.deleted_at.is_(None)))
95+
```
96+
97+
### Join with Column Equality
98+
```python
99+
# Before
100+
query.join(Project, Session.project_id == Project.id)
101+
102+
# After
103+
from sqlalchemy import and_
104+
query.join(Project, and_(Session.project_id == Project.id, Project.deleted_at.is_(None)))
105+
```
106+
107+
### Multiple Joins in Query Chain
108+
```python
109+
# Before
110+
Session.query.join(Project, Session.project)\
111+
.join(Account, Project.account)\
112+
.outerjoin(Proposal, Session.proposal)
113+
114+
# After
115+
from sqlalchemy import and_
116+
Session.query.join(Project, and_(Session.project, Project.deleted_at.is_(None)))\
117+
.join(Account, Project.account)\
118+
.outerjoin(Proposal, and_(Session.proposal, Proposal.deleted_at.is_(None)))
119+
```
120+
121+
## Graph Disable Mode
122+
123+
This codemod includes support for running without the graph feature enabled. This is useful for the faster processing of large codebases and reduced memory usage.
124+
125+
To run in no-graph mode:
126+
```python
127+
codebase = Codebase(
128+
str(repo_path),
129+
programming_language=ProgrammingLanguage.PYTHON,
130+
config=CodebaseConfig(
131+
feature_flags=GSFeatureFlags(disable_graph=True)
132+
)
133+
)
134+
```
135+
136+
## Running the Conversion
137+
138+
```bash
139+
# Install Codegen
140+
pip install codegen
141+
142+
# Run the conversion
143+
python run.py
144+
```
145+
146+
## Learn More
147+
148+
- [SQLAlchemy Documentation](https://docs.sqlalchemy.org/en/20/)
149+
- [Codegen Documentation](https://docs.codegen.com)
150+
151+
## Contributing
152+
153+
Feel free to submit issues and enhancement requests!

‎examples/sqlalchemy_soft_delete/run.py

+105
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
import codegen
2+
from codegen import Codebase
3+
from codegen.sdk.core.detached_symbols.function_call import FunctionCall
4+
from codegen.sdk.enums import ProgrammingLanguage
5+
import shutil
6+
import subprocess
7+
from pathlib import Path
8+
9+
10+
def should_process_join_call(call, soft_delete_models, join_methods):
11+
"""Determine if a function call should be processed for soft delete conditions."""
12+
if str(call.name) not in join_methods:
13+
return False
14+
15+
call_args = list(call.args)
16+
if not call_args:
17+
return False
18+
19+
model_name = str(call_args[0].value)
20+
return model_name in soft_delete_models
21+
22+
23+
def add_deleted_at_check(file, call, model_name):
24+
"""Add the deleted_at check to a join call."""
25+
call_args = list(call.args)
26+
deleted_at_check = f"{model_name}.deleted_at.is_(None)"
27+
28+
if len(call_args) == 1:
29+
print(f"Adding deleted_at check to function call {call.source}")
30+
call_args.append(deleted_at_check)
31+
return
32+
33+
second_arg = call_args[1].value
34+
if second_arg.source == deleted_at_check:
35+
print(f"Skipping {file.filepath} because the deleted_at check is already present")
36+
return
37+
38+
if isinstance(second_arg, FunctionCall) and second_arg.name == "and_":
39+
if deleted_at_check in {str(x) for x in second_arg.args}:
40+
print(f"Skipping {file.filepath} because the deleted_at check is already present")
41+
return
42+
print(f"Adding deleted_at check to and_ call in {file.filepath}")
43+
second_arg.args.append(deleted_at_check)
44+
else:
45+
print(f"Adding deleted_at check to {file.filepath}")
46+
call_args[1].edit(f"and_({second_arg.source}, {deleted_at_check})")
47+
48+
ensure_and_import(file)
49+
50+
51+
def ensure_and_import(file):
52+
"""Ensure the file has the necessary and_ import."""
53+
if not any("and_" in imp.name for imp in file.imports):
54+
print(f"File {file.filepath} does not import and_. Adding import.")
55+
file.add_import_from_import_string("from sqlalchemy import and_")
56+
57+
58+
def clone_repo(repo_url: str, repo_path: Path) -> None:
59+
"""Clone a git repository to the specified path."""
60+
if repo_path.exists():
61+
shutil.rmtree(repo_path)
62+
subprocess.run(["git", "clone", repo_url, str(repo_path)], check=True)
63+
64+
65+
@codegen.function("sqlalchemy-soft-delete")
66+
def process_soft_deletes(codebase):
67+
"""Process soft delete conditions for join methods in the codebase."""
68+
soft_delete_models = {
69+
"User",
70+
"Update",
71+
"Proposal",
72+
"Comment",
73+
"Project",
74+
"Team",
75+
"SavedSession",
76+
}
77+
join_methods = {"join", "outerjoin", "innerjoin"}
78+
79+
for file in codebase.files:
80+
for call in file.function_calls:
81+
if not should_process_join_call(call, soft_delete_models, join_methods):
82+
continue
83+
84+
model_name = str(list(call.args)[0].value)
85+
print(f"Found join method for model {model_name} in file {file.filepath}")
86+
add_deleted_at_check(file, call, model_name)
87+
88+
codebase.commit()
89+
print("commit")
90+
print(codebase.get_diff())
91+
92+
93+
if __name__ == "__main__":
94+
from codegen.sdk.core.codebase import Codebase
95+
from codegen.sdk.codebase.config import CodebaseConfig, GSFeatureFlags
96+
97+
repo_path = Path("/tmp/core")
98+
repo_url = "https://github.com/hasgeek/funnel.git"
99+
100+
try:
101+
clone_repo(repo_url, repo_path)
102+
codebase = Codebase(str(repo_path), programming_language=ProgrammingLanguage.PYTHON, config=CodebaseConfig(feature_flags=GSFeatureFlags(disable_graph=True)))
103+
process_soft_deletes(codebase)
104+
finally:
105+
shutil.rmtree(repo_path)

0 commit comments

Comments
 (0)