Newer
Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
"""Some helper functions for the API controlers."""
from contextlib import contextmanager
from typing import Iterator, Optional, TypeVar
from sqlalchemy import text
from sqlalchemy.engine import Connection
from .db import project_engines
T = TypeVar('T')
@contextmanager
def begin_transaction(project: str) -> Iterator[Connection]:
"""Returns a cursor associated with the project database."""
engine = project_engines[project]
with engine.begin() as connection:
yield connection
def table_column_names(project: str, dataset: str, table: str) -> list[str]:
"""Returns the column names of a table."""
with begin_transaction(project) as connection:
stmt = text(
f"""
SELECT column_name
FROM information_schema.columns
WHERE table_schema = '{dataset}' AND table_name = '{table}'
"""
)
result = connection.execute(stmt)
return result.scalars().all()
def fix_sqlalchemy2_stubs_non_nullable_column(arg: Optional[T]) -> T:
"""Non-nullable model columns are of type Optional[...] for mypy.
It was correctly handled by sqlalchemy-stubs, but it is not yet the case for
sqlalchemy2-stubs==0.0.2a4.
"""
assert arg is not None
return arg