def test_expand_environment_variables_in_connection():
conn_str = "{}://cluster('clustername').database('dbname')".format(TEST_URI_SCHEMA_NAME)
os.environ['KQL_DATABASE'] = conn_str
assert Parser.parse("$KQL_DATABASE {}".format(query1), empty_config) == \
{'connection': conn_str,
'kql': query1,
'options': default_options}
def test_parse_kusto_socket_connection_dsn():
conn_str = "{}://username($USERNAME).password($PASSWORD).cluster('clustername').database('dbname')".format(TEST_URI_SCHEMA_NAME)
result_conn_str = "{}://username('michael').password('michael123').cluster('clustername').database('dbname')".format(TEST_URI_SCHEMA_NAME)
assert Parser.parse("{0} {1}".format(conn_str, query1), empty_config) == \
{'connection': result_conn_str,
'kql': query1,
'options': default_options}
def test_parse_kql_only():
parsed = Parser.parse(query1, empty_config)
print(parsed)
assert parsed == \
{'connection': "",
'kql': query1,
'options': default_options}
def test_parse_kusto_socket_connection_with_credentials():
conn_str = "{}://username('username').password('password').cluster('clustername').database('dbname')".format(
TEST_URI_SCHEMA_NAME)
assert Parser.parse("{0} {1}".format(conn_str, query1), empty_config, TEST_ENGINE, {}) == \
{'connection': conn_str,
'kql': query1,
'options': default_options}
def test_parse_kusto_socket_connection_with_env_credentials():
conn_str = "{}://username($USERNAME).password($PASSWORD).cluster('clustername').database('dbname')".format(
TEST_URI_SCHEMA_NAME)
result_conn_str = "{}://username('michael').password('michael123').cluster('clustername').database('dbname')".format(
TEST_URI_SCHEMA_NAME)
os.environ['USERNAME'] = "******"
os.environ['PASSWORD'] = "******"
assert Parser.parse("{0} {1}".format(conn_str, query1), empty_config, TEST_ENGINE, {}) == \
{'connection': result_conn_str,
'kql': query1,
'options': default_options}
def _parse_common_connection_str(self, conn_str: str, current,
uri_schema_name, mandatory_key: str,
alt_names: list,
valid_keys_combinations: list,
user_ns: dict):
# parse schema part
parts = conn_str.split("://", 1)
if len(parts) != 2 or parts[0] not in alt_names:
raise KqlEngineError(
'invalid connection string, must be prefixed by a valid "<uri schema name>://"'
)
rest = conn_str[len(parts[0]) + 3:].strip()
# get key/values in connection string
parsed_conn_kv = Parser.parse_and_get_kv_string(rest, user_ns)
# In case certificate_pem_file was specified instead of certificate.
pem_file_name = parsed_conn_kv.get(ConnStrKeys.CERTIFICATE_PEM_FILE)
if pem_file_name is not None:
with open(pem_file_name, "r") as pem_file:
parsed_conn_kv[ConnStrKeys.CERTIFICATE] = pem_file.read()
del parsed_conn_kv[ConnStrKeys.CERTIFICATE_PEM_FILE]
matched_keys_set = set(parsed_conn_kv.keys())
# check for unknown keys
all_keys = set(itertools.chain(*valid_keys_combinations))
unknonw_keys_set = matched_keys_set.difference(all_keys)
if len(unknonw_keys_set) > 0:
raise KqlEngineError(
"invalid connection string, detected unknown keys: {0}.".
format(unknonw_keys_set))
# check that mandatory key in matched set
if mandatory_key not in matched_keys_set:
raise KqlEngineError(
"invalid connection strin, mandatory key {0} is missing.".
format(mandatory_key))
# find a valid combination for the set
valid_combinations = [
c for c in valid_keys_combinations if matched_keys_set.issubset(c)
]
# in case of ambiguity, assume it is based on current connection, resolve by copying missing values from current
if len(valid_combinations) > 1:
if current is not None:
for k, v in current._parsed_conn.items():
if k not in matched_keys_set and k not in self._NOT_INHERITABLE_KEYS:
parsed_conn_kv[k] = v
matched_keys_set.add(k)
for k in self._CREDENTIAL_KEYS.intersection(matched_keys_set):
if parsed_conn_kv[k] != current._parsed_conn.get(k):
raise KqlEngineError(
"invalid connection string, not a valid keys set, missing keys."
)
valid_combinations = [
c for c in valid_combinations if matched_keys_set.issubset(c)
]
# only one combination can be accepted
if len(valid_combinations) == 0:
raise KqlEngineError(
"invalid connection string, not a valid keys set, missing keys."
)
conn_keys_list = None
# if still too many choose the shortest
if len(valid_combinations) > 1:
for c in valid_combinations:
if len(c) == 3:
conn_keys_list = c
else:
conn_keys_list = valid_combinations[0]
if conn_keys_list is None:
raise KqlEngineError(
"invalid connection string, not a valid keys set, missing keys."
)
conn_keys_set = set(conn_keys_list)
# in case inheritable fields are missing inherit from current if exist
inherit_keys_set = self._INHERITABLE_KEYS.intersection(
conn_keys_set).difference(matched_keys_set)
if len(inherit_keys_set) > 1:
if current is not None:
for k in inherit_keys_set:
v = current._parsed_conn.get(k)
if v is not None:
parsed_conn_kv[k] = v
matched_keys_set.add(k)
# make sure that all required keys are in set
secret_key_set = self._SECRET_KEYS.intersection(conn_keys_set)
missing_set = conn_keys_set.difference(matched_keys_set).difference(
secret_key_set).difference(self._OPTIONAL_KEYS)
if len(missing_set) > 0:
raise KqlEngineError(
"invalid connection string, missing {0}.".format(missing_set))
# special case although tenant in _OPTIONAL_KEYS
if parsed_conn_kv.get(
ConnStrKeys.TENANT
) is None and ConnStrKeys.CLIENTID in conn_keys_set:
raise KqlEngineError(
"invalid connection string, missing tenant key/value.")
# make sure that all required keys are with proper value
for key in matched_keys_set: # .difference(secret_key_set).difference(self._SHOULD_BE_NULL_KEYS):
if key in self._SHOULD_BE_NULL_KEYS:
if parsed_conn_kv[key] != "":
raise KqlEngineError(
"invalid connection string, key {0} must be empty.".
format(key))
elif key not in self._SECRET_KEYS:
if parsed_conn_kv[key] == "<{0}>".format(
key) or parsed_conn_kv[key] == "":
raise KqlEngineError(
"invalid connection string, key {0} cannot be empty or set to <{1}>."
.format(key, key))
# in case secret is missing, get it from user
if len(secret_key_set) == 1:
s = secret_key_set.pop()
if s not in matched_keys_set or parsed_conn_kv[
s] == "<{0}>".format(s):
parsed_conn_kv[s] = getpass.getpass(
prompt="please enter {0}: ".format(s))
matched_keys_set.add(s)
# set attribuets
self.cluster_name = parsed_conn_kv.get(
ConnStrKeys.CLUSTER) or uri_schema_name
self.database_name = parsed_conn_kv.get(mandatory_key)
self.alias = parsed_conn_kv.get(ConnStrKeys.ALIAS)
bind_url = []
for key in conn_keys_list:
if key not in self._EXCLUDE_FROM_URL_KEYS:
bind_url.append("{0}('{1}')".format(key,
parsed_conn_kv.get(key)))
self.bind_url = "{0}://".format(uri_schema_name) + ".".join(bind_url)
return parsed_conn_kv
def test_parse_kusto_socket_connection():
conn_str = "{}://cluster('clustername').database('dbname')".format(TEST_URI_SCHEMA_NAME)
assert Parser.parse("{0} {1}".format(conn_str, query1), empty_config) == \
{'connection': conn_str,
'kql': query1,
'options': default_options}
def test_parse_with_kql():
assert Parser.parse("dbname@clustername {}".format(query1),
empty_config) == \
{'connection': "dbname@clustername",
'kql': query1,
'options': default_options}
def test_parse_no_kql():
assert Parser.parse("dbname@clustername", empty_config) == \
{'connection': "dbname@clustername",
'kql': '',
'options': default_options}
def execute(self, line, cell="", local_ns={}):
"""Query Kusto or ApplicationInsights using kusto query language (kql). Repository specified by a connect string.
Magic Syntax::
%%kql <connection-string>
<KQL statement>
# Note: establish connection and query.
%%kql <established-connection-reference>
<KQL statemnt>
# Note: query using an established connection.
%%kql
<KQL statement>
# Note: query using current established connection.
%kql <KQL statment>
# Note: single line query using current established connection.
%kql <connection-string>
# Note: established connection only.
Connection string Syntax::
kusto://username('<username>).password(<password>).cluster(<cluster>).database(<database>')
appinsights://appid(<appid>).appkey(<appkey>)
loganalytics://workspace(<workspaceid>).appkey(<appkey>)
%<connectionStringVariable>%
# Note: connection string is taken from the environment variable.
[<sectionName>]
# Note: connection string is built from the dsn file settings, section <sectionName>.
# The dsn filename value is taken from configuartion value Kqlmagic.dsn_filename.
# Note: if password or appkey component is missing, user will be prompted.
# Note: connection string doesn't have to include all components, see examples below.
# Note: substring of the form $name or ${name} in windows also %name%, are replaced by environment variables if exist.
Examples::
%%kql kusto://username('myName').password('myPassword').cluster('myCluster').database('myDatabase')
<KQL statement>
# Note: establish connection to kusto and submit query.
%%kql myDatabase@myCluster
<KQL statement>
# Note: submit query using using an established kusto connection to myDatabase database at cluster myCluster.
%%kql appinsights://appid('myAppid').appkey('myAppkey')
<KQL statement>
# Note: establish connection to ApplicationInsights and submit query.
%%kql myAppid@appinsights
<KQL statement>
# Note: submit query using established ApplicationInsights connection to myAppid.
%%kql loganalytics://workspace('myWorkspaceid').appkey('myAppkey')
<KQL statement>
# Note: establish connection to LogAnalytics and submit query.
%%kql myWorkspaceid@loganalytics
<KQL statement>
# Note: submit query using established LogAnalytics connection to myWorkspaceid.
%%kql
<KQL statement>
# Note: submit query using current established connection.
%kql <KQL statement>
# Note: submit single line query using current established connection.
%%kql kusto://cluster('myCluster').database('myDatabase')
<KQL statement>
# Note: establish connection to kusto using current username and password to form the full connection string and submit query.
%%kql kusto://database('myDatabase')
<KQL statement>
# Note: establish connection to kusto using current username, password and cluster to form the full connection string and submit query.
%kql kusto://username('myName').password('myPassword').cluster('myCluster')
# Note: set current (default) username, passsword and cluster to kusto.
%kql kusto://username('myName').password('myPassword')
# Note: set current (default) username and password to kusto.
%kql kusto://cluster('myCluster')
# Note set current (default) cluster to kusto.
"""
set_logger(Logger(None, create_log_context()))
# save globals and locals so they can be referenced in bind vars
user_ns = self.shell.user_ns.copy()
user_ns.update(local_ns)
logger().debug("To Parsed: \n\rline: {}\n\rcell:\n\r{}".format(
line, cell))
try:
parsed = None
parsed_queries = Parser.parse("%s\n%s" % (line, cell), self)
logger().debug("Parsed: {}".format(parsed_queries))
result = None
for parsed in parsed_queries:
result = self.execute_query(parsed, user_ns)
return result
except Exception as e:
if parsed:
if parsed["options"].get("short_errors", self.short_errors):
Display.showDangerMessage(str(e))
return None
elif self.short_errors:
Display.showDangerMessage(str(e))
return None
raise
