"""Database factory.
This module represents one class `DBFactory` as a start point to any connection with the database.
"""
from __future__ import annotations
import copy
import sys
import inspect
from contextlib import contextmanager, asynccontextmanager, nullcontext
from .exceptions import *
from .db_table import *
from .db_field import *
from .db_types import *
from .db_types import DBTableT, Unassigned
from .translator import Translator
from .helpers import make_async, hybrid_contextmanager
import typing
if typing.TYPE_CHECKING:
from typing import *
from types import SimpleNamespace
from .db_query import DBQuery, DBQueryAsync, DBSQL
from .db_protocol import *
__all__ = ['DBFactory', 'DBFactoryAsync']
[docs]
class DBFactory:
"""Basic database factory class
Use this class to arrange connection to a database, parse table models and run queries.
Note:
It supports a Postgres database only at the moment
Example:
db = DBFactory.postgres(conninfo="postgresql://quazy:quazy@localhost/quazy")
db._debug_mode = True
db.bind_module()
db.clear()
db.create()
...
"""
async_mode: ClassVar[bool] = False
def __init__(self,
connection_pool: DBPoolLike,
translator: type[Translator],
named_factory: type[DBRowFactoryLike],
dict_factory: type[DBRowFactoryLike],
class_factory: type[DBRowFactoryLike]):
"""Basic constructor for database connection
Create a connection via a specific connection pool. Intended to use internally.
Note:
Normally is not intended to run manually, as it intended to use of proxy constructors
like :meth:`postgres()` below
Note:
It is highly appreciated to work with databases supported pools.
Args:
connection_pool: Connection (or pool) to work with the database
translator: Translator class specified to a database type
named_factory: Database row factory class for named results
dict_factory: Database row factory class for dictionary results
class_factory: Database row factory class for class instance results
"""
self._connection_pool: DBPoolLike = connection_pool
self._translator = translator
self._tables: dict[str, type[DBTable]] = dict()
self._named_factory = named_factory
self._dict_factory = dict_factory
self._class_factory = class_factory
self._debug_mode = False
@property
def translator(self):
"""Translator contains all database specific primitives to construct SQL code"""
return self._translator
[docs]
@staticmethod
def postgres(conninfo: str, **kwargs) -> DBFactory | None:
"""Proxy constructor Postgres specific connection
Args:
conninfo: connection string to connect to a database
**kwargs: all keywords arguments passed to `psycopg` constructor
Returns:
DBFactory instance or None
Example::
db = DBFactory.postgres("postgresql://quazy:quazy@localhost/quazy")
"""
import psycopg
from psycopg.rows import namedtuple_row, dict_row, kwargs_row
from .translator_psql import TranslatorPSQL
connection = psycopg.connect(conninfo, **kwargs)
connection.close()
class PsycopgConnection:
conn: psycopg.Connection = None
def connection(self) -> ContextManager[DBConnectionLike] | DBConnectionLike:
if self.conn is None or self.conn.closed:
self.conn = psycopg.connect(conninfo, **kwargs)
return self.conn
@contextmanager
def cursor(self, read_only: bool=False) -> Generator[DBCursorLike]:
with self.connection().cursor(binary=True) as curr:
yield curr
if not read_only:
self.conn.commit()
def close(self):
self.conn.close()
return DBFactory(PsycopgConnection(), TranslatorPSQL, namedtuple_row, dict_row, kwargs_row)
[docs]
@staticmethod
def postgres_pool(conninfo: str, **kwargs) -> DBFactory | None:
"""Proxy constructor Postgres specific pool
Args:
conninfo: connection string to connect to a database
**kwargs: all keywords arguments passed to `psycopg` constructor, except `debug_mode`
Returns:
:class:`DBFactory` instance or None
Example::
db = DBFactory.postgres_pool(conninfo="postgresql://quazy:quazy@localhost/quazy")
"""
from psycopg_pool.pool import ConnectionPool
from psycopg.rows import namedtuple_row, dict_row, kwargs_row
from .translator_psql import TranslatorPSQL
try:
pool = ConnectionPool(conninfo, kwargs=kwargs)
pool.wait()
except Exception as e:
print(str(e))
return None
class PsycopgPoolConnection:
def connection(self) -> ContextManager[DBConnectionLike]:
return pool.connection()
@contextmanager
def cursor(self, read_only:bool=False) -> Generator[DBCursorLike]:
with self.connection() as conn:
with conn.cursor(binary=True) as curr:
yield curr
def close(self):
pool.close()
return DBFactory(PsycopgPoolConnection(), TranslatorPSQL, namedtuple_row, dict_row, kwargs_row)
[docs]
@staticmethod
def sqlite(conn_uri: str, **kwargs) -> DBFactory | None:
"""Proxy constructor SQLite specific connection
Args:
conn_uri: connection string to connect to a database
kwargs: all keywords arguments passed to `sqlite3` constructor, except `debug_mode`
Example::
db = DBFactory.sqlite("file:test.db?mode=rwc")
"""
import sqlite3
from .translator_sqlite import (TranslatorSQLite, namedtuple_row, dict_row, kwargs_row,
ConnectionFactory)
detect_types = kwargs.pop("detect_types", sqlite3.PARSE_DECLTYPES)
connection = sqlite3.connect(conn_uri, uri=True, detect_types=detect_types, factory=ConnectionFactory, **kwargs)
class SQLiteConnection:
@contextmanager
def connection(self) -> Generator[sqlite3.Connection]:
yield connection
def cursor(self, read_only:bool=False) -> ContextManager[DBCursorLike]:
return connection.cursor()
def close(self):
connection.close()
return DBFactory(SQLiteConnection(), TranslatorSQLite, namedtuple_row, dict_row, kwargs_row)
[docs]
def close(self) -> Awaitable[None] | None:
"""Close the database connection and release all resources"""
self._connection_pool.close()
def _bind_subtables(self, cls: type[DBTable]):
for subtab in cls.DB.subtables.values():
subtab.DB.db = self
subtab.DB.source_schema = cls.DB.source_schema
if not subtab.DB.schema:
subtab.DB.schema = cls.DB.schema
[docs]
def bind(self, cls: type[DBTable], schema: str = 'public'):
"""Bind a specific table to the factory instance
Args:
cls: Table to bind
schema: schema name
"""
if cls.__qualname__ in self._tables:
return
cls.DB.db = self
cls.DB.source_schema = schema
if not cls.DB.schema:
cls.DB.schema = schema
self._bind_subtables(cls)
self._tables[cls.__qualname__] = cls
[docs]
def bind_module(self, name: str = None, schema: str = 'public'):
"""Bind a specific module by name to the database factory
Args:
name: module name to bind. If not specified, bind current module
schema: schema name to use
"""
if name:
if name in sys.modules:
globalns = vars(sys.modules[name])
else:
__import__(name)
globalns = vars(sys.modules[name])
else:
globalns = sys._getframe(1).f_globals
if s := globalns.get('_SCHEMA_'):
schema = s
tables: list[type[DBTable]] = list()
for v in globalns.values():
if inspect.isclass(v) and v is not DBTable and issubclass(v, DBTable) and not v.DB.meta:
tables.append(v)
self.bind(v, schema)
for table in tables:
table.resolve_types(globalns)
for table in tables:
self._bind_subtables(table)
for table in tables:
table.resolve_types_many(lambda t: self.bind(t, schema))
for table in tables:
table.setup_validators()
[docs]
def unbind(self, schema: str = "public"):
"""Unbind all tables
Arguments:
schema: schema name to unbind. "Public" by default. If "None", unbind all tables from all schemas.
"""
if schema is None:
self._tables.clear()
else:
for name, table in list(self._tables.items()):
if table.DB.source_schema == schema:
del self._tables[name]
def __contains__(self, item: str | DBTable) -> bool:
"""Check if table exists in the database factory"""
if isinstance(item, str):
return item in self._tables
else:
return item in self._tables.values()
def __getitem__(self, item: str) -> type[DBTable]:
"""Get DBTable class by name"""
return self._tables[item]
[docs]
def query(self, table_class: Optional[type[DBTableT]] = None, name: Optional[str] = None) -> DBQuery[DBTableT]:
"""Create :class:`DBQuery` instance
Create a :class:`DBQuery` instance bound to a specified DBTable or to the whole schema
Args:
table_class: DBTable class to use or None
name: name of the query for subquery request
Returns:
:class:`DBQuery` instance
Example::
q = db.query(Street).select('name')
q = db.query().select(name=lambda s: s.street.name)
"""
from .db_query import DBQuery
return DBQuery[DBTableT](self, table_class, name)
[docs]
def get(self, table_class: type[DBTableT], pk: Any = None, **fields) -> Awaitable[DBTableT] | DBTableT:
"""Request one row from a database table
Hint:
This method is not intended to be used directly. It is shorter to call the `get` method from the `DBTable` instance.
Args:
table_class: DBTable class to use
pk: primary key value to filter row (optional)
**fields: field values to filter row if no pk is specified (optional)
Returns:
:class:`DBTable` instance
"""
query = self.query(table_class)
if pk is not None:
query.filter(pk=pk)
for k, v in fields.items():
query.filter(lambda s: getattr(s, k) == v)
return query.fetch_one()
[docs]
@hybrid_contextmanager
def connection(self) -> AsyncGenerator[DBConnectionLike] | Generator[DBConnectionLike]:
"""Context manager for connection"""
with self._connection_pool.connection() as conn:
yield conn
@hybrid_contextmanager
def cursor(self, read_only: bool=False, _curr: DBCursorLike=None) -> AsyncGenerator[DBCursorLike] | Generator[DBCursorLike]:
with self._connection_pool.cursor(read_only) if _curr is None else nullcontext(_curr) as cursor:
yield cursor
[docs]
def clear(self, schema: str = None) -> Awaitable[None] | None:
"""Drop all known tables in a database
Args:
schema: schema name to use
Warning:
It works without attention. Please double-check before calling.
"""
with self.cursor() as curr: # type: DBConnectionLike
tables = []
for res in curr.execute(self._translator.select_all_tables()):
if self._translator.supports_schema and (not schema or schema == res[0]):
tables.append(f'"{res[0]}"."{res[1]}"')
else:
tables.append(f'"{res[0]}"')
for table_name in tables:
curr.execute(self._translator.drop_table_by_name(table_name))
[docs]
def all_tables(self, schema: str = None, for_stub: bool = False) -> list[type[DBTable]]:
"""Get all known tables in the database
Args:
schema: schema name to use
for_stub: whether to return tables only for stub purposes (uses internally)
Returns:
List of table classes
"""
all_tables = list(self._tables.values())
for table in self._tables.values():
all_tables.extend(table.DB.subtables.values())
for t in all_tables.copy():
if schema and t.DB.source_schema != schema:
all_tables.remove(t)
if for_stub:
def add_bases(t):
if t not in all_tables:
all_tables.insert(0, t)
for base in t.__bases__:
if issubclass(base, DBTable) and base is not DBTable:
add_bases(base)
for t in all_tables.copy():
add_bases(t)
return all_tables
def missed_tables(self, schema: str = None) -> Awaitable[list[type[DBTable]]] | list[type[DBTable]]:
"""Get all tables added as models but not created in the database yet
:meta private:
"""
all_tables = self.all_tables(schema)
with self.select(self._translator.select_all_tables()) as created_tables_query:
#created_tables = [(t.schema, t.table) for t in created_tables_query]
created_tables = created_tables_query.fetchall()
for table in all_tables.copy():
if (table.DB.source_schema, table.DB.table) in created_tables:
all_tables.remove(table)
return all_tables
def check(self, schema: str = None) -> bool:
"""Check all tables are created in the database
:meta private:
"""
return len(self.missed_tables(schema)) == 0
def table_exists(self, table: DBTable) -> Awaitable[bool] | bool:
"""Check if a table exists in the database
:meta private:"""
with self.select(self._translator.is_table_exists(table)) as res:
return res.fetch_one()[0]
[docs]
def create(self, schema: str = None) -> Awaitable[None] | None:
"""Create all added tables in the database"""
all_tables = self.missed_tables(schema)
if not all_tables:
return
all_schemas = set()
for table in all_tables:
if table.DB.schema:
all_schemas.add(table.DB.schema)
with self.cursor() as curr:
# create schemas
if self._translator.supports_schema:
for schema in all_schemas:
curr.execute(self._translator.create_schema(schema))
# create tables
for table in all_tables:
if not table.DB.extendable or table.DB.is_root:
curr.execute(self._translator.create_table(table))
else:
for field_name in table.DB.self_fields:
curr.execute(self._translator.add_field(table, table.DB.fields[field_name]))
# create foreign keys
for table in all_tables:
for field_name, field in table.DB.fields.items():
if table.DB.extendable and field_name not in table.DB.self_fields:
continue
if field.ref and not field.property:
curr.execute(self._translator.add_reference(table, field))
# create indices
for table in all_tables:
for field_name, field in table.DB.fields.items():
if table.DB.extendable and field_name not in table.DB.self_fields:
continue
if field.indexed:
curr.execute(self._translator.create_index(table, field))
[docs]
def insert(self, item: DBTableT, _curr: DBCursorLike=None) -> Awaitable[DBTableT] | DBTableT:
"""Insert item into the database
Args:
item: instance of DBTable to insert
Returns:
updated item, just for chain calls
Note:
Inserted item could contain subtables of subclasses and many-to-many sets.
In this case they will be bulk inserted.
"""
item._before_insert(self)
fields: list[tuple[DBField, Any]] = []
for name, field in item.DB.fields.items():
if field.cid:
fields.append((field, item.DB.discriminator))
continue
elif field.body:
continue
elif field.required and not field.pk and field.default is Unassigned and not field.default_sql:
value = getattr(item, name, None)
if value is None:
raise QuazyMissedField(f"Field `{name}` value is missed for `{item.__class__.__name__}`")
else:
value = getattr(item, name, DefaultValue)
if inspect.isfunction(value):
# materialize values on saving
value = value()
object.__setattr__(item, name, value)
fields.append((field, value))
with self.cursor(_curr=_curr) as curr:
sql, values = getattr(self._translator, "insert")(item, fields)
if self._debug_mode: print(sql)
curr.execute(sql, values)
rows = curr.fetchone()
item.pk = rows[0]
item._modified_fields_.clear()
for field_name, table in item.DB.subtables.items():
for row in getattr(item, field_name):
setattr(row, item.DB.table, item)
self.insert(row, _curr=curr)
if not item.DB.use_slots:
for field_name, field in item.DB.many_fields.items():
for row in getattr(item, field_name):
if getattr(row, field.foreign_field) != item.pk:
setattr(row, field.foreign_field, item.pk)
self.save(row, _curr=curr)
for field_name, field in item.DB.many_to_many_fields.items():
for row in getattr(item, field_name):
if not row.pk:
self.save(row, _curr=curr)
new_indices_sql = self._translator.insert_many_index(field.middle_table, item.DB.table,
field.foreign_table.DB.table)
if self._translator.supports_copy:
with curr.copy(new_indices_sql) as copy:
for row in getattr(item, field_name):
copy.write_row((item.pk, row.pk))
else:
curr.executemany(new_indices_sql, [(item.pk, row.pk) for row in getattr(item, field_name)])
item._after_insert(self)
return item
[docs]
def update(self, item: DBTableT, _curr: DBCursorLike=None) -> Awaitable[DBTableT] | DBTableT:
"""Update item changes to a database
Args:
item: instance of DBTable to insert
Returns:
updated item, just for chain calls
Note:
If an item instance has subtables of subclasses, they will be bulk updated (delete and insert).
"""
item._before_update(self)
fields: list[tuple[DBField, Any]] = []
for name in item._modified_fields_:
field = item.DB.fields[name]
fields.append((field, getattr(item, name, DefaultValue)))
with self.cursor(_curr=_curr) as curr:
if fields:
sql, values = self._translator.update(item.__class__, fields)
values['pk'] = getattr(item, item.DB.pk.name)
curr.execute(sql, values)
item._modified_fields_.clear()
for table in item.DB.subtables.values():
sql = self._translator.delete_related(table, item.DB.table)
curr.execute(sql, (getattr(item, item.DB.pk.name), ))
for row in getattr(item, table.DB.snake_name):
self.insert(row, _curr=curr)
for field_name, field in item.DB.many_fields.items():
for row in getattr(item, field_name):
if getattr(row, field.foreign_field) != item:
setattr(row, field.foreign_field, item)
self.save(row, _curr=curr)
for field_name, field in item.DB.many_to_many_fields.items():
for row in getattr(item, field_name):
if not row.pk:
self.save(row, _curr=curr)
# delete old items, add new items
new_indices = set(row.pk for row in getattr(item, field_name))
old_indices_sql = self._translator.select_many_indices(field.middle_table, item.DB.table, field.foreign_table.DB.table)
curr.execute(old_indices_sql, {"value": item.pk})
results = curr.fetchone()
old_indices = set(results[0]) if results[0] else set()
indices_to_delete = list(old_indices - new_indices)
indices_to_add = list(new_indices - old_indices)
if indices_to_delete:
delete_indices_sql = self._translator.delete_many_indices(field.middle_table, item.DB.table, field.foreign_table.DB.table)
curr.execute(delete_indices_sql, {"value": item.pk, "indices": indices_to_delete})
if indices_to_add:
new_indices_sql = self._translator.insert_many_index(field.middle_table, item.DB.table, field.foreign_table.DB.table)
if self._translator.supports_copy:
with curr.copy(new_indices_sql) as copy:
for index in indices_to_add:
copy.write_row((item.pk, index))
else:
curr.executemany(new_indices_sql, [(item.pk, index) for index in indices_to_add])
item._after_update(self)
return item
[docs]
@hybrid_contextmanager
def select(self, query: Union[DBQuery, str], as_dict: bool = False) -> (
AsyncGenerator[Iterator[DBTable | SimpleNamespace | dict[str, Any]]] |
Generator[Iterator[DBTable | SimpleNamespace | dict[str, Any]]]):
"""Select items from the database
It performs a prepared query to a database and yields results. Is not intended for direct calls.
Use preferably `select` method of DBQuery instance.
Result type depends on query type:
if a query is based on a specific DBTable, the result is an instance of DBTable
if a query is based on a whole schema, the result is SimpleNamespace
Args:
query: instance of DBQuery or string
as_dict: results yield as dict instead of instance of :class:`DBTable` or SimpleNamespace or named tuple
Warning:
This method is intended only for data reading. To make modification request, use :meth:`execute` instead.
Yields:
instance of :class:`DBTable` or SimpleNamespace or dict
"""
with self.cursor(read_only=True) as curr:
if not isinstance(query, str):
if not query.is_frozen:
sql = self._translator.select(query)
if self._debug_mode: print(sql)
else:
sql = query.frozen_sql
if as_dict:
row_factory = self._dict_factory
elif query.fetch_objects:
row_factory = self._class_factory(lambda **kwargs: query.table_class.raw(_db_=self, **kwargs))
else:
row_factory = self._named_factory
curr.row_factory = row_factory
yield curr.execute(sql, query.args)
else:
curr.row_factory = self._dict_factory if as_dict else self._named_factory
yield curr.execute(query)
[docs]
@hybrid_contextmanager
def execute(self, query: str, as_dict: bool = False) -> (
AsyncGenerator[Iterator[DBTable | SimpleNamespace | dict[str, Any]]] |
Generator[Iterator[DBTable | SimpleNamespace | dict[str, Any]]]):
"""Execute text query directly into database.
Similar to :meth:`select` method, but allow changes.
Args:
query: string
as_dict: results yield as dict instead named tuple
"""
with self.cursor() as curr:
curr.row_factory = self._dict_factory if as_dict else self._named_factory
yield curr.execute(query)
[docs]
def update_many(self, query: DBQuery[DBTableT], **values) -> Awaitable[None] | None:
"""Update items in the database by a query
Arguments:
query: instance of DBQuery bound to a table
**values: dict of values to update
"""
fields: list[tuple[DBField, Any]] = []
for name, value in values.items():
field = query.table_class.DB.fields[name]
fields.append((field, value))
with self.cursor() as curr:
sql, values = self._translator.update(query.table_class, fields, query)
if self._debug_mode: print(sql)
curr.execute(sql, query.args | values)
[docs]
def describe(self, query: Union[DBQuery[DBTableT], str]) -> list[DBField]:
"""Describe query result fields information
It performs a prepared query to a database requesting zero rows just to collect information about fields.
Args:
query: instance of DBQuery or string
Returns:
list of :class:`DBField` instances
"""
from quazy.db_query import DBQuery
if typing.TYPE_CHECKING:
from psycopg.cursor import BaseCursor, RowMaker
from psycopg.rows import DictRow
def extract_description(cols_info: Sequence[tuple[str, ...]]) -> list[DBField]:
res = []
has_fields = isinstance(query, DBQuery) and query.table_class is not None
for col in cols_info:
if has_fields and (field := query.table_class.DB.fields.get(col[0])) is not None:
res.append(field)
elif (has_fields and col[0].endswith("__view")
and (field := query.table_class.DB.fields.get(col[0][:-6])) is not None):
field = copy.deepcopy(field)
field.name += "__view"
field.ux.name += "__view"
field.type = str
field.ref = False
res.append(field)
else:
field = DBField(col[0])
field.prepare(col[0])
field.type = self._translator.TYPES_BY_OID[col[1]]
res.append(field)
return res
with self.cursor() as curr:
if isinstance(query, DBQuery):
query = query.copy().set_window(limit=0)
sql = self._translator.select(query)
if self._debug_mode: print(sql)
curr.execute(sql, query.args)
return extract_description(curr.description)
else:
curr.execute(query)
return extract_description(curr.description)
[docs]
def save(self, _item: DBTableT, _lookup_field: str | None = None, _curr: DBCursorLike=None, **kwargs) -> Awaitable[DBTableT] | DBTableT:
"""Save item to the database
It checks whether an item needs to be inserted or updated. Specify `lookup_field` to avoid searching by a primary key.
Args:
_item: instance of DBTable
_lookup_field: field name or None
kwargs: additional values to update item fields before saving it to the database
Returns:
updated instance, just for chain calls
"""
for name, value in kwargs.items():
setattr(_item, name, value)
pk_name = _item.__class__.DB.pk.name
if _lookup_field:
row_id = self.query(_item.__class__)\
.filter(lambda x: x[_lookup_field] == getattr(_item, _lookup_field))\
.set_window(limit=1)\
.select(pk_name)\
.fetch_value()
if row_id:
setattr(_item, pk_name, row_id)
self.update(_item, _curr=_curr)
else:
self.insert(_item, _curr=_curr)
else:
if getattr(_item, pk_name):
self.update(_item, _curr=_curr)
else:
self.insert(_item, _curr=_curr)
return _item
[docs]
def delete(self, table: type[DBTable] = None, *,
item: DBTableT = None,
id: Any = None,
items: typing.Iterator[DBTableT] = None,
query: DBQuery[DBTableT] = None,
filter: Callable[[DBTableT], DBSQL] = None,
_curr: DBCursorLike = None) -> Awaitable[None] | None:
"""Delete an item or items from the database
It has many possible ways to delete expected items:
* by item
* by table and id (for a primary key)
* by item list
* by query based on DBTable
* by table and lamba filter
Args:
table: DBTable class
item: instance of DBTable
id: instance primary key id
items: iterable of DBTable instances
query: instance of DBQuery based on DBTable
filter: callable lamba added to a query
Raises:
QuazyWrongOperation: wrong arguments usage
"""
if id is not None:
if table is None:
raise QuazyWrongOperation("Both `id` and `table` should be specified")
with self.cursor(_curr=_curr) as curr:
sql = self._translator.delete_related(table, table.DB.pk.column)
curr.execute(sql, (id, ))
elif item is not None:
item._before_delete(self)
with self.cursor(_curr=_curr) as curr:
sql = self._translator.delete_related(type(item), item.DB.pk.column)
curr.execute(sql, (item.pk, ))
item._after_delete(self)
elif items is not None:
with self.cursor(_curr=_curr) as curr:
for item in items:
self.delete(item=item, _curr=curr)
elif query is not None:
if not query.fetch_objects:
raise QuazyWrongOperation('Query should be objects related')
builder = self.query(query.table_class)
sub = builder.with_query(query.select("pk"), not_materialized=True)
with self.cursor(_curr=_curr) as curr:
sql = self._translator.delete_selected(builder, sub)
if self._debug_mode: print(sql)
curr.execute(sql, builder.args)
elif filter is not None:
if table is None:
raise QuazyWrongOperation("Both `filter` and `table` should be specified")
query = self.query(table).filter(filter)
self.delete(query=query, _curr=_curr)
elif table is not None:
with self.cursor(_curr=_curr) as curr:
sql = getattr(self._translator, "delete")(table)
curr.execute(sql)
else:
raise QuazyWrongOperation('Specify one of the arguments')
ASYNC_COMMON_FUNCS = ('cursor', 'execute', 'executemany', 'fetchone', 'fetchall', 'copy', 'write_row')
class DBFactoryAsync(DBFactory):
async_mode: ClassVar[bool] = True
@staticmethod
def postgres(conninfo: str, **kwargs) -> DBFactoryAsync | None:
import psycopg
from psycopg.rows import namedtuple_row, dict_row, kwargs_row
from .translator_psql import TranslatorPSQL
debug_mode = kwargs.pop("debug_mode", False)
connection = psycopg.connect(conninfo, **kwargs)
connection.close()
class PsycopgConnection:
conn: psycopg.AsyncConnection = None
async def connection(self) -> AsyncContextManager[psycopg.AsyncConnection] | psycopg.AsyncConnection:
if self.conn is None or self.conn.closed:
self.conn = await psycopg.AsyncConnection.connect(conninfo, **kwargs)
return self.conn
@asynccontextmanager
async def cursor(self, read_only: bool=False) -> AsyncGenerator[DBCursorLike]:
conn = await self.connection()
async with conn.cursor(binary=True) as curr:
yield curr
if not read_only:
await self.conn.commit()
async def close(self):
await self.conn.close()
return DBFactoryAsync(PsycopgConnection(), TranslatorPSQL, namedtuple_row, dict_row, kwargs_row, debug_mode)
@staticmethod
def postgres_pool(conninfo: str, **kwargs) -> DBFactoryAsync | None:
import psycopg
from psycopg_pool.pool_async import AsyncConnectionPool
from psycopg.rows import namedtuple_row, dict_row, kwargs_row
from .translator_psql import TranslatorPSQL
connection = psycopg.connect(conninfo, **kwargs)
connection.close()
class PsycopgPoolConnection:
pool: AsyncConnectionPool = None
@asynccontextmanager
async def connection(self) -> AsyncGenerator[DBConnectionLike]:
if self.pool is None:
self.pool = AsyncConnectionPool(conninfo, kwargs=kwargs)
await self.pool.open()
await self.pool.wait()
async with self.pool.connection() as conn:
yield conn
@asynccontextmanager
async def cursor(self, read_only:bool=False) -> AsyncGenerator[DBCursorLike]:
async with self.connection() as conn:
async with conn.cursor(binary=True) as curr:
yield curr
async def close(self):
await self.pool.close()
return DBFactoryAsync(PsycopgPoolConnection(), TranslatorPSQL, namedtuple_row, dict_row, kwargs_row)
@staticmethod
def sqlite(conn_uri: str, **kwargs) -> DBFactoryAsync | None:
try:
import aiosqlite
except ImportError:
raise NotImplementedError('SQLite async is not implemented. Install `aiosqlite` module')
import sqlite3
from .translator_sqlite import (TranslatorSQLite, namedtuple_row, dict_row, kwargs_row,
ConnectionFactoryAsync)
detect_types = kwargs.pop("detect_types", sqlite3.PARSE_DECLTYPES)
class SQLiteConnection:
conn: aiosqlite.Connection = None
async def connection(self) -> DBConnectionLike:
if self.conn is None:
connector = lambda: sqlite3.connect(conn_uri, uri=True, detect_types=detect_types, **kwargs)
self.conn = await ConnectionFactoryAsync(connector, iter_chunk_size=64)
await self.conn.execute("PRAGMA foreign_keys = ON")
return self.conn
@asynccontextmanager
async def cursor(self, read_only:bool=False) -> AsyncGenerator[DBCursorLike]:
conn = await self.connection()
async with conn.cursor() as curr:
yield curr
async def close(self):
await self.conn.close()
return DBFactoryAsync(SQLiteConnection(), TranslatorSQLite, namedtuple_row, dict_row, kwargs_row)
def query(self, table_class: Optional[type[DBTableT]] = None, name: Optional[str] = None) -> DBQueryAsync[DBTableT]:
from .db_query import DBQueryAsync
return DBQueryAsync[DBTableT](self, table_class, name)
connection = make_async(DBFactory.connection, ('connection',))
cursor = make_async(DBFactory.cursor, ('cursor',))
close = make_async(DBFactory.close, ('close',))
get = make_async(DBFactory.get, ('fetch_one',))
clear = make_async(DBFactory.clear, ASYNC_COMMON_FUNCS + ('results',))
missed_tables = make_async(DBFactory.missed_tables, ('select', 'fetchall'))
create = make_async(DBFactory.create, ASYNC_COMMON_FUNCS + ('missed_tables', ))
table_exists = make_async(DBFactory.table_exists, ASYNC_COMMON_FUNCS + ('select',))
insert = make_async(DBFactory.insert, ASYNC_COMMON_FUNCS + ('insert', 'save'))
update = make_async(DBFactory.update, ASYNC_COMMON_FUNCS + ('insert', 'save'))
update_many = make_async(DBFactory.update_many, ASYNC_COMMON_FUNCS)
select = make_async(DBFactory.select, ASYNC_COMMON_FUNCS)
execute = make_async(DBFactory.execute, ASYNC_COMMON_FUNCS)
save = make_async(DBFactory.save, ASYNC_COMMON_FUNCS + ('fetch_value', 'update', 'insert'))
delete = make_async(DBFactory.delete, ASYNC_COMMON_FUNCS + ('delete',))