from metaflow import FlowSpec, step, pypi, pypi_base, kubernetes, environment

@pypi_base(
    python='3.10',
    extra_indices=[f"<https://us-python.pkg.dev/.../simple>"],
    packages={
        'keyrings.google-artifactregistry-auth': '>=1.1.1' ,
        'my-private-package': '3.4.1',
    }
)
class PrivateLibFlow(FlowSpec):
    @environment(vars={"METAFLOW_DEBUG_CONDA": "1"})
    @step
    def start(self):
        import pkg_resources
        installed_packages = pkg_resources.working_set
        installed_packages_list = sorted(["%s==%s" % (i.key, i.version) for i in installed_packages])
        print(installed_packages_list)
        from my_private_package import Something
        self.next(self.next_step)

    @pypi(
        packages={
            "gcsfs": "2024.6.0"
        }
    )
    @step
    def next_step(self):
        print("next_step")
        self.next(self.end)

    @step
    def end(self):
        print("yay!")

if __name__ == '__main__':
    PrivateLibFlow()

GOOGLE_APPLICATION_CREDENTIALS=/Users/maja/.config/gcloud/application_default_credentials.json  METAFLOW_DEBUG_CONDA=1 python3 workflows/minimal.py  --datastore=gs --environment=conda run