|
1 | | -import json |
2 | | -import http.client |
3 | | -import datetime |
4 | | -import re |
| 1 | +import sqlparse |
5 | 2 |
|
6 | 3 | from django.db import DatabaseError, Error, DataError, OperationalError, \ |
7 | 4 | IntegrityError, InternalError, ProgrammingError, NotSupportedError, InterfaceError |
8 | | -from django.conf import settings |
9 | | -from django.utils import timezone |
10 | 5 |
|
11 | 6 | class D1Result: |
12 | 7 | lastrowid = None |
@@ -121,81 +116,68 @@ def process_query(self, query, params=None): |
121 | 116 | def run_query(self, query, params=None): |
122 | 117 | proc_query, params = self.process_query(query, params) |
123 | 118 |
|
124 | | - # print(query) |
125 | | - # print(params) |
126 | | - |
127 | 119 | cf_workers = self.import_from_javascript("cloudflare:workers") |
128 | | - # print(dir(cf_workers.env)) |
129 | 120 | db = getattr(cf_workers.env, self.binding) |
130 | 121 |
|
131 | 122 | if params: |
132 | 123 | stmt = db.prepare(proc_query).bind(*params); |
133 | 124 | else: |
134 | 125 | stmt = db.prepare(proc_query); |
135 | 126 |
|
| 127 | + read_only = is_read_only_query(proc_query) |
| 128 | + |
136 | 129 | try: |
137 | | - resp = self.run_sync(stmt.all()) |
| 130 | + if read_only: |
| 131 | + resp = self.run_sync(stmt.raw()) |
| 132 | + else: |
| 133 | + resp = self.run_sync(stmt.all()) |
138 | 134 | except: |
139 | 135 | from js import Error |
140 | 136 | Error.stackTraceLimit = 1e10 |
141 | 137 | raise Error(Error.new().stack) |
142 | 138 |
|
143 | | - results = self._convert_results(resp.results.to_py()) |
144 | | - |
145 | | - # print(results) |
146 | | - # print(f'rowsRead: {resp.meta.rows_read}') |
147 | | - # print(f'rowsWritten: {resp.meta.rows_written}') |
148 | | - # print('---') |
| 139 | + # this is a hack, because D1 Raw method (required for reading rows) doesn't return metadata |
| 140 | + if read_only: |
| 141 | + results = self._convert_results_list(resp.to_py()) |
| 142 | + rows_read = len(results) |
| 143 | + rows_written = 0 |
| 144 | + else: |
| 145 | + results = self._convert_results_dict(resp.results.to_py()) |
| 146 | + rows_read = resp.meta.rows_read |
| 147 | + rows_written = resp.meta.rows_written |
149 | 148 |
|
150 | 149 | return results, { |
151 | | - "rows_read": resp.meta.rows_read, |
152 | | - "rows_written": resp.meta.rows_written, |
| 150 | + "rows_read": rows_read, |
| 151 | + "rows_written": rows_written, |
153 | 152 | } |
154 | 153 |
|
155 | | - def _convert_results(self, data): |
156 | | - """ |
157 | | - Convert any datetime strings in the result set to actual timezone-aware datetime objects. |
158 | | - """ |
159 | | - # print('before') |
160 | | - # print(data) |
| 154 | + def _convert_results_dict(self, data): |
161 | 155 | result = [] |
162 | 156 |
|
163 | 157 | for row in data: |
164 | 158 | row_items = () |
165 | 159 | for k, v in row.items(): |
166 | | - if isinstance(v, str): |
167 | | - v = self._parse_datetime(v) |
168 | 160 | row_items += (v,) |
169 | 161 |
|
170 | 162 | result.append(row_items) |
171 | 163 |
|
172 | | - # print('after') |
173 | | - # print(result) |
| 164 | + return result |
| 165 | + |
| 166 | + def _convert_results_list(self, data): |
| 167 | + result = [] |
| 168 | + |
| 169 | + for row in data: |
| 170 | + row_items = () |
| 171 | + for v in row: |
| 172 | + row_items += (v,) |
| 173 | + |
| 174 | + result.append(row_items) |
| 175 | + |
174 | 176 | return result |
175 | 177 |
|
176 | 178 | query = None |
177 | 179 | params = None |
178 | 180 |
|
179 | | - def _parse_datetime(self, value): |
180 | | - """ |
181 | | - Parse the string value to a timezone-aware datetime object, if applicable. |
182 | | - Handles both datetime strings with and without milliseconds. |
183 | | - Uses Django's timezone utilities for proper conversion. |
184 | | - """ |
185 | | - datetime_formats = ["%Y-%m-%d %H:%M:%S.%f", "%Y-%m-%d %H:%M:%S"] |
186 | | - |
187 | | - for dt_format in datetime_formats: |
188 | | - try: |
189 | | - naive_dt = datetime.datetime.strptime(value, dt_format) |
190 | | - # If Django is using timezones, convert to an aware datetime object |
191 | | - if timezone.is_naive(naive_dt): |
192 | | - return timezone.make_aware(naive_dt, timezone.get_default_timezone()) |
193 | | - return naive_dt |
194 | | - except (ValueError, TypeError): |
195 | | - continue # Try the next format if parsing fails |
196 | | - |
197 | | - return value # If it's not a datetime string, return the original value |
198 | | - |
199 | 181 | def execute(self, query, params=None): |
200 | 182 | if params: |
201 | 183 | newParams = [] |
@@ -252,3 +234,22 @@ def fetchmany(self, size=1): |
252 | 234 |
|
253 | 235 | def close(self): |
254 | 236 | return |
| 237 | + |
| 238 | + |
| 239 | +def is_read_only_query(query: str) -> bool: |
| 240 | + parsed = sqlparse.parse(query.strip()) |
| 241 | + |
| 242 | + if not parsed: |
| 243 | + return False # Invalid or empty query |
| 244 | + |
| 245 | + # Get the first statement |
| 246 | + statement = parsed[0] |
| 247 | + |
| 248 | + # Check if the statement is a SELECT query |
| 249 | + if statement.get_type().upper() == "SELECT": |
| 250 | + return True |
| 251 | + |
| 252 | + # List of modifying query types |
| 253 | + modifying_types = {"INSERT", "UPDATE", "DELETE", "CREATE", "ALTER", "DROP", "REPLACE"} |
| 254 | + |
| 255 | + return statement.get_type().upper() not in modifying_types |
0 commit comments