1 files changed, 17 insertions, 8 deletions

diff --git a/pyenc/enumerate_classes.py b/pyenc/enumerate_classes.py
index dfeb1dc..29cf2f7 100644
--- a/pyenc/enumerate_classes.py
+++ b/pyenc/enumerate_classes.py
@@ -114,6 +114,8 @@ def enumerate_files(path_base, environment):
     """
     path = os.path.join(path_base, environment.name)
     files = find(path, type='f', name='*.pp')
+    files = [f for f in find(path, type='f', name='*.pp')
+             if os.path.basename(os.path.dirname(f)) == 'manifests']
 
     try:
         for puppet_file in (model.PuppetFile(path=file) for file in files):
@@ -177,28 +179,35 @@ def run(path_base: Path = '/etc/puppetlabs/code/environments',
     enumerate_files(path_base, environment)
 
     ### Find all puppet files which we haven't parsed
-    base = model.db.session \
+    subexpr = model.db.session \
                    .query(model.PuppetFile.path,
                           model.PuppetFile.checksum,
-                          # Selects any of the availably environmentns. Since the checksum is the same
-                          # the file should also be the same, regardles of which environment we chose
-                          model.db.func.min(model.PuppetFile.environment_id)) \
+                          # Selects any of the availably environmentns. Since the checksum
+                          # is the same the file should also be the same, regardles of
+                          # which environment we chose
+                          model.db.func.min(model.PuppetFile.environment_id).label('env_id')) \
                    .outerjoin(model.PuppetFileContent,
                               model.PuppetFile.checksum == model.PuppetFileContent.checksum) \
                    .where(model.PuppetFileContent.json == None) \
                    .group_by(model.PuppetFile.checksum,
-                             model.PuppetFile.path)
+                             model.PuppetFile.path) \
+                   .cte()
+
+    base = model.db.session \
+                   .query(subexpr.c.path,
+                          subexpr.c.checksum,
+                          model.Environment.name) \
+                   .join(model.Environment,
+                         model.Environment.id == subexpr.c.env_id)
 
     files = base.all()
     count = base.count()
-    environments = {e.id: e.name for e in model.Environment.query.all()}
 
     db.session.commit()
 
     # Parse all puppet files, and store their output into pupet_file_content
     try:
-        for (i, (path, checksum, env_id)) in enumerate(files):
-            env = environments[env_id]
+        for (i, (path, checksum, env)) in enumerate(files):
             print(f'\x1b[2K{env} {path}')
             print(f'{i} / {count}', end='\r')