Fix podcast queries

Fix param names in create_pod
AtomikLabs · May 24, 2024 · a6905b6 · a6905b6
1 parent 36746eb
commit a6905b6
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 10 deletions.
diff --git a/orchestration/airflow/dags/publishing/tasks/create_pod.py b/orchestration/airflow/dags/publishing/tasks/create_pod.py
@@ -374,17 +374,17 @@ def create_pod_node(
             pod.relate(
                 driver=driver,
                 label=CATEGORIZED_BY,
-                start_label=Podcast.label,
+                start_label=Podcast.LABEL,
                 start_uuid=pod.uuid,
-                end_label=ArxivCategory.label,
+                end_label=ArxivCategory.LABEL,
                 end_uuid=category_node.uuid,
             )
             category_node.relate(
                 driver=driver,
                 label=CATEGORIZES,
-                start_label=ArxivCategory.label,
+                start_label=ArxivCategory.LABEL,
                 start_uuid=category_node.uuid,
-                end_label=Podcast.label,
+                end_label=Podcast.LABEL,
                 end_uuid=pod.uuid,
             )
 
@@ -399,17 +399,17 @@ def create_pod_node(
                     pod.relate(
                         driver=driver,
                         label=PUBLISHES,
-                        start_label=Podcast.label,
+                        start_label=Podcast.LABEL,
                         start_uuid=pod.uuid,
-                        end_label=ArxivRecord.label,
+                        end_label=ArxivRecord.LABEL,
                         end_uuid=record_node.uuid,
                     )
                     record_node.relate(
                         driver=driver,
                         label=PUBLISHED_IN,
-                        start_label=ArxivRecord.label,
+                        start_label=ArxivRecord.LABEL,
                         start_uuid=record_node.uuid,
-                        end_label=Podcast.label,
+                        end_label=Podcast.LABEL,
                         end_uuid=pod.uuid,
                     )
                 except Exception as e:

diff --git a/orchestration/airflow/dags/shared/models/podcast.py b/orchestration/airflow/dags/shared/models/podcast.py
@@ -162,7 +162,8 @@ def create(
                     episode_date: $episode_date,
                     script_key: $script_key,
                     audio_key: $audio_key,
-                    created_at: $created_at
+                    created: $created,
+                    last_modified: $last_modified
                 })
             """
             parameters = {
@@ -177,7 +178,7 @@ def create(
                 "created": now,
                 "last_modified": now,
             }
-            self.execute_query(query, parameters)
+            self.driver.execute_query(query, parameters)
         except Exception as e:
             self.logger.error(f"Failed to create podcast: {e}", method=self.create.__name__)
             raise e
@@ -238,6 +239,35 @@ def find_all(cls, driver: Driver):
             logger.error("Error finding all podcasts", method=cls.findall.__name__, error=str(e))
             raise e
 
+    @classmethod
+    def load(self) -> bool:
+        try:
+            if not self.title or not self.season or not self.episode or not self.part or not self.episode_date:
+                raise ValueError("title, season, episode, part, and episode_date are required for loading")
+            records, _, _ = self.driver.execute_query(
+                "MATCH (p:Podcast {title: $title, season: $season, episode: $episode, part: $part, episode_date: $episode_date}) RETURN p",
+                title=self.title,
+                season=self.season,
+                episode=self.episode,
+                part=self.part,
+                episode_date=self.episode_date,
+                database_=self.db,
+            )
+            if records:
+                data = records[0].data().get("p", {})
+                self.uuid = data.get("uuid", "")
+                self.created = data.get("created", "")
+                self.last_modified = data.get("last_modified", "")
+                self.script_key = data.get("script_key", "")
+                self.audio_key = data.get("audio_key", "")
+
+                if not self.script_key or not self.audio_key:
+                    raise ValueError("Failed to properly load podcast")
+            return True if records else False
+        except Exception as e:
+            logger.error("Error loading podcast", pod=self, method=self.load.__name__, error=str(e))
+            raise e
+
     def relate(
         self,
         driver: Driver,