def patch_connection(connection_id, session, update_mask=None):
"""
Update a connection entry
"""
try:
data = connection_schema.load(request.json, partial=True)
except ValidationError as err:
# If validation get to here, it is extra field validation.
raise BadRequest(detail=str(err.messages))
non_update_fields = ['connection_id', 'conn_id']
connection = session.query(Connection).filter_by(conn_id=connection_id).first()
if connection is None:
raise NotFound("Connection not found")
if data.get('conn_id', None) and connection.conn_id != data['conn_id']:
raise BadRequest("The connection_id cannot be updated.")
if update_mask:
update_mask = [i.strip() for i in update_mask]
data_ = {}
for field in update_mask:
if field in data and field not in non_update_fields:
data_[field] = data[field]
else:
raise BadRequest(f"'{field}' is unknown or cannot be updated.")
data = data_
for key in data:
setattr(connection, key, data[key])
session.add(connection)
session.commit()
return connection_schema.dump(connection)
def post_connection(*, session: Session = NEW_SESSION) -> APIResponse:
"""Create connection entry"""
body = request.json
try:
data = connection_schema.load(body)
except ValidationError as err:
raise BadRequest(detail=str(err.messages))
conn_id = data['conn_id']
query = session.query(Connection)
connection = query.filter_by(conn_id=conn_id).first()
if not connection:
connection = Connection(**data)
session.add(connection)
session.commit()
return connection_schema.dump(connection)
raise AlreadyExists(detail=f"Connection already exist. ID: {conn_id}")
def post_connection(session):
"""
Create connection entry
"""
body = request.json
try:
result = connection_schema.load(body)
except ValidationError as err:
raise BadRequest(detail=str(err.messages))
data = result.data
conn_id = data['conn_id']
query = session.query(Connection)
connection = query.filter_by(conn_id=conn_id).first()
if not connection:
connection = Connection(**data)
session.add(connection)
session.commit()
return connection_schema.dump(connection)
raise AlreadyExists("Connection already exist. ID: %s" % conn_id)
def test_deserialize(self):
den = {
'connection_id': "mysql_default",
'conn_type': 'mysql',
'host': 'mysql',
'login': '******',
'schema': 'testschema',
'port': 80,
'extra': "{'key':'string'}",
}
result = connection_schema.load(den)
assert result == {
'conn_id': "mysql_default",
'conn_type': 'mysql',
'host': 'mysql',
'login': '******',
'schema': 'testschema',
'port': 80,
'extra': "{'key':'string'}",
}
def test_connection() -> APIResponse:
"""
To test a connection, this method first creates an in-memory dummy conn_id & exports that to an
env var, as some hook classes tries to find out the conn from their __init__ method & errors out
if not found. It also deletes the conn id env variable after the test.
"""
body = request.json
dummy_conn_id = get_random_string()
conn_env_var = f'{CONN_ENV_PREFIX}{dummy_conn_id.upper()}'
try:
data = connection_schema.load(body)
data['conn_id'] = dummy_conn_id
conn = Connection(**data)
os.environ[conn_env_var] = conn.get_uri()
status, message = conn.test_connection()
return connection_test_schema.dump({"status": status, "message": message})
except ValidationError as err:
raise BadRequest(detail=str(err.messages))
finally:
if conn_env_var in os.environ:
del os.environ[conn_env_var]
