コード例 #1
0
 def popup_schema(file_path, conn):
     if file_path:
         conn_name = conn.kql_engine.get_conn_name() if isinstance(
             conn, CacheEngine) else conn.get_conn_name()
         button_text = "popup schema " + conn_name
         window_name = "_" + conn_name.replace("@", "_at_") + "_schema"
         Display.show_window(window_name,
                             file_path,
                             button_text=button_text,
                             onclick_visibility="visible")
コード例 #2
0
    def acquire_token(self, **options):
        """Acquire tokens from AAD."""
        token = self._adal_context.acquire_token(self._resource,
                                                 self._username,
                                                 self._client_id)
        if token is not None:
            expiration_date = dateutil.parser.parse(
                token[TokenResponseFields.EXPIRES_ON])
            if expiration_date > datetime.now() + timedelta(minutes=1):
                return self._get_header(token)
            if TokenResponseFields.REFRESH_TOKEN in token:
                token = self._adal_context.acquire_token_with_refresh_token(
                    token[TokenResponseFields.REFRESH_TOKEN], self._client_id,
                    self._resource)
                if token is not None:
                    return self._get_header(token)

        if self._authentication_method is AuthenticationMethod.aad_username_password:
            token = self._adal_context.acquire_token_with_username_password(
                self._resource, self._username, self._password,
                self._client_id)
        elif self._authentication_method is AuthenticationMethod.aad_application_key:
            token = self._adal_context.acquire_token_with_client_credentials(
                self._resource, self._client_id, self._client_secret)
        elif self._authentication_method is AuthenticationMethod.aad_device_login:
            # print(code[OAuth2DeviceCodeResponseParameters.MESSAGE])
            # webbrowser.open(code[OAuth2DeviceCodeResponseParameters.VERIFICATION_URL])
            # token = self._adal_context.acquire_token_with_device_code(
            #     self._resource, code, self._client_id
            # )
            code: dict = self._adal_context.acquire_user_code(
                self._resource, self._client_id)
            url = code[OAuth2DeviceCodeResponseParameters.VERIFICATION_URL]
            device_code = code[
                OAuth2DeviceCodeResponseParameters.USER_CODE].strip()

            html_str = ("""<!DOCTYPE html>
                <html><body>

                <!-- h1 id="user_code_p"><b>""" + device_code +
                        """</b><br></h1-->

                <input  id="kql_MagicCodeAuthInput" type="text" readonly style="font-weight: bold; border: none;" size = '"""
                        + str(len(device_code)) + """' value='""" +
                        device_code + """'>

                <button id='kql_MagicCodeAuth_button', onclick="this.style.visibility='hidden';kql_MagicCodeAuthFunction()">Copy code to clipboard and authenticate</button>

                <script>
                var kql_MagicUserCodeAuthWindow = null
                function kql_MagicCodeAuthFunction() {
                    /* Get the text field */
                    var copyText = document.getElementById("kql_MagicCodeAuthInput");

                    /* Select the text field */
                    copyText.select();

                    /* Copy the text inside the text field */
                    document.execCommand("copy");

                    /* Alert the copied text */
                    // alert("Copied the text: " + copyText.value);

                    var w = screen.width / 2;
                    var h = screen.height / 2;
                    params = 'width='+w+',height='+h
                    kql_MagicUserCodeAuthWindow = window.open('""" + url +
                        """', 'kql_MagicUserCodeAuthWindow', params);

                    // TODO: save selected cell index, so that the clear will be done on the lince cell
                }
                </script>

                </body></html>""")

            if options.get("notebook_app") in ["visualstudiocode", "ipython"]:
                Display.show_window('verification_url', url, **options)
                # Display.showInfoMessage("Code: {0}".format(device_code))
                Display.showInfoMessage(
                    "Copy code: {0} to verification url: {1} and authenticate".
                    format(device_code, url), **options)
            else:
                Display.show_html(html_str)

            try:
                token = self._adal_context.acquire_token_with_device_code(
                    self._resource, code, self._client_id)
            finally:
                html_str = """<!DOCTYPE html>
                    <html><body><script>

                        // close authentication window
                        if (kql_MagicUserCodeAuthWindow && kql_MagicUserCodeAuthWindow.opener != null && !kql_MagicUserCodeAuthWindow.closed) {
                            kql_MagicUserCodeAuthWindow.close()
                        }
                        // TODO: make sure, you clear the right cell. BTW, not sure it is a must to do any clearing

                        // clear output cell
                        Jupyter.notebook.clear_output(Jupyter.notebook.get_selected_index())

                        // TODO: if in run all mode, move to last cell, otherwise move to next cell
                        // move to next cell

                    </script></body></html>"""

                Display.show_html(html_str)
        elif self._authentication_method is AuthenticationMethod.aad_application_certificate:
            token = self._adal_context.acquire_token_with_client_certificate(
                self._resource, self._client_id, self._certificate,
                self._thumbprint)
        else:
            raise AuthenticationError("Unknown authentication method.")
        return self._get_header(token)
コード例 #3
0
    def execute_query(self, parsed, user_ns, result_set=None):
        if Help_html.showfiles_base_url is None:
            window_location = user_ns.get("NOTEBOOK_URL")
            if window_location is not None:
                Help_html.flush(window_location,
                                notebook_app=self.notebook_app)
            else:
                self.submit_get_notebook_url()

        query = parsed["query"].strip()
        options = parsed["options"]
        suppress_results = options.get(
            "suppress_results", False) and options.get(
                "enable_suppress_result", self.enable_suppress_result)
        connection_string = parsed["connection"]

        special_info = False
        if options.get("version"):
            print("{0} version: {1}".format(Constants.MAGIC_PACKAGE_NAME,
                                            VERSION))
            special_info = True

        if options.get("palette"):
            palette = Palette(
                palette_name=options.get("palette_name", self.palette_name),
                n_colors=options.get("palette_colors", self.palette_colors),
                desaturation=options.get("palette_desaturation",
                                         self.palette_desaturation),
                to_reverse=options.get("palette_reverse", False),
            )
            html_str = palette._repr_html_()
            Display.show_html(html_str)
            special_info = True

        if options.get("popup_palettes"):
            n_colors = options.get("palette_colors", self.palette_colors)
            desaturation = options.get("palette_desaturation",
                                       self.palette_desaturation)
            palettes = Palettes(n_colors=n_colors, desaturation=desaturation)
            html_str = palettes._repr_html_()
            button_text = "popup {0} colors palettes".format(n_colors)
            file_name = "{0}_colors_palettes".format(n_colors)
            if desaturation is not None and desaturation != 1.0 and desaturation != 0:
                file_name += "_desaturation{0}".format(str(desaturation))
                button_text += " (desaturation {0})".format(str(desaturation))
            file_path = Display._html_to_file_path(html_str, file_name,
                                                   **options)
            Display.show_window(file_name,
                                file_path,
                                button_text=button_text,
                                onclick_visibility="visible")
            special_info = True

        if options.get("popup_help"):
            help_url = "http://aka.ms/kdocs"
            # 'https://docs.loganalytics.io/docs/Language-Reference/Tabular-operators'
            # 'http://aka.ms/kdocs'
            # 'https://kusdoc2.azurewebsites.net/docs/queryLanguage/query-essentials/readme.html'
            # import requests
            # f = requests.get(help_url)
            # html = f.text.replace('width=device-width','width=500')
            # Display.show(html, **{"popup_window" : True, 'name': 'KustoQueryLanguage'})
            button_text = "popup kql help "
            Display.show_window("KustoQueryLanguage",
                                help_url,
                                button_text,
                                onclick_visibility="visible")
            special_info = True

        if special_info and not query and not connection_string:
            return None

        try:
            #
            # set connection
            #
            conn = Connection.get_connection(connection_string, **options)

        # parse error
        except KqlEngineError as e:
            if options.get("short_errors", self.short_errors):
                msg = Connection.tell_format(connection_string)
                Display.showDangerMessage(str(e))
                Display.showInfoMessage(msg)
                return None
            else:
                raise

        # parse error
        except ConnectionError as e:
            if options.get("short_errors", self.short_errors):
                Display.showDangerMessage(str(e))
                self._show_connection_info(show_conn_info="list")
                return None
            else:
                raise

        try:
            # validate connection
            if not conn.options.get(
                    "validate_connection_string_done") and options.get(
                        "validate_connection_string",
                        self.validate_connection_string):
                retry_with_code = False
                try:
                    conn.validate(**options)
                    conn.set_validation_result(True)
                except Exception as e:
                    msg = str(e)
                    if msg.find("AADSTS50079") > 0 and msg.find(
                            "multi-factor authentication") > 0 and isinstance(
                                conn, KustoEngine):
                        Display.showDangerMessage(str(e))
                        retry_with_code = True
                    else:
                        raise e

                if retry_with_code:
                    Display.showInfoMessage(
                        "replaced connection with code authentication")
                    database_name = conn.get_database()
                    cluster_name = conn.get_cluster()
                    uri_schema_name = conn._URI_SCHEMA_NAME
                    connection_string = "{0}://code().cluster('{1}').database('{2}')".format(
                        uri_schema_name, cluster_name, database_name)
                    conn = Connection.get_connection(connection_string,
                                                     **options)
                    conn.validate(**options)
                    conn.set_validation_result(True)

            conn.options["validate_connection_string_done"] = True

            schema_file_path = None
            if options.get("popup_schema") or (
                    not conn.options.get("auto_popup_schema_done") and
                    options.get("auto_popup_schema", self.auto_popup_schema)):
                schema_file_path = Database_html.get_schema_file_path(
                    conn, **options)
                Database_html.popup_schema(schema_file_path, conn)

            conn.options["auto_popup_schema_done"] = True
            if not conn.options.get("add_schema_to_help_done") and options.get(
                    "add_schema_to_help"):
                schema_file_path = schema_file_path or Database_html.get_schema_file_path(
                    conn, **options)
                Help_html.add_menu_item(conn.get_conn_name(), schema_file_path,
                                        **options)
                conn.options["add_schema_to_help_done"] = True

            if not query:
                #
                # If NO  kql query, just return the current connection
                #
                if not connection_string and Connection.connections and not suppress_results:
                    self._show_connection_info(**options)
                return None
            #
            # submit query
            #
            start_time = time.time()

            params_dict_name = options.get('params_dict')
            dictionary = user_ns.get(
                params_dict_name) if params_dict_name is not None and len(
                    params_dict_name) > 0 else user_ns
            parametrized_query = Parameterizer(dictionary).expand(
                query) if result_set is None else result_set.parametrized_query
            raw_query_result = conn.execute(parametrized_query, user_ns,
                                            **options)

            end_time = time.time()

            #
            # model query results
            #
            if result_set is None:
                fork_table_id = 0
                saved_result = ResultSet(raw_query_result,
                                         parametrized_query,
                                         fork_table_id=0,
                                         fork_table_resultSets={},
                                         metadata={},
                                         options=options)
                saved_result.metadata["magic"] = self
                saved_result.metadata["parsed"] = parsed
                saved_result.metadata["connection"] = conn.get_conn_name()
            else:
                fork_table_id = result_set.fork_table_id
                saved_result = result_set.fork_result(0)
                saved_result.feedback_info = []
                saved_result._update(raw_query_result)

            result = saved_result

            if not connection_string and Connection.connections:
                saved_result.metadata["conn_info"] = self._get_connection_info(
                    **options)
            else:
                saved_result.metadata["conn_info"] = []

            saved_result.metadata["start_time"] = start_time
            saved_result.metadata["end_time"] = end_time

            if options.get("feedback", self.feedback):
                minutes, seconds = divmod(end_time - start_time, 60)
                saved_result.feedback_info.append(
                    "Done ({:0>2}:{:06.3f}): {} records".format(
                        int(minutes), seconds, saved_result.records_count))

            if options.get("columns_to_local_vars",
                           self.columns_to_local_vars):
                # Instead of returning values, set variables directly in the
                # users namespace. Variable names given by column names

                if options.get("feedback", self.feedback):
                    saved_result.feedback_info.append(
                        "Returning raw data to local variables")

                self.shell.user_ns.update(saved_result.to_dict())
                result = None

            if options.get("auto_dataframe", self.auto_dataframe):
                if options.get("feedback", self.feedback):
                    saved_result.feedback_info.append(
                        "Returning data converted to pandas dataframe")
                result = saved_result.to_dataframe()

            if options.get("result_var") and result_set is None:
                result_var = options["result_var"]
                if options.get("feedback", self.feedback):
                    saved_result.feedback_info.append(
                        "Returning data to local variable {}".format(
                            result_var))
                self.shell.user_ns.update({
                    result_var:
                    result if result is not None else saved_result
                })
                result = None

            if options.get('cache') and not options.get(
                    'use_cache') and not isinstance(conn, CacheEngine):
                file_path = CacheClient().save(raw_query_result,
                                               conn.get_database(),
                                               conn.get_cluster(),
                                               parametrized_query, **options)
                if options.get("feedback", self.feedback):
                    saved_result.feedback_info.append("query results cached")

            if options.get('save_as') is not None:
                file_path = CacheClient().save(raw_query_result,
                                               conn.get_database(),
                                               conn.get_cluster(),
                                               parametrized_query,
                                               filepath=options.get('save_as'),
                                               **options)
                if options.get("feedback", self.feedback):
                    saved_result.feedback_info.append(
                        "query results saved as {0}".format(file_path))

            saved_result.suppress_result = False
            saved_result.display_info = False
            if result is not None:
                if suppress_results:
                    saved_result.suppress_result = True
                elif options.get("auto_dataframe", self.auto_dataframe):
                    Display.showSuccessMessage(saved_result.feedback_info)
                else:
                    saved_result.display_info = True

            if result_set is None:
                saved_result._create_fork_results()
            else:
                saved_result._update_fork_results()

            # Return results into the default ipython _ variable
            self.shell.user_ns.update({
                options.get("last_raw_result_var", self.last_raw_result_var):
                saved_result
            })

            if result == saved_result:
                result = saved_result.fork_result(fork_table_id)
            return result

        except Exception as e:
            if not connection_string and Connection.connections and not suppress_results:
                # display list of all connections
                self._show_connection_info(**options)

            if options.get("short_errors", self.short_errors):
                Display.showDangerMessage(e)
                return None
            else:
                raise e