def _get_kusto_database_engine(cls, database_name, cluster_name, alias, user_ns: dict):
if cluster_name in cls._ENGINE_MAP.keys():
raise KqlEngineError(
'invalid connection_str, connection_str pattern "database@cluster" cannot be used for "appinsights", "loganalytics" and "cache"'
)
cluster_conn_name = "@" + cluster_name
cluster_conn = cls.connections.get(cluster_conn_name)
if cluster_conn is None:
raise KqlEngineError(
'invalid connection_str, connection_str pattern "database@cluster" can be used only after a previous connection was established to a cluster'
)
details = {ConnStrKeys.DATABASE: database_name, ConnStrKeys.CLUSTER: cluster_name, ConnStrKeys.ALIAS: alias}
return KustoEngine(details, user_ns, conn_class=Connection)
def __init__(self, connect_str, user_ns:dict, **kwargs):
engine = self._find_engine(connect_str)
# wasn't found in connection list, but maybe a kusto database connection
if engine is None:
if "@" in connect_str:
engine = KustoEngine
else:
valid_prefixes_str = ", ".join(["{0}://".format(s) for s in self._ENGINE_MAP.keys()])
raise KqlEngineError("invalid connection_str, unknown <uri schema name>. valid uri schemas are: {0}".format(valid_prefixes_str))
last_current = self.last_current_by_engine.get(engine.__name__)
if engine != KustoEngine:
conn_engine = engine(connect_str, user_ns, last_current)
else:
if "://" in connect_str:
if last_current:
last_cluster_name = last_current.get_cluster()
last_current = self.connections.get("@" + last_cluster_name)
cluster_conn_engine = engine(connect_str, user_ns, last_current)
cluster_name = cluster_conn_engine.get_cluster()
Connection._set_current(cluster_conn_engine, conn_name="@" + cluster_name)
database_name = cluster_conn_engine.get_database()
alias = cluster_conn_engine.get_alias()
else:
database_name, cluster_name = connect_str.split("@")
alias = None
conn_engine = Connection._get_kusto_database_engine(database_name, cluster_name, alias, user_ns)
if kwargs.get("use_cache") and engine != CacheEngine:
conn_engine = CacheEngine(conn_engine, user_ns, last_current, cache_name=kwargs.get("use_cache"))
Connection._set_current(conn_engine)
def get_client(self):
if self.client is None:
cluster_connection = self.conn_class.get_connection_by_name("@" + self.cluster_name)
if cluster_connection is None:
raise KqlEngineError("connection to cluster not set.")
return cluster_connection.get_client()
else:
return self.client
def validate(self, **kwargs):
client = self.get_client()
if not client:
raise KqlEngineError("Client is not defined.")
# query = "range c from 1 to 10 step 1 | count"
filename = self._VALIDATION_FILE_NAME
database = self.get_database()
response = client.execute(database,
filename,
accept_partial_results=False,
timeout=None)
# print(response.json_response)
table = KqlResponse(response, **kwargs).tables[0]
if table.rowcount() != 1 or table.colcount() != 1 or [
r for r in table.fetchall()
][0][0] != 10:
raise KqlEngineError("Client failed to validate connection.")
def validate(self, **kwargs):
ip = get_ipython()
root_path = ip.starting_dir.replace("\\", "/")
client = self.get_client()
if not client:
raise KqlEngineError("Client is not defined.")
# query = "range c from 1 to 10 step 1 | count"
filename = 'validation_file.json'
database = self.get_database()
response = client.execute(database,
filename,
accept_partial_results=False,
timeout=None)
# print(response.json_response)
table = KqlResponse(response, **kwargs).tables[0]
if table.rowcount() != 1 or table.colcount() != 1 or [
r for r in table.fetchall()
][0][0] != 10:
raise KqlEngineError("Client failed to validate connection.")
def _run_magic(magic, content):
if magic == "reload_ext":
return None
if magic == "config":
if "=" in content:
return "dummy_setting"
return True
check.equal(magic, "kql")
if "KqlErrorUnk" in content:
resp = '{"error": {"code": "UnknownError"}}'
raise KqlError(http_response=resp, message=resp)
if "KqlErrorWS" in content:
resp = '{"error": {"code": "WorkspaceNotFoundError"}}'
raise KqlError(http_response=resp, message=resp)
if "KqlEngineError" in content:
raise KqlEngineError("Test Error")
if "AdalErrorUnk" in content:
resp = {"error_description": "unknown error"}
raise AdalError("Test Error", error_response=resp)
if "AdalErrorNR" in content:
raise AdalError("Test Error")
if "AdalErrorPoll" in content:
raise AdalError("Unexpected polling state code_expired")
if "AuthenticationError" in content:
raise AuthenticationError("Test Error")
if content == "--schema":
return {
"table1": {
"field1": int,
"field2": str
},
"table2": {
"field1": int,
"field2": str
},
}
if "query_partial" in content:
return KqlResultTest(code=0, partial=True, status="partial")
if "query_failed" in content:
return KqlResultTest(code=1, partial=False, status="failed")
return KqlResultTest(code=0, partial=False, status="success")
