another-ayx-python-sdk 2.4.2__py3-none-any.whl

another_ayx_python_sdk/providers/amp_provider/builders/record_packet_builder.py

@@ -0,0 +1,133 @@
+# Copyright (C) 2022 Alteryx, Inc. All rights reserved.
+#
+# Licensed under the ALTERYX SDK AND API LICENSE AGREEMENT;
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.alteryx.com/alteryx-sdk-and-api-license-agreement
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Record packet building utilities for converting between core and protobuf."""
+from typing import List, Tuple
+
+from another_ayx_python_sdk.core.constants import NULL_VALUE_PLACEHOLDER
+from another_ayx_python_sdk.core.metadata import Metadata
+from another_ayx_python_sdk.core.record_packet_base import RecordPacketBase
+from another_ayx_python_sdk.providers.amp_provider.amp_record_packet import AMPRecordPacket
+from another_ayx_python_sdk.providers.amp_provider.builders.record_builder import RecordBuilder
+from another_ayx_python_sdk.providers.amp_provider.resources.generated.record_packet_pb2 import (
+    RecordPacket as ProtobufRecordPacket,
+)
+
+
+class RecordPacketBuilder:
+    """Utilities for converting record packets between protobuf and core objects."""
+
+    record_builder = RecordBuilder()
+
+    @classmethod
+    def from_protobuf(
+        cls, protobuf_record_packet: ProtobufRecordPacket, metadata: Metadata
+    ) -> Tuple[RecordPacketBase, int, float]:
+        """
+        Convert a protobuf to a record packet.
+
+        Parameters
+        ----------
+        protobuf_record_packet
+            Protobuf representation of a record packet.
+        metadata
+            Metadata associated with the record packet.
+
+        Returns
+        -------
+        Tuple[RecordPacketBase, int, float]
+            AMPRecordPacket representation of the protobuf record packet, sequence, and progress
+        """
+        import pandas as pd
+
+        parsed_records = [
+            cls.record_builder.from_protobuf(protobuf_record)
+            for protobuf_record in protobuf_record_packet.records
+        ]
+
+        return (
+            AMPRecordPacket(
+                metadata=metadata,
+                df=pd.DataFrame(
+                    parsed_records, columns=[field.name for field in metadata]
+                ),
+            ),
+            protobuf_record_packet.sequence,
+            protobuf_record_packet.progress,
+        )
+
+    @classmethod
+    def to_protobuf(
+        cls, amp_record_packet: RecordPacketBase, sequence: int, progress: float
+    ) -> List["ProtobufRecordPacket"]:
+        """
+        Convert a record packet to a protobuf.
+
+        Parameters
+        ----------
+        amp_record_packet
+            Protobuf representation of a record packet.
+        sequence
+            Position that the current record packet would be in,
+            out of all record packets. (First, second, third, etc.)
+        progress
+            Overall progress (ranging from 0.0 to 1.0)
+
+        Returns
+        -------
+        List[ProtobufRecordPacket]
+            A list of ProtobufRecordPackets. If the passed in AMPRecordPacket would surpass 64MB,
+            then there will be more than one element in list
+        """
+        df = amp_record_packet.to_dataframe()
+        df.fillna(NULL_VALUE_PLACEHOLDER, inplace=True)
+
+        metadata = amp_record_packet.metadata
+
+        protobuf_records = [
+            cls.record_builder.to_protobuf(record, metadata)
+            for record in df.itertuples(index=False)
+        ]
+
+        size_limit = 64 * (2**20)  # 64 MB
+        curr_size = 0
+
+        pb_record_packets = []
+        current_records = []
+
+        for record in protobuf_records:
+            record_size = len(record.data)
+            if curr_size + record_size > size_limit:
+                pb_record_packets.append(
+                    ProtobufRecordPacket(
+                        sequence=sequence,
+                        progress=progress,
+                        records=current_records,
+                    )
+                )
+                sequence += 1
+                current_records = []
+                curr_size = 0
+
+            current_records.append(record)
+            curr_size += record_size
+
+        if current_records:
+            pb_record_packets.append(
+                ProtobufRecordPacket(
+                    sequence=sequence,
+                    progress=progress,
+                    records=current_records,
+                )
+            )
+        return pb_record_packets

another_ayx_python_sdk/providers/amp_provider/builders/tool_config_builder.py

@@ -0,0 +1,59 @@
+# Copyright (C) 2022 Alteryx, Inc. All rights reserved.
+#
+# Licensed under the ALTERYX SDK AND API LICENSE AGREEMENT;
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.alteryx.com/alteryx-sdk-and-api-license-agreement
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tool configuration builder file."""
+from typing import Any, Dict
+
+import xmltodict
+
+
+class ToolConfigBuilder:
+    """Builder class for converting the tool configuration XML to and from a dict."""
+
+    @staticmethod
+    def to_xml(config_dict: Dict[str, Any]) -> str:
+        """
+        Convert a tool configuration dictionary to the expected XML formatted string.
+
+        Parameters
+        ----------
+        config_dict
+            The dictionary representation of the tool configuration XML
+
+        Returns
+        -------
+        str
+            Tool configuration unparsed into an XML string
+        """
+        return xmltodict.unparse(
+            {"Configuration": config_dict}, short_empty_elements=True
+        )
+
+    @staticmethod
+    def from_xml(xml_string: str) -> Dict[str, Any]:
+        """
+        Convert a tool configuration XML to the expected dictionary.
+
+        Parameters
+        ----------
+        xml_string
+            The tool configuration XML string to be converted into a Python dict
+
+        Returns
+        -------
+        dict
+            The dictionary representation of the passed in XML
+        """
+        return dict(
+            xmltodict.parse(xml_string, strip_whitespace=False)["Configuration"] or {}
+        )

another_ayx_python_sdk/providers/amp_provider/cng_certs.py

@@ -0,0 +1,197 @@
+# Copyright (C) 2022 Alteryx, Inc. All rights reserved.
+#
+# Licensed under the ALTERYX SDK AND API LICENSE AGREEMENT;
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.alteryx.com/alteryx-sdk-and-api-license-agreement
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# type: ignore
+"""READ cert/private key from windows STORE."""
+
+
+from ctypes import (
+    POINTER,
+    Structure,
+    WinDLL,
+    c_char_p,
+    c_ulong,
+    c_void_p,
+    cast,
+    create_string_buffer,
+    pointer,
+    string_at,
+)
+from ctypes.wintypes import BYTE, DWORD, LPCWSTR
+from typing import Optional, Tuple
+
+import wincertstore as wcs
+
+if wcs.USE_LAST_ERROR:
+    Ncrypt = WinDLL("Ncrypt.dll", use_last_error=True)
+else:
+    Ncrypt = WinDLL("Ncrypt.dll")
+
+NCRYPT_KEY_HANDLE = c_void_p
+
+NCRYPT_PROV_HANDLE = c_void_p
+
+NCryptOpenStorageProvider = Ncrypt.NCryptOpenStorageProvider
+NCryptOpenStorageProvider.argtypes = [POINTER(NCRYPT_PROV_HANDLE), LPCWSTR, DWORD]
+NCryptOpenStorageProvider.restype = DWORD
+
+MS_KEY_STORAGE_PROVIDER = "Microsoft Software Key Storage Provider"
+
+NCryptOpenKey = Ncrypt.NCryptOpenKey
+NCryptOpenKey.argtypes = [
+    NCRYPT_PROV_HANDLE,
+    POINTER(NCRYPT_KEY_HANDLE),
+    LPCWSTR,
+    DWORD,
+    DWORD,
+]
+NCryptOpenKey.restype = DWORD
+
+
+BCRYPT_RSAFULLPRIVATE_BLOB = "RSAFULLPRIVATEBLOB"
+CNG_RSA_PRIVATE_KEY_BLOB = 83
+SZOID_RSA_RSA = "1.2.840.113549.1.1.1"
+X509_ASN_ENCODING = 1
+CRYPT_ENCODE_ALLOC_FLAG = 0x8000
+PKCS_PRIVATE_KEY_INFO = 44
+
+NCryptExportKey = Ncrypt.NCryptExportKey
+NCryptExportKey.argtypes = [
+    NCRYPT_KEY_HANDLE,
+    NCRYPT_KEY_HANDLE,
+    LPCWSTR,
+    c_void_p,
+    POINTER(BYTE),
+    DWORD,
+    POINTER(DWORD),
+    DWORD,
+]
+NCryptExportKey.restype = DWORD
+
+
+class CryptPrivateKeyInfo(Structure):
+    """Docstring."""
+
+    __slots__ = ()
+    _fields_ = [
+        ("Version", DWORD),
+        # CRYPT_ALGORITHM_IDENTIFIER/Algorithm
+        ("pszObjId", c_char_p),
+        # CRYPT_ALGORITHM_IDENTIFIER.CRYPT_OBJID_BLOB
+        ("algoCbData", DWORD),
+        ("algoPbData", POINTER(BYTE)),
+        # CRYPT_DER_BLOB/PrivateKey
+        ("pkeyCbData", DWORD),
+        ("pkeyPbData", POINTER(BYTE)),
+        ("pAttributes", c_void_p),
+    ]
+
+
+CryptEncodeObjectEx = wcs.crypt32.CryptEncodeObjectEx
+CryptEncodeObjectEx.argtypes = [
+    DWORD,
+    c_void_p,
+    c_void_p,
+    DWORD,
+    c_void_p,
+    POINTER(c_void_p),
+    POINTER(DWORD),
+]
+CryptEncodeObjectEx.restype = DWORD
+
+
+class PrivateKeyPem(wcs.ContextStruct):
+    """Docstring."""
+
+    cert_type = "PRIVATE KEY"
+    # WARNING: please DO NOT CHANGE THIS VARIABLE NAME: it is used in wincertstore
+    dwCertEncodingType = 0  # noqa
+    # End warning
+
+    def __init__(self, pb_key: c_void_p, cb_key: c_ulong) -> None:
+        self.pb_key = pb_key
+        self.cb_key = cb_key
+
+    def get_encoded(self) -> bytes:
+        """Docstring."""
+        return string_at(self.pb_key, self.cb_key.value)
+
+
+def read_windows_store_chain(cert_name: str) -> Optional[Tuple[bytes, bytes]]:
+    """Docstring."""
+    with wcs.CertSystemStore("MY") as store:
+        for cert in store.itercerts(usage=wcs.SERVER_AUTH):
+            if (
+                cert.get_name() == "localhost"
+                and cert.get_name(wcs.CERT_NAME_FRIENDLY_DISPLAY_TYPE) == cert_name
+            ):
+                certpem = cert.get_pem()
+                provider_handle = NCRYPT_PROV_HANDLE()
+                NCryptOpenStorageProvider(provider_handle, MS_KEY_STORAGE_PROVIDER, 0)
+                key_handle = NCRYPT_KEY_HANDLE()
+                NCryptOpenKey(provider_handle, key_handle, cert_name, 0, 0)
+                cb_key = DWORD()
+                NCryptExportKey(
+                    key_handle,
+                    None,
+                    BCRYPT_RSAFULLPRIVATE_BLOB,
+                    None,
+                    None,
+                    0,
+                    cb_key,
+                    0,
+                )
+                pb_key = create_string_buffer(cb_key.value)
+                NCryptExportKey(
+                    key_handle,
+                    None,
+                    BCRYPT_RSAFULLPRIVATE_BLOB,
+                    None,
+                    cast(pb_key, POINTER(BYTE)),
+                    cb_key.value,
+                    cb_key,
+                    0,
+                )
+
+                private_key_info = CryptPrivateKeyInfo()
+                private_key_info.Version = 0
+                private_key_info.pszObjId = SZOID_RSA_RSA.encode("utf-8")
+
+                enc1_cb_key = DWORD()
+                enc1_pb_key = c_void_p()
+                CryptEncodeObjectEx(
+                    X509_ASN_ENCODING,
+                    cast(CNG_RSA_PRIVATE_KEY_BLOB, c_void_p),
+                    pb_key,
+                    CRYPT_ENCODE_ALLOC_FLAG,
+                    None,
+                    enc1_pb_key,
+                    enc1_cb_key,
+                )
+                private_key_info.pkeyCbData = enc1_cb_key
+                private_key_info.pkeyPbData = cast(enc1_pb_key, POINTER(BYTE))
+                enc_cb_key = DWORD()
+                enc_pb_key = c_void_p()
+                CryptEncodeObjectEx(
+                    X509_ASN_ENCODING,
+                    cast(PKCS_PRIVATE_KEY_INFO, c_void_p),
+                    pointer(private_key_info),
+                    CRYPT_ENCODE_ALLOC_FLAG,
+                    None,
+                    enc_pb_key,
+                    enc_cb_key,
+                )
+                pkp = PrivateKeyPem(enc_pb_key, enc_cb_key)
+                keypem = pkp.get_pem()
+                return bytes(keypem, "utf-8"), bytes(certpem, "utf-8")
+    return None

another_ayx_python_sdk/providers/amp_provider/data_transport/__init__.py

@@ -0,0 +1,18 @@
+# Copyright (C) 2022 Alteryx, Inc. All rights reserved.
+#
+# Licensed under the ALTERYX SDK AND API LICENSE AGREEMENT;
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.alteryx.com/alteryx-sdk-and-api-license-agreement
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Implementations of TransportBase."""
+from .amp_transport import AmpStreamTransport
+from .transport_base import TransportBase
+
+__all__ = ["TransportBase", "AmpStreamTransport"]

another_ayx_python_sdk/providers/amp_provider/data_transport/amp_transport.py

@@ -0,0 +1,127 @@
+# Copyright (C) 2022 Alteryx, Inc. All rights reserved.
+#
+# Licensed under the ALTERYX SDK AND API LICENSE AGREEMENT;
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.alteryx.com/alteryx-sdk-and-api-license-agreement
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Arrows stream formatted transport interface."""
+import logging
+from collections import defaultdict
+from typing import Generator, TYPE_CHECKING, Tuple
+
+from another_ayx_python_sdk.providers.amp_provider.data_transport.transport_base import (
+    TransportBase,
+)
+from another_ayx_python_sdk.providers.amp_provider.grpc_helpers.record_transfer_msgs import (
+    new_record_transfer_out_msg,
+)
+
+import pyarrow as pa
+
+if TYPE_CHECKING:
+    from another_ayx_python_sdk.providers.amp_provider.resources.generated.sdk_tool_service_v2_pb2 import (  # noqa: F401
+        RecordTransferIn,
+        RecordTransferOut,
+    )
+
+# Size of stream chunks, in n bytes.
+# Desired 1MB default is causing some failures. Fairly sure this is on the slicing
+# logic, and just needs better bounds to check so arrow doesn't freak out.
+# Boxing that for another ticket after bigger stuff is done, though.
+DEFAULT_CHUNK_SIZE = 1 * (10**6)
+
+
+class AmpStreamTransport(TransportBase):
+    """Transport interface to deliver data in arrows stream format."""
+
+    logger = logging.getLogger()
+
+    def __init__(self) -> None:
+        self.stream_buffer: dict = defaultdict(dict)
+        self.chunk_size = DEFAULT_CHUNK_SIZE
+        self.incomplete_streams = 0
+
+    # TODO: Rename this to something better for Transport interface.
+    def send_record(
+        self,
+        record_batch: "pa.RecordBatch",
+        anchor_name: str,
+        connection_name: str = "",
+    ) -> Generator["RecordTransferOut", None, None]:
+        """Convert a given record batch into a sequence of RecordTransferOut msgs."""
+        stream = self.get_rec_stream(record_batch)
+        for chunk, end_of_chunk in self.iter_stream(stream):
+            record_out = {
+                "anchor_name": anchor_name,
+                "data": chunk,
+                "connection_name": connection_name,
+                "end_of_chunk": end_of_chunk,
+            }
+            yield new_record_transfer_out_msg(record_out, "outgoing_records")
+
+    def get_rec_stream(self, record_batch: "pa.RecordBatch") -> "pa.Buffer":
+        """Get an arrows formatted stream buffer."""
+        try:
+            sink = pa.BufferOutputStream()
+            writer = pa.ipc.new_stream(sink, record_batch.schema)
+            writer.write(record_batch)
+            writer.close()
+        except Exception as e:
+            self.logger.error(repr(e))
+            raise e
+        return sink.getvalue()
+
+    def iter_stream(
+        self, stream: "pa.Buffer"
+    ) -> Generator[Tuple[bytes, bool], None, None]:
+        """Break up a given stream by set chunk size."""
+        buf = pa.BufferReader(stream)
+        end_of_chunk = False
+        while not end_of_chunk:
+            end_of_chunk = buf.tell() + self.chunk_size >= buf.size()
+            chunk = buf.read() if end_of_chunk else buf.read(self.chunk_size)
+            yield (chunk, end_of_chunk)
+
+    def _unpack_and_append_rec_in_payload(self, record_in: "RecordTransferIn") -> None:
+        anchor_name = record_in.incoming_records.anchor_name
+        connection_name = record_in.incoming_records.connection_name
+        data = record_in.incoming_records.data
+        if not self.stream_buffer[anchor_name].get(connection_name):
+            self.incomplete_streams += 1
+            self.stream_buffer[anchor_name][connection_name] = pa.BufferOutputStream()
+        self.stream_buffer[anchor_name][connection_name].write(data)
+
+    def _handle_completed_stream(self, record_in: "RecordTransferIn") -> None:
+        self.logger.debug("Stream terminater received, completing batch.")
+        anchor_name = record_in.incoming_records.anchor_name
+        connection_name = record_in.incoming_records.connection_name
+        data = self.stream_buffer[anchor_name][connection_name]
+        try:
+            reader = pa.ipc.open_stream(data.getvalue())
+            self.stream_buffer[anchor_name][connection_name] = pa.BufferOutputStream()
+            self.incomplete_streams -= 1
+            return reader.read_all()
+        except pa.ArrowException as ae:
+            self.logger.error(
+                "Exception caught reading completed batch chunk: %s", repr(ae)
+            )
+            raise ae
+
+    def receive_record(self, record_in_msg: "RecordTransferIn") -> "pa.Table":
+        """
+        Receive RecordTransferIn messages.
+
+        Returns None if the data received is only a chunk of a full record.
+        Returns pa.Table if `end_of_chunk` is True.
+        """
+        self._unpack_and_append_rec_in_payload(record_in_msg)
+        if record_in_msg.incoming_records.end_of_chunk:
+            return self._handle_completed_stream(record_in_msg)
+        return None

another_ayx_python_sdk/providers/amp_provider/data_transport/transport_base.py

@@ -0,0 +1,35 @@
+# Copyright (C) 2022 Alteryx, Inc. All rights reserved.
+#
+# Licensed under the ALTERYX SDK AND API LICENSE AGREEMENT;
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.alteryx.com/alteryx-sdk-and-api-license-agreement
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Transport implementation.
+
+Transports provide the functionality for data management and IO.
+An implementation should provide an independant way to receive, parse,
+serialize, and send data to another service.
+"""
+from abc import ABC, abstractmethod
+
+
+class TransportBase(ABC):
+    """Transport Manager for the tool."""
+
+    @abstractmethod
+    def send_record(self) -> None:
+        """Push data to destination."""
+        raise NotImplementedError()
+
+    @abstractmethod
+    def receive_record(self) -> None:
+        """Pull data from source."""
+        raise NotImplementedError()

another_ayx_python_sdk/providers/amp_provider/grpc_helpers/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (C) 2022 Alteryx, Inc. All rights reserved.
+#
+# Licensed under the ALTERYX SDK AND API LICENSE AGREEMENT;
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.alteryx.com/alteryx-sdk-and-api-license-agreement
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""gRPC Helper Functions."""