Module pggm_datalab_utils.db

Functions

def cursor(db_server: str, db_name: str) ‑> Iterator[Union[pyodbc.Cursor, sqlite3.Cursor]]

Obtain a cursor for a certain database server and database name. Internally uses get_db_connection(). Use this as a context manager, which will handle closing the cursor and the connection. NOTE: this will not handle transaction support: most of the time that means you need to commit your transactions yourself! Example usage:

with cursor('my_server.net', 'test') as c:
    my_data = c.execute('select * from test_database').fetchall()

Expand source code

@contextmanager
def cursor(db_server: str, db_name: str) -> Iterator[Cursor]:
    """
    Obtain a cursor for a certain database server and database name. Internally uses `get_db_connection`. Use this as
    a context manager, which will handle closing the cursor and the connection. NOTE: this will not handle transaction
    support: most of the time that means you need to commit your transactions yourself!
    Example usage:
    ```
    with cursor('my_server.net', 'test') as c:
        my_data = c.execute('select * from test_database').fetchall()
    ```
    """
    conn = get_db_connection(db_server, db_name)
    c = conn.cursor()
    try:
        yield c
    finally:
        c.close()
        conn.close()

def get_all(c: Union[pyodbc.Cursor, sqlite3.Cursor], table_name: str) ‑> List[Dict[str, Any]]

Get all current data from table table_name.

IMPORTANT WARNING: table_name is not sanitized. Don't pass untrusted table names to this function!

Expand source code

def get_all(c: Cursor, table_name: str) -> RecordList:
    """
    Get all current data from table `table_name`.

    IMPORTANT WARNING: `table_name` is not sanitized. Don't pass untrusted table names to this function!
    """
    return query(c, f'select * from {table_name}')

def get_db_connection(server: str, database: str) ‑> Union[pyodbc.Connection, sqlite3.Connection]

Initiate pyodbc connection to a SQL Server database. This function is intended to be suitable for cloud development: in the cloud you use environment variables to log in under a certain username and password, whereas locally you simply log in using your AAD credentials. You may specify DB_DRIVER to a pyodbc-compatible driver name for your system. It defaults to {ODBC Driver 17 for SQL Server}. If you specify the value SQLITE this routine will use the built-in sqlite3 library to connect instead. By default, this connection will try to log into the database with the user account it's executing under. This is compatible with AAD login and suitable for local development. If you want to log into a database from another environment, you will have to use Windows credentials. Save the username in the environment variable DB_UID, the password in DB_PASSWORD.

Expand source code

def get_db_connection(server: str, database: str) -> Connection:
    """
    Initiate pyodbc connection to a SQL Server database. This function is intended to be suitable for cloud development:
    in the cloud you use environment variables to log in under a certain username and password, whereas locally you
    simply log in using your AAD credentials.
    You may specify `DB_DRIVER` to a pyodbc-compatible driver name for your system. It defaults
    to `{ODBC Driver 17 for SQL Server}`. If you specify the value `SQLITE` this routine will use the built-in sqlite3
    library to connect instead.
    By default, this connection will try to log into the database with the user account it's executing under. This is
    compatible with AAD login and suitable for local development. If you want to log into a database from another
    environment, you will have to use Windows credentials. Save the username in the environment variable `DB_UID`,
    the password in `DB_PASSWORD`.
    """
    driver = os.environ.get('DB_DRIVER', '{ODBC Driver 17 for SQL Server}')
    if driver == 'SQLITE':
        logging.info(f'Connecting to SQLite database {database} because driver=SQLITE. Ignoring server {server}.')

        # Make sqlite3 somewhat well-behaved.
        sqlite3.register_converter('datetime', lambda b: datetime.fromisoformat(b.decode()))
        sqlite3.register_converter('json', json.loads)
        sqlite3.register_adapter(list, json.dumps)
        sqlite3.register_adapter(dict, json.dumps)

        return sqlite3.connect(database, detect_types=sqlite3.PARSE_DECLTYPES)
    elif user := os.environ.get('DB_UID', False):
        logging.info(f'Logging into {database}/{server} as {user}.')
        pwd = os.environ['DB_PASSWORD']
        return pyodbc.connect(
            f'Driver={driver};'
            f'Server={server};'
            f'Database={database};'
            f'Uid={os.environ["DB_UID"]};'
            f'Pwd={pwd};')
    else:
        logging.info(f'Logging into {database}/{server} as program user.')
        return pyodbc.connect(
            f'Driver={driver};'
            f'Server={server};'
            f'Database={database};'
            f'trusted_connection=yes;')

def insert_with_return(c: Union[pyodbc.Cursor, sqlite3.Cursor], table_name: str, data: Dict[str, Any], return_columns: Union[str, tuple, ForwardRef(None)] = None) ‑> Dict[str, Any]

Insert data into the database, returning a set of return columns. The primary use for this is if you have columns generated by your database, like an identity. Returns input record with returned columns added (if any).

Expand source code

def insert_with_return(
        c: Cursor, table_name: str, data: Record, return_columns: Optional[Union[str, tuple]] = None
) -> Record:
    """
    Insert data into the database, returning a set of return columns. The primary use for this is if you have columns
    generated by your database, like an identity. Returns input record with returned columns added (if any).
    """
    validate([data])
    return_columns = (return_columns,) if isinstance(return_columns, str) else return_columns

    columns = data.keys()
    insert_data = tuple(data[col] for col in columns)
    text_columns = ', '.join(columns)
    placeholders = ', '.join('?' for _ in columns)
    # Dispatch on cursor type for now, pyodbc type is for MSSQL only
    if return_columns is None:
        sql = f'insert into {table_name}({text_columns}) values ({placeholders})'
        c.execute(sql)
        return data
    elif isinstance(c, sqlite3.Cursor):
        text_return_columns = ', '.join(return_columns)
        sql = f'insert into {table_name}({text_columns}) values ({placeholders}) returning {text_return_columns}'
        output = query(c, sql, insert_data)[0]
        return {**data, **output}
    else:
        text_return_columns = ', '.join(f'Inserted.{col}' for col in return_columns)
        sql = f'insert into {table_name}({text_columns}) output {text_return_columns} values ({placeholders})'
        output = query(c, sql, insert_data)[0]
        return {**data, **output}

def query(c: Union[pyodbc.Cursor, sqlite3.Cursor], sql: str, data: Optional[tuple] = None) ‑> List[Dict[str, Any]]

Call c.execute(sql, data).fetchall() and format the resulting rowset a list of records of the form [{colname: value}].

Expand source code

def query(c: Cursor, sql: str, data: Optional[tuple] = None) -> RecordList:
    """
    Call `c.execute(sql, data).fetchall()` and format the resulting rowset a list of records of the form
    [{colname: value}].
    """
    if data is None:
        result = c.execute(sql).fetchall()
    else:
        result = c.execute(sql, data).fetchall()
    headers = [name for name, *_ in c.description]
    return [dict(zip(headers, r)) for r in result]

def validate(data: List[Dict[str, Any]])

Expand source code

def validate(data: RecordList):
    assert len(unique := set(tuple(sorted(r.keys())) for r in data)) == 1, \
        f'Non-uniform list of dictionaries passed, got differing keys {unique}.'
    assert not any(non_str := {k: type(k) for k in data[0].keys() if not isinstance(k, str)}), \
        f'Non-string keys in data, got keys with types {non_str}.'

def write(c: Union[pyodbc.Cursor, sqlite3.Cursor], table_name: str, data: List[Dict[str, Any]], primary_key: Union[str, tuple, ForwardRef(None)] = None, *, update=True, insert=True, delete=True)

Update data in database table. We check identity based on the keys of the IndexedPyFrame. update, insert, and delete control which actions to take. By default, this function emits the correct update, insert, and delete queries to make the database table equal to the in-memory table. - update=True means rows already in the database will be updated with the in-memory data - insert=True means rows not already in the database will be added from the in-memory data - delete=True means rows present in the database but not in the in-memory database will be deleted

If primary_key is None, only inserting is supported.

IMPORTANT WARNING: table_name is not sanitized. Don't pass untrusted table names to this function!

Expand source code

def write(c: Cursor, table_name: str, data: RecordList, primary_key: Optional[Union[str, tuple]] = None, *,
          update=True, insert=True, delete=True):
    """
    Update data in database table. We check identity based on the keys of the IndexedPyFrame.
    `update`, `insert`, and `delete` control which actions to take. By default, this function emits the correct update,
    insert, and delete queries to make the database table equal to the in-memory table.
    - `update=True` means rows already in the database will be updated with the in-memory data
    - `insert=True` means rows not already in the database will be added from the in-memory data
    - `delete=True` means rows present in the database but not in the in-memory database will be deleted

    If primary_key is None, only inserting is supported.

    IMPORTANT WARNING: `table_name` is not sanitized. Don't pass untrusted table names to this function!
    """
    validate(data)

    # Deal with primary key, list of writeable columns, indexed data, data in db
    if primary_key is None:
        assert not update and not delete, 'updating and deleting without specifying a primary key not supported'
        primary_key = tuple()
        data = {i: r for i, r in enumerate(data)}
        columns = tuple(k for k in data[0].keys())
        in_db = set()
    else:
        primary_key = (primary_key,) if isinstance(primary_key, str) else tuple(primary_key)
        assert all(isinstance(r[k], Hashable) for r in data for k in primary_key)
        if any(empty_strings := [name for name in data[0].keys() if any(r[name] == '' for r in data)]):
            logging.warning(f'Columns {empty_strings} contain empty strings. '
                            f'Generally inserting empty strings into a database is a bad idea.')

        # List of writeable columns (for updates we don't try to overwrite the primary key)
        columns = tuple(k for k in data[0].keys() if k not in primary_key)

        # Indexed data on primary key
        data = {tuple(r[i] for i in primary_key): r for r in data}

        # Data present in database
        sql = f'select {", ".join(primary_key)} from {table_name}'
        in_db = {tuple(r[k] for k in primary_key) for r in query(c, sql)}

    if update and (update_keys := data.keys() & in_db):
        update_data = [
            (tuple(data[k][col] for col in columns) + tuple(data[k][col] for col in primary_key)) for k in update_keys
        ]

        # Cannot use keyword placeholders because pyodbc doesn't support named paramstyle. Would be better.
        assignment = ', '.join(f'{col}=?' for col in columns)
        pk_cols = ' AND '.join(f'{col}=?' for col in primary_key)
        sql = f'update {table_name} set {assignment} where {pk_cols}'

        c.executemany(sql, update_data)

    if insert and (insert_keys := data.keys() - in_db):
        insert_data = [tuple(data[k][col] for col in columns + primary_key) for k in insert_keys]

        placeholders = ', '.join(f'?' for _ in columns + primary_key)
        text_columns = ', '.join(columns + primary_key)
        sql = f'insert into {table_name}({text_columns}) VALUES ({placeholders})'

        c.executemany(sql, insert_data)

    if delete and (delete_keys := in_db - data.keys()):
        condition = ' AND '.join(f'{k}=?' for k in primary_key)
        sql = f'delete from {table_name} where {condition}'

        c.executemany(sql, list(delete_keys))