|
6 | 6 |
|
7 | 7 | TODO: Add support for multiple database backends? |
8 | 8 | """ |
| 9 | +import os |
9 | 10 | from pathlib import Path |
10 | 11 | from typing import Iterable |
11 | 12 |
|
12 | | -from playhouse import sqlite_ext |
| 13 | +from playhouse import postgres_ext as ext |
13 | 14 |
|
14 | | -db_proxy = sqlite_ext.DatabaseProxy() |
| 15 | +db_proxy = ext.DatabaseProxy() |
15 | 16 |
|
16 | 17 |
|
17 | | -class Document(sqlite_ext.Model): |
| 18 | +class Document(ext.Model): |
18 | 19 | """Document main model.""" |
19 | 20 |
|
20 | | - page = sqlite_ext.TextField(null=False, unique=True) |
21 | | - title = sqlite_ext.TextField(null=False) |
| 21 | + page = ext.TextField(null=False, unique=True) |
| 22 | + title = ext.TextField(null=False) |
22 | 23 |
|
23 | 24 | class Meta: # noqa: D106 |
24 | 25 | database = db_proxy |
25 | 26 |
|
26 | 27 |
|
27 | | -class Section(sqlite_ext.Model): |
| 28 | +class Section(ext.Model): |
28 | 29 | """Section unit of document.""" |
29 | 30 |
|
30 | | - document = sqlite_ext.ForeignKeyField(Document) |
31 | | - root = sqlite_ext.BooleanField(default=False, null=False) |
32 | | - ref = sqlite_ext.TextField(null=False) |
33 | | - title = sqlite_ext.TextField(null=False) |
34 | | - body = sqlite_ext.TextField(null=False) |
| 31 | + document = ext.ForeignKeyField(Document) |
| 32 | + root = ext.BooleanField(default=False, null=False) |
| 33 | + ref = ext.TextField(null=False) |
| 34 | + title = ext.TextField(null=False) |
| 35 | + body = ext.TextField(null=False) |
35 | 36 |
|
36 | 37 | class Meta: # noqa: D106 |
37 | 38 | database = db_proxy |
38 | 39 |
|
39 | 40 |
|
40 | | -class Content(sqlite_ext.FTS5Model): |
| 41 | +class Content(ext.Model): |
41 | 42 | """Searching model.""" |
42 | 43 |
|
43 | | - rowid = sqlite_ext.RowIDField() |
44 | | - title = sqlite_ext.SearchField() |
45 | | - body = sqlite_ext.SearchField() |
| 44 | + rowid = ext.IntegerField() |
| 45 | + title = ext.TextField() |
| 46 | + body = ext.TextField() |
46 | 47 |
|
47 | 48 | class Meta: # noqa: D106 |
48 | 49 | database = db_proxy |
49 | | - options = {"tokenize": "trigram"} |
| 50 | + # TODO: This is an option from SQLite, it does not work on other DBMS. |
| 51 | + # options = {"tokenize": "trigram"} |
50 | 52 |
|
51 | 53 |
|
52 | 54 | def store_document(document: Document, sections: Iterable[Section]): |
@@ -74,16 +76,25 @@ def search_documents(keyword: str) -> Iterable[Section]: |
74 | 76 | ) |
75 | 77 |
|
76 | 78 |
|
77 | | -def bind(db_path: Path): |
| 79 | +def bind(db_type: str, db_path: Path): |
78 | 80 | """Bind connection. |
79 | 81 |
|
80 | 82 | This works only set db into proxy, not included creating tables. |
81 | 83 | """ |
82 | | - db = sqlite_ext.SqliteExtDatabase(db_path) |
| 84 | + if db_type == "sqlite": |
| 85 | + db = ext.SqliteExtDatabase(db_path) |
| 86 | + elif db_type == "postgresql": |
| 87 | + db = ext.PostgresqlExtDatabase(db_path) |
| 88 | + if "POSTGRES_LOG_STATEMENT" in os.environ: |
| 89 | + db.execute_sql( |
| 90 | + f"SET log_statement='{os.environ['POSTGRES_LOG_STATEMENT']}';" |
| 91 | + ) |
| 92 | + else: |
| 93 | + raise ValueError(f"Unknown database type: {db_type}") |
83 | 94 | db_proxy.initialize(db) |
84 | 95 |
|
85 | 96 |
|
86 | | -def initialize(db_path: Path): |
| 97 | +def initialize(db_type: str, db_path: Path): |
87 | 98 | """Bind connection and create tables.""" |
88 | | - bind(db_path) |
| 99 | + bind(db_type, db_path) |
89 | 100 | db_proxy.create_tables([Document, Section, Content]) |
0 commit comments