|
| 1 | +from typing import Literal |
| 2 | +from pydantic import BaseModel, ConfigDict, Field, model_validator |
| 3 | +from pydantic.types import StrictStr, SecretStr, StrictInt |
| 4 | + |
| 5 | + |
| 6 | +class SQLDBConfig(BaseModel): |
| 7 | + """Configuration validator for SQL databases for PostgreSQL and MySQL""" |
| 8 | + |
| 9 | + provider: Literal["postgresql", "mysql"] = "postgresql" |
| 10 | + """Database provider, postgresql or mysql""" |
| 11 | + host: StrictStr = "localhost" |
| 12 | + """PostgreSQL server host""" |
| 13 | + port: StrictInt = 5432 |
| 14 | + """server port, default 5432""" |
| 15 | + database: StrictStr = "hololinked" |
| 16 | + """database name, default hololinked""" |
| 17 | + user: StrictStr = "hololinked" |
| 18 | + """user name, default hololinked, recommended not to use admin user""" |
| 19 | + password: SecretStr = Field(default=SecretStr(""), repr=False) |
| 20 | + """user password, default empty""" |
| 21 | + dialect: Literal["asyncpg", "psycopg", "asyncmy", "mysqldb", ""] = "" |
| 22 | + """dialect to use, default psycopg for postgresql""" |
| 23 | + uri: SecretStr = "" |
| 24 | + """Full database URI, overrides other settings, default empty""" |
| 25 | + |
| 26 | + model_config = ConfigDict(extra="forbid") |
| 27 | + |
| 28 | + @property |
| 29 | + def URL(self) -> str: |
| 30 | + if self.uri: |
| 31 | + return self.uri.get_secret_value() |
| 32 | + if self.provider == "postgresql": |
| 33 | + return f"postgresql{'+' + self.dialect if self.dialect else ''}://{self.user}:{self.password.get_secret_value()}@{self.host}:{self.port}/{self.database}" |
| 34 | + return f"mysql{'+' + self.dialect if self.dialect else ''}://{self.user}:{self.password.get_secret_value()}@{self.host}:{self.port}/{self.database}" |
| 35 | + |
| 36 | + @model_validator(mode="after") |
| 37 | + def _at_least_one(self): |
| 38 | + if not self.uri and (not self.host or not self.port or not self.database or not self.user or not self.password): |
| 39 | + raise ValueError("Provide either database URI or all of 'host', 'port', 'database', 'user', and 'password'") |
| 40 | + return self |
| 41 | + |
| 42 | + |
| 43 | +class SQLiteConfig(BaseModel): |
| 44 | + """Configuration validator for SQLite database""" |
| 45 | + |
| 46 | + provider: Literal["sqlite"] = "sqlite" |
| 47 | + """Database provider, only sqlite is supported""" |
| 48 | + dialect: SecretStr = "pysqlite" |
| 49 | + """dialect to use, aiosqlite for async, pysqlite for sync""" |
| 50 | + file: str = "" |
| 51 | + """SQLite database file, default is empty string, which leads to an DB with name of thing ID""" |
| 52 | + in_memory: bool = False |
| 53 | + """Use in-memory SQLite database, default False as it is not persistent""" |
| 54 | + uri: SecretStr = "" |
| 55 | + """Full database URI, overrides other settings, default empty""" |
| 56 | + |
| 57 | + model_config = ConfigDict(extra="forbid") |
| 58 | + |
| 59 | + @property |
| 60 | + def URL(self) -> str: |
| 61 | + if self.uri: |
| 62 | + return self.uri.get_secret_value() |
| 63 | + if self.in_memory: |
| 64 | + return f"sqlite+{self.dialect}:///:memory:" |
| 65 | + elif self.file: |
| 66 | + return f"sqlite+{self.dialect}:///{self.file}" |
| 67 | + raise NotImplementedError("Either 'uri' or 'file' or 'in_memory' must be provided for SQLiteConfig") |
| 68 | + |
| 69 | + |
| 70 | +class MongoDBConfig(BaseModel): |
| 71 | + """Configuration validator for MongoDB database""" |
| 72 | + |
| 73 | + provider: Literal["mongo"] = "mongo" |
| 74 | + """Database provider, only mongo is supported""" |
| 75 | + host: StrictStr = "localhost" |
| 76 | + """MongoDB server host""" |
| 77 | + port: StrictInt = 27017 |
| 78 | + """server port, default 27017""" |
| 79 | + database: StrictStr = "hololinked" |
| 80 | + """database name, default hololinked""" |
| 81 | + user: StrictStr = "" |
| 82 | + """user name, default empty, recommended not to use admin user""" |
| 83 | + password: SecretStr = "" |
| 84 | + """user password, default empty""" |
| 85 | + authSource: StrictStr = "" |
| 86 | + """authentication source database, default empty""" |
| 87 | + uri: SecretStr = "" |
| 88 | + """Full database URI, overrides other settings, default empty""" |
| 89 | + |
| 90 | + model_config = ConfigDict(extra="forbid") |
| 91 | + |
| 92 | + @property |
| 93 | + def URL(self) -> str: |
| 94 | + if self.uri: |
| 95 | + return self.uri.get_secret_value() |
| 96 | + if self.user and self.password: |
| 97 | + if self.authSource: |
| 98 | + return f"mongodb://{self.user}:{self.password.get_secret_value()}@{self.host}:{self.port}/?authSource={self.authSource}" |
| 99 | + return f"mongodb://{self.user}:{self.password.get_secret_value()}@{self.host}:{self.port}/" |
| 100 | + return f"mongodb://{self.host}:{self.port}/" |
| 101 | + |
| 102 | + @model_validator(mode="after") |
| 103 | + def _at_least_one(self): |
| 104 | + if not self.uri and (not self.host or not self.port or not self.database or not self.user or not self.password): |
| 105 | + raise ValueError("Provide either database URI or all of 'host', 'port', 'database', 'user', and 'password'") |
| 106 | + return self |
0 commit comments