refactor: make timestamp in local tz

TEAMSchools · Jan 11, 2024 · a8c4dde · a8c4dde
1 parent 11706b1
commit a8c4dde
Show file tree

Hide file tree

Showing 12 changed files with 40 additions and 16 deletions.
diff --git a/src/teamster/core/powerschool/assets.py b/src/teamster/core/powerschool/assets.py
@@ -19,6 +19,7 @@
 def build_powerschool_table_asset(
     code_location,
     asset_name,
+    local_timezone,
     partitions_def: (
         FiscalYearPartitionsDefinition | MonthlyPartitionsDefinition | None
     ) = None,
@@ -40,7 +41,7 @@ def _asset(
         ssh_powerschool: SSHResource,
         db_powerschool: OracleResource,
     ):
-        now = pendulum.now().start_of("hour")
+        now = pendulum.now(tz=local_timezone).start_of("hour")
 
         asset_metadata = context.assets_def.metadata_by_key[context.assets_def.key]
 

diff --git a/src/teamster/core/powerschool/schedules.py b/src/teamster/core/powerschool/schedules.py
@@ -74,10 +74,8 @@ def _schedule(
                         "partition_column"
                     ]
 
-                    latest_materialization_fmt = (
-                        latest_materialization_datetime.in_timezone(
-                            execution_timezone
-                        ).format("YYYY-MM-DDTHH:mm:ss.SSSSSS")
+                    latest_materialization_fmt = latest_materialization_datetime.format(
+                        "YYYY-MM-DDTHH:mm:ss.SSSSSS"
                     )
 
                     [(count,)] = db_powerschool.engine.execute_query(

diff --git a/src/teamster/core/powerschool/sensors.py b/src/teamster/core/powerschool/sensors.py
@@ -61,7 +61,7 @@ def _sensor(
                     )  # type: ignore
                 )
 
-                window_start_fmt = latest_materialization_datetime.format(
+                latest_materialization_fmt = latest_materialization_datetime.format(
                     "YYYY-MM-DDTHH:mm:ss.SSSSSS"
                 )
 
@@ -71,7 +71,7 @@ def _sensor(
                         "SELECT COUNT(*) "
                         f"FROM {asset.key.path[-1]} "
                         f"WHERE {partition_column} >= "
-                        f"TO_TIMESTAMP('{window_start_fmt}', "
+                        f"TO_TIMESTAMP('{latest_materialization_fmt}', "
                         "'YYYY-MM-DD\"T\"HH24:MI:SS.FF6')"
                     ),
                     partition_size=1,

diff --git a/src/teamster/kippcamden/__init__.py b/src/teamster/kippcamden/__init__.py
@@ -1,10 +1,11 @@
 import pendulum
+from pendulum.tz import timezone
 
 from teamster.core.utils.classes import FiscalYear
 
 CODE_LOCATION = "kippcamden"
 
-LOCAL_TIMEZONE = pendulum.timezone(name="America/New_York")
+LOCAL_TIMEZONE = timezone(name="America/New_York")
 NOW = pendulum.now(tz=LOCAL_TIMEZONE)
 TODAY = NOW.start_of(unit="day")
 CURRENT_FISCAL_YEAR = FiscalYear(datetime=TODAY, start_month=7)
diff --git a/src/teamster/kippcamden/powerschool/assets.py b/src/teamster/kippcamden/powerschool/assets.py
@@ -16,6 +16,7 @@
     build_powerschool_table_asset(
         **asset,
         code_location=CODE_LOCATION,
+        local_timezone=LOCAL_TIMEZONE,
         op_tags={MAX_RUNTIME_SECONDS_TAG: (60 * 10)},
     )
     for asset in config_from_files([(f"{config_dir}/assets-full.yaml")])["assets"]
@@ -25,6 +26,7 @@
     build_powerschool_table_asset(
         **asset,
         code_location=CODE_LOCATION,
+        local_timezone=LOCAL_TIMEZONE,
         op_tags={MAX_RUNTIME_SECONDS_TAG: (60 * 10)},
     )
     for asset in config_from_files([(f"{config_dir}/assets-nonpartition.yaml")])[
@@ -36,6 +38,7 @@
     build_powerschool_table_asset(
         **asset,
         code_location=CODE_LOCATION,
+        local_timezone=LOCAL_TIMEZONE,
         partitions_def=FiscalYearPartitionsDefinition(
             start_date=pendulum.datetime(year=2016, month=7, day=1),
             start_month=7,
@@ -54,6 +57,7 @@
     build_powerschool_table_asset(
         **asset,
         code_location=CODE_LOCATION,
+        local_timezone=LOCAL_TIMEZONE,
         partitions_def=MonthlyPartitionsDefinition(
             start_date=pendulum.datetime(year=2016, month=7, day=1),
             timezone=LOCAL_TIMEZONE.name,

diff --git a/src/teamster/kippmiami/__init__.py b/src/teamster/kippmiami/__init__.py
@@ -1,10 +1,11 @@
 import pendulum
+from pendulum.tz import timezone
 
 from teamster.core.utils.classes import FiscalYear
 
 CODE_LOCATION = "kippmiami"
 
-LOCAL_TIMEZONE = pendulum.timezone(name="America/New_York")
+LOCAL_TIMEZONE = timezone(name="America/New_York")
 NOW = pendulum.now(tz=LOCAL_TIMEZONE)
 TODAY = NOW.start_of(unit="day")
 CURRENT_FISCAL_YEAR = FiscalYear(datetime=TODAY, start_month=7)
diff --git a/src/teamster/kippmiami/powerschool/assets.py b/src/teamster/kippmiami/powerschool/assets.py
@@ -13,13 +13,18 @@
 config_dir = f"src/teamster/{CODE_LOCATION}/powerschool/config"
 
 full_assets = [
-    build_powerschool_table_asset(**cfg, code_location=CODE_LOCATION)
+    build_powerschool_table_asset(
+        **cfg, code_location=CODE_LOCATION, local_timezone=LOCAL_TIMEZONE
+    )
     for cfg in config_from_files([(f"{config_dir}/assets-full.yaml")])["assets"]
 ]
 
 nonpartition_assets = [
     build_powerschool_table_asset(
-        **cfg, code_location=CODE_LOCATION, op_tags={MAX_RUNTIME_SECONDS_TAG: (60 * 10)}
+        **cfg,
+        code_location=CODE_LOCATION,
+        local_timezone=LOCAL_TIMEZONE,
+        op_tags={MAX_RUNTIME_SECONDS_TAG: (60 * 10)},
     )
     for cfg in config_from_files([(f"{config_dir}/assets-nonpartition.yaml")])["assets"]
 ]
@@ -28,6 +33,7 @@
     build_powerschool_table_asset(
         **asset,
         code_location=CODE_LOCATION,
+        local_timezone=LOCAL_TIMEZONE,
         partitions_def=FiscalYearPartitionsDefinition(
             start_date=pendulum.datetime(year=2018, month=7, day=1),
             start_month=7,
@@ -46,6 +52,7 @@
     build_powerschool_table_asset(
         **asset,
         code_location=CODE_LOCATION,
+        local_timezone=LOCAL_TIMEZONE,
         partitions_def=MonthlyPartitionsDefinition(
             start_date=pendulum.datetime(year=2018, month=7, day=1),
             timezone=LOCAL_TIMEZONE.name,

diff --git a/src/teamster/kippnewark/__init__.py b/src/teamster/kippnewark/__init__.py
@@ -1,10 +1,11 @@
 import pendulum
+from pendulum.tz import timezone
 
 from teamster.core.utils.classes import FiscalYear
 
 CODE_LOCATION = "kippnewark"
 
-LOCAL_TIMEZONE = pendulum.timezone(name="America/New_York")
+LOCAL_TIMEZONE = timezone(name="America/New_York")
 NOW = pendulum.now(tz=LOCAL_TIMEZONE)
 TODAY = NOW.start_of(unit="day")
 CURRENT_FISCAL_YEAR = FiscalYear(datetime=TODAY, start_month=7)
diff --git a/src/teamster/kippnewark/powerschool/assets.py b/src/teamster/kippnewark/powerschool/assets.py
@@ -14,14 +14,20 @@
 
 full_assets = [
     build_powerschool_table_asset(
-        **cfg, code_location=CODE_LOCATION, op_tags={MAX_RUNTIME_SECONDS_TAG: (60 * 10)}
+        **cfg,
+        code_location=CODE_LOCATION,
+        local_timezone=LOCAL_TIMEZONE,
+        op_tags={MAX_RUNTIME_SECONDS_TAG: (60 * 10)},
     )
     for cfg in config_from_files([(f"{config_dir}/assets-full.yaml")])["assets"]
 ]
 
 nonpartition_assets = [
     build_powerschool_table_asset(
-        **cfg, code_location=CODE_LOCATION, op_tags={MAX_RUNTIME_SECONDS_TAG: (60 * 10)}
+        **cfg,
+        code_location=CODE_LOCATION,
+        local_timezone=LOCAL_TIMEZONE,
+        op_tags={MAX_RUNTIME_SECONDS_TAG: (60 * 10)},
     )
     for cfg in config_from_files([(f"{config_dir}/assets-nonpartition.yaml")])["assets"]
 ]
@@ -30,6 +36,7 @@
     build_powerschool_table_asset(
         **asset,
         code_location=CODE_LOCATION,
+        local_timezone=LOCAL_TIMEZONE,
         partitions_def=FiscalYearPartitionsDefinition(
             start_date=pendulum.datetime(year=2016, month=7, day=1),
             start_month=7,
@@ -48,6 +55,7 @@
     build_powerschool_table_asset(
         **asset,
         code_location=CODE_LOCATION,
+        local_timezone=LOCAL_TIMEZONE,
         partitions_def=MonthlyPartitionsDefinition(
             start_date=pendulum.datetime(year=2016, month=7, day=1),
             timezone=LOCAL_TIMEZONE.name,

diff --git a/src/teamster/kipptaf/__init__.py b/src/teamster/kipptaf/__init__.py
@@ -1,10 +1,11 @@
 import pendulum
+from pendulum.tz import timezone
 
 from teamster.core.utils.classes import FiscalYear
 
 CODE_LOCATION = "kipptaf"
 
-LOCAL_TIMEZONE = pendulum.timezone(name="America/New_York")
+LOCAL_TIMEZONE = timezone(name="America/New_York")
 NOW = pendulum.now(tz=LOCAL_TIMEZONE)
 TODAY = NOW.start_of(unit="day")
 CURRENT_FISCAL_YEAR = FiscalYear(datetime=TODAY, start_month=7)
diff --git a/src/teamster/staging/__init__.py b/src/teamster/staging/__init__.py
@@ -1,10 +1,11 @@
 import pendulum
+from pendulum.tz import timezone
 
 from teamster.core.utils.classes import FiscalYear
 
 CODE_LOCATION = "staging"
 
-LOCAL_TIMEZONE = pendulum.timezone(name="America/New_York")
+LOCAL_TIMEZONE = timezone(name="America/New_York")
 NOW = pendulum.now(tz=LOCAL_TIMEZONE)
 TODAY = NOW.start_of(unit="day")
 CURRENT_FISCAL_YEAR = FiscalYear(datetime=TODAY, start_month=7)
diff --git a/tests/assets/test_powerschool_assets.py b/tests/assets/test_powerschool_assets.py
@@ -15,6 +15,7 @@
 def _test_asset(asset_name, partitions_def=None, partition_column=None):
     asset = build_powerschool_table_asset(
         code_location="staging",
+        local_timezone=LOCAL_TIMEZONE,
         asset_name=asset_name,
         partitions_def=partitions_def,
         partition_column=partition_column,