|
1 | 1 | from __future__ import with_statement |
2 | 2 |
|
3 | | -import sqlite3 |
4 | 3 | import os |
| 4 | +from os import path |
5 | 5 |
|
| 6 | +from code42cli.errors import Code42CLIError |
6 | 7 | from code42cli.util import get_user_project_path |
7 | 8 |
|
8 | 9 |
|
| 10 | +class Cursor(object): |
| 11 | + def __init__(self, location): |
| 12 | + self._location = location |
| 13 | + self._name = path.basename(location) |
| 14 | + |
| 15 | + @property |
| 16 | + def name(self): |
| 17 | + return self._name |
| 18 | + |
| 19 | + @property |
| 20 | + def value(self): |
| 21 | + with open(self._location) as checkpoint: |
| 22 | + return checkpoint.read() |
| 23 | + |
| 24 | + |
9 | 25 | class BaseCursorStore(object): |
10 | | - _PRIMARY_KEY_COLUMN_NAME = u"cursor_id" |
11 | | - _timestamp_column_name = u"OVERRIDE" |
12 | | - _primary_key = u"OVERRIDE" |
13 | | - |
14 | | - def __init__(self, db_table_name, db_file_path=None): |
15 | | - self._table_name = db_table_name |
16 | | - if db_file_path is None: |
17 | | - db_path = get_user_project_path(u"db") |
18 | | - db_file = u"file_event_checkpoints.db" |
19 | | - db_file_path = os.path.join(db_path, db_file) |
20 | | - |
21 | | - self._connection = sqlite3.connect(db_file_path) |
22 | | - if self._is_empty(): |
23 | | - self._init_table() |
24 | | - |
25 | | - def _get(self, columns, primary_key): |
26 | | - query = u"SELECT {0} FROM {1} WHERE {2}=?" |
27 | | - query = query.format(columns, self._table_name, self._PRIMARY_KEY_COLUMN_NAME) |
28 | | - with self._connection as conn: |
29 | | - cursor = conn.cursor() |
30 | | - cursor.execute(query, (primary_key,)) |
31 | | - return cursor.fetchall() |
32 | | - |
33 | | - def _set(self, column_name, new_value, primary_key): |
34 | | - query = u"UPDATE {0} SET {1}=? WHERE {2}=?".format( |
35 | | - self._table_name, column_name, self._PRIMARY_KEY_COLUMN_NAME |
36 | | - ) |
37 | | - with self._connection as conn: |
38 | | - conn.execute(query, (new_value, primary_key)) |
39 | | - |
40 | | - def _delete(self, primary_key): |
41 | | - query = u"DELETE FROM {0} WHERE {1}=?".format( |
42 | | - self._table_name, self._PRIMARY_KEY_COLUMN_NAME |
43 | | - ) |
44 | | - with self._connection as conn: |
45 | | - conn.execute(query, (primary_key,)) |
46 | | - |
47 | | - def _row_exists(self, primary_key): |
48 | | - query = u"SELECT * FROM {0} WHERE {1}=?" |
49 | | - query = query.format(self._table_name, self._PRIMARY_KEY_COLUMN_NAME) |
50 | | - with self._connection as conn: |
51 | | - cursor = conn.cursor() |
52 | | - cursor.execute(query, (primary_key,)) |
53 | | - query_result = cursor.fetchone() |
54 | | - if not query_result: |
55 | | - return False |
56 | | - return True |
57 | | - |
58 | | - def _drop_table(self): |
59 | | - drop_query = u"DROP TABLE {0}".format(self._table_name) |
60 | | - with self._connection as conn: |
61 | | - conn.execute(drop_query) |
62 | | - |
63 | | - def _is_empty(self): |
64 | | - table_count_query = u""" |
65 | | - SELECT COUNT(name) |
66 | | - FROM sqlite_master |
67 | | - WHERE type='table' AND name=? |
68 | | - """ |
69 | | - with self._connection as conn: |
70 | | - cursor = conn.cursor() |
71 | | - cursor.execute(table_count_query, (self._table_name,)) |
72 | | - query_result = cursor.fetchone() |
73 | | - if query_result: |
74 | | - return int(query_result[0]) <= 0 |
75 | | - |
76 | | - def _init_table(self): |
77 | | - columns = u"{0}, {1}".format(self._PRIMARY_KEY_COLUMN_NAME, self._timestamp_column_name) |
78 | | - create_table_query = u"CREATE TABLE {0} ({1})".format(self._table_name, columns) |
79 | | - with self._connection as conn: |
80 | | - conn.execute(create_table_query) |
81 | | - |
82 | | - def _insert_new_row(self): |
83 | | - insert_query = u"INSERT INTO {0} VALUES(?, null)".format(self._table_name) |
84 | | - with self._connection as conn: |
85 | | - conn.execute(insert_query, (self._primary_key,)) |
86 | | - |
87 | | - def get_stored_cursor_timestamp(self): |
88 | | - """Gets the last stored date observed timestamp.""" |
89 | | - rows = self._get(self._timestamp_column_name, self._primary_key) |
90 | | - if rows and rows[0]: |
91 | | - return rows[0][0] |
| 26 | + def __init__(self, dir_path): |
| 27 | + self._dir_path = dir_path |
92 | 28 |
|
93 | | - def replace_stored_cursor_timestamp(self, new_date_observed_timestamp): |
| 29 | + def get(self, cursor_name): |
| 30 | + """Gets the last stored date observed timestamp.""" |
| 31 | + try: |
| 32 | + location = path.join(self._dir_path, cursor_name) |
| 33 | + with open(location) as checkpoint: |
| 34 | + return float(checkpoint.read()) |
| 35 | + except FileNotFoundError: |
| 36 | + return None |
| 37 | + |
| 38 | + def replace(self, cursor_name, new_timestamp): |
94 | 39 | """Replaces the last stored date observed timestamp with the given one.""" |
95 | | - self._set( |
96 | | - column_name=self._timestamp_column_name, |
97 | | - new_value=new_date_observed_timestamp, |
98 | | - primary_key=self._primary_key, |
99 | | - ) |
| 40 | + location = path.join(self._dir_path, cursor_name) |
| 41 | + with open(location, "w") as checkpoint: |
| 42 | + return checkpoint.write(str(new_timestamp)) |
| 43 | + |
| 44 | + def delete(self, cursor_name): |
| 45 | + """Removes a single cursor from the store.""" |
| 46 | + try: |
| 47 | + location = path.join(self._dir_path, cursor_name) |
| 48 | + os.remove(location) |
| 49 | + except FileNotFoundError: |
| 50 | + msg = "No checkpoint named {0} exists for this profile.".format(cursor_name) |
| 51 | + raise Code42CLIError(msg) |
100 | 52 |
|
101 | 53 | def clean(self): |
102 | | - """Removes profile cursor data from store.""" |
103 | | - self._delete(self._primary_key) |
| 54 | + """Removes all cursors from this store.""" |
| 55 | + cursors = self.get_all_cursors() |
| 56 | + for cursor in cursors: |
| 57 | + self.delete(cursor.name) |
104 | 58 |
|
| 59 | + def get_all_cursors(self): |
| 60 | + """Returns a list of all cursors stored in this directory (which istypically scoped to a profile).""" |
| 61 | + dir_contents = os.listdir(self._dir_path) |
| 62 | + return [Cursor(f) for f in dir_contents if self._is_file(f)] |
105 | 63 |
|
106 | | -class FileEventCursorStore(BaseCursorStore): |
107 | | - _timestamp_column_name = u"insertionTimestamp" |
| 64 | + def _is_file(self, node_name): |
| 65 | + return path.isfile(path.join(self._dir_path, node_name)) |
108 | 66 |
|
109 | | - def __init__(self, profile_name, db_file_path=None): |
110 | | - self._primary_key = profile_name |
111 | | - super(FileEventCursorStore, self).__init__(u"file_event_checkpoints", db_file_path) |
112 | | - if not self._row_exists(self._primary_key): |
113 | | - self._insert_new_row() |
114 | 67 |
|
| 68 | +class FileEventCursorStore(BaseCursorStore): |
| 69 | + def __init__(self, profile_name): |
| 70 | + dir_path = get_user_project_path(u"file_event_checkpoints", profile_name) |
| 71 | + super(FileEventCursorStore, self).__init__(dir_path) |
115 | 72 |
|
116 | | -class AlertCursorStore(BaseCursorStore): |
117 | | - _timestamp_column_name = u"createdAt" |
118 | 73 |
|
119 | | - def __init__(self, profile_name, db_file_path=None): |
120 | | - self._primary_key = profile_name |
121 | | - super(AlertCursorStore, self).__init__(u"alert_checkpoints", db_file_path) |
122 | | - if not self._row_exists(self._primary_key): |
123 | | - self._insert_new_row() |
| 74 | +class AlertCursorStore(BaseCursorStore): |
| 75 | + def __init__(self, profile_name): |
| 76 | + dir_path = get_user_project_path(u"alert_checkpoints", profile_name) |
| 77 | + super(AlertCursorStore, self).__init__(dir_path) |
124 | 78 |
|
125 | 79 |
|
126 | 80 | def get_file_event_cursor_store(profile_name): |
|
0 commit comments