def data():
from sqlalchemy import create_engine, text
from sqlalchemy.engine.url import URL
from scraper import settings
from app.config import database
#database_settings = settings.DB_SETTINGS
#engine = create_engine(URL.create(**database_settings))
engine = create_engine(URL.create(**database.default))
query = text("""
with news as (
select distinct
base_url as outlet
,parent_url as link
from news
where published_at between '2016-01-01'::date and '2017-01-01'::date
and parent_url is not null
)
select
outlet
,string_agg(link, ' ') as links
,count(1) as total
from news
where link is not null
group by outlet
order by total desc
""")
with engine.connect() as conn:
df = pd.read_sql_query(query, conn)
return df
def postgres_base_uri(request):
"""The base URI connection string to postgres.
This should contain all necessary connection information except the database name.
"""
uri = request.config.getoption('postgres_uri')
if not uri:
# Not set, return None; indicates the tests are not to use postgres
return None
# If the URI contains a database name, we need to remove it
from sqlalchemy.engine.url import make_url, URL
url = make_url(uri)
url_kwargs = {
'drivername': url.drivername,
'username': url.username,
'password': url.password,
'host': url.host,
'port': url.port,
# Purposely remove database and query.
# 'database': None,
# 'query': None,
}
base_uri = str(URL.create(**url_kwargs))
return base_uri
def __init__(self,
db_path: str = '',
db_type: str = ''): # type (str) -> ()
"""
Connect to database.
:param db_path: Path to database including database name
:param db_type: Database type: sqlite, Postgresql, mysql, etc.
"""
self.database_name = DATABASE_NAME
self.database_type = db_type if db_type else 'sqlite'
self.database_path = db_path if db_path else os.path.join(
utilities.path(), self.database_name)
db_credentials = {
'username': None,
'password': None,
'host': None,
'database': str(self.database_path),
'port': None
}
self.database_url = URL.create(
drivername=self.database_type,
username=db_credentials['username'],
password=db_credentials['password'],
host=db_credentials['host'],
port=db_credentials['port'],
database=db_credentials['database'],
)
def engine_postgis():
"""
Initiaties a connection engine to a postGIS database that must already exist.
"""
sqlalchemy = pytest.importorskip("sqlalchemy")
from sqlalchemy.engine.url import URL
user = os.environ.get("PGUSER")
password = os.environ.get("PGPASSWORD")
host = os.environ.get("PGHOST")
port = os.environ.get("PGPORT")
dbname = "test_geopandas"
try:
con = sqlalchemy.create_engine(
URL.create(
drivername="postgresql+psycopg2",
username=user,
database=dbname,
password=password,
host=host,
port=port,
))
con.begin()
except Exception:
pytest.skip("Cannot connect with postgresql database")
yield con
con.dispose()
def open_db(dbstr,
echo=False,
serial=False,
poolclass=sqlalchemy.pool.NullPool,
**kwargs):
"""Create an engine to access a DbHandle.
"""
url = parse_dbstring(dbstr)
query = dict(url.query)
tablename = query.pop('table')
dbname = query.pop('dbname', None)
if dbname is None:
dbname = url.database
# mutate url
url = url._asdict()
url['query'] = query
url = URL.create(**url)
if serial:
engine = create_engine(url,
echo=echo,
poolclass=poolclass,
isolation_level='SERIALIZABLE')
else:
engine = create_engine(url, echo=echo, poolclass=poolclass)
return db_from_engine(engine,
table_prefix=tablename,
dbname=dbname,
**kwargs)
def ensure_mssql_ready_for_tests(config):
"""Test connection to MSSQL instance
and check the existence of database AHJO_TEST.
"""
try:
if not config.getoption('mssql_host'):
raise Exception('MSSQL Server not given')
connection_url = URL.create(
drivername="mssql+pyodbc",
username=config.getoption('mssql_username'),
password=config.getoption('mssql_password'),
host=config.getoption('mssql_host'),
port=config.getoption('mssql_port'),
database='master',
query={'driver': 'SQL Server'})
engine = create_engine(connection_url)
with engine.connect() as connection:
query = "SELECT name FROM sys.databases WHERE UPPER(name) = ?"
result = connection.execute(query, (TEST_DB_NAME, ))
if result.fetchall():
raise Exception(
f"There already exists a database with name '{TEST_DB_NAME}'"
)
return True
except:
return False
def __init__(
self,
config: Optional[Dict[str, Any]] = None,
configuration_section: Optional[str] = None,
uri: Optional[str] = None,
use_ssh_tunnel: bool = False,
sshtunnel_params: Optional[Dict[str, Any]] = None,
**credentials,
):
"""
:param use_ssh_tunnel: boolean - default false. Whether to tunnel sqlalchemy connection
through an ssh tunnel or not
:param sshtunnel_params: Dict of ssh params - specify according to sshtunnel project
https://github.com/pahaz/sshtunnel/ - direct passthrough
"""
self.use_ssh_tunnel = use_ssh_tunnel
# Sort out which credentials are the final ones -- default to remaining passed params
if configuration_section is not None:
if config is None:
raise SimpleMLError(
"Cannot use config section without a config file")
# Default to credentials in config file
credentials = dict(config[configuration_section])
elif uri is not None:
# Deconstruct URI into credentials
url = make_url(uri)
credentials = {
"drivername": url.drivername,
"username": url.username,
"password": url.password,
"host": url.host,
"port": url.port,
"database": url.database,
"query": url.query,
}
# Reconfigure credentials if SSH tunnel specified
if self.use_ssh_tunnel:
LOGGER.warning("""
Usage: call Database.open_tunnel() before Database.initialize() and
end script with Database.close_tunnel()
Configure connection with supported parameters passed via
`sshtunnel_params={**configs}`. Binding and routing through local
port is automatically handled, but other parameters like `set_keepalive`
may be interesting. https://sshtunnel.readthedocs.io/en/latest/
""")
# Overwrite passed ports and hosts to route localhost port to the
# original destination via tunnel
if sshtunnel_params is None:
sshtunnel_params = {}
credentials, self.ssh_config = self.configure_ssh_tunnel(
credentials, sshtunnel_params)
try:
# New syntax (1.4.2+)
self.url = URL.create(**credentials)
except AttributeError:
# Old syntax
self.url = URL(**credentials)
def connect(force_migrate=False):
"""Module initialisation
The connection with the underlying storage is initialised.
Meta information is available via the Base variable.
Data retrieval is facilitated via the session object
:return: True when the connection has been established
"""
global session, engine
try:
engine = create_engine(URL.create(**GOB_MGMT_DB),
connect_args={'sslmode': 'require'})
migrate_storage(force_migrate)
# Declarative base model to create database tables and classes
Base.metadata.bind = engine
session = Session(engine)
except DBAPIError as e:
# Catch any connection errors
print(f"Connect failed: {str(e)}")
disconnect() # Cleanup
return is_connected()
def connect():
"""Module initialisation
The connection with the underlying storage is initialised.
Meta information is available via the Base variale.
Data retrieval is facilitated via the session object
:return:
"""
global session, _Base, metadata
engine = create_engine(URL.create(**GOB_DB),
connect_args={'sslmode': 'require'})
session = scoped_session(
sessionmaker(autocommit=True,
autoflush=False,
bind=engine,
query_cls=AuthorizedQuery))
with warnings.catch_warnings():
# Ignore warnings for unsupported reflection for expression-based indexes
warnings.simplefilter("ignore", category=sa_exc.SAWarning)
_Base = automap_base()
_Base.prepare(engine, reflect=True) # Long running statement !
Base.metadata.bind = engine # Bind engine to metadata of the base class
Base.query = session.query_property() # Used by graphql to execute queries
metadata = MetaData(engine)
set_session(session)
profiled_query.activate()
def get_url(db_name=None):
"""Return an sqlalchemy URL for database
Return database options, only the database name db_name can be overloaded::
url = get_url(db_name='Mydb')
..note::
Since 0.5.3, an URL can be define by the configuration file.
The *username*, *password* and *database* if overwrite by the
options if they are filled::
# db_url = 'postgresql:///db'
get_url()
==> 'postgresql:///db'
# db_user_name = 'jssuzanne'
# db_password = '******'
get_url()
==> 'postgresql://*****:*****@/db'
# db_name = 'db1'
get_url()
==> 'postgresql://*****:*****@/db1'
get_url(db_name='Mydb')
==> 'postgresql://*****:*****@/Mydb'
:param db_name: database name
:rtype: SqlAlchemy URL
:exception: ConfigurationException
"""
url = Configuration.get('db_url', None)
drivername = Configuration.get('db_driver_name', None)
username = Configuration.get('db_user_name', None)
password = Configuration.get('db_password', None)
host = Configuration.get('db_host', None)
port = Configuration.get('db_port', None)
database = Configuration.get('db_name', None)
query = Configuration.get('db_query', {})
if db_name is not None:
database = db_name
if url:
url = make_url(url)
if username:
url = url.set(username=username)
if password:
url = url.set(password=password)
if database:
url = url.set(database=database)
return url
if drivername is None:
raise ConfigurationException('No driver name defined')
return URL.create(
drivername, username=username, password=password, host=host,
port=port, database=database, query=query)
def __init__(
self,
*,
url: str | URL | dict[str, Any] | None = None,
bind: Connection | Engine | AsyncConnection | AsyncEngine
| None = None,
engine_args: dict[str, Any] | None = None,
session_args: dict[str, Any] | None = None,
commit_executor_workers: int = 5,
ready_callback: Callable[[Engine, sessionmaker], Any] | str
| None = None,
poolclass: str | type[Pool] | None = None,
resource_name: str = "default",
):
self.resource_name = resource_name
self.commit_executor_workers = commit_executor_workers
self.ready_callback = resolve_reference(ready_callback)
engine_args = engine_args or {}
session_args = session_args or {}
session_args["expire_on_commit"] = False
session_args["future"] = True
if bind:
self.bind = bind
self.engine = cast("Engine | AsyncEngine", bind.engine)
else:
if isinstance(url, dict):
url = URL.create(**url)
elif isinstance(url, str):
url = make_url(url)
elif url is None:
raise TypeError('both "url" and "bind" cannot be None')
# This is a hack to get SQLite to play nice with asphalt-sqlalchemy's
# juggling of connections between multiple threads. The same connection
# should, however, never be used in multiple threads at once.
if url.get_dialect().name == "sqlite":
connect_args = engine_args.setdefault("connect_args", {})
connect_args.setdefault("check_same_thread", False)
if isinstance(poolclass, str):
poolclass = resolve_reference(poolclass)
pool_class = cast("type[Pool]", poolclass)
try:
self.engine = self.bind = create_async_engine(
url, poolclass=pool_class, **engine_args)
except InvalidRequestError:
self.engine = self.bind = create_engine(url,
poolclass=pool_class,
**engine_args)
if url.get_dialect().name == "sqlite":
apply_sqlite_hacks(self.engine)
if isinstance(self.engine, AsyncEngine):
session_args.setdefault("class_", AsyncSession)
self.sessionmaker = sessionmaker(bind=self.bind, **session_args)
def plot():
x = request.args.get('x', 'first')
y = request.args.get('y', 'second')
static = request.args.get('static')
cache = request.args.get('cache', True, type=json.loads)
filename = 'plot_{}_{}.png'.format(x, y)
if not cache or not os.path.exists(resource_path(filename)):
label = request.args.get('label', 'outlet')
database_settings = settings.DB_SETTINGS
engine = create_engine(URL.create(**database_settings))
query = text("""
select
outlet
,labels
,case :x
when 'first'
then first
when 'second'
then second
else third
end as x
,case :y
when 'first'
then first
when 'second'
then second
else third
end as y
,total as total
from news.outlet_svd
where outlet is not null
""")
with engine.connect() as conn:
params = {'x': x, 'y': y}
df = pd.read_sql_query(query, conn, params=params)
#df.set_index('label')
fig, ax = plt.subplots()
ax.scatter(x=df['x'],
y=df['y'],
s=df['total'],
c=df['labels'],
cmap=plt.cm.RdBu)
df['range'] = np.abs((df['y'] - np.mean(df['y'])) / np.std(df['y']))
labels = df[df.range > 2.]
for i, row in labels.iterrows():
ax.text(x=row.loc['x'] + 0.3,
y=row.loc['y'] + 0.3,
s=row.loc['outlet'],
fontdict=dict(color='red', size=10))
ax.set(xlabel=x,
ylabel=y,
title='plot of {0} vs {1} component of svd'.format(x, y))
ax.grid()
fig.savefig(resource_path(filename), format='png')
return app.send_static_file(filename)
def test_simple_url(self):
url = URL.create(
drivername='clickhouse+native',
host='localhost',
database='default',
)
connect_args = self.dialect.create_connect_args(url)
self.assertEqual(str(connect_args[0][0]),
'clickhouse://localhost/default')
def DB_DSN(self) -> URL:
return URL.create(
self.DB_DRIVER,
self.DB_USER,
self.DB_PASSWORD,
self.DB_HOST,
self.DB_PORT,
self.DB_DATABASE,
)
def create_connection_obj():
"""Creates engine object from database.ini configuration"""
filename = "database.ini"
parser = ConfigParser()
parser.read(filename)
params = {k: v for k, v in parser.items("postgresql")}
conn = create_engine(URL.create(**params))
return conn
def DB_DSN(self) -> URL: # pylint: disable=invalid-name
return URL.create(
self.DB_DRIVER,
self.DB_USER,
self.DB_PASSWORD,
self.DB_HOST,
self.DB_PORT,
self.DB_DATABASE,
)
def refresh():
engine = create_engine(URL.create(**database.default))
data = svd.model(svd.data())
with engine.connect() as conn:
data.to_sql('outlet_svd',
conn,
schema='news',
if_exists='replace',
index=False)
return redirect(url_for('index'))
async def test_component_start_async():
"""Test that the component creates all the expected (asynchronous) resources."""
url = URL.create("sqlite+aiosqlite", database=":memory:")
component = SQLAlchemyComponent(url=url)
async with Context() as ctx:
await component.start(ctx)
ctx.require_resource(AsyncEngine)
ctx.require_resource(sessionmaker)
ctx.require_resource(AsyncSession)
def create_sqlalchemy_url(conn_info: dict, use_master_db: bool = False) -> URL:
"""Create url for sqlalchemy/alembic.
If use_master_db flag is on, pass 'master' database to url.
Arguments
---------
conn_info
Dictionary holding information needed to establish database connection.
use_master_db
Indicator to connect to 'master' database.
Returns
-------
sqlalchemy.engine.url.URL
Connection url for sqlalchemy/alembic.
"""
if use_master_db is True:
database = MASTER_DB.get(conn_info.get('dialect'))
else:
database = conn_info.get('database')
query = {}
# Add optional driver to query-dictionary
if conn_info.get('driver') is not None:
query['driver'] = conn_info.get('driver')
# sqlalchemy does not have full url support for different Azure authentications
# ODBC connection string must be added to query
azure_auth = conn_info.get('azure_auth')
if azure_auth is not None:
odbc = ''
if azure_auth.lower() == 'activedirectorypassword':
authentication = 'ActiveDirectoryPassword'
odbc = f"Pwd{{{conn_info.get('password')}}};"
elif azure_auth.lower() == 'activedirectoryintegrated':
authentication = 'ActiveDirectoryIntegrated'
elif azure_auth.lower() == 'activedirectoryinteractive':
authentication = 'ActiveDirectoryInteractive'
else:
raise Exception(
"Unknown Azure authentication type! Check variable 'azure_authentication'.")
query['odbc_connect'] = odbc + "Driver={{{driver}}};Server={server};Database={database};Uid={{{uid}}};Encrypt=yes;TrustServerCertificate=no;Authentication={auth}".format(
driver=conn_info.get('driver'),
server=conn_info.get('server'),
database=database,
uid=conn_info.get('username'),
auth=authentication
)
return URL.create(
drivername=conn_info.get('dialect'),
username=conn_info.get('username'),
password=conn_info.get('password'),
host=conn_info.get('host'),
port=conn_info.get('port'),
database=database,
query=query
)
def get_engine(self):
comp = env.postgis
# Need to check connection params to see if
# they changed for each connection request
credhash = (self.hostname, self.port, self.database, self.username,
self.password)
if self.id in comp._engine:
engine = comp._engine[self.id]
if engine._credhash == credhash:
return engine
else:
del comp._engine[self.id]
connect_timeout = int(comp.options['connect_timeout'].total_seconds())
statement_timeout_ms = int(
comp.options['statement_timeout'].total_seconds()) * 1000
args = dict(client_encoding='utf-8',
connect_args=dict(connect_timeout=connect_timeout,
options='-c statement_timeout=%d' %
statement_timeout_ms))
engine_url = make_engine_url(
EngineURL.create('postgresql+psycopg2',
host=self.hostname,
port=self.port,
database=self.database,
username=self.username,
password=self.password))
engine = db.create_engine(engine_url, **args)
resid = self.id
@db.event.listens_for(engine, 'connect')
def _connect(dbapi, record):
logger.debug("Resource #%d, pool 0x%x, connection 0x%x created",
resid, id(dbapi), id(engine))
@db.event.listens_for(engine, 'checkout')
def _checkout(dbapi, record, proxy):
logger.debug("Resource #%d, pool 0x%x, connection 0x%x retrieved",
resid, id(dbapi), id(engine))
@db.event.listens_for(engine, 'checkin')
def _checkin(dbapi, record):
logger.debug("Resource #%d, pool 0x%x, connection 0x%x returned",
resid, id(dbapi), id(engine))
engine._credhash = credhash
comp._engine[self.id] = engine
return engine
def test_secure_false(self):
url = URL.create(drivername='clickhouse+native',
username='******',
password='******',
host='localhost',
port=9001,
database='default',
query={'secure': 'False'})
connect_args = self.dialect.create_connect_args(url)
self.assertEqual(
str(connect_args[0][0]),
'clickhouse://*****:*****@localhost:9001/default?secure=False')
def _service_from_path(self, file):
if file.suffix == '.json':
with open(file, 'r', encoding='utf-8') as f:
try:
return json.load(f)
except Exception as e:
warnings.warn(
f'*** Failed to load service JSON: {file}: {e}')
return None
elif file.suffix == '.sqlite':
s = URL.create(drivername='sqlite', database=str(file.absolute()))
e = create_engine(s)
return e
def get_db():
if "CLOUD_DEPLOYED" in os.environ.keys():
if "db" not in g:
db_user = os.environ["DB_USER"]
db_pass = os.environ["DB_PASS"]
db_name = os.environ["DB_NAME"]
db_socket_dir = os.environ.get("DB_SOCKET_DIR", "/cloudsql")
cloud_sql_connection_name = os.environ["CLOUD_SQL_CONNECTION_NAME"]
g.db = create_engine(
URL.create(
drivername="postgresql+pg8000",
username=db_user,
password=db_pass,
database=db_name,
query={
"unix_sock": "{}/{}/.s.PGSQL.5432".format(
db_socket_dir, cloud_sql_connection_name
)
},
)
)
metadata.reflect(bind=g.db)
return g.db
else:
if "db" not in g:
config = ConfigParser()
config.read("config.ini")
g.db = create_engine(
URL.create(
drivername="postgresql+pg8000",
username=config["localdev"]["user"],
password=config["localdev"]["password"],
database=config["localdev"]["database"],
port=config["localdev"]["port"],
host=config["localdev"]["host"],
)
)
metadata.reflect(bind=g.db)
return g.db
def pg_dsn() -> str:
"""
:return: DSN url suitable for sqlalchemy and aiopg.
"""
return str(
URL.create(
database='hello_world',
password=os.getenv('PGPASS', 'benchmarkdbpass'),
host='tfb-database',
port='5432',
username=os.getenv('PGUSER', 'benchmarkdbuser'),
drivername='postgresql',
))
def create_db_engine(config: Dict[str, str]) -> Engine:
"""Create database engine
:param config: Dictionary with connection parameters
:returns: Database engine object
"""
try:
conn_url = URL.create(**config)
except TypeError:
logging.error("Database connection parameter error")
raise
return create_engine(conn_url)
def mssql_engine(request, ahjo_config, mssql_sample):
"""Create engine for MSSQL server test database.
"""
config = ahjo_config(mssql_sample)
connection_url = URL.create(
drivername="mssql+pyodbc",
username=request.config.getoption('mssql_username'),
password=request.config.getoption('mssql_password'),
host=config['target_server_hostname'],
port=config['sql_port'],
database=config['target_database_name'],
query={'driver': config['sql_driver']})
return create_engine(connection_url)
def test_dynamic_password():
url = URL.create('postgresql',
host="fake_host",
database="fake_database",
port=6543,
username="******",
password="******")
engine = PostgresDb._create_engine(url)
counter[0] = 0
last_base[0] = None
handle_dynamic_token_authentication(engine, next_token, base="password")
with pytest.raises(OperationalError):
conn = engine.connect()
assert counter[0] == 1
assert last_base[0] == "password"
def test_obtain_new_iam_token(monkeypatch, without_aws_env):
import moto
from sqlalchemy.engine.url import URL
url = URL.create(
'postgresql',
host="fakehost",
database="fake_db",
port=5432,
username="******",
password="******",
)
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "fake-key-id")
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "fake-secret")
with moto.mock_iam():
token = obtain_new_iam_auth_token(url, region_name='us-west-1')
assert isinstance(token, str)
def create_engine():
"""
Create a SQL Alchemy engine using config.DatabaseConfig.
:return: the engine
"""
connection_url = URL.create(
drivername=config.DB_DRIVER,
username=config.DB_USERNAME,
password=config.DB_PASSWORD,
host=config.DB_HOST,
port=config.DB_PORT,
database=config.DB_DATABASE,
)
return sqlalchemy.create_engine(
connection_url,
executemany_mode="values",
executemany_values_page_size=1000,
executemany_batch_page_size=200,
)
def async_engine_to_sync_engine(engine):
"""
If an engine is async, creates a new engine that is synchronous.
It does that by switching to a synchronous driver for given database:
* asyncpg -> psycopg2
* aiosqlite -> sqlite3
* ...
See variable `async_to_sync_drivers_dict` in this module.
This is probably not super elegant but it saves a lot of hassle for
everything that "surround" tests like setups, verifications, cleaning up...
For the "proper" testing part though this should obviously not be used
as it would defeat the purpose! E.g. in the presence of an engine using
`asyncpg` we will want to test it with `aupsert`. We are not going to convert
it to a synchronous engine and pass it to `upsert`...
"""
u = engine.url
params = {
attr: getattr(u, attr)
for attr in ('drivername', 'username', 'password', 'host', 'port',
'database', 'query')
}
# case where it is already sync
if not is_async_sqla_obj(engine):
return engine
# case where it is async but we don't recognize the driver
if not any(s in params['drivername'] for s in async_to_sync_drivers_dict):
raise NotImplementedError(
f'Cannot create a sync engine from this unknown async engine: {engine}'
)
# convert to a sync engine
for s in async_to_sync_drivers_dict:
if s in params['drivername']:
params['drivername'] = async_to_sync_drivers_dict[s]
break
new_u = URL.create(**params)
return create_engine(str(new_u))
