Merge pull request #235 from atlanhq/DVX-186

Correct defects found while testing custom package generation

HISTORY.md

@@ -1,3 +1,15 @@
+## 1.9.3 (January 30, 2024)
+
+### Bug fixes
+
+- raise exception if tag in query in FluentSearch does not exist
+- (Experimental) fix defects custom package generation
+
+### QOL Improvements
+
+- add api to assign a token to a purpose
+- add api to enable running sql queries
+
 ## 1.9.2 (January 24, 2024)
 
 ### Bug fixes

pyatlan/pkg/create_package_files.py

@@ -0,0 +1,174 @@
+import re
+import sys
+from pathlib import Path
+from typing import Optional
+
+FILE_NAME = "package_config.py"
+OUTPUT_FILE = Path("package_config.py")
+CODE = """from pathlib import Path
+from pyatlan.pkg.models import CustomPackage, generate, PullPolicy
+from pyatlan.pkg.ui import UIConfig, UIRule, UIStep
+from pyatlan.pkg.widgets import (
+    APITokenSelector,
+    BooleanInput,
+    ConnectionCreator,
+    ConnectionSelector,
+    ConnectorTypeSelector,
+    DateInput,
+    DropDown,
+    FileUploader,
+    KeygenInput,
+    MultipleGroups,
+    MultipleUsers,
+    NumericInput,
+    PasswordInput,
+    Radio,
+    SingleGroup,
+    SingleUser,
+    TextInput,
+)
+
+PARENT = Path(__file__).parent
+TOP = PARENT.parent
+
+def create_package() -> CustomPackage:
+    \"\"\"Create the custom package\"\"\"
+    return CustomPackage(
+        package_id="{package_id}",
+
+        container_image="ghcr.io/atlanhq/{image}",
+        container_command=["python", "'-m'", "{package_name}.main"],
+        outputs={{
+            "debug-logs": "/tmp/debug.log",
+        }}
+    )
+
+if __name__ == "__main__":
+    package = create_package()
+    generate(pkg=package, path=TOP / "generated_packages", operation="package")
+    generate(pkg=package, path=PARENT, operation="config")
+"""
+
+LOGGING_CONF = """[loggers]
+keys=root,pyatlan,urllib3
+
+[handlers]
+keys=consoleHandler,fileHandler,jsonHandler
+
+[formatters]
+keys=simpleFormatter,jsonFormatter
+
+[logger_root]
+level=INFO
+handlers=consoleHandler,fileHandler
+
+[logger_pyatlan]
+level=DEBUG
+handlers=fileHandler,jsonHandler
+qualname=pyatlan
+propagate=0
+
+[logger_urllib3]
+level=DEBUG
+handlers=fileHandler,jsonHandler
+qualname=urllib3
+propagate=0
+
+[handler_consoleHandler]
+class=StreamHandler
+formatter=simpleFormatter
+args=(sys.stdout,)
+
+[handler_fileHandler]
+class=FileHandler
+level=DEBUG
+formatter=simpleFormatter
+args=('/tmp/debug.log',)
+
+[handler_jsonHandler]
+class=FileHandler
+level=DEBUG
+formatter=jsonFormatter
+args=('/tmp/pyatlan.json',)
+
+[formatter_simpleFormatter]
+format=%(asctime)s - %(name)s - %(levelname)s - %(message)s
+
+[formatter_jsonFormatter]
+format=%(asctime)s - %(name)s - %(levelname)s - %(message)s
+class=pyatlan.utils.JsonFormatter
+"""
+REQUIREMENTS = """pyatlan
+"""
+DOCKER_FILE = """FROM python:3.9-bookworm
+
+
+RUN wget -O /usr/local/bin/dumb-init https://github.com/Yelp/dumb-init/releases/download/v1.2.5/dumb-init_1.2.5_x86_64
+RUN chmod +x /usr/local/bin/dumb-init
+
+COPY requirements.txt requirements.txt
+
+RUN pip3 install -r requirements.txt
+
+WORKDIR /app
+
+ADD {package_name} /app/{package_name}
+
+ENTRYPOINT ["/usr/local/bin/dumb-init", "--"]
+"""
+MAIN = """import logging
+from {package_name}.{package_name}_cfg import RuntimeConfig
+from pyatlan.pkg.utils import set_package_ops
+
+LOGGER = logging.getLogger(__name__)
+
+def main():
+    \"\"\""Main logic\"\"\"
+    runtime_config = RuntimeConfig()
+    custom_config = runtime_config.custom_config
+    client = set_package_ops(runtime_config)
+
+if __name__ == '__main__':
+    main()
+"""
+
+
+def write_file(output_file: Path, data: Optional[str] = None):
+    if output_file.exists():
+        print(f"{output_file.absolute()} already exists. Leaving unchanged.")
+        return
+    with output_file.open("w") as output:
+        if data:
+            output.write(data)
+
+
+def main(package_name: str):
+    output_dir = Path(package_name)
+    package_id = f"@csa/{sys.argv[1].replace('_', '-')}"
+    output_file = output_dir / FILE_NAME
+    output_dir.mkdir(exist_ok=True)
+    write_file(Path("requirements.txt"), REQUIREMENTS)
+    write_file(Path("Dockerfile"), DOCKER_FILE.format(package_name=package_name))
+    write_file((output_dir / "__init__.py"))
+    write_file((output_dir / "logging.conf"), LOGGING_CONF)
+    write_file(
+        output_file,
+        CODE.format(
+            package_id=package_id,
+            package_name=package_name,
+            image=f"csa-{package_name.replace('_', '-')}",
+        ),
+    )
+    write_file(Path(output_dir / "main.py"), MAIN.format(package_name=package_name))
+
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print("Please specify the python package name for the package")
+        exit(1)
+    if not re.fullmatch(r"\w+", sys.argv[1], re.ASCII):
+        print(
+            "The package name can only consist of alphanumeric characters and the underscore"
+        )
+
+    main(sys.argv[1])

pyatlan/pkg/models.py

@@ -2,7 +2,6 @@
 # Copyright 2023 Atlan Pte. Ltd.
 import json
 import logging
-import re
 import textwrap
 from enum import Enum
 from importlib import resources
@@ -146,7 +145,7 @@ class CustomPackage(BaseModel):
     :ivar keywords list[str]: (optional) list of any keyword labels to apply to the package
     :ivar container_image str: container image to run the logic of the custom package
     :ivar container_image_pull_policy PullPolicy: (optional) override the default IfNotPresent policy
-    :ivar container_command list[str[: the full command to run in the container image, as a list rather than spaced
+    :ivar container_command list[str]: the full command to run in the container image, as a list rather than spaced
     (must be provided if you have not specified the class above)
     :ivar allow_schedule bool: (optional) whether to allow the package to be scheduled (default, true) or only run
     immediately (false)
@@ -248,7 +247,6 @@ def create_package(self):
 
     def create_config(self):
         self.create_config_class()
-        self.create_logging_conf()
 
     def create_templates(self):
         self._templates_dir.mkdir(parents=True, exist_ok=True)
@@ -267,17 +265,9 @@ def create_configmaps(self):
     def create_config_class(self):
         template = self._env.get_template("package_config.jinja2")
         content = template.render({"pkg": self.pkg})
-        file_name = re.sub(r"\W+", "_", self.pkg.package_name).lower()
+        file_name = f'{self.pkg.package_id[5:].replace("-","_")}_cfg.py'
 
-        with (
-            self.path / f"{file_name}{'' if file_name.endswith('_') else '_'}cfg.py"
-        ).open("w") as script:
-            script.write(content)
-
-    def create_logging_conf(self):
-        template = self._env.get_template("logging_conf.jinja2")
-        content = template.render({"pkg": self.pkg})
-        with (self.path / "logging.conf").open("w") as script:
+        with (self.path / file_name).open("w") as script:
             script.write(content)
 
 
@@ -297,3 +287,8 @@ def generate(pkg: CustomPackage, path: Path, operation: Literal["package", "conf
         writer.create_package()
     else:
         writer.create_config()
+
+
+class ConnectorAndConnection(BaseModel):
+    source: AtlanConnectorType
+    connections: list[str]

pyatlan/pkg/templates/default_template.jinja2

@@ -18,14 +18,14 @@ spec:
 {%- set ns = namespace(has_s3=false) %}
 {%-  for name, property in pkg.ui_config.properties.items()  %}
     {%- if property.ui.s3_artifact %}
-        {%- if loop.first %}
+        {%- if not ns.has_s3 %}
         artifacts:
         {%- endif %}
         {%-  set ns.has_s3 = true %}
-          - name: {{ name }}_s3
-            path: "/tmp/{{ name }}/{% raw %}{{inputs.parameters.{% endraw %}{{ name }}{% raw %}}}{% endraw %}"
-            s3:
-              key: "{% raw %}{{inputs.parameters.{% endraw %}{{ name }}{% raw %}}}{% endraw %}"
+            - name: {{ name }}_s3
+              path: "/tmp/{{ name }}/{% raw %}{{inputs.parameters.{% endraw %}{{ name }}{% raw %}}}{% endraw %}"
+              s3:
+                key: "{% raw %}{{inputs.parameters.{% endraw %}{{ name }}{% raw %}}}{% endraw %}"
     {%- endif  %}
 {%- endfor %}
 {%- if not ns.has_s3 %}

pyatlan/pkg/templates/package_config.jinja2

@@ -1,12 +1,11 @@
 from datetime import datetime
 from pathlib import Path
-from pydantic import BaseModel, BaseSettings, Field, parse_obj_as, validator
+from pydantic import BaseModel, BaseSettings, Field, validator
 from pyatlan.model.assets import Connection
-from pyatlan.model.enums import AtlanConnectorType
+from pyatlan.pkg.models import ConnectorAndConnection
+from pyatlan.pkg.utils import validate_connection, validate_multiselect, validate_connector_and_connection
 from typing import Any, Optional, Union
-import json
 import logging.config
-import os
 
 PARENT = Path(__file__).parent
 LOGGING_CONF = PARENT / "logging.conf"
@@ -16,29 +15,6 @@ LOGGER = logging.getLogger(__name__)
 
 ENV = 'env'
 
-
-def validate_multiselect(v, values, **kwargs):
-    if isinstance(v, str):
-        if v.startswith('['):
-            data = json.loads(v)
-            v = parse_obj_as(list[str], data)
-        else:
-            v = [v]
-    return v
-
-
-def validate_connection(v, values, config, field, **kwargs):
-    v = Connection.parse_raw(v)
-
-
-class ConnectorAndConnection(BaseModel):
-    source: AtlanConnectorType
-    connections: list[str]
-
-
-def validate_connector_and_connection(v, values, config, field, **kwargs):
-    return ConnectorAndConnection.parse_raw(v)
-
 class CustomConfig(BaseModel):
     """"""""
 {%- for key, value in pkg.ui_config.properties.items() %}
@@ -79,3 +55,16 @@ class RuntimeConfig(BaseSettings):
             if field_name == 'custom_config':
                 return CustomConfig.parse_raw(raw_value)
             return cls.json_loads(raw_value)
+
+    @property
+    def envars_as_dict(self) -> dict[str, Any]:
+        """
+        :return dict: returns a dict of the environment variable names and values consumed by this RuntimeConfig.
+        the name of an environment variable will be the key and the value will be the value. This is provided
+        to facilitate testing
+        """
+        ret_val: dict[str, Any] = {}
+        for key, value in RuntimeConfig.Config.fields.items():
+            if field_value := getattr(self, key):
+                ret_val[value["env"]] = field_value.json()
+        return ret_val