-
Notifications
You must be signed in to change notification settings - Fork 64
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
- Loading branch information
1 parent
5b63cf6
commit b5ffd31
Showing
3 changed files
with
58 additions
and
26 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -8,26 +8,26 @@ | |
# Copyright (c) 2016, Shay Palachy <[email protected]> | ||
|
||
# standard library imports | ||
import datetime | ||
import pickle | ||
import time | ||
import datetime | ||
|
||
pyodbc = None | ||
# third party imports | ||
with suppress(ImportError): | ||
import pyodbc | ||
|
||
# local imports | ||
from .base import _BaseCore, RecalculationNeeded | ||
from .base import RecalculationNeeded, _BaseCore | ||
|
||
class _OdbcCore(_BaseCore): | ||
|
||
class _OdbcCore(_BaseCore): | ||
def __init__( | ||
self, | ||
hash_func, | ||
wait_for_calc_timeout, | ||
connection_string, | ||
table_name, | ||
self, | ||
hash_func, | ||
wait_for_calc_timeout, | ||
connection_string, | ||
table_name, | ||
): | ||
if "pyodbc" not in sys.modules: | ||
warnings.warn( | ||
|
@@ -43,7 +43,8 @@ def __init__( | |
def ensure_table_exists(self): | ||
with pyodbc.connect(self.connection_string) as conn: | ||
cursor = conn.cursor() | ||
cursor.execute(f""" | ||
cursor.execute( | ||
f""" | ||
IF NOT EXISTS (SELECT * FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_NAME = N'{self.table_name}') | ||
BEGIN | ||
CREATE TABLE {self.table_name} ( | ||
|
@@ -54,13 +55,17 @@ def ensure_table_exists(self): | |
PRIMARY KEY (key) | ||
); | ||
END | ||
""") | ||
""" | ||
) | ||
conn.commit() | ||
|
||
def get_entry_by_key(self, key): | ||
with pyodbc.connect(self.connection_string) as conn: | ||
cursor = conn.cursor() | ||
cursor.execute(f"SELECT value, time, being_calculated FROM {self.table_name} WHERE key = ?", key) | ||
cursor.execute( | ||
f"SELECT value, time, being_calculated FROM {self.table_name} WHERE key = ?", | ||
key, | ||
) | ||
row = cursor.fetchone() | ||
if row: | ||
return { | ||
|
@@ -73,34 +78,48 @@ def get_entry_by_key(self, key): | |
def set_entry(self, key, func_res): | ||
with pyodbc.connect(self.connection_string) as conn: | ||
cursor = conn.cursor() | ||
cursor.execute(f""" | ||
cursor.execute( | ||
f""" | ||
MERGE INTO {self.table_name} USING (SELECT 1 AS dummy) AS src ON (key = ?) | ||
WHEN MATCHED THEN | ||
UPDATE SET value = ?, time = GETDATE(), being_calculated = 0 | ||
WHEN NOT MATCHED THEN | ||
INSERT (key, value, time, being_calculated) VALUES (?, ?, GETDATE(), 0); | ||
""", key, pickle.dumps(func_res), key, pickle.dumps(func_res)) | ||
""", | ||
key, | ||
pickle.dumps(func_res), | ||
key, | ||
pickle.dumps(func_res), | ||
) | ||
conn.commit() | ||
|
||
def mark_entry_being_calculated(self, key): | ||
with pyodbc.connect(self.connection_string) as conn: | ||
cursor = conn.cursor() | ||
cursor.execute(f"UPDATE {self.table_name} SET being_calculated = 1 WHERE key = ?", key) | ||
cursor.execute( | ||
f"UPDATE {self.table_name} SET being_calculated = 1 WHERE key = ?", | ||
key, | ||
) | ||
conn.commit() | ||
|
||
def mark_entry_not_calculated(self, key): | ||
with pyodbc.connect(self.connection_string) as conn: | ||
cursor = conn.cursor() | ||
cursor.execute(f"UPDATE {self.table_name} SET being_calculated = 0 WHERE key = ?", key) | ||
cursor.execute( | ||
f"UPDATE {self.table_name} SET being_calculated = 0 WHERE key = ?", | ||
key, | ||
) | ||
conn.commit() | ||
|
||
def wait_on_entry_calc(self, key): | ||
start_time = datetime.datetime.now() | ||
while True: | ||
entry = self.get_entry_by_key(key) | ||
if entry and not entry['being_calculated']: | ||
return entry['value'] | ||
if (datetime.datetime.now() - start_time).total_seconds() > self.wait_for_calc_timeout: | ||
if entry and not entry["being_calculated"]: | ||
return entry["value"] | ||
if ( | ||
datetime.datetime.now() - start_time | ||
).total_seconds() > self.wait_for_calc_timeout: | ||
raise RecalculationNeeded() | ||
time.sleep(1) | ||
|
||
|
@@ -113,5 +132,7 @@ def clear_cache(self): | |
def clear_being_calculated(self): | ||
with pyodbc.connect(self.connection_string) as conn: | ||
cursor = conn.cursor() | ||
cursor.execute(f"UPDATE {self.table_name} SET being_calculated = 0") | ||
cursor.execute( | ||
f"UPDATE {self.table_name} SET being_calculated = 0" | ||
) | ||
conn.commit() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters