1 files changed, 3 insertions, 3 deletions
diff --git a/import_yaml.py b/import_yaml.py
index d602cad..0007e75 100755
--- a/import_yaml.py
+++ b/import_yaml.py
@@ -19,16 +19,16 @@ for fqdn, val in data.items():
     h = model.Host.query.where(model.Host.fqdn == fqdn).first()
     if not h:
         h = model.Host(fqdn=fqdn)
-    h.environment = model.PuppetEnvironment \
+    h.environment = model.Environment \
         .query \
-        .where(model.PuppetEnvironment.name == val.get('environment', 'production')) \
+        .where(model.Environment.name == val.get('environment', 'production')) \
         .one()
 
     classes = val['classes']
     if type(classes) == dict:
         classes = classes.keys()
     cls = model.PuppetClass.query \
-               .where(model.PuppetClass.class_name.in_(classes)).all()
+               .where(model.PuppetClass.name.in_(classes)).all()
     print(cls)
 
     h.classes.extend(cls)