database.py 6.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253
  1. import datetime
  2. import json
  3. import settings
  4. import sqlite3
  5. # Open (or create database)
  6. db = sqlite3.connect(settings.database_path)
  7. # Returns SQLite rows as dictionaries instead of tuples.
  8. # https://docs.python.org/3/library/sqlite3.html#sqlite3.Connection.row_factory
  9. db.row_factory = sqlite3.Row
  10. # Create database tables for federation if they do not exist, for example if
  11. # it's a new database.
  12. with db:
  13. # Actors' private keys used for signing Activities
  14. db.execute("""
  15. CREATE TABLE IF NOT EXISTS private_key (
  16. actor TEXT NOT NULL,
  17. key BLOB NOT NULL,
  18. PRIMARY KEY (actor, key)
  19. )""")
  20. # List of remote actors that follow a local actor
  21. db.execute("""
  22. CREATE TABLE IF NOT EXISTS follower (
  23. local_actor TEXT NOT NULL,
  24. remote_actor TEXT NOT NULL,
  25. PRIMARY KEY (local_actor, remote_actor)
  26. )""")
  27. db.execute("CREATE INDEX IF NOT EXISTS remote_actor_idx ON follower (remote_actor ASC)")
  28. db.execute("CREATE INDEX IF NOT EXISTS local_actor_idx ON follower (local_actor ASC)")
  29. # List of local actors that follow a remote actor
  30. db.execute("""
  31. CREATE TABLE IF NOT EXISTS following (
  32. local_actor TEXT NOT NULL,
  33. remote_actor TEXT NOT NULL,
  34. PRIMARY KEY (local_actor, remote_actor)
  35. )""")
  36. db.execute("CREATE INDEX IF NOT EXISTS local_actor_idx ON following (local_actor ASC)")
  37. db.execute("CREATE INDEX IF NOT EXISTS remote_actor_idx ON following (remote_actor ASC)")
  38. # INBOX table used to store incoming messages
  39. db.execute("""
  40. CREATE TABLE IF NOT EXISTS inbox (
  41. datetime TEXT NOT NULL,
  42. actor TEXT NOT NULL,
  43. message TEXT NOT NULL
  44. )""")
  45. db.execute("CREATE INDEX IF NOT EXISTS dt_idx ON inbox (datetime ASC)")
  46. db.execute("CREATE INDEX IF NOT EXISTS actor_idx ON inbox (actor ASC)")
  47. # OUTBOX table used to store outgoing messages
  48. db.execute("""
  49. CREATE TABLE IF NOT EXISTS outbox (
  50. activity_id TEXT UNIQUE NOT NULL,
  51. datetime TEXT NOT NULL,
  52. actor TEXT NOT NULL,
  53. message TEXT NOT NULL
  54. )""")
  55. db.execute("CREATE INDEX IF NOT EXISTS dt_idx ON outbox (datetime ASC)")
  56. db.execute("CREATE INDEX IF NOT EXISTS actor_idx ON outbox (actor ASC)")
  57. # Create database tables for a demo forge if they do not exist
  58. with db:
  59. # Tickets table. This forge only keeps track of tickets number, so it can
  60. # generate new IDs.
  61. db.execute("""
  62. CREATE TABLE IF NOT EXISTS ticket (
  63. repository TEXT NOT NULL,
  64. counter INTEGER NOT NULL DEFAULT 0,
  65. PRIMARY KEY (repository)
  66. )""")
  67. # Get the private key of a local actor
  68. def get_private_key(local_actor):
  69. with db:
  70. cursor = db.execute (
  71. """
  72. SELECT key
  73. FROM private_key
  74. WHERE actor = ?
  75. """,
  76. [ local_actor ]
  77. )
  78. results = cursor.fetchone()
  79. return None if not results else results['key']
  80. # Set the private key of a local actor
  81. def set_private_key(actor, key):
  82. with db:
  83. db.execute (
  84. """
  85. INSERT OR IGNORE INTO private_key (actor, key)
  86. VALUES (?, ?)
  87. """,
  88. [ actor, key ]
  89. )
  90. # Add a remote follower to a local actor
  91. def add_follower(local_actor, remote_actor):
  92. with db:
  93. db.execute (
  94. """
  95. INSERT OR IGNORE INTO follower (local_actor, remote_actor)
  96. VALUES (?, ?)
  97. """,
  98. [ local_actor, remote_actor ]
  99. )
  100. # Add a local actor that is following a remote actor
  101. def add_following(local_actor, remote_actor):
  102. with db:
  103. db.execute (
  104. """
  105. INSERT OR IGNORE INTO following (local_actor, remote_actor)
  106. VALUES (?, ?)
  107. """,
  108. [ local_actor, remote_actor ]
  109. )
  110. # Get the list of remote actors that follow a local actor
  111. def get_followers(local_actor):
  112. with db:
  113. cursor = db.execute (
  114. """
  115. SELECT remote_actor
  116. FROM follower
  117. WHERE local_actor = ?
  118. """,
  119. [ local_actor ]
  120. )
  121. results = cursor.fetchall()
  122. followers = [ follower['remote_actor'] for follower in results ]
  123. return followers
  124. # Store a new message in the INBOX
  125. def store_inbox(actor, message):
  126. with db:
  127. db.execute (
  128. """
  129. INSERT OR IGNORE INTO inbox (datetime, actor, message)
  130. VALUES (?, ?, ?)
  131. """,
  132. [ datetime.datetime.utcnow().replace(tzinfo=datetime.timezone.utc).isoformat(),
  133. actor, json.dumps(message) ]
  134. )
  135. # Store a new message in the OUTBOX
  136. def store_outbox(actor, message):
  137. with db:
  138. db.execute (
  139. """
  140. INSERT OR IGNORE INTO outbox (activity_id, datetime, actor, message)
  141. VALUES (?, ?, ?, ?)
  142. """,
  143. [ message['id'],
  144. datetime.datetime.utcnow().replace(tzinfo=datetime.timezone.utc).isoformat(),
  145. actor, json.dumps(message) ]
  146. )
  147. # Get all INBOX messages for an actor
  148. def get_inbox_messages(actor):
  149. with db:
  150. cursor = db.execute (
  151. """
  152. SELECT datetime, message
  153. FROM inbox
  154. WHERE actor = ?
  155. ORDER BY datetime DESC
  156. """,
  157. [ actor ]
  158. )
  159. results = cursor.fetchall()
  160. messages = [ json.loads(result['message']) for result in results ]
  161. return messages
  162. # Get an Activity that was generated by a local actor
  163. def get_activity(id):
  164. with db:
  165. cursor = db.execute (
  166. """
  167. SELECT message
  168. FROM outbox
  169. WHERE activity_id = ?
  170. """,
  171. [ id ]
  172. )
  173. results = cursor.fetchone()
  174. return None if not results else results['message']
  175. def add_ticket(repository):
  176. with db:
  177. db.execute('BEGIN TRANSACTION')
  178. # Create a new row if doesn't exist
  179. db.execute ("INSERT OR IGNORE INTO ticket (repository) VALUES(?)", [ repository ])
  180. # Increment counter
  181. db.execute (
  182. """
  183. UPDATE ticket
  184. SET counter = counter + 1
  185. WHERE repository = ?
  186. """,
  187. [ repository ]
  188. )
  189. cursor = db.execute (
  190. """
  191. SELECT counter
  192. FROM ticket
  193. WHERE repository = ?
  194. """,
  195. [ repository ]
  196. )
  197. return cursor.fetchone()['counter']